Improve preventing concurrent calls to handler

Author: viceroypenguin

src/Immediate.Cache.Shared/ApplicationCacheBase.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics.CodeAnalysis;
 using Immediate.Handlers.Shared;
 using Microsoft.Extensions.Caching.Memory;
 
@@ -113,22 +114,20 @@ public void SetValue(TRequest request, TResponse value) =>
 	public void RemoveValue(TRequest request) =>
 		GetCacheValue(request).RemoveValue();
 
+	[SuppressMessage("Design", "CA1001:Types that own disposable fields should be disposable", Justification = "CancellationTokenSource does not need to be disposed here.")]
 	private sealed class CacheValue(
 		TRequest request,
 		Owned<IHandler<TRequest, TResponse>> handler
 	)
 	{
-		private TResponse? _response;
 		private CancellationTokenSource? _tokenSource;
+		private TaskCompletionSource<TResponse>? _responseSource;
 		private readonly object _lock = new();
 
 		public async ValueTask<TResponse> GetValue()
 		{
-			if (_response is not null)
-				return _response;
-
-			lock (_lock)
-				_tokenSource ??= new CancellationTokenSource();
+			if (!TryAcquireResponseSource())
+				return await _responseSource.Task.ConfigureAwait(false);
 
 			var token = _tokenSource.Token;
 
@@ -146,20 +145,43 @@ public async ValueTask<TResponse> GetValue()
 						.ConfigureAwait(false);
 
 					lock (_lock)
-						return _response ??= response;
+					{
+						_responseSource.SetResult(response);
+						return response;
+					}
 				}
 			}
-			catch (OperationCanceledException) when (_response 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()
+		{
+			if (_responseSource is not null)
+				return false;
+
+			lock (_lock)
 			{
-				return _response;
+				if (_responseSource is not null)
+					return false;
+
+				_tokenSource = new CancellationTokenSource();
+				_responseSource = new TaskCompletionSource<TResponse>();
+				return true;
 			}
 		}
 
 		public void SetValue(TResponse response)
 		{
 			lock (_lock)
 			{
-				_response = response;
+				_responseSource = new TaskCompletionSource<TResponse>();
+				_responseSource.SetResult(response);
 				_tokenSource?.Cancel();
 			}
 		}
@@ -168,7 +190,7 @@ public void RemoveValue()
 		{
 			lock (_lock)
 			{
-				_response = null;
+				_responseSource = null;
 				_tokenSource?.Cancel();
 			}
 		}

tests/Immediate.Cache.FunctionalTests/ApplicationCacheTests.cs

@@ -13,6 +13,7 @@ public ApplicationCacheTests()
 		var services = new ServiceCollection();
 		_ = services.AddHandlers();
 		_ = services.AddSingleton<GetValueCache>();
+		_ = services.AddSingleton<DelayGetValueCache>();
 		_ = services.AddSingleton(typeof(Owned<>));
 		_ = services.AddMemoryCache();
 
@@ -71,4 +72,54 @@ public async Task RemoveValueRemovesValue()
 		Assert.Equal(3, response.Value);
 		Assert.True(response.ExecutedHandler);
 	}
+
+	[Test]
+	public async Task SimultaneousAccessIsSerialized()
+	{
+		var request1 = new DelayGetValue.Query()
+		{
+			Value = 1,
+			Name = "Request1",
+			CompletionSource = new(),
+		};
+
+		var request2 = new DelayGetValue.Query()
+		{
+			Value = 1,
+			Name = "Request2",
+			CompletionSource = new(),
+		};
+
+		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
+		var response1Task = cache.GetValue(request1);
+		var response2Task = cache.GetValue(request2);
+
+		// both waiting until tcs triggered
+		Assert.False(response1Task.IsCompleted);
+		Assert.False(response2Task.IsCompleted);
+
+		Assert.Equal(0, request1.TimesExecuted);
+		Assert.Equal(0, request2.TimesExecuted);
+
+		// request2 does nothing at this point
+		request2.CompletionSource.SetResult();
+
+		Assert.False(response1Task.IsCompleted);
+		Assert.False(response2Task.IsCompleted);
+
+		Assert.Equal(0, request1.TimesExecuted);
+		Assert.Equal(0, request2.TimesExecuted);
+
+		// trigger request1, which should run exactly once
+		request1.CompletionSource.SetResult();
+
+		var response1 = await response1Task;
+		var response2 = await response2Task;
+
+		Assert.Equal(1, request1.TimesExecuted);
+		Assert.Equal(0, request2.TimesExecuted);
+
+		// ensure both responses get the same response back
+		Assert.Equal(response1.RandomValue, response2.RandomValue);
+	}
 }

tests/Immediate.Cache.FunctionalTests/DelayGetValue.cs

@@ -0,0 +1,31 @@
+using Immediate.Handlers.Shared;
+
+namespace Immediate.Cache.FunctionalTests;
+
+[Handler]
+public static partial class DelayGetValue
+{
+	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 sealed record Response(int Value, bool ExecutedHandler, Guid RandomValue);
+
+	private static readonly object s_lock = new();
+
+	private static async ValueTask<Response> HandleAsync(
+		Query query,
+		CancellationToken _
+	)
+	{
+		await query.CompletionSource.Task;
+		lock (s_lock)
+			query.TimesExecuted++;
+
+		return new(query.Value, ExecutedHandler: true, RandomValue: Guid.NewGuid());
+	}
+}

tests/Immediate.Cache.FunctionalTests/DelayGetValueCache.cs

@@ -0,0 +1,22 @@
+using System.Diagnostics.CodeAnalysis;
+using Immediate.Handlers.Shared;
+using Microsoft.Extensions.Caching.Memory;
+
+namespace Immediate.Cache.FunctionalTests;
+
+public sealed class DelayGetValueCache(
+	IMemoryCache memoryCache,
+	Owned<IHandler<DelayGetValue.Query, DelayGetValue.Response>> ownedHandler
+) : ApplicationCacheBase<DelayGetValue.Query, DelayGetValue.Response>(
+	memoryCache,
+	ownedHandler
+)
+{
+	[SuppressMessage(
+		"Design",
+		"CA1062:Validate arguments of public methods",
+		Justification = "Not a public method"
+	)]
+	protected override string TransformKey(DelayGetValue.Query request) =>
+		$"GetValue(query: {request.Value})";
+}