Prevent external-event loss after canceled waits in isolated worker

src/Worker/Core/Shims/TaskOrchestrationContextWrapper.EventSource.cs

@@ -30,10 +30,14 @@ interface IEventSource
         /// Tries to set the result on tcs.
         /// </summary>
         /// <param name="result">The result.</param>
-        void TrySetResult(object result);
+        /// <returns>
+        /// <c>true</c> if the result was set successfully, meaning this waiter is a live consumer of the event;
+        /// <c>false</c> if the waiter was already completed, canceled, or abandoned and cannot accept the result.
+        /// </returns>
+        bool TrySetResult(object? result);
     }
 
-    class EventTaskCompletionSource<T> : TaskCompletionSource<T>, IEventSource
+    sealed class EventTaskCompletionSource<T> : TaskCompletionSource<T>, IEventSource
     {
         /// <inheritdoc/>
         public Type EventType => typeof(T);
@@ -42,7 +46,7 @@ class EventTaskCompletionSource<T> : TaskCompletionSource<T>, IEventSource
         public IEventSource? Next { get; set; }
 
         /// <inheritdoc/>
-        void IEventSource.TrySetResult(object result) => this.TrySetResult((T)result);
+        bool IEventSource.TrySetResult(object? result) => this.TrySetResult((T)result!);
     }
 
     class NamedQueue<TValue>

src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs

@@ -473,7 +473,8 @@
     /// <param name="rawEventPayload">The serialized event payload.</param>
     internal void CompleteExternalEvent(string eventName, string rawEventPayload)
     {
-        if (this.externalEventSources.TryGetValue(eventName, out IEventSource? waiter))
+        Dictionary<Type, object?>? deserializedValues = null;
+        while (this.externalEventSources.TryGetValue(eventName, out IEventSource? waiter))
         {
             // Get the waiter at the top of the stack (most recent waiter)
             // If we're going to raise an event we should remove it from the pending collection
@@ -489,35 +490,47 @@
                 this.externalEventSources[eventName] = next;
             }
 
-            object? value;
-            if (waiter.EventType == typeof(OperationResult))
+            deserializedValues ??= new Dictionary<Type, object?>();
+            if (!deserializedValues.TryGetValue(waiter.EventType, out object? value))
             {
-                // use the framework-defined deserialization for entity responses, not the application-defined data converter,
-                // because we are just unwrapping the entity response without yet deserializing any application-defined data.
-                value = this.entityFeature!.EntityContext.DeserializeEntityResponseEvent(rawEventPayload);
+                if (waiter.EventType == typeof(OperationResult))
+                {
+                    // use the framework-defined deserialization for entity responses, not the application-defined data converter,
+                    // because we are just unwrapping the entity response without yet deserializing any application-defined data.
+                    value = this.entityFeature!.EntityContext.DeserializeEntityResponseEvent(rawEventPayload);
+                }
+                else
+                {
+                    value = this.DataConverter.Deserialize(rawEventPayload, waiter.EventType);
+                }
+
+                deserializedValues[waiter.EventType] = value;
             }
-            else
+
+            if (waiter.TrySetResult(value))
             {
-                value = this.DataConverter.Deserialize(rawEventPayload, waiter.EventType);
+                // The event was delivered to a live waiter. We're done.
+                return;
             }
 
-            waiter.TrySetResult(value);
+            // This waiter was already completed, canceled, or abandoned (e.g. the losing side of a
+            // Task.WhenAny) and cannot accept the event. It has already been popped off the stack above,
+            // so continue the loop to try the next waiter underneath it, if any. If there are no more
+            // waiters, fall through to the buffering/forwarding logic below as if no one was listening.
+        }
+
+        if (this.preserveUnprocessedEventsOnContinueAsNew)
+        {
+            // ContinueAsNew has already been scheduled with event preservation enabled.
+            // Forward late-arriving events directly to the next execution instead of buffering
+            // them on the current wrapper instance, which is about to be discarded.
+            this.ForwardRawExternalEvent(eventName, rawEventPayload);
         }
         else
         {
-            if (this.preserveUnprocessedEventsOnContinueAsNew)
-            {
-                // ContinueAsNew has already been scheduled with event preservation enabled.
-                // Forward late-arriving events directly to the next execution instead of buffering
-                // them on the current wrapper instance, which is about to be discarded.
-                this.ForwardRawExternalEvent(eventName, rawEventPayload);
-            }
-            else
-            {
-                // The orchestrator isn't waiting for this event (yet?). Save it in case
-                // the orchestrator wants it later.
-                this.externalEventBuffer.Add(eventName, rawEventPayload);
-            }
+            // The orchestrator isn't waiting for this event (yet?). Save it in case
+            // the orchestrator wants it later.
+            this.externalEventBuffer.Add(eventName, rawEventPayload);
         }
     }
 
