Improve behavior for concurrent access

Author: viceroypenguin

src/Immediate.Cache.Shared/ApplicationCacheBase.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using Immediate.Handlers.Shared;
 using Microsoft.Extensions.Caching.Memory;
@@ -129,11 +130,42 @@ Owned<IHandler<TRequest, TResponse>> handler
 
 		public async ValueTask<TResponse> GetValue(CancellationToken cancellationToken)
 		{
-			if (!TryAcquireResponseSource(cancellationToken))
-				return await _responseSource.Task.ConfigureAwait(false);
+			try
+			{
+				return await GetHandlerTask().WaitAsync(cancellationToken).ConfigureAwait(false);
+			}
+			catch (OperationCanceledException) when (
+				!cancellationToken.IsCancellationRequested
+				&& _responseSource?.Task is { IsCompletedSuccessfully: true } task
+			)
+			{
+				return await task.ConfigureAwait(false);
+			}
+		}
 
-			var token = _tokenSource.Token;
+		[SuppressMessage("Maintainability", "CA1508:Avoid dead conditional code", Justification = "Double-checked lock pattern")]
+		private Task<TResponse> GetHandlerTask()
+		{
+			if (_responseSource is not null)
+				return _responseSource.Task;
+
+			lock (_lock)
+			{
+				if (_responseSource is not null)
+					return _responseSource.Task;
+
+				_tokenSource = new();
+				_responseSource = new TaskCompletionSource<TResponse>();
 
+				return Task.Run(() => RunHandler(_tokenSource, _responseSource));
+			}
+		}
+
+		private async Task<TResponse> RunHandler(
+			CancellationTokenSource tokenSource,
+			TaskCompletionSource<TResponse> responseSource
+		)
+		{
 			try
 			{
 				var scope = handler.GetScope();
@@ -143,42 +175,25 @@ public async ValueTask<TResponse> GetValue(CancellationToken cancellationToken)
 					var response = await scope.Service
 						.HandleAsync(
 							request,
-							cancellationToken: token
+							tokenSource.Token
 						)
 						.ConfigureAwait(false);
 
 					lock (_lock)
 					{
-						_responseSource.SetResult(response);
+						Debug.Assert(_responseSource == responseSource);
+						responseSource.SetResult(response);
+
 						return response;
 					}
 				}
 			}
-			catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && _responseSource is not null)
+			catch (OperationCanceledException) when (_responseSource is not null)
 			{
 				return await _responseSource.Task.ConfigureAwait(false);
 			}
 		}
 
-		[MemberNotNull(nameof(_responseSource))]
-		[MemberNotNullWhen(true, nameof(_tokenSource))]
-		[SuppressMessage("Maintainability", "CA1508:Avoid dead conditional code", Justification = "Double-checked lock pattern")]
-		private bool TryAcquireResponseSource(CancellationToken cancellationToken)
-		{
-			if (_responseSource is not null)
-				return false;
-
-			lock (_lock)
-			{
-				if (_responseSource is not null)
-					return false;
-
-				_tokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-				_responseSource = new TaskCompletionSource<TResponse>();
-				return true;
-			}
-		}
-
 		public void SetValue(TResponse response)
 		{
 			lock (_lock)

tests/Immediate.Cache.FunctionalTests/ApplicationCacheTests.cs

@@ -138,8 +138,91 @@ public async Task ProperlyUsesCancellationToken()
 		var responseTask = cache.GetValue(request, tcs.Token);
 
 		await tcs.CancelAsync();
-		request.CompletionSource.SetResult();
 
 		Assert.True(responseTask.IsCanceled);
 	}
+
+	[Test]
+	public async Task CancellingFirstAccessOperatesCorrectly()
+	{
+		using var cts1 = new CancellationTokenSource();
+		var request1 = new DelayGetValue.Query()
+		{
+			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>();
+		var response1Task = cache.GetValue(request1, cts1.Token);
+		var response2Task = cache.GetValue(request2, cts2.Token);
+
+		// both waiting until cancellation triggered
+		Assert.False(response1Task.IsCompleted);
+		Assert.False(response2Task.IsCompleted);
+
+		Assert.Equal(0, request1.TimesExecuted);
+		Assert.Equal(0, request2.TimesExecuted);
+
+		// cancel query1; query2 should remain uncancelled
+		await cts1.CancelAsync();
+
+		Assert.True(response1Task.IsCanceled);
+		Assert.False(response2Task.IsCanceled);
+
+		await cts2.CancelAsync();
+
+		Assert.True(response1Task.IsCanceled);
+		Assert.True(response2Task.IsCanceled);
+	}
+
+	[Test]
+	public async Task CancellingSecondAccessOperatesCorrectly()
+	{
+		using var cts1 = new CancellationTokenSource();
+		var request1 = new DelayGetValue.Query()
+		{
+			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>();
+		var response1Task = cache.GetValue(request1, cts1.Token);
+		var response2Task = cache.GetValue(request2, cts2.Token);
+
+		// both waiting until cancellation triggered
+		Assert.False(response1Task.IsCompleted);
+		Assert.False(response2Task.IsCompleted);
+
+		Assert.Equal(0, request1.TimesExecuted);
+		Assert.Equal(0, request2.TimesExecuted);
+
+		// cancel query2; query1 should remain uncancelled
+		await cts2.CancelAsync();
+
+		Assert.False(response1Task.IsCanceled);
+		Assert.True(response2Task.IsCanceled);
+
+		await cts1.CancelAsync();
+
+		Assert.True(response1Task.IsCanceled);
+		Assert.True(response2Task.IsCanceled);
+	}
 }

tests/Immediate.Cache.FunctionalTests/DelayGetValue.cs

@@ -11,7 +11,6 @@ public sealed class Query
 		public required string Name { get; init; }
 		public required TaskCompletionSource CompletionSource { get; init; }
 		public int TimesExecuted { get; set; }
-		public CancellationToken CancellationToken { get; set; }
 	}
 
 	public sealed record Response(int Value, bool ExecutedHandler, Guid RandomValue);
@@ -23,9 +22,7 @@ private static async ValueTask<Response> HandleAsync(
 		CancellationToken token
 	)
 	{
-		query.CancellationToken = token;
-		await query.CompletionSource.Task;
-		token.ThrowIfCancellationRequested();
+		await query.CompletionSource.Task.WaitAsync(token);
 
 		lock (s_lock)
 			query.TimesExecuted++;