Address concurrency bug (#23)

Author: viceroypenguin

src/Immediate.Cache.Shared/ApplicationCacheBase.cs

@@ -133,27 +133,42 @@ public async ValueTask<TResponse> GetValue(CancellationToken cancellationToken)
 		[SuppressMessage("Maintainability", "CA1508:Avoid dead conditional code", Justification = "Double-checked lock pattern")]
 		private Task<TResponse> GetHandlerTask()
 		{
-			if (_responseSource is not null)
+			if (_responseSource is { Task.Status: not (TaskStatus.Faulted or TaskStatus.Canceled) })
 				return _responseSource.Task;
 
 			lock (_lock)
 			{
-				if (_responseSource is not null)
+				if (_responseSource is { Task.Status: not (TaskStatus.Faulted or TaskStatus.Canceled) })
 					return _responseSource.Task;
 
-				var ts = _tokenSource = new();
-				_responseSource = new();
+				// escape current sync context
+				_ = Task.Factory.StartNew(
+					RunHandler,
+					CancellationToken.None,
+					TaskCreationOptions.PreferFairness,
+					TaskScheduler.Current
+				);
 
-				return Task.Run(() => RunHandler(ts));
+				return (_responseSource = new()).Task;
 			}
 		}
 
-		private async Task<TResponse> RunHandler(CancellationTokenSource tokenSource)
+		private async Task RunHandler()
 		{
 			while (true)
 			{
-				if (_responseSource?.Task is { IsCompletedSuccessfully: true } task)
-					return await task.ConfigureAwait(false);
+				CancellationTokenSource tokenSource;
+
+				lock (_lock)
+				{
+					if (_responseSource?.Task is { IsCompletedSuccessfully: true })
+						return;
+
+					if (_tokenSource is null or { IsCancellationRequested: true })
+						_tokenSource = new();
+
+					tokenSource = _tokenSource;
+				}
 
 				try
 				{
@@ -171,26 +186,25 @@ private async Task<TResponse> RunHandler(CancellationTokenSource tokenSource)
 
 						lock (_lock)
 						{
-							if (!token.IsCancellationRequested)
-							{
-								var rs = _responseSource ??= new();
-								rs.SetResult(response);
-
-								return response;
-							}
+							if (!tokenSource.IsCancellationRequested)
+								_responseSource!.SetResult(response);
 						}
 					}
 				}
 				catch (OperationCanceledException) when (tokenSource.IsCancellationRequested)
 				{
 				}
-
-				lock (_lock)
+#pragma warning disable CA1031 // Do not catch general exception types
+				// no one is listening to `RunHandler`; return the exception via `SetException`
+				catch (Exception ex)
+#pragma warning restore CA1031
 				{
-					if (_tokenSource is null or { IsCancellationRequested: true })
-						_tokenSource = new();
-
-					tokenSource = _tokenSource;
+					lock (_lock)
+					{
+						if (!tokenSource.IsCancellationRequested)
+							_responseSource?.SetException(ex);
+						return;
+					}
 				}
 			}
 		}
