TrackPublishHandlerAsActiveAsync closure and synchronous invocation h…

…int (#26986) * TrackPublishHandlerAsActiveAsync closure and synchronous invocation hint * Update sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubBufferedProducerClient.cs Co-authored-by: Jesse Squire <jesse.squire@gmail.com> * Update sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubBufferedProducerClient.cs Co-authored-by: Jesse Squire <jesse.squire@gmail.com> * Properly fix whitespace, sigh Co-authored-by: Jesse Squire <jesse.squire@gmail.com>
Azure · Feb 14, 2022 · e781c3a · e781c3a
1 parent c04cc51
commit e781c3a
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubBufferedProducerClient.cs b/sdk/eventhub/Azure.Messaging.EventHubs/src/Producer/EventHubBufferedProducerClient.cs
@@ -1966,8 +1966,14 @@ private Task TrackPublishHandlerAsActiveAsync(Task handlerTask)
 
             _activePublishingHandlers.TryAdd(handlerTask, 0);
 
-            var continationTask = handlerTask.ContinueWith((runTask, trackedTask) => _activePublishingHandlers.TryRemove((Task)trackedTask, out _), handlerTask, TaskScheduler.Default);
-            return continationTask;
+            var continuationTask = handlerTask.ContinueWith(static (runTask, state) =>
+            {
+                var (trackedTask, activeHandlers) = (Tuple<Task, ConcurrentDictionary<Task, byte>>)state;
+                return activeHandlers.TryRemove(trackedTask, out _);
+            },
+            Tuple.Create(handlerTask, _activePublishingHandlers), CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
+
+            return continuationTask;
         }
 
         /// <summary>