From 4f655c30ccceee64526d75c7411e9cd9bb19f450 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Thu, 23 Apr 2026 14:13:12 -0700 Subject: [PATCH 1/2] Preserve late events after continue-as-new Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Shims/TaskOrchestrationContextWrapper.cs | 48 ++++++++++++++----- .../TaskOrchestrationContextWrapperTests.cs | 36 +++++++++++++- 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs index bfbf1e47..e91ba426 100644 --- a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs +++ b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs @@ -30,6 +30,7 @@ sealed partial class TaskOrchestrationContextWrapper : TaskOrchestrationContext int newGuidCounter; object? customStatus; + bool preserveUnprocessedEventsOnContinueAsNew; TaskOrchestrationEntityContext? entityFeature; /// @@ -349,24 +350,31 @@ public override void ContinueAsNew(ContinueAsNewOptions options) { 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.ForwardBufferedExternalEvent(eventName, eventPayload); } } } @@ -477,12 +485,30 @@ internal void CompleteExternalEvent(string eventName, string 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); + 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.ForwardBufferedExternalEvent(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 ForwardBufferedExternalEvent(string eventName, string eventPayload) + { + OrchestrationInstance instance = new() { InstanceId = this.InstanceId }; +#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 + } + /// /// Gets the serialized custom status. /// diff --git a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs index 06df4317..fa894e9b 100644 --- a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Reflection; using DurableTask.Core; using Microsoft.Extensions.Logging.Abstractions; @@ -8,6 +9,9 @@ namespace Microsoft.DurableTask.Worker.Shims; public class TaskOrchestrationContextWrapperTests { + static readonly MethodInfo CompleteExternalEventMethod = typeof(TaskOrchestrationContextWrapper) + .GetMethod("CompleteExternalEvent", BindingFlags.Instance | BindingFlags.NonPublic)!; + [Fact] public void Ctor_NullParent_Populates() { @@ -103,6 +107,30 @@ 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.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 +146,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 +179,9 @@ public override Task ScheduleTask(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 +242,4 @@ public override void SendEvent(OrchestrationInstance orchestrationInstance, stri throw new NotImplementedException(); } } -} \ No newline at end of file +} From 57409250008ce4a476081b266b64eb7d56d9e121 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Thu, 23 Apr 2026 14:35:41 -0700 Subject: [PATCH 2/2] Tighten wrapper regression test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs | 8 ++++---- .../Shims/TaskOrchestrationContextWrapperTests.cs | 5 ++++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs index e91ba426..646858c4 100644 --- a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs +++ b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs @@ -374,7 +374,7 @@ public override void ContinueAsNew(ContinueAsNewOptions options) // Send all the buffered external events to ourself. foreach ((string eventName, string eventPayload) in this.externalEventBuffer.TakeAll()) { - this.ForwardBufferedExternalEvent(eventName, eventPayload); + this.ForwardRawExternalEvent(eventName, eventPayload); } } } @@ -490,7 +490,7 @@ internal void CompleteExternalEvent(string eventName, string rawEventPayload) // 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.ForwardBufferedExternalEvent(eventName, rawEventPayload); + this.ForwardRawExternalEvent(eventName, rawEventPayload); } else { @@ -501,11 +501,11 @@ internal void CompleteExternalEvent(string eventName, string rawEventPayload) } } - void ForwardBufferedExternalEvent(string eventName, string eventPayload) + 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(eventPayload)); + this.innerContext.SendEvent(instance, eventName, new RawInput(rawEventPayload)); #pragma warning restore CS0618 // Type or member is obsolete } diff --git a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs index fa894e9b..17ed6f63 100644 --- a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs @@ -3,6 +3,7 @@ using System.Reflection; using DurableTask.Core; +using DurableTask.Core.Serializing.Internal; using Microsoft.Extensions.Logging.Abstractions; namespace Microsoft.DurableTask.Worker.Shims; @@ -10,7 +11,8 @@ namespace Microsoft.DurableTask.Worker.Shims; public class TaskOrchestrationContextWrapperTests { static readonly MethodInfo CompleteExternalEventMethod = typeof(TaskOrchestrationContextWrapper) - .GetMethod("CompleteExternalEvent", BindingFlags.Instance | BindingFlags.NonPublic)!; + .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() @@ -123,6 +125,7 @@ public void ContinueAsNew_WithPreserveUnprocessedEvents_ForwardsLateArrivingEven 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().Which.Value.Should().Be("\"payload\""); innerContext.LastContinueAsNewInput.Should().Be("new-input"); }