Preserve late events after continue-as-new

src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs

@@ -30,6 +30,7 @@
 
     int newGuidCounter;
     object? customStatus;
+    bool preserveUnprocessedEventsOnContinueAsNew;
     TaskOrchestrationEntityContext? entityFeature;
 
     /// <summary>
@@ -349,24 +350,31 @@
     {
         Check.NotNull(options);
 
-        if (!string.IsNullOrWhiteSpace(options.NewVersion))
+        this.preserveUnprocessedEventsOnContinueAsNew = options.PreserveUnprocessedEvents;
+
+        try
         {
-            this.innerContext.ContinueAsNew(options.NewVersion, options.NewInput);
+            if (!string.IsNullOrWhiteSpace(options.NewVersion))
+            {
+                this.innerContext.ContinueAsNew(options.NewVersion, options.NewInput);
+            }
+            else
+            {
+                this.innerContext.ContinueAsNew(options.NewInput);
+            }
         }
-        else
+        catch
         {
-            this.innerContext.ContinueAsNew(options.NewInput);
+            this.preserveUnprocessedEventsOnContinueAsNew = false;
+            throw;
         }
 
         if (options.PreserveUnprocessedEvents)
         {
             // Send all the buffered external events to ourself.
-            OrchestrationInstance instance = new() { InstanceId = this.InstanceId };
             foreach ((string eventName, string eventPayload) in this.externalEventBuffer.TakeAll())
             {
-#pragma warning disable CS0618 // Type or member is obsolete -- 'internal' usage.
-                this.innerContext.SendEvent(instance, eventName, new RawInput(eventPayload));
-#pragma warning restore CS0618 // Type or member is obsolete
+                this.ForwardRawExternalEvent(eventName, eventPayload);
             }
         }
     }
@@ -477,17 +485,35 @@
         }
         else
         {
-            // The orchestrator isn't waiting for this event (yet?). Save it in case
-            // the orchestrator wants it later.
-            this.externalEventBuffer.Add(eventName, rawEventPayload);
+            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);
+            }
         }
     }
 
+    void ForwardRawExternalEvent(string eventName, string rawEventPayload)
+    {
+        OrchestrationInstance instance = new() { InstanceId = this.InstanceId };
+#pragma warning disable CS0618 // Type or member is obsolete -- 'internal' usage.
+        this.innerContext.SendEvent(instance, eventName, new RawInput(rawEventPayload));
+#pragma warning restore CS0618 // Type or member is obsolete
+    }
+
     /// <summary>
     /// 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

@@ -1,13 +1,19 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Reflection;
 using DurableTask.Core;
+using DurableTask.Core.Serializing.Internal;
 using Microsoft.Extensions.Logging.Abstractions;
 
 namespace Microsoft.DurableTask.Worker.Shims;
 
 public class TaskOrchestrationContextWrapperTests
 {
+    static readonly MethodInfo CompleteExternalEventMethod = typeof(TaskOrchestrationContextWrapper)
+        .GetMethod(nameof(TaskOrchestrationContextWrapper.CompleteExternalEvent), BindingFlags.Instance | BindingFlags.NonPublic)
+        ?? throw new InvalidOperationException($"{nameof(TaskOrchestrationContextWrapper)}.{nameof(TaskOrchestrationContextWrapper.CompleteExternalEvent)} was not found.");
+
     [Fact]
     public void Ctor_NullParent_Populates()
     {
@@ -103,6 +109,31 @@ public void ContinueAsNew_WithOptionsNoVersion_CallsInnerContextWithoutVersion()
         innerContext.LastContinueAsNewVersion.Should().BeNull();
     }
 
+    [Fact]
+    public void ContinueAsNew_WithPreserveUnprocessedEvents_ForwardsLateArrivingEventsToNextExecution()
+    {
+        // Arrange
+        TrackingOrchestrationContext innerContext = new();
+        OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null);
+        TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input");
+
+        // Act
+        wrapper.ContinueAsNew("new-input", preserveUnprocessedEvents: true);
+        InvokeCompleteExternalEvent(wrapper, "Event", "\"payload\"");
+
+        // Assert
+        innerContext.SentEvents.Should().ContainSingle();
+        innerContext.SentEvents[0].InstanceId.Should().Be(wrapper.InstanceId);
+        innerContext.SentEvents[0].EventName.Should().Be("Event");
+        innerContext.SentEvents[0].EventData.Should().BeOfType<RawInput>().Which.Value.Should().Be("\"payload\"");
+        innerContext.LastContinueAsNewInput.Should().Be("new-input");
+    }
+
+    static void InvokeCompleteExternalEvent(TaskOrchestrationContextWrapper wrapper, string eventName, string rawEventPayload)
+    {
+        CompleteExternalEventMethod.Invoke(wrapper, [eventName, rawEventPayload]);
+    }
+
     sealed class TrackingOrchestrationContext : OrchestrationContext
     {
         public TrackingOrchestrationContext()
@@ -118,6 +149,8 @@ public TrackingOrchestrationContext()
 
         public string? LastContinueAsNewVersion { get; private set; }
 
+        public List<(string InstanceId, string EventName, object EventData)> SentEvents { get; } = [];
+
         public override void ContinueAsNew(object input)
         {
             this.LastContinueAsNewInput = input;
@@ -149,7 +182,9 @@ public override Task<TResult> ScheduleTask<TResult>(string name, string version,
             => throw new NotImplementedException();
 
         public override void SendEvent(OrchestrationInstance orchestrationInstance, string eventName, object eventData)
-            => throw new NotImplementedException();
+        {
+            this.SentEvents.Add((orchestrationInstance.InstanceId, eventName, eventData));
+        }
     }
 
     class TestOrchestrationContext : OrchestrationContext
@@ -210,4 +245,4 @@ public override void SendEvent(OrchestrationInstance orchestrationInstance, stri
             throw new NotImplementedException();
         }
     }
-}
+}