Validate Using CancellationToken

Author: viceroypenguin

src/Immediate.Cache.Shared/ApplicationCacheBase.cs

@@ -86,11 +86,14 @@ private CacheValue GetCacheValue(TRequest request)
 	/// <param name="request">
 	///		The request payload to be cached.
 	/// </param>
+	/// <param name="cancellationToken">
+	///		The token to monitor for cancellation requests. The default value is <see cref="CancellationToken.None"/>.
+	/// </param>
 	/// <returns>
 	///		The response payload from executing the handler.
 	/// </returns>
-	public ValueTask<TResponse> GetValue(TRequest request) =>
-		GetCacheValue(request).GetValue();
+	public ValueTask<TResponse> GetValue(TRequest request, CancellationToken cancellationToken = default) =>
+		GetCacheValue(request).GetValue(cancellationToken);
 
 	/// <summary>
 	///		Sets the value for a particular cache entry, bypassing the execution of the handler.
@@ -124,9 +127,9 @@ Owned<IHandler<TRequest, TResponse>> handler
 		private TaskCompletionSource<TResponse>? _responseSource;
 		private readonly Lock _lock = new();
 
-		public async ValueTask<TResponse> GetValue()
+		public async ValueTask<TResponse> GetValue(CancellationToken cancellationToken)
 		{
-			if (!TryAcquireResponseSource())
+			if (!TryAcquireResponseSource(cancellationToken))
 				return await _responseSource.Task.ConfigureAwait(false);
 
 			var token = _tokenSource.Token;
@@ -151,7 +154,7 @@ public async ValueTask<TResponse> GetValue()
 					}
 				}
 			}
-			catch (OperationCanceledException) when (_responseSource is not null)
+			catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && _responseSource is not null)
 			{
 				return await _responseSource.Task.ConfigureAwait(false);
 			}
@@ -160,7 +163,7 @@ public async ValueTask<TResponse> GetValue()
 		[MemberNotNull(nameof(_responseSource))]
 		[MemberNotNullWhen(true, nameof(_tokenSource))]
 		[SuppressMessage("Maintainability", "CA1508:Avoid dead conditional code", Justification = "Double-checked lock pattern")]
-		private bool TryAcquireResponseSource()
+		private bool TryAcquireResponseSource(CancellationToken cancellationToken)
 		{
 			if (_responseSource is not null)
 				return false;
@@ -170,7 +173,7 @@ private bool TryAcquireResponseSource()
 				if (_responseSource is not null)
 					return false;
 
-				_tokenSource = new CancellationTokenSource();
+				_tokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 				_responseSource = new TaskCompletionSource<TResponse>();
 				return true;
 			}

tests/Immediate.Cache.FunctionalTests/ApplicationCacheTests.cs

@@ -122,4 +122,24 @@ public async Task SimultaneousAccessIsSerialized()
 		// ensure both responses get the same response back
 		Assert.Equal(response1.RandomValue, response2.RandomValue);
 	}
+
+	[Test]
+	public async Task ProperlyUsesCancellationToken()
+	{
+		var request = new DelayGetValue.Query()
+		{
+			Value = 4,
+			Name = "Request4",
+			CompletionSource = new(),
+		};
+
+		using var tcs = new CancellationTokenSource();
+		var cache = _serviceProvider.GetRequiredService<DelayGetValueCache>();
+		var responseTask = cache.GetValue(request, tcs.Token);
+
+		await tcs.CancelAsync();
+		request.CompletionSource.SetResult();
+
+		Assert.True(responseTask.IsCanceled);
+	}
 }

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 CancellationToken CancellationToken { get; set; }
 	}
 
 	public sealed record Response(int Value, bool ExecutedHandler, Guid RandomValue);
@@ -19,10 +20,13 @@ public sealed record Response(int Value, bool ExecutedHandler, Guid RandomValue)
 
 	private static async ValueTask<Response> HandleAsync(
 		Query query,
-		CancellationToken _
+		CancellationToken token
 	)
 	{
+		query.CancellationToken = token;
 		await query.CompletionSource.Task;
+		token.ThrowIfCancellationRequested();
+
 		lock (s_lock)
 			query.TimesExecuted++;