@@ -533,7 +546,7 @@
     /// Gets the serialized custom status.
     /// </summary>
     /// <returns>The custom status serialized to a string, or <c>null</c> if there is not custom status.</returns>
     internal string? GetSerializedCustomStatus()
     {
         return this.DataConverter.Serialize(this.customStatus);
     }

test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs

@@ -391,6 +391,123 @@ await wrapper.CallSubOrchestratorAsync<string>(
         innerContext.LastSubOrchestrationVersion.Should().Be(string.Empty);
     }
 
+    [Fact]
+    public void CompleteExternalEvent_CanceledWaiter_BuffersEventForNextWaiter()
+    {
+        // Arrange: reproduces the sequence from Azure/azure-functions-durable-extension#1676.
+        // Two differently-named external event waits are registered with a shared cancellation
+        // token, similar to a Task.WhenAny loop. The waiter for "event_1" is canceled (e.g.
+        // because "event_0" won the race), but "event_1" is later raised while the canceled
+        // waiter is still the only registered listener for that name.
+        TrackingOrchestrationContext innerContext = new();
+        OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
+        TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
+
+        using CancellationTokenSource cts = new();
+        Task<string> canceledWait = wrapper.WaitForExternalEvent<string>("event_1", cts.Token);
+        cts.Cancel();
+
+        // Act: the event arrives after the only waiter for it has been canceled/abandoned.
+        InvokeCompleteExternalEvent(wrapper, "event_1", "\"payload\"");
+
+        // Assert: the canceled waiter must not "consume" the event. It should be buffered so
+        // that a subsequent WaitForExternalEvent call (e.g. in the next ContinueAsNew generation)
+        // can still observe it.
+        canceledWait.IsCanceled.Should().BeTrue();
+        Task<string> nextWait = wrapper.WaitForExternalEvent<string>("event_1");
+        nextWait.IsCompletedSuccessfully.Should().BeTrue();
+        nextWait.Result.Should().Be("payload");
+    }
+
+    [Fact]
+    public void CompleteExternalEvent_CanceledWaiter_WithContinueAsNewPreserved_ForwardsEventToNextExecution()
+    {
+        // Arrange: same as above, but the event arrives after ContinueAsNew(preserveUnprocessedEvents: true)
+        // has already been called (e.g. while an activity scheduled before ContinueAsNew is still
+        // completing). The event should be forwarded to the next execution, not silently dropped.
+        TrackingOrchestrationContext innerContext = new();
+        OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
+        TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
+
+        using CancellationTokenSource cts = new();
+        Task<string> canceledWait = wrapper.WaitForExternalEvent<string>("event_1", cts.Token);
+        cts.Cancel();
+
+        wrapper.ContinueAsNew("new-input", preserveUnprocessedEvents: true);
+
+        // Act
+        InvokeCompleteExternalEvent(wrapper, "event_1", "\"payload\"");
+
+        // Assert
+        canceledWait.IsCanceled.Should().BeTrue();
+        innerContext.SentEvents.Should().ContainSingle();
+        innerContext.SentEvents[0].InstanceId.Should().Be(wrapper.InstanceId);
+        innerContext.SentEvents[0].EventName.Should().Be("event_1");
+        innerContext.SentEvents[0].EventData.Should().BeOfType<RawInput>().Which.Value.Should().Be("\"payload\"");
+    }
+
+    [Fact]
+    public void CompleteExternalEvent_MultipleCanceledWaiters_BuffersEventAfterSkippingAll()
+    {
+        // Arrange: several waiters for the same event name are canceled/abandoned in a row
+        // (e.g. a loop that repeatedly races and abandons the same event). None of them should
+        // be able to consume the event.
+        TrackingOrchestrationContext innerContext = new();
+        CountingDataConverter converter = new();
+        OrchestrationInvocationContext invocationContext = new(
+            "Test",
+            new DurableTaskWorkerOptions { DataConverter = converter },
+            NullLoggerFactory.Instance,
+            null);
+        TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
+
+        using CancellationTokenSource cts1 = new();
+        using CancellationTokenSource cts2 = new();
+        using CancellationTokenSource cts3 = new();
+        Task<string> wait1 = wrapper.WaitForExternalEvent<string>("event_1", cts1.Token);
+        Task<string> wait2 = wrapper.WaitForExternalEvent<string>("event_1", cts2.Token);
+        Task<string> wait3 = wrapper.WaitForExternalEvent<string>("event_1", cts3.Token);
+        cts1.Cancel();
+        cts2.Cancel();
+        cts3.Cancel();
+
+        // Act
+        InvokeCompleteExternalEvent(wrapper, "event_1", "\"payload\"");
+
+        // Assert
+        converter.DeserializeCallCounts[typeof(string)].Should().Be(1);
+        wait1.IsCanceled.Should().BeTrue();
+        wait2.IsCanceled.Should().BeTrue();
+        wait3.IsCanceled.Should().BeTrue();
+        Task<string> nextWait = wrapper.WaitForExternalEvent<string>("event_1");
+        nextWait.IsCompletedSuccessfully.Should().BeTrue();
+        nextWait.Result.Should().Be("payload");
+    }
+
+    [Fact]
+    public void CompleteExternalEvent_ActiveWaiterAboveCanceledWaiter_DeliversToActiveWaiter()
+    {
+        // Arrange: the most recent (top-of-stack) waiter is still live, but an older waiter
+        // underneath it for the same event name was already canceled. The event should go to
+        // the live waiter, and the canceled one underneath should remain untouched (not consumed).
+        TrackingOrchestrationContext innerContext = new();
+        OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
+        TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
+
+        using CancellationTokenSource cts1 = new();
+        Task<string> canceledWait = wrapper.WaitForExternalEvent<string>("event_1", cts1.Token);
+        cts1.Cancel();
+        Task<string> activeWait = wrapper.WaitForExternalEvent<string>("event_1");
+
+        // Act
+        InvokeCompleteExternalEvent(wrapper, "event_1", "\"payload\"");
+
+        // Assert
+        canceledWait.IsCanceled.Should().BeTrue();
+        activeWait.IsCompletedSuccessfully.Should().BeTrue();
+        activeWait.Result.Should().Be("payload");
+    }
+
     static IReadOnlyDictionary<string, string> GetLastScheduledTaskTags(TrackingOrchestrationContext innerContext)
     {
         PropertyInfo tagsProperty = innerContext.LastScheduledTaskOptions!.GetType().GetProperty("Tags")!;
@@ -402,6 +519,19 @@ static void InvokeCompleteExternalEvent(TaskOrchestrationContextWrapper wrapper,
         CompleteExternalEventMethod.Invoke(wrapper, [eventName, rawEventPayload]);
     }
 
+    sealed class CountingDataConverter : DataConverter
+    {
+        public Dictionary<Type, int> DeserializeCallCounts { get; } = [];
+
+        public override object? Deserialize(string? data, Type targetType)
+        {
+            this.DeserializeCallCounts[targetType] = this.DeserializeCallCounts.GetValueOrDefault(targetType) + 1;
+            return targetType == typeof(string) ? "payload" : null;
+        }
+
+        public override string? Serialize(object? value) => throw new NotSupportedException();
+    }
+
     sealed class TrackingOrchestrationContext : OrchestrationContext
     {
         public TrackingOrchestrationContext(string? version = null)