Remove _tokenRegistration

Azure · Dec 11, 2024 · 15bbc50 · 15bbc50
1 parent 70644de
commit 15bbc50
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 29 deletions.
diff --git a/extensions/Worker.Extensions.Http.AspNetCore/src/Coordinator/ContextReference.cs b/extensions/Worker.Extensions.Http.AspNetCore/src/Coordinator/ContextReference.cs
@@ -12,16 +12,14 @@ internal class ContextReference : IDisposable
     {
         private readonly TaskCompletionSource<bool> _functionStartTask = new();
         private readonly TaskCompletionSource<bool> _functionCompletionTask = new();
+        private readonly CancellationToken _invocationCancellationToken;
 
         private TaskCompletionSource<HttpContext> _httpContextValueSource = new();
         private TaskCompletionSource<FunctionContext> _functionContextValueSource = new();
 
-        private CancellationToken _token;
-        private CancellationToken _invocationToken;
-        private CancellationTokenRegistration _tokenRegistration;
-
-        public ContextReference(string invocationId)
+        public ContextReference(string invocationId, CancellationToken invocationCancellationToken = default)
         {
+            _invocationCancellationToken = invocationCancellationToken;
             InvocationId = invocationId;
         }
 
@@ -33,20 +31,6 @@ public ContextReference(string invocationId)
 
         public TaskCompletionSource<FunctionContext> FunctionContextValueSource { get => _functionContextValueSource; set => _functionContextValueSource = value; }
 
-        internal void SetCancellationToken(CancellationToken contextToken, CancellationToken invocationToken)
-        {
-            _token = contextToken;
-            _invocationToken = invocationToken; // Only to check if the invocation is canceled
-
-            _tokenRegistration = _token.Register(() => // We don't want these to depend on the invocation CT as it should be up to the function to decide how to handle it
-            {
-                _functionStartTask.TrySetCanceled();
-                _functionCompletionTask.TrySetCanceled();
-                _functionContextValueSource.TrySetCanceled();
-                _httpContextValueSource.TrySetCanceled();
-            });
-        }
-
         internal Task InvokeFunctionAsync()
         {
             _functionStartTask.SetResult(true);
@@ -62,9 +46,10 @@ internal void CompleteFunction()
 
             if (_httpContextValueSource.Task.IsCompleted)
             {
-                if (_httpContextValueSource.Task.IsCanceled || _invocationToken.IsCancellationRequested || _token.IsCancellationRequested)
+                if (_httpContextValueSource.Task.IsCanceled || _invocationCancellationToken.IsCancellationRequested)
                 {
                     _functionCompletionTask.TrySetCanceled();
+                    _functionContextValueSource.TrySetCanceled();
                 }
                 else
                 {
@@ -79,10 +64,10 @@ internal void CompleteFunction()
 
         public void Dispose()
         {
-            if (_tokenRegistration != default)
-            {
-                _tokenRegistration.Dispose();
-            }
+            _functionStartTask?.TrySetCanceled();
+            _functionCompletionTask?.TrySetCanceled();
+            _httpContextValueSource?.TrySetCanceled();
+            _functionContextValueSource?.TrySetCanceled();
         }
     }
 }
diff --git a/extensions/Worker.Extensions.Http.AspNetCore/src/Coordinator/DefaultHttpCoordinator.cs b/extensions/Worker.Extensions.Http.AspNetCore/src/Coordinator/DefaultHttpCoordinator.cs
@@ -18,7 +18,6 @@ internal class DefaultHttpCoordinator : IHttpCoordinator
 
         private readonly ConcurrentDictionary<string, ContextReference> _contextReferenceList;
         private readonly ExtensionTrace _logger;
-        private readonly CancellationTokenSource _cts = new(); // I am not sure if introducing a new cts here is the best idea
 
         public DefaultHttpCoordinator(ExtensionTrace extensionTrace)
         {
@@ -45,8 +44,7 @@ public async Task<FunctionContext> SetHttpContextAsync(string invocationId, Http
 
         public async Task<HttpContext> SetFunctionContextAsync(string invocationId, FunctionContext context)
         {
-            var contextRef = _contextReferenceList.GetOrAdd(invocationId, static id => new ContextReference(id));
-            contextRef.SetCancellationToken(_cts.Token, context.CancellationToken);
+            var contextRef = _contextReferenceList.GetOrAdd(invocationId, id => new ContextReference(id, context.CancellationToken));
             contextRef.FunctionContextValueSource.SetResult(context);
 
             _logger.FunctionContextSet(invocationId);
@@ -93,8 +91,6 @@ public void CompleteFunctionInvocation(string invocationId)
             {
                 // do something here?
             }
-
-            _cts?.Dispose();
         }
     }