@@ -199,8 +213,11 @@ public void SetValue(TResponse response)
 		{
 			lock (_lock)
 			{
-				_responseSource = new TaskCompletionSource<TResponse>();
+				if (_responseSource is null or { Task.IsCompleted: true })
+					_responseSource = new TaskCompletionSource<TResponse>();
+
 				_responseSource.SetResult(response);
+
 				_tokenSource?.Cancel();
 				_tokenSource = null;
 			}
@@ -210,7 +227,9 @@ public void RemoveValue()
 		{
 			lock (_lock)
 			{
-				_responseSource = null;
+				if (_responseSource is { Task.IsCompleted: true })
+					_responseSource = null;
+
 				_tokenSource?.Cancel();
 				_tokenSource = null;
 			}

tests/Immediate.Cache.FunctionalTests/ApplicationCacheTests.cs

@@ -80,14 +80,12 @@ public async Task SimultaneousAccessIsSerialized()
 		{
 			Value = 1,
 			Name = "Request1",
-			CompletionSource = new(),
 		};
 
 		var request2 = new DelayGetValue.Query()
 		{
 			Value = 1,
 			Name = "Request2",
-			CompletionSource = new(),
 		};
 
 		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
@@ -102,7 +100,7 @@ public async Task SimultaneousAccessIsSerialized()
 		Assert.Equal(0, request2.TimesExecuted);
 
 		// request2 does nothing at this point
-		request2.CompletionSource.SetResult();
+		request2.WaitForTestToContinueOperation.SetResult();
 
 		Assert.False(response1Task.IsCompleted);
 		Assert.False(response2Task.IsCompleted);
@@ -111,7 +109,7 @@ public async Task SimultaneousAccessIsSerialized()
 		Assert.Equal(0, request2.TimesExecuted);
 
 		// trigger request1, which should run exactly once
-		request1.CompletionSource.SetResult();
+		request1.WaitForTestToContinueOperation.SetResult();
 
 		var response1 = await response1Task;
 		var response2 = await response2Task;
@@ -130,7 +128,6 @@ public async Task ProperlyUsesCancellationToken()
 		{
 			Value = 1,
 			Name = "Request1",
-			CompletionSource = new(),
 		};
 
 		using var tcs = new CancellationTokenSource();
@@ -144,14 +141,13 @@ public async Task ProperlyUsesCancellationToken()
 		Assert.False(request.CancellationToken.IsCancellationRequested);
 
 		// actual handler will continue executing in spite of no remaining callers
-		request.CompletionSource.SetResult();
+		request.WaitForTestToContinueOperation.SetResult();
 
 		// check that value is now properly in cache
 		var request2 = new DelayGetValue.Query()
 		{
 			Value = 1,
 			Name = "Request2",
-			CompletionSource = new(),
 		};
 
 		var response = await cache.GetValue(request2);
@@ -168,15 +164,13 @@ public async Task CancellingFirstAccessOperatesCorrectly()
 		{
 			Value = 1,
 			Name = "Request1",
-			CompletionSource = new(),
 		};
 
 		using var cts2 = new CancellationTokenSource();
 		var request2 = new DelayGetValue.Query()
 		{
 			Value = 1,
 			Name = "Request2",
-			CompletionSource = new(),
 		};
 
 		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
@@ -210,15 +204,13 @@ public async Task CancellingSecondAccessOperatesCorrectly()
 		{
 			Value = 1,
 			Name = "Request1",
-			CompletionSource = new(),
 		};
 
 		using var cts2 = new CancellationTokenSource();
 		var request2 = new DelayGetValue.Query()
 		{
 			Value = 1,
 			Name = "Request2",
-			CompletionSource = new(),
 		};
 
 		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
@@ -251,19 +243,18 @@ public async Task RemovingValueCancelsExistingOperation()
 		{
 			Value = 1,
 			Name = "Request1",
-			CompletionSource = new(),
 		};
 
 		using var tcs = new CancellationTokenSource();
 		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
 		var responseTask = cache.GetValue(request, tcs.Token);
 
-		cache.RemoveValue(request);
+		await request.WaitForTestToStartExecuting.Task;
 
-		// allow IC task to be run
-		await Task.Delay(10);
+		cache.RemoveValue(request);
+		await request.WaitForTestToFinalize.Task;
 
-		request.CompletionSource.SetResult();
+		request.WaitForTestToContinueOperation.SetResult();
 
 		var response = await responseTask;
 
@@ -281,22 +272,39 @@ public async Task SettingValueCancelsExistingOperation()
 		{
 			Value = 1,
 			Name = "Request1",
-			CompletionSource = new(),
 		};
 
 		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
 		var responseTask = cache.GetValue(request, default);
 
-		// allow IC task to be run
-		await Task.Delay(10);
+		await request.WaitForTestToStartExecuting.Task;
 
 		cache.SetValue(request, new(5, ExecutedHandler: false, Guid.NewGuid()));
 
 		var response = await responseTask;
 		Assert.Equal(5, response.Value);
 		Assert.False(response.ExecutedHandler);
 
+		await request.WaitForTestToFinalize.Task;
+
 		Assert.Equal(0, request.TimesExecuted);
 		Assert.Equal(1, request.TimesCancelled);
 	}
+
+	[Test]
+	public async Task ExceptionGetsPropagatedCorrectly()
+	{
+		var request = new DelayGetValue.Query()
+		{
+			Value = 1,
+			Name = "Request1",
+			ThrowException = true,
+		};
+
+		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
+		var responseTask = cache.GetValue(request, default);
+
+		var ex = await Assert.ThrowsAsync<InvalidOperationException>(async () => await responseTask);
+		Assert.Equal("Test Exception 1", ex.Message);
+	}
 }

tests/Immediate.Cache.FunctionalTests/DelayGetValue.cs

@@ -9,10 +9,14 @@ public sealed class Query
 	{
 		public required int Value { get; init; }
 		public required string Name { get; init; }
-		public required TaskCompletionSource CompletionSource { get; init; }
 		public int TimesExecuted { get; set; }
 		public int TimesCancelled { get; set; }
 		public CancellationToken CancellationToken { get; set; }
+		public bool ThrowException { get; init; }
+
+		public TaskCompletionSource WaitForTestToContinueOperation { get; } = new();
+		public TaskCompletionSource WaitForTestToStartExecuting { get; } = new();
+		public TaskCompletionSource WaitForTestToFinalize { get; } = new();
 	}
 
 	public sealed record Response(int Value, bool ExecutedHandler, Guid RandomValue);
@@ -24,10 +28,14 @@ private static async ValueTask<Response> HandleAsync(
 		CancellationToken token
 	)
 	{
+		if (query.ThrowException)
+			throw new InvalidOperationException($"Test Exception {query.Value}");
+
 		try
 		{
 			query.CancellationToken = token;
-			await query.CompletionSource.Task.WaitAsync(token);
+			_ = query.WaitForTestToStartExecuting.TrySetResult();
+			await query.WaitForTestToContinueOperation.Task.WaitAsync(token);
 
 			lock (s_lock)
 				query.TimesExecuted++;
@@ -40,5 +48,9 @@ CancellationToken token
 				query.TimesCancelled++;
 			throw;
 		}
+		finally
+		{
+			_ = query.WaitForTestToFinalize.TrySetResult();
+		}
 	}
 }