Improve remove behavior (#17)

Author: viceroypenguin

src/Immediate.Cache.Shared/ApplicationCacheBase.cs

@@ -1,4 +1,3 @@
-using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using Immediate.Handlers.Shared;
 using Microsoft.Extensions.Caching.Memory;
@@ -128,20 +127,8 @@ Owned<IHandler<TRequest, TResponse>> handler
 		private TaskCompletionSource<TResponse>? _responseSource;
 		private readonly Lock _lock = new();
 
-		public async ValueTask<TResponse> GetValue(CancellationToken cancellationToken)
-		{
-			try
-			{
-				return await GetHandlerTask().WaitAsync(cancellationToken).ConfigureAwait(false);
-			}
-			catch (OperationCanceledException) when (
-				!cancellationToken.IsCancellationRequested
-				&& _responseSource?.Task is { IsCompletedSuccessfully: true } task
-			)
-			{
-				return await task.ConfigureAwait(false);
-			}
-		}
+		public async ValueTask<TResponse> GetValue(CancellationToken cancellationToken) =>
+			await GetHandlerTask().WaitAsync(cancellationToken).ConfigureAwait(false);
 
 		[SuppressMessage("Maintainability", "CA1508:Avoid dead conditional code", Justification = "Double-checked lock pattern")]
 		private Task<TResponse> GetHandlerTask()
@@ -154,43 +141,57 @@ private Task<TResponse> GetHandlerTask()
 				if (_responseSource is not null)
 					return _responseSource.Task;
 
-				_tokenSource = new();
-				_responseSource = new TaskCompletionSource<TResponse>();
+				var ts = _tokenSource = new();
+				_responseSource = new();
 
-				return Task.Run(() => RunHandler(_tokenSource, _responseSource));
+				return Task.Run(() => RunHandler(ts));
 			}
 		}
 
-		private async Task<TResponse> RunHandler(
-			CancellationTokenSource tokenSource,
-			TaskCompletionSource<TResponse> responseSource
-		)
+		private async Task<TResponse> RunHandler(CancellationTokenSource tokenSource)
 		{
-			try
+			while (true)
 			{
-				var scope = handler.GetScope();
+				if (_responseSource?.Task is { IsCompletedSuccessfully: true } task)
+					return await task.ConfigureAwait(false);
 
-				await using (scope.ConfigureAwait(false))
+				try
 				{
-					var response = await scope.Service
-						.HandleAsync(
-							request,
-							tokenSource.Token
-						)
-						.ConfigureAwait(false);
-
-					lock (_lock)
-					{
-						Debug.Assert(_responseSource == responseSource);
-						responseSource.SetResult(response);
+					var token = tokenSource.Token;
+					var scope = handler.GetScope();
 
-						return response;
+					await using (scope.ConfigureAwait(false))
+					{
+						var response = await scope.Service
+							.HandleAsync(
+								request,
+								token
+							)
+							.ConfigureAwait(false);
+
+						lock (_lock)
+						{
+							if (!token.IsCancellationRequested)
+							{
+								var rs = _responseSource ??= new();
+								rs.SetResult(response);
+
+								return response;
+							}
+						}
 					}
 				}
-			}
-			catch (OperationCanceledException) when (_responseSource is not null)
-			{
-				return await _responseSource.Task.ConfigureAwait(false);
+				catch (OperationCanceledException) when (tokenSource.IsCancellationRequested)
+				{
+				}
+
+				lock (_lock)
+				{
+					if (_tokenSource is null or { IsCancellationRequested: true })
+						_tokenSource = new();
+
+					tokenSource = _tokenSource;
+				}
 			}
 		}
 
@@ -201,6 +202,7 @@ public void SetValue(TResponse response)
 				_responseSource = new TaskCompletionSource<TResponse>();
 				_responseSource.SetResult(response);
 				_tokenSource?.Cancel();
+				_tokenSource = null;
 			}
 		}
 
@@ -210,6 +212,7 @@ public void RemoveValue()
 			{
 				_responseSource = null;
 				_tokenSource?.Cancel();
+				_tokenSource = null;
 			}
 		}
 	}

tests/Immediate.Cache.FunctionalTests/ApplicationCacheTests.cs

@@ -260,13 +260,18 @@ public async Task RemovingValueCancelsExistingOperation()
 
 		cache.RemoveValue(request);
 
-		_ = await Assert.ThrowsAnyAsync<OperationCanceledException>(
-			async () => await responseTask
-		);
+		// allow IC task to be run
+		await Task.Delay(10);
 
-		Assert.True(responseTask.IsCanceled);
-		Assert.Equal(0, request.TimesExecuted);
-		Assert.True(request.CancellationToken.IsCancellationRequested);
+		request.CompletionSource.SetResult();
+
+		var response = await responseTask;
+
+		Assert.Equal(1, request.TimesExecuted);
+		Assert.Equal(1, request.TimesCancelled);
+
+		Assert.Equal(1, response.Value);
+		Assert.True(response.ExecutedHandler);
 	}
 
 	[Test]
@@ -279,9 +284,11 @@ public async Task SettingValueCancelsExistingOperation()
 			CompletionSource = new(),
 		};
 
-		using var tcs = new CancellationTokenSource();
 		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
-		var responseTask = cache.GetValue(request, tcs.Token);
+		var responseTask = cache.GetValue(request, default);
+
+		// allow IC task to be run
+		await Task.Delay(10);
 
 		cache.SetValue(request, new(5, ExecutedHandler: false, Guid.NewGuid()));
 
@@ -290,6 +297,6 @@ public async Task SettingValueCancelsExistingOperation()
 		Assert.False(response.ExecutedHandler);
 
 		Assert.Equal(0, request.TimesExecuted);
-		Assert.True(request.CancellationToken.IsCancellationRequested);
+		Assert.Equal(1, request.TimesCancelled);
 	}
 }

tests/Immediate.Cache.FunctionalTests/DelayGetValue.cs

@@ -11,6 +11,7 @@ public sealed class Query
 		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; }
 	}
 
@@ -23,12 +24,21 @@ private static async ValueTask<Response> HandleAsync(
 		CancellationToken token
 	)
 	{
-		query.CancellationToken = token;
-		await query.CompletionSource.Task.WaitAsync(token);
+		try
+		{
+			query.CancellationToken = token;
+			await query.CompletionSource.Task.WaitAsync(token);
 
-		lock (s_lock)
-			query.TimesExecuted++;
+			lock (s_lock)
+				query.TimesExecuted++;
 
-		return new(query.Value, ExecutedHandler: true, RandomValue: Guid.NewGuid());
+			return new(query.Value, ExecutedHandler: true, RandomValue: Guid.NewGuid());
+		}
+		catch
+		{
+			lock (s_lock)
+				query.TimesCancelled++;
+			throw;
+		}
 	}
 }