From 6485e9fbe510514b48fd6598a96d6d882326b89e Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 15:05:22 -0700 Subject: [PATCH 01/10] Reduce per-call SHA1 allocation in deterministic NewGuid() TaskOrchestrationContextWrapper.NewGuid() previously created and disposed a new SHA1 instance on every call. Since a single wrapper instance is used for the duration of one orchestration execution, and orchestrator code within that execution runs sequentially (never concurrently), the SHA1 instance can safely be cached on the wrapper and reused across calls via HashAlgorithm.Initialize(), avoiding a per-call allocation. The hashed name, DNS namespace bytes, byte-swap order, and RFC 4122 version/variant bit handling are all unchanged, so generated GUIDs are byte-for-byte identical to before this change for the same inputs. Adds regression tests asserting: - stable, golden-value deterministic GUIDs for fixed inputs - distinct GUIDs across repeated calls with the same instance/timestamp - identical GUID sequences across two independent wrapper instances given identical inputs (simulated replay) - the cached SHA1 instance is reused (not reallocated) across calls Fixes #778 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- .../Shims/TaskOrchestrationContextWrapper.cs | 23 ++- .../TaskOrchestrationContextWrapperTests.cs | 131 +++++++++++++++++- 2 files changed, 146 insertions(+), 8 deletions(-) diff --git a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs index c92ee5d6..e20c2d77 100644 --- a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs +++ b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs @@ -33,6 +33,13 @@ sealed partial class TaskOrchestrationContextWrapper : TaskOrchestrationContext bool preserveUnprocessedEventsOnContinueAsNew; TaskOrchestrationEntityContext? entityFeature; + // Cached and reused across NewGuid() calls (instead of creating and disposing a new SHA1 instance + // per call) to reduce per-call allocation overhead. A single TaskOrchestrationContextWrapper is used + // for the duration of a single orchestration execution, and orchestrator code executes sequentially + // (never concurrently) within that execution, so reusing this instance is safe. HashAlgorithm.Initialize() + // resets all internal state before each use, so the computed hash is identical to using a fresh instance. + SHA1? cachedHashAlgorithm; + /// /// Initializes a new instance of the class. /// @@ -434,16 +441,18 @@ static void SwapByteArrayElements(byte[] byteArray, int left, int right) byte[] namespaceValueByteArray = namespaceValueGuid.ToByteArray(); SwapByteArrayValues(namespaceValueByteArray); - byte[] hashByteArray; #pragma warning disable CA5350 // Do Not Use Weak Cryptographic Algorithms -- not for cryptography - using (HashAlgorithm hashAlgorithm = SHA1.Create()) /* CodeQL [SM02196] Suppressed: SHA1 is not used for cryptographic purposes here. The information being hashed is not sensitive, + SHA1 hashAlgorithm = this.cachedHashAlgorithm ??= SHA1.Create(); /* CodeQL [SM02196] Suppressed: SHA1 is not used for cryptographic purposes here. The information being hashed is not sensitive, and the goal is to generate a deterministic Guid. We cannot update to SHA2-based algorithms without breaking customers' inflight orchestrations. */ - { - hashAlgorithm.TransformBlock(namespaceValueByteArray, 0, namespaceValueByteArray.Length, null, 0); - hashAlgorithm.TransformFinalBlock(nameByteArray, 0, nameByteArray.Length); - hashByteArray = hashAlgorithm.Hash; - } + + // Reset internal state before each use since this instance is cached and reused across calls + // rather than being created and disposed per call. This produces byte-for-byte identical hashes + // to constructing a new SHA1 instance for every call. + hashAlgorithm.Initialize(); + hashAlgorithm.TransformBlock(namespaceValueByteArray, 0, namespaceValueByteArray.Length, null, 0); + hashAlgorithm.TransformFinalBlock(nameByteArray, 0, nameByteArray.Length); + byte[] hashByteArray = hashAlgorithm.Hash; #pragma warning restore CA5350 // Do Not Use Weak Cryptographic Algorithms -- not for cryptography byte[] newGuidByteArray = new byte[16]; diff --git a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs index e8335e79..497603f5 100644 --- a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs @@ -2,6 +2,7 @@ // Licensed under the MIT License. using System.Collections.Generic; +using System.Globalization; using System.Reflection; using DurableTask.Core; using DurableTask.Core.Serializing.Internal; @@ -391,6 +392,119 @@ await wrapper.CallSubOrchestratorAsync( innerContext.LastSubOrchestrationVersion.Should().Be(string.Empty); } + [Fact] + public void NewGuid_FixedInputs_ProducesStableDeterministicValue() + { + // Arrange — these golden values were computed independently (offline, using the documented + // algorithm: SHA1("9e952958-5e33-4daf-827f-2fa12937b875" bytes + name bytes), with the RFC 4122 + // byte swaps and version/variant bits applied) for the given instance ID, timestamp, and counter. + // This regression test protects replay compatibility: it must keep producing these exact GUIDs. + TestOrchestrationContext innerContext = new( + "fixed-instance-id", + DateTime.Parse("2023-05-06T07:08:09.1234567Z", null, DateTimeStyles.RoundtripKind)); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + // Act + Guid first = wrapper.NewGuid(); + Guid second = wrapper.NewGuid(); + Guid third = wrapper.NewGuid(); + + // Assert + first.Should().Be(Guid.Parse("0f353f85-75d2-56f8-89b5-a7773ace7605")); + second.Should().Be(Guid.Parse("b0fd1465-f3d8-5a7e-98b1-f34137b15060")); + third.Should().Be(Guid.Parse("12bec829-d5e1-563c-ac70-9806cad148c1")); + } + + [Fact] + public void NewGuid_DifferentInstanceId_ProducesDifferentStableDeterministicValue() + { + // Arrange — same timestamp and counter as the other golden-value test, but a different + // instance ID, computed independently the same way. Confirms the instance ID is still part of + // the hashed name and that the namespace/algorithm/byte-ordering were not altered. + TestOrchestrationContext innerContext = new( + "other-instance-id", + DateTime.Parse("2023-05-06T07:08:09.1234567Z", null, DateTimeStyles.RoundtripKind)); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + // Act + Guid result = wrapper.NewGuid(); + + // Assert + result.Should().Be(Guid.Parse("258d445c-0c1e-594c-a4a1-0a837e4ebe92")); + } + + [Fact] + public void NewGuid_CalledRepeatedly_ProducesDistinctValuesEachTime() + { + // Arrange — the internal counter advances on every call, so repeated calls with the same + // instance ID and timestamp must still yield distinct GUIDs. + TestOrchestrationContext innerContext = new( + "repeat-instance-id", + DateTime.Parse("2024-01-01T00:00:00.0000000Z", null, DateTimeStyles.RoundtripKind)); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + // Act + List results = new(); + for (int i = 0; i < 5; i++) + { + results.Add(wrapper.NewGuid()); + } + + // Assert — all five results are distinct from one another. + results.Distinct().Should().HaveCount(5); + } + + [Fact] + public void NewGuid_ReplayingSameHistory_ProducesIdenticalGuidSequence() + { + // Arrange — simulates replay: two independent wrapper instances (as would be created for two + // separate replay passes over the same orchestration history) observe the same instance ID and + // the same sequence of CurrentUtcDateTime values as history is replayed. + string instanceId = "replay-instance-id"; + DateTime timestamp = DateTime.Parse("2022-11-11T11:11:11.1111111Z", null, DateTimeStyles.RoundtripKind); + + TestOrchestrationContext innerContext1 = new(instanceId, timestamp); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper1 = new(innerContext1, invocationContext, "input"); + + TestOrchestrationContext innerContext2 = new(instanceId, timestamp); + TaskOrchestrationContextWrapper wrapper2 = new(innerContext2, invocationContext, "input"); + + // Act — generate the same number of GUIDs from both "replay passes". + Guid[] pass1 = [wrapper1.NewGuid(), wrapper1.NewGuid(), wrapper1.NewGuid()]; + Guid[] pass2 = [wrapper2.NewGuid(), wrapper2.NewGuid(), wrapper2.NewGuid()]; + + // Assert — replay must produce an identical sequence of GUIDs given identical inputs. + pass2.Should().Equal(pass1); + } + + [Fact] + public void NewGuid_MultipleCalls_ReuseCachedHashAlgorithmInstance() + { + // Arrange — verifies the optimization from + // https://github.com/microsoft/durabletask-dotnet/issues/778: the SHA1 instance backing + // NewGuid() is created once and reused across calls, rather than being constructed and + // disposed on every call. + TestOrchestrationContext innerContext = new(); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + FieldInfo cachedHashAlgorithmField = typeof(TaskOrchestrationContextWrapper) + .GetField("cachedHashAlgorithm", BindingFlags.Instance | BindingFlags.NonPublic)!; + + // Act + wrapper.NewGuid(); + object? afterFirstCall = cachedHashAlgorithmField.GetValue(wrapper); + wrapper.NewGuid(); + object? afterSecondCall = cachedHashAlgorithmField.GetValue(wrapper); + + // Assert — the same underlying instance is reused rather than a new one being allocated. + afterFirstCall.Should().NotBeNull(); + afterSecondCall.Should().BeSameAs(afterFirstCall); + } + static IReadOnlyDictionary GetLastScheduledTaskTags(TrackingOrchestrationContext innerContext) { PropertyInfo tagsProperty = innerContext.LastScheduledTaskOptions!.GetType().GetProperty("Tags")!; @@ -508,15 +622,30 @@ public override void SendEvent(OrchestrationInstance orchestrationInstance, stri class TestOrchestrationContext : OrchestrationContext { + // Only set when a fixed value is supplied via the constructor overload below; otherwise the + // base class's (internally-set) value is used, preserving prior behavior for existing callers. + readonly DateTime? fixedCurrentUtcDateTime; + public TestOrchestrationContext() + : this(Guid.NewGuid().ToString(), currentUtcDateTime: null) + { + } + + // Allows tests to pin the InstanceId and CurrentUtcDateTime that feed into NewGuid(), since + // OrchestrationContext.CurrentUtcDateTime's setter is internal to DurableTask.Core and cannot + // be assigned directly from this assembly. + public TestOrchestrationContext(string instanceId, DateTime? currentUtcDateTime) { this.OrchestrationInstance = new() { - InstanceId = Guid.NewGuid().ToString(), + InstanceId = instanceId, ExecutionId = Guid.NewGuid().ToString(), }; + this.fixedCurrentUtcDateTime = currentUtcDateTime; } + public override DateTime CurrentUtcDateTime => this.fixedCurrentUtcDateTime ?? base.CurrentUtcDateTime; + public override void ContinueAsNew(object input) { throw new NotImplementedException(); From de8d9f1a89d709e0b0a7b96e52e59f89b6967517 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 16:23:46 -0700 Subject: [PATCH 02/10] Fix SHA1 disposal leak; document per-TFM native-allocation scope Address blocking review feedback on PR #780 (issue #778): 1. Cached SHA1 instance is now disposed. TaskOrchestrationContextWrapper implements IDisposable, and TaskOrchestrationShim disposes the previous wrapper right before replacing it with a new one on each replay/decision task, so the cached hash algorithm no longer accumulates undisposed instances. The base TaskOrchestration type has no disposal hook, so a documented CA1001 suppression covers the final instance per orchestration execution's lifetime, which is left for the GC/finalizer. 2. Corrected the code comment to honestly scope the optimization per TFM, verified against reference source: on .NET Framework, SHA1CryptoServiceProvider.Initialize() disposes/recreates its native CAPI handle every call, so native-handle churn is not eliminated there, only managed-side allocation. On modern .NET (net5.0+) on Windows, the CNG-based implementation can reuse and reset its native handle, so this caching also avoids native allocation on those runtimes. Added regression tests verifying Dispose() releases the cached SHA1 instance, is idempotent/safe to call without a prior NewGuid() call, and that GUID generation remains byte-identical after a mid-sequence dispose. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- .../Shims/TaskOrchestrationContextWrapper.cs | 22 ++++- .../Core/Shims/TaskOrchestrationShim.cs | 14 +++ .../TaskOrchestrationContextWrapperTests.cs | 85 +++++++++++++++++++ 3 files changed, 120 insertions(+), 1 deletion(-) diff --git a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs index e20c2d77..42e83f07 100644 --- a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs +++ b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs @@ -16,7 +16,7 @@ namespace Microsoft.DurableTask.Worker.Shims; /// /// A wrapper to go from to . /// -sealed partial class TaskOrchestrationContextWrapper : TaskOrchestrationContext +sealed partial class TaskOrchestrationContextWrapper : TaskOrchestrationContext, IDisposable { // We use a stack (a custom implementation using a single-linked list) to make it easier for users // to abandon external events that they no longer care about. The common case is a Task.WhenAny in a loop. @@ -38,6 +38,15 @@ sealed partial class TaskOrchestrationContextWrapper : TaskOrchestrationContext // for the duration of a single orchestration execution, and orchestrator code executes sequentially // (never concurrently) within that execution, so reusing this instance is safe. HashAlgorithm.Initialize() // resets all internal state before each use, so the computed hash is identical to using a fresh instance. + // This instance is disposed via Dispose() (see TaskOrchestrationShim, which disposes the previous + // wrapper before replacing it with a new one on the next replay/decision task). + // + // Note: on .NET Framework, the underlying SHA1CryptoServiceProvider.Initialize() disposes and + // recreates its native CAPI hash handle on every call, so the native-handle churn is not eliminated + // there -- only the managed-side allocation (the HashAlgorithm object itself and SHA1.Create()'s + // provider lookup) is avoided. On modern .NET (net5.0+) running on Windows, the CNG-based + // implementation can use a reusable hash handle (BCRYPT_HASH_REUSABLE_FLAG) and reset it in place, + // so this caching also avoids native-handle churn on those runtimes. SHA1? cachedHashAlgorithm; /// @@ -467,6 +476,17 @@ and the goal is to generate a deterministic Guid. We cannot update to SHA2-based return new Guid(newGuidByteArray); } + /// + /// Releases the resources cached by this instance, including the instance used by + /// . This should be called once this wrapper is no longer needed, i.e. once the + /// orchestration execution that owns it has completed and it is being replaced or discarded. + /// + public void Dispose() + { + this.cachedHashAlgorithm?.Dispose(); + this.cachedHashAlgorithm = null; + } + /// /// exits the critical section, if currently within a critical section. Otherwise, this has no effect. /// diff --git a/src/Worker/Core/Shims/TaskOrchestrationShim.cs b/src/Worker/Core/Shims/TaskOrchestrationShim.cs index eb7a179b..82fc6ec7 100644 --- a/src/Worker/Core/Shims/TaskOrchestrationShim.cs +++ b/src/Worker/Core/Shims/TaskOrchestrationShim.cs @@ -14,6 +14,12 @@ namespace Microsoft.DurableTask.Worker.Shims; /// This class is intended for use with alternate .NET-based durable task runtimes. It's not intended for use /// in application code. /// +#pragma warning disable CA1001 // Types that own disposable fields should be disposable -- the base + // TaskOrchestration type (defined in DurableTask.Core) has no disposal + // hook, so this type cannot implement IDisposable in a way that would ever + // be invoked. Instead, wrapperContext is disposed explicitly before each + // replacement in Execute(); only the final instance for a given orchestration + // execution's lifetime is left for the garbage collector/finalizer to reclaim. partial class TaskOrchestrationShim : TaskOrchestration { readonly ITaskOrchestrator implementation; @@ -64,6 +70,13 @@ public TaskOrchestrationShim( innerContext.ErrorDataConverter = converterShim; object? input = this.DataConverter.Deserialize(rawInput, this.implementation.InputType); + + // Dispose the previous execution's wrapper (if any) before replacing it. Each call to Execute + // represents a new replay/decision task with a fresh wrapper; orchestrator code always runs + // synchronously within a single Execute call, so the previous wrapper is guaranteed to no longer + // be in use once we reach this point, and releasing its cached resources (e.g. the SHA1 instance + // used by NewGuid) here avoids accumulating undisposed instances across replays. + this.wrapperContext?.Dispose(); this.wrapperContext = new(innerContext, this.invocationContext, input, this.properties); string instanceId = innerContext.OrchestrationInstance.InstanceId; @@ -119,3 +132,4 @@ public override void RaiseEvent(OrchestrationContext context, string name, strin this.wrapperContext?.CompleteExternalEvent(name, input); } } +#pragma warning restore CA1001 // Types that own disposable fields should be disposable diff --git a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs index 497603f5..09d8e1ab 100644 --- a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs @@ -4,6 +4,7 @@ using System.Collections.Generic; using System.Globalization; using System.Reflection; +using System.Security.Cryptography; using DurableTask.Core; using DurableTask.Core.Serializing.Internal; using Microsoft.Extensions.Logging.Abstractions; @@ -505,6 +506,90 @@ public void NewGuid_MultipleCalls_ReuseCachedHashAlgorithmInstance() afterSecondCall.Should().BeSameAs(afterFirstCall); } + [Fact] + public void Dispose_ReleasesCachedHashAlgorithm() + { + // Arrange + TestOrchestrationContext innerContext = new(); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + FieldInfo cachedHashAlgorithmField = typeof(TaskOrchestrationContextWrapper) + .GetField("cachedHashAlgorithm", BindingFlags.Instance | BindingFlags.NonPublic)!; + + wrapper.NewGuid(); + SHA1 cachedInstance = (SHA1)cachedHashAlgorithmField.GetValue(wrapper)!; + + // Act + wrapper.Dispose(); + + // Assert — the field is cleared, and the underlying instance was actually disposed (not merely + // dereferenced), confirmed by it throwing when used afterwards. + cachedHashAlgorithmField.GetValue(wrapper).Should().BeNull(); + Action useAfterDispose = () => cachedInstance.ComputeHash([1, 2, 3]); + useAfterDispose.Should().Throw(); + } + + [Fact] + public void Dispose_CalledMultipleTimes_DoesNotThrow() + { + // Arrange + TestOrchestrationContext innerContext = new(); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + wrapper.NewGuid(); + + // Act + Action dispose = () => + { + wrapper.Dispose(); + wrapper.Dispose(); + }; + + // Assert — disposing an already-disposed (or never-used) wrapper is safe. + dispose.Should().NotThrow(); + } + + [Fact] + public void Dispose_WithoutPriorNewGuidCall_DoesNotThrow() + { + // Arrange — the cached SHA1 instance is lazily created, so Dispose() must tolerate the case + // where NewGuid() was never called. + TestOrchestrationContext innerContext = new(); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + // Act + Action dispose = () => wrapper.Dispose(); + + // Assert + dispose.Should().NotThrow(); + } + + [Fact] + public void NewGuid_AfterDispose_StillProducesStableDeterministicValue() + { + // Arrange — Dispose() releases the cached SHA1 instance, but the wrapper lazily creates a new + // one on the next NewGuid() call (via the `??=` pattern). This must still produce byte-identical + // GUIDs to the ones computed with a fresh instance, proving disposal does not affect correctness. + TestOrchestrationContext innerContext = new( + "fixed-instance-id", + DateTime.Parse("2023-05-06T07:08:09.1234567Z", null, DateTimeStyles.RoundtripKind)); + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + // Act + wrapper.Dispose(); // dispose before any use is allowed (no-op, since nothing was cached yet) + Guid first = wrapper.NewGuid(); + wrapper.Dispose(); // dispose the now-cached instance mid-sequence + Guid second = wrapper.NewGuid(); // should transparently create a new instance and continue correctly + Guid third = wrapper.NewGuid(); + + // Assert — identical to the golden values in NewGuid_FixedInputs_ProducesStableDeterministicValue. + first.Should().Be(Guid.Parse("0f353f85-75d2-56f8-89b5-a7773ace7605")); + second.Should().Be(Guid.Parse("b0fd1465-f3d8-5a7e-98b1-f34137b15060")); + third.Should().Be(Guid.Parse("12bec829-d5e1-563c-ac70-9806cad148c1")); + } + static IReadOnlyDictionary GetLastScheduledTaskTags(TrackingOrchestrationContext innerContext) { PropertyInfo tagsProperty = innerContext.LastScheduledTaskOptions!.GetType().GetProperty("Tags")!; From 2f31ad6cc3a833bb8b9aeeb508251b97b87568fb Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 17:09:21 -0700 Subject: [PATCH 03/10] Fix cached SHA1 lifecycle leak at shim/executor/session boundaries TaskOrchestrationShim.Execute() runs exactly once per shim instance in every real call site; extended-session resumption reuses the cached executor via TaskOrchestrationExecutor.ExecuteNewEvents() and never calls Execute() again. This meant the round-2 fix (disposing the *previous* wrapper inside Execute() before replacing it) never actually freed the *final* cached SHA1 for normal execution or for extended sessions, leaking a native hash handle per orchestration instance until GC/finalization. This introduces deterministic ownership-based disposal instead: - TaskOrchestrationShim now implements IDisposable directly (the defensive dispose-previous-wrapper logic in Execute() is kept as a safety net but documented as effectively unreachable for current callers). - GrpcDurableTaskWorker's processor wraps shim construction/execution in try/finally and disposes the shim once its single work item completes. - GrpcOrchestrationRunner disposes the shim immediately when it is not handed off to the extended-session cache (completed on first execution, or the request isn't part of an extended session), and registers a PostEvictionCallbackRegistration on cached entries so the shim is disposed exactly once whenever the cache entry is evicted for any reason (explicit Remove, sliding-expiration timeout, capacity eviction, or cache disposal) -- never while the orchestration may still resume. No generated GUID bytes, algorithm, namespace, endianness, version bits, target frameworks, or public API are changed. Tests added: - TaskOrchestrationShimTests: Dispose() forwards to and actually releases the wrapper's cached SHA1 (ObjectDisposedException on reuse), is a no-op when Execute() was never called, and is idempotent across repeated calls. - GrpcOrchestrationRunnerTests: extended-session cache eviction via explicit removal and via sliding-expiration timeout both dispose the cached shim's SHA1, using reflection (no InternalsVisibleTo from this assembly) to reach the private wrapperContext/cachedHashAlgorithm fields via the public ExtendedSessionState.TaskOrchestration property. Verified: Worker.Tests 145/145 passed, Worker.Grpc.Tests 139/139 passed (including new eviction tests). Fixes #778 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- .../Core/Shims/TaskOrchestrationShim.cs | 39 +++-- .../Grpc/GrpcDurableTaskWorker.Processor.cs | 27 +++- src/Worker/Grpc/GrpcOrchestrationRunner.cs | 32 +++- .../Shims/TaskOrchestrationShimTests.cs | 139 ++++++++++++++++++ .../GrpcOrchestrationRunnerTests.cs | 127 ++++++++++++++++ 5 files changed, 341 insertions(+), 23 deletions(-) create mode 100644 test/Worker/Core.Tests/Shims/TaskOrchestrationShimTests.cs diff --git a/src/Worker/Core/Shims/TaskOrchestrationShim.cs b/src/Worker/Core/Shims/TaskOrchestrationShim.cs index 82fc6ec7..b7c8f1de 100644 --- a/src/Worker/Core/Shims/TaskOrchestrationShim.cs +++ b/src/Worker/Core/Shims/TaskOrchestrationShim.cs @@ -13,14 +13,16 @@ namespace Microsoft.DurableTask.Worker.Shims; /// /// This class is intended for use with alternate .NET-based durable task runtimes. It's not intended for use /// in application code. +/// +/// The base type (defined in DurableTask.Core) has no disposal hook of its +/// own, so the framework will never call automatically. Callers that construct a +/// directly (e.g. the gRPC worker processor and the orchestration +/// runner) own its lifetime and are responsible for disposing it once they are done with it -- typically +/// immediately after the single call completes, or (for extended sessions) when the +/// cached shim is evicted/removed. +/// /// -#pragma warning disable CA1001 // Types that own disposable fields should be disposable -- the base - // TaskOrchestration type (defined in DurableTask.Core) has no disposal - // hook, so this type cannot implement IDisposable in a way that would ever - // be invoked. Instead, wrapperContext is disposed explicitly before each - // replacement in Execute(); only the final instance for a given orchestration - // execution's lifetime is left for the garbage collector/finalizer to reclaim. -partial class TaskOrchestrationShim : TaskOrchestration +partial class TaskOrchestrationShim : TaskOrchestration, IDisposable { readonly ITaskOrchestrator implementation; readonly OrchestrationInvocationContext invocationContext; @@ -71,11 +73,12 @@ public TaskOrchestrationShim( object? input = this.DataConverter.Deserialize(rawInput, this.implementation.InputType); - // Dispose the previous execution's wrapper (if any) before replacing it. Each call to Execute - // represents a new replay/decision task with a fresh wrapper; orchestrator code always runs - // synchronously within a single Execute call, so the previous wrapper is guaranteed to no longer - // be in use once we reach this point, and releasing its cached resources (e.g. the SHA1 instance - // used by NewGuid) here avoids accumulating undisposed instances across replays. + // Defensively dispose any previous wrapper before replacing it, in case this shim instance is + // ever reused across more than one Execute call. Current callers construct a fresh shim per + // execution and call Execute exactly once, so the actual resource cleanup for this shim's wrapper + // happens via Dispose() (see the class remarks); this is still safe to do if it ever runs, since + // orchestrator code always runs synchronously within a single Execute call, so a previous wrapper + // is guaranteed to no longer be in use once we reach this point. this.wrapperContext?.Dispose(); this.wrapperContext = new(innerContext, this.invocationContext, input, this.properties); @@ -131,5 +134,15 @@ public override void RaiseEvent(OrchestrationContext context, string name, strin { this.wrapperContext?.CompleteExternalEvent(name, input); } + + /// + /// Releases the resources (e.g. the cached instance + /// backing ) held by this shim's current wrapper. Callers + /// that construct this shim directly are responsible for calling this once they are finished with it, + /// since the base type provides no framework-invoked disposal hook. + /// + public void Dispose() + { + this.wrapperContext?.Dispose(); + } } -#pragma warning restore CA1001 // Types that own disposable fields should be disposable diff --git a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs index 6d63f6af..20105510 100644 --- a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs +++ b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs @@ -774,14 +774,25 @@ await this.ExecuteWithRetryAsync( }; TaskOrchestration shim = this.shimFactory.CreateOrchestration(name, orchestrator, parent); - TaskOrchestrationExecutor executor = new( - runtimeState, - shim, - BehaviorOnContinueAsNew.Carryover, - request.EntityParameters.ToCore(), - ErrorPropagationMode.UseFailureDetails, - this.exceptionPropertiesProvider); - result = executor.Execute(); + try + { + TaskOrchestrationExecutor executor = new( + runtimeState, + shim, + BehaviorOnContinueAsNew.Carryover, + request.EntityParameters.ToCore(), + ErrorPropagationMode.UseFailureDetails, + this.exceptionPropertiesProvider); + result = executor.Execute(); + } + finally + { + // This worker (unlike the extended-session path in GrpcOrchestrationRunner) never + // reuses a shim across work items, so it owns and must dispose it once execution + // of this single work item completes (e.g. to release the SHA1 instance cached by + // NewGuid). + (shim as IDisposable)?.Dispose(); + } } else { diff --git a/src/Worker/Grpc/GrpcOrchestrationRunner.cs b/src/Worker/Grpc/GrpcOrchestrationRunner.cs index 5fbe2228..0b9c273e 100644 --- a/src/Worker/Grpc/GrpcOrchestrationRunner.cs +++ b/src/Worker/Grpc/GrpcOrchestrationRunner.cs @@ -207,15 +207,31 @@ public static string LoadAndRun( if (addToExtendedSessions && !executor.IsCompleted) { - // addToExtendedSessions can only be set to true if extendedSessions is not null + // addToExtendedSessions can only be set to true if extendedSessions is not null. + // The shim is now owned by the cache; it must not be disposed here since the + // orchestration may resume via ExecuteNewEvents() on a future call. Instead, register + // an eviction callback so it's disposed exactly once, whenever this entry is removed + // for any reason (explicit Remove, sliding-expiration timeout, capacity eviction, or + // cache disposal). + MemoryCacheEntryOptions cacheEntryOptions = new() + { + SlidingExpiration = TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds), + }; + cacheEntryOptions.RegisterPostEvictionCallback(DisposeEvictedExtendedSession); extendedSessions!.Set( request.InstanceId, new(runtimeState, shim, executor), - new MemoryCacheEntryOptions { SlidingExpiration = TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds) }); + cacheEntryOptions); } else { extendedSessions?.Remove(request.InstanceId); + + // This execution either isn't part of an extended session, or it completed on its + // first execution, so the shim was never handed off to the cache above. Nothing else + // will ever use it again, so it must be disposed here to release its resources (e.g. + // the SHA1 instance cached by NewGuid). + (shim as IDisposable)?.Dispose(); } } } @@ -232,4 +248,16 @@ public static string LoadAndRun( byte[] responseBytes = response.ToByteArray(); return Convert.ToBase64String(responseBytes); } + + // Invoked by the extended-sessions MemoryCache whenever a cached ExtendedSessionState entry is + // evicted, for any reason (explicit Remove, sliding-expiration timeout, capacity eviction, or the + // cache itself being disposed). Disposes the cached shim's resources (e.g. the SHA1 instance used by + // NewGuid) exactly once, at the point where the orchestration can no longer resume via this entry. + static void DisposeEvictedExtendedSession(object key, object? value, EvictionReason reason, object? state) + { + if (value is ExtendedSessionState sessionState && sessionState.TaskOrchestration is IDisposable disposable) + { + disposable.Dispose(); + } + } } diff --git a/test/Worker/Core.Tests/Shims/TaskOrchestrationShimTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationShimTests.cs new file mode 100644 index 00000000..20167272 --- /dev/null +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationShimTests.cs @@ -0,0 +1,139 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Reflection; +using System.Security.Cryptography; +using DurableTask.Core; +using Microsoft.Extensions.Logging.Abstractions; + +namespace Microsoft.DurableTask.Worker.Shims; + +public class TaskOrchestrationShimTests +{ + static readonly FieldInfo ShimWrapperContextField = typeof(TaskOrchestrationShim) + .GetField("wrapperContext", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException($"{nameof(TaskOrchestrationShim)}.wrapperContext was not found."); + + static readonly FieldInfo CachedHashAlgorithmField = typeof(TaskOrchestrationContextWrapper) + .GetField("cachedHashAlgorithm", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException( + $"{nameof(TaskOrchestrationContextWrapper)}.cachedHashAlgorithm was not found."); + + [Fact] + public void Dispose_ForwardsToWrapperContext_ReleasingCachedHashAlgorithm() + { + // Arrange. We inject the wrapper context the same way Execute() would, so we can prove that + // TaskOrchestrationShim.Dispose() forwards to TaskOrchestrationContextWrapper.Dispose(), + // actually releasing the cached SHA1 instance backing the deterministic NewGuid() + // optimization from issue #778 (not merely dereferencing it). + TaskOrchestrationShim shim = CreateShim(); + TaskOrchestrationContextWrapper wrapperContext = CreateWrapperContext(); + wrapperContext.NewGuid(); // Populate the cached SHA1 instance. + ShimWrapperContextField.SetValue(shim, wrapperContext); + + SHA1 cachedInstance = (SHA1)CachedHashAlgorithmField.GetValue(wrapperContext)!; + + // Act + shim.Dispose(); + + // Assert + CachedHashAlgorithmField.GetValue(wrapperContext).Should().BeNull(); + Action useAfterDispose = () => cachedInstance.ComputeHash(new byte[] { 1, 2, 3 }); + useAfterDispose.Should().Throw(); + } + + [Fact] + public void Dispose_WithoutExecute_DoesNotThrow() + { + // Arrange. Execute() was never called, so the shim's wrapperContext field is still null. This + // must remain a safe no-op (e.g. eviction/teardown paths may run before the shim ever executes). + TaskOrchestrationShim shim = CreateShim(); + + // Act + Action dispose = shim.Dispose; + + // Assert + dispose.Should().NotThrow(); + } + + [Fact] + public void Dispose_CalledMultipleTimes_DoesNotThrow() + { + // Arrange. Both the processor's try/finally and, defensively, an eviction callback could end up + // disposing the same shim; disposal must be idempotent. + TaskOrchestrationShim shim = CreateShim(); + TaskOrchestrationContextWrapper wrapperContext = CreateWrapperContext(); + wrapperContext.NewGuid(); + ShimWrapperContextField.SetValue(shim, wrapperContext); + + // Act + Action dispose = () => + { + shim.Dispose(); + shim.Dispose(); + }; + + // Assert + dispose.Should().NotThrow(); + } + + static TaskOrchestrationShim CreateShim() + { + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + return new TaskOrchestrationShim(invocationContext, new NoOpOrchestrator()); + } + + static TaskOrchestrationContextWrapper CreateWrapperContext() + { + OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); + TestOrchestrationContext innerContext = new(); + return new TaskOrchestrationContextWrapper(innerContext, invocationContext, deserializedInput: null); + } + + sealed class NoOpOrchestrator : ITaskOrchestrator + { + public Type InputType => typeof(object); + + public Type OutputType => typeof(object); + + public Task RunAsync(TaskOrchestrationContext context, object? input) => Task.FromResult(input); + } + + sealed class TestOrchestrationContext : OrchestrationContext + { + public TestOrchestrationContext() + { + this.OrchestrationInstance = new() + { + InstanceId = Guid.NewGuid().ToString(), + ExecutionId = Guid.NewGuid().ToString(), + }; + } + + public override void ContinueAsNew(object input) => throw new NotImplementedException(); + + public override void ContinueAsNew(string newVersion, object input) => throw new NotImplementedException(); + + public override Task CreateSubOrchestrationInstance(string name, string version, object input) + => throw new NotImplementedException(); + + public override Task CreateSubOrchestrationInstance( + string name, string version, string instanceId, object input) + => throw new NotImplementedException(); + + public override Task CreateSubOrchestrationInstance( + string name, string version, string instanceId, object input, IDictionary tags) + => throw new NotImplementedException(); + + public override Task CreateTimer(DateTime fireAt, T state) => throw new NotImplementedException(); + + public override Task CreateTimer(DateTime fireAt, T state, CancellationToken cancelToken) + => throw new NotImplementedException(); + + public override Task ScheduleTask(string name, string version, params object[] parameters) + => throw new NotImplementedException(); + + public override void SendEvent(OrchestrationInstance orchestrationInstance, string eventName, object eventData) + => throw new NotImplementedException(); + } +} diff --git a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs index 0bb6be24..0d9c628b 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Reflection; +using System.Security.Cryptography; using Google.Protobuf; using Google.Protobuf.Collections; using Google.Protobuf.WellKnownTypes; @@ -455,6 +457,92 @@ public void PastEventIncluded_Means_ExtendedSession_Evicted() Assert.True(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out extendedSession)); } + [Fact] + public void ExternallyEndedExtendedSession_Evicted_DisposesCachedShimResources() + { + // Regression test for the round-3 lifecycle fix: the extended-session MemoryCache must dispose + // the cached shim (and, transitively, its wrapper's cached SHA1 backing NewGuid()) whenever an + // entry is evicted -- not just when the shim is replaced within a single Execute() call. + using var extendedSessions = new ExtendedSessionsCache(); + var historyEvent = new Protobuf.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + ExecutionStarted = new Protobuf.ExecutionStartedEvent() + { + OrchestrationInstance = new Protobuf.OrchestrationInstance + { + InstanceId = TestInstanceId, + ExecutionId = TestExecutionId, + }, + } + }; + Protobuf.OrchestratorRequest orchestratorRequest = CreateOrchestratorRequest([historyEvent]); + orchestratorRequest.Properties.Add(new MapField() { + { "IncludeState", Value.ForBool(true) }, + { "IsExtendedSession", Value.ForBool(true) }, + { "ExtendedSessionIdleTimeoutInSeconds", Value.ForNumber(DefaultExtendedSessionIdleTimeoutInSeconds) } }); + byte[] requestBytes = orchestratorRequest.ToByteArray(); + string requestString = Convert.ToBase64String(requestBytes); + GrpcOrchestrationRunner.LoadAndRun(requestString, new NewGuidThenCallSubOrchestrationOrchestrator(), extendedSessions); + Assert.True(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out object? extendedSession)); + SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(extendedSession!); + + // Now set the extended session flag to false for this instance, which removes/evicts the cache + // entry and should synchronously invoke the eviction callback that disposes the cached shim. + orchestratorRequest.Properties.Clear(); + orchestratorRequest.Properties.Add(new MapField() { + { "IncludeState", Value.ForBool(true) }, + { "IsExtendedSession", Value.ForBool(false) }, + { "ExtendedSessionIdleTimeoutInSeconds", Value.ForNumber(DefaultExtendedSessionIdleTimeoutInSeconds) } }); + requestBytes = orchestratorRequest.ToByteArray(); + requestString = Convert.ToBase64String(requestBytes); + GrpcOrchestrationRunner.LoadAndRun(requestString, new NewGuidThenCallSubOrchestrationOrchestrator(), extendedSessions); + Assert.False(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out _)); + + Action useAfterDispose = () => cachedHashAlgorithm.ComputeHash(new byte[] { 1, 2, 3 }); + useAfterDispose.Should().Throw(); + } + + [Fact] + public async Task Stale_ExtendedSession_Evicted_DisposesCachedShimResources_Async() + { + // Regression test for the round-3 lifecycle fix: sliding-expiration eviction of a stale + // extended session must also dispose the cached shim's resources. + using var extendedSessions = new ExtendedSessionsCache(); + int extendedSessionIdleTimeout = 5; + var historyEvent = new Protobuf.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + ExecutionStarted = new Protobuf.ExecutionStartedEvent() + { + OrchestrationInstance = new Protobuf.OrchestrationInstance + { + InstanceId = TestInstanceId, + ExecutionId = TestExecutionId, + }, + } + }; + Protobuf.OrchestratorRequest orchestratorRequest = CreateOrchestratorRequest([historyEvent]); + orchestratorRequest.Properties.Add(new MapField() { + { "IncludeState", Value.ForBool(true) }, + { "IsExtendedSession", Value.ForBool(true) }, + { "ExtendedSessionIdleTimeoutInSeconds", Value.ForNumber(extendedSessionIdleTimeout) } }); + byte[] requestBytes = orchestratorRequest.ToByteArray(); + string requestString = Convert.ToBase64String(requestBytes); + GrpcOrchestrationRunner.LoadAndRun(requestString, new NewGuidThenCallSubOrchestrationOrchestrator(), extendedSessions); + Assert.True(extendedSessions.GetOrInitializeCache(extendedSessionIdleTimeout).TryGetValue(TestInstanceId, out object? extendedSession)); + SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(extendedSession!); + + // Wait for longer than the timeout to account for finite cache scan for stale items frequency + await Task.Delay(extendedSessionIdleTimeout * 1000 * 2); + Assert.False(extendedSessions.GetOrInitializeCache(extendedSessionIdleTimeout).TryGetValue(TestInstanceId, out _)); + + Action useAfterDispose = () => cachedHashAlgorithm.ComputeHash(new byte[] { 1, 2, 3 }); + useAfterDispose.Should().Throw(); + } + [Fact] public void Null_ExtendedSessionsCache_IsOkay() { @@ -498,6 +586,32 @@ public void Null_ExtendedSessionsCache_IsOkay() Assert.Equal(Protobuf.OrchestrationStatus.Completed, response.Actions[0].CompleteOrchestration.OrchestrationStatus); } + // TaskOrchestrationShim and TaskOrchestrationContextWrapper are internal to the Worker.Core assembly + // and not visible to this test assembly via InternalsVisibleTo, so reflection is used to reach into + // the cached shim (exposed only as the public ExtendedSessionState.TaskOrchestration property, typed + // as the public base class TaskOrchestration) and pull out its wrapper's cached SHA1 instance. + static SHA1 GetCachedHashAlgorithm(object extendedSessionState) + { + PropertyInfo taskOrchestrationProperty = extendedSessionState.GetType() + .GetProperty("TaskOrchestration", BindingFlags.Instance | BindingFlags.Public) + ?? throw new InvalidOperationException("ExtendedSessionState.TaskOrchestration was not found."); + object shim = taskOrchestrationProperty.GetValue(extendedSessionState) + ?? throw new InvalidOperationException("ExtendedSessionState.TaskOrchestration was null."); + + FieldInfo wrapperContextField = shim.GetType() + .GetField("wrapperContext", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException("TaskOrchestrationShim.wrapperContext was not found."); + object wrapperContext = wrapperContextField.GetValue(shim) + ?? throw new InvalidOperationException("TaskOrchestrationShim.wrapperContext was null."); + + FieldInfo cachedHashAlgorithmField = wrapperContext.GetType() + .GetField("cachedHashAlgorithm", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException( + "TaskOrchestrationContextWrapper.cachedHashAlgorithm was not found."); + return (SHA1)(cachedHashAlgorithmField.GetValue(wrapperContext) + ?? throw new InvalidOperationException("cachedHashAlgorithm was null; NewGuid() may not have run.")); + } + static Protobuf.OrchestratorRequest CreateOrchestratorRequest(IEnumerable newEvents) { var orchestratorRequest = new Protobuf.OrchestratorRequest() @@ -529,4 +643,17 @@ public override async Task RunAsync(TaskOrchestrationContext context, st return input; } } + + // Same shape as CallSubOrchestrationOrchestrator (so the orchestration is left pending in the + // extended-session cache) but also calls NewGuid() before awaiting, so the cached shim's wrapper + // has a live SHA1 instance whose disposal we can observe once the extended session is evicted. + class NewGuidThenCallSubOrchestrationOrchestrator : TaskOrchestrator + { + public override async Task RunAsync(TaskOrchestrationContext context, string input) + { + context.NewGuid(); + await context.CallSubOrchestratorAsync(nameof(SimpleOrchestrator)); + return input; + } + } } From 38d424dfec9504ef52f93636845cc4c227a2810b Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 17:22:35 -0700 Subject: [PATCH 04/10] Fix remaining MemoryCache disposal lifecycle gaps Addresses three remaining Medium lifecycle issues from round-4 review of PR #780: 1. ExtendedSessionsCache.Dispose() now calls MemoryCache.Clear() before Dispose(). Confirmed against the pinned Microsoft.Extensions.Caching.Memory 8.0.1 source that MemoryCache.Dispose() alone does NOT invoke post-eviction callbacks for entries still present in the cache -- only Clear(), Remove(), and the internal capacity/expiration removal paths do. Without this fix, any extended session still cached at worker shutdown would leak its cached shim (and SHA1 instance). 2. GrpcOrchestrationRunner's direct-execution path now wraps TaskOrchestrationExecutor construction/Execute()/cache Set() in try/finally, using a transferredShimToCache flag to gate disposal. If Execute() (or the subsequent Set() call) throws before the shim's ownership is transferred to the extended-sessions cache, the shim is now still deterministically disposed in finally instead of leaking. 3. The round-3 eviction-disposal tests asserted disposal immediately after triggering an eviction/removal, but MemoryCache dispatches post-eviction callbacks via Task.Factory.StartNew (i.e. asynchronously, on a background thread), so this was a latent race. Added a bounded polling helper (WaitUntilDisposedAsync) used by both existing tests, and added a new regression test (ExtendedSessionsCache_Dispose_DisposesCachedShimResources) proving that disposing the cache itself while an extended session is still pending disposes the cached shim's resources. Verified: - Worker.Tests: 145/145 passed - Worker.Grpc.Tests: 140/140 passed (139 + 1 new) Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- src/Worker/Core/ExtendedSessionsCache.cs | 8 ++ src/Worker/Grpc/GrpcOrchestrationRunner.cs | 77 +++++++++------- .../GrpcOrchestrationRunnerTests.cs | 89 +++++++++++++++++-- 3 files changed, 138 insertions(+), 36 deletions(-) diff --git a/src/Worker/Core/ExtendedSessionsCache.cs b/src/Worker/Core/ExtendedSessionsCache.cs index 59df2536..ef7f4661 100644 --- a/src/Worker/Core/ExtendedSessionsCache.cs +++ b/src/Worker/Core/ExtendedSessionsCache.cs @@ -24,6 +24,14 @@ public class ExtendedSessionsCache : IDisposable /// public void Dispose() { + // MemoryCache.Dispose() does NOT invoke post-eviction callbacks for entries that are still + // present in the cache -- it merely tears down the cache's internal state. Any entries + // (e.g. cached extended-session state holding an IDisposable shim) that are still cached at + // shutdown would therefore never be disposed. Calling Clear() first forces every remaining + // entry to be removed via the normal removal path, which does invoke eviction callbacks + // synchronously-scheduled (via Task.Factory.StartNew) for each entry, ensuring deterministic + // cleanup of any cached shim/SHA1 resources before the cache itself is torn down. + this.extendedSessions?.Clear(); this.extendedSessions?.Dispose(); GC.SuppressFinalize(this); } diff --git a/src/Worker/Grpc/GrpcOrchestrationRunner.cs b/src/Worker/Grpc/GrpcOrchestrationRunner.cs index 0b9c273e..3d085da3 100644 --- a/src/Worker/Grpc/GrpcOrchestrationRunner.cs +++ b/src/Worker/Grpc/GrpcOrchestrationRunner.cs @@ -197,41 +197,56 @@ public static string LoadAndRun( : ActivatorUtilities.GetServiceOrCreateInstance(services); TaskOrchestration shim = factory.CreateOrchestration(orchestratorName, implementation, properties, parent); - TaskOrchestrationExecutor executor = new( - runtimeState, - shim, - BehaviorOnContinueAsNew.Carryover, - request.EntityParameters.ToCore(), - ErrorPropagationMode.UseFailureDetails); - result = executor.Execute(); - - if (addToExtendedSessions && !executor.IsCompleted) + // Tracks whether ownership of the shim has been successfully transferred to the + // extended-sessions cache. Execute() (or the subsequent cache Set() call) could throw; + // in that case ownership is never transferred, and the finally block below must dispose + // the shim itself rather than leaking it. + bool transferredShimToCache = false; + try { - // addToExtendedSessions can only be set to true if extendedSessions is not null. - // The shim is now owned by the cache; it must not be disposed here since the - // orchestration may resume via ExecuteNewEvents() on a future call. Instead, register - // an eviction callback so it's disposed exactly once, whenever this entry is removed - // for any reason (explicit Remove, sliding-expiration timeout, capacity eviction, or - // cache disposal). - MemoryCacheEntryOptions cacheEntryOptions = new() + TaskOrchestrationExecutor executor = new( + runtimeState, + shim, + BehaviorOnContinueAsNew.Carryover, + request.EntityParameters.ToCore(), + ErrorPropagationMode.UseFailureDetails); + result = executor.Execute(); + + if (addToExtendedSessions && !executor.IsCompleted) + { + // addToExtendedSessions can only be set to true if extendedSessions is not null. + // The shim is now owned by the cache; it must not be disposed here since the + // orchestration may resume via ExecuteNewEvents() on a future call. Instead, register + // an eviction callback so it's disposed exactly once, whenever this entry is removed + // for any reason (explicit Remove, sliding-expiration timeout, capacity eviction, or + // cache disposal). + MemoryCacheEntryOptions cacheEntryOptions = new() + { + SlidingExpiration = TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds), + }; + cacheEntryOptions.RegisterPostEvictionCallback(DisposeEvictedExtendedSession); + extendedSessions!.Set( + request.InstanceId, + new(runtimeState, shim, executor), + cacheEntryOptions); + transferredShimToCache = true; + } + else { - SlidingExpiration = TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds), - }; - cacheEntryOptions.RegisterPostEvictionCallback(DisposeEvictedExtendedSession); - extendedSessions!.Set( - request.InstanceId, - new(runtimeState, shim, executor), - cacheEntryOptions); + extendedSessions?.Remove(request.InstanceId); + } } - else + finally { - extendedSessions?.Remove(request.InstanceId); - - // This execution either isn't part of an extended session, or it completed on its - // first execution, so the shim was never handed off to the cache above. Nothing else - // will ever use it again, so it must be disposed here to release its resources (e.g. - // the SHA1 instance cached by NewGuid). - (shim as IDisposable)?.Dispose(); + if (!transferredShimToCache) + { + // This execution either isn't part of an extended session, it completed on its + // first execution, or an exception was thrown before ownership could be + // transferred to the cache above. In every one of these cases nothing else will + // ever use the shim again, so it must be disposed here to release its resources + // (e.g. the SHA1 instance cached by NewGuid). + (shim as IDisposable)?.Dispose(); + } } } } diff --git a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs index 0d9c628b..978b4b1a 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Diagnostics; using System.Reflection; using System.Security.Cryptography; using Google.Protobuf; @@ -458,11 +459,16 @@ public void PastEventIncluded_Means_ExtendedSession_Evicted() } [Fact] - public void ExternallyEndedExtendedSession_Evicted_DisposesCachedShimResources() + public async Task ExternallyEndedExtendedSession_Evicted_DisposesCachedShimResources() { // Regression test for the round-3 lifecycle fix: the extended-session MemoryCache must dispose // the cached shim (and, transitively, its wrapper's cached SHA1 backing NewGuid()) whenever an // entry is evicted -- not just when the shim is replaced within a single Execute() call. + // + // Note: MemoryCache invokes post-eviction callbacks via Task.Factory.StartNew (i.e. + // asynchronously, on a background thread), so disposal is not guaranteed to have happened by + // the time Remove()/TryGetValue() returns. WaitUntilDisposedAsync polls with a bounded timeout + // instead of asserting disposal immediately, to avoid flakiness under load. using var extendedSessions = new ExtendedSessionsCache(); var historyEvent = new Protobuf.HistoryEvent { @@ -500,8 +506,7 @@ public void ExternallyEndedExtendedSession_Evicted_DisposesCachedShimResources() GrpcOrchestrationRunner.LoadAndRun(requestString, new NewGuidThenCallSubOrchestrationOrchestrator(), extendedSessions); Assert.False(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out _)); - Action useAfterDispose = () => cachedHashAlgorithm.ComputeHash(new byte[] { 1, 2, 3 }); - useAfterDispose.Should().Throw(); + await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); } [Fact] @@ -509,6 +514,11 @@ public async Task Stale_ExtendedSession_Evicted_DisposesCachedShimResources_Asyn { // Regression test for the round-3 lifecycle fix: sliding-expiration eviction of a stale // extended session must also dispose the cached shim's resources. + // + // Note: MemoryCache invokes post-eviction callbacks via Task.Factory.StartNew (i.e. + // asynchronously, on a background thread), so disposal is not guaranteed to have happened + // immediately after the scan removes the entry. WaitUntilDisposedAsync polls with a bounded + // timeout instead of asserting disposal immediately, to avoid flakiness under load. using var extendedSessions = new ExtendedSessionsCache(); int extendedSessionIdleTimeout = 5; var historyEvent = new Protobuf.HistoryEvent @@ -539,8 +549,47 @@ public async Task Stale_ExtendedSession_Evicted_DisposesCachedShimResources_Asyn await Task.Delay(extendedSessionIdleTimeout * 1000 * 2); Assert.False(extendedSessions.GetOrInitializeCache(extendedSessionIdleTimeout).TryGetValue(TestInstanceId, out _)); - Action useAfterDispose = () => cachedHashAlgorithm.ComputeHash(new byte[] { 1, 2, 3 }); - useAfterDispose.Should().Throw(); + await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); + } + + [Fact] + public async Task ExtendedSessionsCache_Dispose_DisposesCachedShimResources() + { + // Regression test for round-4: MemoryCache.Dispose() alone does NOT invoke post-eviction + // callbacks for entries that are still cached (confirmed against the pinned + // Microsoft.Extensions.Caching.Memory 8.0.1 source), so ExtendedSessionsCache.Dispose() must + // explicitly Clear() the cache before disposing it. Otherwise a still-pending extended session + // at worker shutdown would leak its cached shim's resources (e.g. the SHA1 instance backing + // NewGuid()) for the remaining lifetime of the process. + var extendedSessions = new ExtendedSessionsCache(); + var historyEvent = new Protobuf.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + ExecutionStarted = new Protobuf.ExecutionStartedEvent() + { + OrchestrationInstance = new Protobuf.OrchestrationInstance + { + InstanceId = TestInstanceId, + ExecutionId = TestExecutionId, + }, + } + }; + Protobuf.OrchestratorRequest orchestratorRequest = CreateOrchestratorRequest([historyEvent]); + orchestratorRequest.Properties.Add(new MapField() { + { "IncludeState", Value.ForBool(true) }, + { "IsExtendedSession", Value.ForBool(true) }, + { "ExtendedSessionIdleTimeoutInSeconds", Value.ForNumber(DefaultExtendedSessionIdleTimeoutInSeconds) } }); + byte[] requestBytes = orchestratorRequest.ToByteArray(); + string requestString = Convert.ToBase64String(requestBytes); + GrpcOrchestrationRunner.LoadAndRun(requestString, new NewGuidThenCallSubOrchestrationOrchestrator(), extendedSessions); + Assert.True(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out object? extendedSession)); + SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(extendedSession!); + + // Simulate a worker shutting down while an extended session is still pending in the cache. + extendedSessions.Dispose(); + + await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); } [Fact] @@ -612,6 +661,36 @@ static SHA1 GetCachedHashAlgorithm(object extendedSessionState) ?? throw new InvalidOperationException("cachedHashAlgorithm was null; NewGuid() may not have run.")); } + // Eviction callbacks on the extended-sessions MemoryCache are dispatched via + // Task.Factory.StartNew (i.e. asynchronously, on a background thread pool task) rather than + // synchronously on the calling thread -- see Microsoft.Extensions.Caching.Memory's + // CacheEntryTokens.InvokeEvictionCallbacks. This helper polls with a bounded timeout for the given + // SHA1 instance to become disposed, instead of asserting disposal immediately after triggering an + // eviction, to avoid flakiness from that inherent async scheduling. + static async Task WaitUntilDisposedAsync(SHA1 hashAlgorithm, TimeSpan timeout) + { + Stopwatch stopwatch = Stopwatch.StartNew(); + while (true) + { + try + { + hashAlgorithm.ComputeHash([1, 2, 3]); + } + catch (ObjectDisposedException) + { + return; + } + + if (stopwatch.Elapsed >= timeout) + { + throw new TimeoutException( + $"SHA1 instance was not disposed within {timeout} of the triggering eviction/disposal."); + } + + await Task.Delay(20); + } + } + static Protobuf.OrchestratorRequest CreateOrchestratorRequest(IEnumerable newEvents) { var orchestratorRequest = new Protobuf.OrchestratorRequest() From 8c8be3904acef65db3cd4f9d6f0e5faf3ae4b284 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 17:40:10 -0700 Subject: [PATCH 05/10] Fix ExtendedSessionsCache.Dispose() idempotency Round-5 fix: MemoryCache.Clear() (added in round 4 to force eviction callbacks to run before shutdown) throws ObjectDisposedException if called on an already-disposed MemoryCache. Since Dispose() was not guarded against repeat calls, invoking it twice (e.g. duplicate shutdown hooks, or concurrent shutdown paths) would throw instead of being a safe no-op. - Add an `int disposed` field guarded via `Interlocked.Exchange(ref this.disposed, 1) != 0` at the top of Dispose(), making it both idempotent (repeat calls short-circuit) and thread-safe (concurrent calls only let one caller proceed). - Clarify the disposal comment: Clear() only guarantees eviction callbacks are scheduled (via Task.Factory.StartNew) before Dispose() returns, not that they have completed. - Add ExtendedSessionsCache_Dispose_CalledMultipleTimes_DoesNotThrow (sequential double-dispose) and ExtendedSessionsCache_Dispose_CalledConcurrently_DoesNotThrow (8-thread concurrent dispose) regression tests. Verified: targeted GrpcOrchestrationRunnerTests 17/17, full Worker.Grpc.Tests 142/142, full Worker.Tests 145/145. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- src/Worker/Core/ExtendedSessionsCache.cs | 18 +++- .../GrpcOrchestrationRunnerTests.cs | 85 +++++++++++++++++++ 2 files changed, 100 insertions(+), 3 deletions(-) diff --git a/src/Worker/Core/ExtendedSessionsCache.cs b/src/Worker/Core/ExtendedSessionsCache.cs index ef7f4661..024b9da8 100644 --- a/src/Worker/Core/ExtendedSessionsCache.cs +++ b/src/Worker/Core/ExtendedSessionsCache.cs @@ -12,6 +12,7 @@ namespace Microsoft.DurableTask.Worker; public class ExtendedSessionsCache : IDisposable { MemoryCache? extendedSessions; + int disposed; /// /// Gets a value indicating whether returns whether or not the cache has been initialized. @@ -24,13 +25,24 @@ public class ExtendedSessionsCache : IDisposable /// public void Dispose() { + if (System.Threading.Interlocked.Exchange(ref this.disposed, 1) != 0) + { + // Already disposed (or a concurrent Dispose() call won the race). MemoryCache.Clear() + // and MemoryCache.Dispose() are not safe to call more than once -- Clear() throws + // ObjectDisposedException if the cache has already been disposed -- so this guard makes + // Dispose() idempotent and safe under concurrent callers. + return; + } + // MemoryCache.Dispose() does NOT invoke post-eviction callbacks for entries that are still // present in the cache -- it merely tears down the cache's internal state. Any entries // (e.g. cached extended-session state holding an IDisposable shim) that are still cached at // shutdown would therefore never be disposed. Calling Clear() first forces every remaining - // entry to be removed via the normal removal path, which does invoke eviction callbacks - // synchronously-scheduled (via Task.Factory.StartNew) for each entry, ensuring deterministic - // cleanup of any cached shim/SHA1 resources before the cache itself is torn down. + // entry to be removed via the normal removal path, which does invoke eviction callbacks for + // each entry, ensuring they are triggered instead of silently skipped. Note that eviction + // callbacks are queued asynchronously (via Task.Factory.StartNew), so Clear() does not + // guarantee those callbacks have completed by the time Dispose() returns -- it only + // guarantees they are scheduled before the cache itself is torn down. this.extendedSessions?.Clear(); this.extendedSessions?.Dispose(); GC.SuppressFinalize(this); diff --git a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs index 978b4b1a..f8a676d3 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -592,6 +592,91 @@ public async Task ExtendedSessionsCache_Dispose_DisposesCachedShimResources() await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); } + [Fact] + public async Task ExtendedSessionsCache_Dispose_CalledMultipleTimes_DoesNotThrow() + { + // Regression test: ExtendedSessionsCache.Dispose() calls MemoryCache.Clear() before + // MemoryCache.Dispose() (see above). MemoryCache.Clear() throws ObjectDisposedException if + // the cache was already disposed, so without an idempotency guard, a second Dispose() call + // (e.g. from a duplicate shutdown-hook invocation) would throw instead of being a safe no-op. + var extendedSessions = new ExtendedSessionsCache(); + var historyEvent = new Protobuf.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + ExecutionStarted = new Protobuf.ExecutionStartedEvent() + { + OrchestrationInstance = new Protobuf.OrchestrationInstance + { + InstanceId = TestInstanceId, + ExecutionId = TestExecutionId, + }, + } + }; + Protobuf.OrchestratorRequest orchestratorRequest = CreateOrchestratorRequest([historyEvent]); + orchestratorRequest.Properties.Add(new MapField() { + { "IncludeState", Value.ForBool(true) }, + { "IsExtendedSession", Value.ForBool(true) }, + { "ExtendedSessionIdleTimeoutInSeconds", Value.ForNumber(DefaultExtendedSessionIdleTimeoutInSeconds) } }); + byte[] requestBytes = orchestratorRequest.ToByteArray(); + string requestString = Convert.ToBase64String(requestBytes); + GrpcOrchestrationRunner.LoadAndRun(requestString, new NewGuidThenCallSubOrchestrationOrchestrator(), extendedSessions); + Assert.True(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out object? extendedSession)); + SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(extendedSession!); + + // Act: dispose the cache twice in a row. + Exception? firstDisposeException = Record.Exception(() => extendedSessions.Dispose()); + Exception? secondDisposeException = Record.Exception(() => extendedSessions.Dispose()); + + // Assert: neither call throws, and the cached shim resources are still disposed exactly once. + Assert.Null(firstDisposeException); + Assert.Null(secondDisposeException); + await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); + } + + [Fact] + public async Task ExtendedSessionsCache_Dispose_CalledConcurrently_DoesNotThrow() + { + // Regression test: guards the Dispose() idempotency fix above against a race between two + // threads calling Dispose() at (approximately) the same time -- e.g. overlapping shutdown + // paths -- rather than only the simpler sequential double-dispose case above. + var extendedSessions = new ExtendedSessionsCache(); + var historyEvent = new Protobuf.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + ExecutionStarted = new Protobuf.ExecutionStartedEvent() + { + OrchestrationInstance = new Protobuf.OrchestrationInstance + { + InstanceId = TestInstanceId, + ExecutionId = TestExecutionId, + }, + } + }; + Protobuf.OrchestratorRequest orchestratorRequest = CreateOrchestratorRequest([historyEvent]); + orchestratorRequest.Properties.Add(new MapField() { + { "IncludeState", Value.ForBool(true) }, + { "IsExtendedSession", Value.ForBool(true) }, + { "ExtendedSessionIdleTimeoutInSeconds", Value.ForNumber(DefaultExtendedSessionIdleTimeoutInSeconds) } }); + byte[] requestBytes = orchestratorRequest.ToByteArray(); + string requestString = Convert.ToBase64String(requestBytes); + GrpcOrchestrationRunner.LoadAndRun(requestString, new NewGuidThenCallSubOrchestrationOrchestrator(), extendedSessions); + Assert.True(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out object? extendedSession)); + SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(extendedSession!); + + // Act: dispose the cache concurrently from several threads. + Task[] disposeTasks = Enumerable.Range(0, 8) + .Select(_ => Task.Run(() => extendedSessions.Dispose())) + .ToArray(); + Exception? concurrentDisposeException = await Record.ExceptionAsync(() => Task.WhenAll(disposeTasks)); + + // Assert: none of the concurrent calls throw, and the cached shim resources are still + // disposed exactly once. + Assert.Null(concurrentDisposeException); + await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); + } + [Fact] public void Null_ExtendedSessionsCache_IsOkay() { From 057ac9e9f3cb8fd272c2cca417d2279fbc595fdc Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 17:53:38 -0700 Subject: [PATCH 06/10] Fix ExtendedSessionsCache Dispose/GetOrInitializeCache race (round 6 review) - Replace the round-5 Interlocked.Exchange disposed flag with a shared syncRoot lock guarding both Dispose() and GetOrInitializeCache(). This closes a race where Dispose() could observe an uninitialized cache field, mark itself disposed, and return -- while a concurrent GetOrInitializeCache() call then lazily constructed a brand-new MemoryCache that would never be disposed (since disposed was already permanently true), permanently leaking it. - GetOrInitializeCache() now throws ObjectDisposedException immediately if called after Dispose(), instead of possibly returning an already-disposed MemoryCache or racing to create an unreachable one. - Dispose() clears the extendedSessions field to null under the lock before tearing down the captured local reference outside the lock (avoids holding the lock during Clear()/Dispose(), which is safe since the eviction callback never calls back into ExtendedSessionsCache). - Add ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_NeverLeaksCache, racing Dispose() against GetOrInitializeCache() across many iterations and asserting the returned cache (if any) is provably disposed, and that a repeated Dispose() call remains a safe no-op. - Fix a stale comment in ExternallyEndedExtendedSession_Evicted_DisposesCachedShimResources claiming the eviction callback fires synchronously; it is dispatched asynchronously and the test already awaits WaitUntilDisposedAsync. - Use CultureInfo.InvariantCulture instead of null for the DateTime.Parse format provider in the NewGuid() golden-value regression tests, removing any dependency on the test runner's current culture. Verified: targeted GrpcOrchestrationRunnerTests 18/18, full Worker.Grpc.Tests 143/143, full Worker.Tests 145/145 (golden GUID values unchanged after the culture fix). Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- src/Worker/Core/ExtendedSessionsCache.cs | 57 +++++++++++++---- .../TaskOrchestrationContextWrapperTests.cs | 10 +-- .../GrpcOrchestrationRunnerTests.cs | 64 ++++++++++++++++++- 3 files changed, 112 insertions(+), 19 deletions(-) diff --git a/src/Worker/Core/ExtendedSessionsCache.cs b/src/Worker/Core/ExtendedSessionsCache.cs index 024b9da8..3b9a6dd0 100644 --- a/src/Worker/Core/ExtendedSessionsCache.cs +++ b/src/Worker/Core/ExtendedSessionsCache.cs @@ -11,8 +11,18 @@ namespace Microsoft.DurableTask.Worker; /// public class ExtendedSessionsCache : IDisposable { + // Guards both the lazy-initialization of `extendedSessions` in GetOrInitializeCache() and the + // disposal state transition in Dispose(). Without this shared lock, Dispose() could observe + // `extendedSessions` as null (because it hasn't been lazily created yet), mark itself disposed, + // and return -- while a concurrent GetOrInitializeCache() call races in and constructs a brand + // new MemoryCache immediately afterwards. That cache would never be disposed (Dispose() has + // already run and is now a permanent no-op), leaking it and any entries added to it. The lock + // makes initialization and disposal mutually exclusive, so there's no window where a cache can + // be created after (or concurrently with) disposal. + readonly object syncRoot = new(); + MemoryCache? extendedSessions; - int disposed; + bool disposed; /// /// Gets a value indicating whether returns whether or not the cache has been initialized. @@ -25,13 +35,25 @@ public class ExtendedSessionsCache : IDisposable /// public void Dispose() { - if (System.Threading.Interlocked.Exchange(ref this.disposed, 1) != 0) + MemoryCache? cacheToDispose; + lock (this.syncRoot) { - // Already disposed (or a concurrent Dispose() call won the race). MemoryCache.Clear() - // and MemoryCache.Dispose() are not safe to call more than once -- Clear() throws - // ObjectDisposedException if the cache has already been disposed -- so this guard makes - // Dispose() idempotent and safe under concurrent callers. - return; + if (this.disposed) + { + // Already disposed by a previous (or concurrent, now-completed) call. MemoryCache.Clear() + // and MemoryCache.Dispose() are not safe to call more than once -- Clear() throws + // ObjectDisposedException if the cache has already been disposed -- so this guard makes + // Dispose() idempotent and safe under concurrent callers. + return; + } + + this.disposed = true; + + // Clear the field (under the same lock used by GetOrInitializeCache()) so that no caller + // can observe or lazily recreate a cache after this point; GetOrInitializeCache() checks + // `this.disposed` under the lock and throws ObjectDisposedException instead. + cacheToDispose = this.extendedSessions; + this.extendedSessions = null; } // MemoryCache.Dispose() does NOT invoke post-eviction callbacks for entries that are still @@ -43,8 +65,8 @@ public void Dispose() // callbacks are queued asynchronously (via Task.Factory.StartNew), so Clear() does not // guarantee those callbacks have completed by the time Dispose() returns -- it only // guarantees they are scheduled before the cache itself is torn down. - this.extendedSessions?.Clear(); - this.extendedSessions?.Dispose(); + cacheToDispose?.Clear(); + cacheToDispose?.Dispose(); GC.SuppressFinalize(this); } @@ -56,13 +78,22 @@ public void Dispose() /// This specifies how often the cache checks for stale items, and evicts them. /// /// The IMemoryCache that holds the cached . + /// The cache has already been disposed. public MemoryCache GetOrInitializeCache(double expirationScanFrequencyInSeconds) { - this.extendedSessions ??= new MemoryCache(new MemoryCacheOptions + lock (this.syncRoot) { - ExpirationScanFrequency = TimeSpan.FromSeconds(expirationScanFrequencyInSeconds / 5), - }); + if (this.disposed) + { + throw new ObjectDisposedException(nameof(ExtendedSessionsCache)); + } - return this.extendedSessions; + this.extendedSessions ??= new MemoryCache(new MemoryCacheOptions + { + ExpirationScanFrequency = TimeSpan.FromSeconds(expirationScanFrequencyInSeconds / 5), + }); + + return this.extendedSessions; + } } } diff --git a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs index 09d8e1ab..5544470c 100644 --- a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs @@ -402,7 +402,7 @@ public void NewGuid_FixedInputs_ProducesStableDeterministicValue() // This regression test protects replay compatibility: it must keep producing these exact GUIDs. TestOrchestrationContext innerContext = new( "fixed-instance-id", - DateTime.Parse("2023-05-06T07:08:09.1234567Z", null, DateTimeStyles.RoundtripKind)); + DateTime.Parse("2023-05-06T07:08:09.1234567Z", CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind)); OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); @@ -425,7 +425,7 @@ public void NewGuid_DifferentInstanceId_ProducesDifferentStableDeterministicValu // the hashed name and that the namespace/algorithm/byte-ordering were not altered. TestOrchestrationContext innerContext = new( "other-instance-id", - DateTime.Parse("2023-05-06T07:08:09.1234567Z", null, DateTimeStyles.RoundtripKind)); + DateTime.Parse("2023-05-06T07:08:09.1234567Z", CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind)); OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); @@ -443,7 +443,7 @@ public void NewGuid_CalledRepeatedly_ProducesDistinctValuesEachTime() // instance ID and timestamp must still yield distinct GUIDs. TestOrchestrationContext innerContext = new( "repeat-instance-id", - DateTime.Parse("2024-01-01T00:00:00.0000000Z", null, DateTimeStyles.RoundtripKind)); + DateTime.Parse("2024-01-01T00:00:00.0000000Z", CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind)); OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); @@ -465,7 +465,7 @@ public void NewGuid_ReplayingSameHistory_ProducesIdenticalGuidSequence() // separate replay passes over the same orchestration history) observe the same instance ID and // the same sequence of CurrentUtcDateTime values as history is replayed. string instanceId = "replay-instance-id"; - DateTime timestamp = DateTime.Parse("2022-11-11T11:11:11.1111111Z", null, DateTimeStyles.RoundtripKind); + DateTime timestamp = DateTime.Parse("2022-11-11T11:11:11.1111111Z", CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind); TestOrchestrationContext innerContext1 = new(instanceId, timestamp); OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); @@ -573,7 +573,7 @@ public void NewGuid_AfterDispose_StillProducesStableDeterministicValue() // GUIDs to the ones computed with a fresh instance, proving disposal does not affect correctness. TestOrchestrationContext innerContext = new( "fixed-instance-id", - DateTime.Parse("2023-05-06T07:08:09.1234567Z", null, DateTimeStyles.RoundtripKind)); + DateTime.Parse("2023-05-06T07:08:09.1234567Z", CultureInfo.InvariantCulture, DateTimeStyles.RoundtripKind)); OrchestrationInvocationContext invocationContext = new("Test", new(), NullLoggerFactory.Instance, null); TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); diff --git a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs index f8a676d3..e4d96456 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -495,7 +495,9 @@ public async Task ExternallyEndedExtendedSession_Evicted_DisposesCachedShimResou SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(extendedSession!); // Now set the extended session flag to false for this instance, which removes/evicts the cache - // entry and should synchronously invoke the eviction callback that disposes the cached shim. + // entry and queues the eviction callback that disposes the cached shim. The callback runs + // asynchronously (see the note above), which is why this test awaits WaitUntilDisposedAsync + // below instead of asserting disposal immediately. orchestratorRequest.Properties.Clear(); orchestratorRequest.Properties.Add(new MapField() { { "IncludeState", Value.ForBool(true) }, @@ -677,6 +679,66 @@ public async Task ExtendedSessionsCache_Dispose_CalledConcurrently_DoesNotThrow( await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); } + [Fact] + public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_NeverLeaksCache() + { + // Regression test: guards against a race between Dispose() and GetOrInitializeCache() where + // Dispose() could observe the `extendedSessions` field as still null (not yet lazily + // created), mark itself disposed, and return having done nothing -- while a concurrent + // GetOrInitializeCache() call then constructs a brand-new MemoryCache that Dispose() has + // already finished running and will never see again, permanently leaking it (all future + // Dispose() calls short-circuit once `disposed` is true). + // + // The fix synchronizes both methods under a single shared lock, so the two operations are + // always fully serialized -- never interleaved -- for any given ExtendedSessionsCache + // instance: + // * If GetOrInitializeCache() completes (and returns a cache) strictly before Dispose() + // acquires the lock, Dispose() is then guaranteed to observe and dispose that exact + // cache. + // * If Dispose() completes strictly first, GetOrInitializeCache() must throw + // ObjectDisposedException instead of creating a now-unreachable cache. + // Run many iterations, racing the two calls from separate threads, to exercise both + // orderings rather than relying on a single attempt. + for (int iteration = 0; iteration < 50; iteration++) + { + var extendedSessions = new ExtendedSessionsCache(); + + Task getOrInitTask = Task.Run(() => + { + try + { + return extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds); + } + catch (ObjectDisposedException) + { + // Losing the race to a Dispose() that ran first is an expected, safe outcome. + return null; + } + }); + Task disposeTask = Task.Run(() => extendedSessions.Dispose()); + + await Task.WhenAll(getOrInitTask, disposeTask); + MemoryCache? cache = await getOrInitTask; + + if (cache is not null) + { + // GetOrInitializeCache() won the race and returned a cache. Because both methods are + // mutually exclusive under the shared lock, and Task.WhenAll has already awaited the + // Dispose() call to completion, Dispose() must have run strictly after + // initialization -- so it is guaranteed to have already captured and disposed this + // exact cache instance. Verify it is genuinely disposed (not merely leaked out of + // reach) by asserting further use throws ObjectDisposedException. + Assert.Throws(() => cache.TryGetValue("any-key", out _)); + } + + // Regardless of which call won the race, a repeated Dispose() call must remain a safe, + // idempotent no-op -- proving the cache reached a single, well-defined disposed state + // with no lingering, undisposed MemoryCache left behind. + Exception? repeatDisposeException = Record.Exception(() => extendedSessions.Dispose()); + Assert.Null(repeatDisposeException); + } + } + [Fact] public void Null_ExtendedSessionsCache_IsOkay() { From 24630f9f7725821aa8122a528b79e6ecfefbebce Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 18:09:26 -0700 Subject: [PATCH 07/10] Round 7: strengthen ExtendedSessionsCache race/dispose test coverage - Add two deterministic (non-racy) tests asserting GetOrInitializeCache() throws ObjectDisposedException after Dispose(), covering both the never-initialized and already-initialized cache scenarios. - Rewrite the Dispose()/GetOrInitializeCache() race stress test to use a Barrier(2) so both competing tasks start racing at (approximately) the same instant, instead of relying on Task.Run queue-order bias which let the 'init wins' ordering dominate and never genuinely exercise the 'dispose wins' path. Track and assert that both orderings actually occurred across the 50 iterations. - Add a new deterministic ExtendedSessionsCache_Dispose_DisposesCachedEntryExactlyOnce test using a CountingDisposable spy wired up with the same PostEvictionCallbackRegistration pattern GrpcOrchestrationRunner uses for real cached shims, proving exact-once disposal of cached content tied to eviction (not just that the owning MemoryCache object becomes unusable). This is intentionally separate from the racing stress test: racing a Set() call concurrently against Dispose()'s internal Clear()/Dispose() split introduces a narrow window where an entry added in between is never evicted by that cache instance -- a test-harness artifact, not a production scenario. No production code changes in this round; NewGuid determinism, cache handoff exception cleanup, and eviction/shutdown lifecycle already passed final review. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- .../GrpcOrchestrationRunnerTests.cs | 134 +++++++++++++++++- 1 file changed, 131 insertions(+), 3 deletions(-) diff --git a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs index e4d96456..30e5c5d4 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -679,6 +679,38 @@ public async Task ExtendedSessionsCache_Dispose_CalledConcurrently_DoesNotThrow( await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); } + [Fact] + public void GetOrInitializeCache_AfterDisposeWithoutPriorInitialization_ThrowsObjectDisposedException() + { + // Deterministic (non-racy) regression test for the exact bug scenario that motivated the + // shared-lock fix: Dispose() runs while the cache has never been lazily initialized (the + // `extendedSessions` field is still null). Without the fix, Dispose() would simply mark + // itself disposed and return, and a *subsequent* GetOrInitializeCache() call would happily + // construct a brand-new MemoryCache that nothing would ever dispose again (since `disposed` + // is now permanently true). GetOrInitializeCache() must instead throw immediately. + var extendedSessions = new ExtendedSessionsCache(); + + extendedSessions.Dispose(); + + Assert.Throws( + () => extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds)); + } + + [Fact] + public void GetOrInitializeCache_AfterDisposeOfInitializedCache_ThrowsObjectDisposedException() + { + // Deterministic (non-racy) regression test for the same post-dispose contract, but covering + // the case where the cache *was* already lazily initialized (and thus disposed/torn down by + // Dispose()) before the later GetOrInitializeCache() call is made. + var extendedSessions = new ExtendedSessionsCache(); + extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds); + + extendedSessions.Dispose(); + + Assert.Throws( + () => extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds)); + } + [Fact] public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_NeverLeaksCache() { @@ -697,14 +729,32 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve // cache. // * If Dispose() completes strictly first, GetOrInitializeCache() must throw // ObjectDisposedException instead of creating a now-unreachable cache. - // Run many iterations, racing the two calls from separate threads, to exercise both - // orderings rather than relying on a single attempt. + // + // A Barrier coordinates the two competing threads to start racing at (approximately) the + // same instant on every iteration -- without it, Task.Run scheduling order alone tends to + // let whichever task was queued first win essentially every time, which would mean the test + // never actually exercises the "Dispose() wins" code path. After the race, this test asserts + // both orderings were actually observed at least once across the iterations below, so the + // test itself is proven to have genuinely stressed both branches rather than trivially + // passing on one alone. + // + // Exact-once disposal of cached *content* tied to eviction is verified separately and + // deterministically by ExtendedSessionsCache_Dispose_DisposesCachedEntryExactlyOnce below -- + // racing an entry Set() call concurrently against this same Dispose() would itself introduce + // a spurious window (between Dispose()'s internal Clear() and its subsequent Dispose() call) + // where an entry added in between would never be evicted-and-disposed by *this* cache + // instance, which is a test-harness artifact rather than anything a real caller does. + int getOrInitWonCount = 0; + int disposeWonCount = 0; + for (int iteration = 0; iteration < 50; iteration++) { var extendedSessions = new ExtendedSessionsCache(); + using var barrier = new Barrier(2); Task getOrInitTask = Task.Run(() => { + barrier.SignalAndWait(); try { return extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds); @@ -715,13 +765,19 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve return null; } }); - Task disposeTask = Task.Run(() => extendedSessions.Dispose()); + Task disposeTask = Task.Run(() => + { + barrier.SignalAndWait(); + extendedSessions.Dispose(); + }); await Task.WhenAll(getOrInitTask, disposeTask); MemoryCache? cache = await getOrInitTask; if (cache is not null) { + getOrInitWonCount++; + // GetOrInitializeCache() won the race and returned a cache. Because both methods are // mutually exclusive under the shared lock, and Task.WhenAll has already awaited the // Dispose() call to completion, Dispose() must have run strictly after @@ -730,6 +786,10 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve // reach) by asserting further use throws ObjectDisposedException. Assert.Throws(() => cache.TryGetValue("any-key", out _)); } + else + { + disposeWonCount++; + } // Regardless of which call won the race, a repeated Dispose() call must remain a safe, // idempotent no-op -- proving the cache reached a single, well-defined disposed state @@ -737,6 +797,44 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve Exception? repeatDisposeException = Record.Exception(() => extendedSessions.Dispose()); Assert.Null(repeatDisposeException); } + + Assert.True( + getOrInitWonCount > 0, + "Expected at least one iteration where GetOrInitializeCache() won the race against Dispose(); " + + "the coordinated race did not actually exercise this ordering."); + Assert.True( + disposeWonCount > 0, + "Expected at least one iteration where Dispose() won the race against GetOrInitializeCache(); " + + "the coordinated race did not actually exercise this ordering."); + } + + [Fact] + public async Task ExtendedSessionsCache_Dispose_DisposesCachedEntryExactlyOnce() + { + // Deterministic (non-racing) regression test proving that Dispose() drives exact-once + // disposal of *cached content* via the eviction-callback path, not merely that the owning + // MemoryCache object itself becomes unusable afterwards. A CountingDisposable spy is + // registered with a post-eviction callback wired up exactly like GrpcOrchestrationRunner + // does for real cached shims, so the assertion reflects genuine production disposal wiring. + var extendedSessions = new ExtendedSessionsCache(); + MemoryCache cache = extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds); + + var spy = new CountingDisposable(); + var options = new MemoryCacheEntryOptions(); + options.RegisterPostEvictionCallback( + (key, value, reason, state) => ((CountingDisposable)value!).Dispose()); + cache.Set("spy", spy, options); + + Assert.Equal(0, spy.DisposeCount); + + extendedSessions.Dispose(); + + await WaitUntilDisposedAsync(spy, TimeSpan.FromSeconds(10)); + Assert.Equal(1, spy.DisposeCount); + + // A repeated Dispose() call must not trigger a second eviction/disposal of the same entry. + extendedSessions.Dispose(); + Assert.Equal(1, spy.DisposeCount); } [Fact] @@ -838,6 +936,24 @@ static async Task WaitUntilDisposedAsync(SHA1 hashAlgorithm, TimeSpan timeout) } } + // Same bounded-polling shape as the SHA1 overload above, but for the CountingDisposable spy used + // by the Dispose()/GetOrInitializeCache() race test, since MemoryCache eviction callbacks (and + // thus disposal of a cache entry's content) are likewise dispatched asynchronously. + static async Task WaitUntilDisposedAsync(CountingDisposable spy, TimeSpan timeout) + { + Stopwatch stopwatch = Stopwatch.StartNew(); + while (spy.DisposeCount == 0) + { + if (stopwatch.Elapsed >= timeout) + { + throw new TimeoutException( + $"CountingDisposable spy was not disposed within {timeout} of the triggering eviction/disposal."); + } + + await Task.Delay(20); + } + } + static Protobuf.OrchestratorRequest CreateOrchestratorRequest(IEnumerable newEvents) { var orchestratorRequest = new Protobuf.OrchestratorRequest() @@ -882,4 +998,16 @@ public override async Task RunAsync(TaskOrchestrationContext context, st return input; } } + + // Minimal disposable spy used by the Dispose()/GetOrInitializeCache() race test to verify + // exact-once disposal semantics precisely -- via a real, observable Dispose() call count -- rather + // than only inferring disposal indirectly through the owning MemoryCache object becoming unusable. + sealed class CountingDisposable : IDisposable + { + int disposeCount; + + public int DisposeCount => Volatile.Read(ref this.disposeCount); + + public void Dispose() => Interlocked.Increment(ref this.disposeCount); + } } From acfe249fdefe5a12cc7b086817c0f70a5e56fe5a Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 18:15:38 -0700 Subject: [PATCH 08/10] Round 8: remove probabilistic both-winners assertion from race test The Dispose()/GetOrInitializeCache() race stress test asserted that both the 'init wins' and 'dispose wins' orderings occurred at least once across 50 Barrier-coordinated iterations. A Barrier release does not guarantee fair or varying thread-pool scheduling, so correct, race-free code could legitimately let the same side win every iteration -- making this assertion's pass/fail outcome probabilistic and CI-flaky rather than a genuine correctness check. Remove the win-count tracking and both-orderings-required assertions. The test still asserts the invariants that must hold under every possible interleaving: a cache handed back by GetOrInitializeCache() before a concurrent Dispose() completes is provably disposed (TryGetValue throws ObjectDisposedException), and repeated Dispose() calls remain a safe, idempotent no-op. Also give each Barrier.SignalAndWait() a bounded timeout instead of waiting indefinitely, so a hung/stalled participant surfaces as a test failure (TimeoutException) rather than hanging the test run. The deterministic post-dispose contract tests (GetOrInitializeCache_AfterDispose*) and the deterministic exact-once disposal test (ExtendedSessionsCache_Dispose_DisposesCachedEntryExactlyOnce) are unchanged and continue to directly catch the original leak. No production code changes in this round. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- .../GrpcOrchestrationRunnerTests.cs | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs index 30e5c5d4..9c04b3bb 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -732,11 +732,19 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve // // A Barrier coordinates the two competing threads to start racing at (approximately) the // same instant on every iteration -- without it, Task.Run scheduling order alone tends to - // let whichever task was queued first win essentially every time, which would mean the test - // never actually exercises the "Dispose() wins" code path. After the race, this test asserts - // both orderings were actually observed at least once across the iterations below, so the - // test itself is proven to have genuinely stressed both branches rather than trivially - // passing on one alone. + // let whichever task was queued first win essentially every time. This is used only to + // encourage genuine contention on the shared lock; it does not (and cannot) guarantee that + // both the "GetOrInitializeCache() wins" and "Dispose() wins" orderings occur across the + // iterations below -- a Barrier release is not a scheduling guarantee, and correct, + // race-free code may legitimately let the same side win every single iteration depending on + // thread-pool scheduling. Asserting that both orderings must occur would therefore make this + // test's pass/fail outcome probabilistic (and CI-flaky) rather than a genuine correctness + // check. Instead, this test asserts only outcomes that must hold under *every* possible + // interleaving: whichever side wins, no cache is ever leaked or double-disposed. + // + // Each SignalAndWait() call uses a bounded timeout rather than waiting indefinitely, so a + // hung/stalled participant surfaces as a test failure (via TimeoutException) instead of the + // test run hanging. // // Exact-once disposal of cached *content* tied to eviction is verified separately and // deterministically by ExtendedSessionsCache_Dispose_DisposesCachedEntryExactlyOnce below -- @@ -744,8 +752,7 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve // a spurious window (between Dispose()'s internal Clear() and its subsequent Dispose() call) // where an entry added in between would never be evicted-and-disposed by *this* cache // instance, which is a test-harness artifact rather than anything a real caller does. - int getOrInitWonCount = 0; - int disposeWonCount = 0; + TimeSpan barrierTimeout = TimeSpan.FromSeconds(10); for (int iteration = 0; iteration < 50; iteration++) { @@ -754,7 +761,12 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve Task getOrInitTask = Task.Run(() => { - barrier.SignalAndWait(); + if (!barrier.SignalAndWait(barrierTimeout)) + { + throw new TimeoutException( + "Barrier synchronization timed out waiting for both racing tasks to start."); + } + try { return extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds); @@ -767,7 +779,12 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve }); Task disposeTask = Task.Run(() => { - barrier.SignalAndWait(); + if (!barrier.SignalAndWait(barrierTimeout)) + { + throw new TimeoutException( + "Barrier synchronization timed out waiting for both racing tasks to start."); + } + extendedSessions.Dispose(); }); @@ -776,8 +793,6 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve if (cache is not null) { - getOrInitWonCount++; - // GetOrInitializeCache() won the race and returned a cache. Because both methods are // mutually exclusive under the shared lock, and Task.WhenAll has already awaited the // Dispose() call to completion, Dispose() must have run strictly after @@ -786,10 +801,6 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve // reach) by asserting further use throws ObjectDisposedException. Assert.Throws(() => cache.TryGetValue("any-key", out _)); } - else - { - disposeWonCount++; - } // Regardless of which call won the race, a repeated Dispose() call must remain a safe, // idempotent no-op -- proving the cache reached a single, well-defined disposed state @@ -797,15 +808,6 @@ public async Task ExtendedSessionsCache_DisposeRaceWithGetOrInitializeCache_Neve Exception? repeatDisposeException = Record.Exception(() => extendedSessions.Dispose()); Assert.Null(repeatDisposeException); } - - Assert.True( - getOrInitWonCount > 0, - "Expected at least one iteration where GetOrInitializeCache() won the race against Dispose(); " + - "the coordinated race did not actually exercise this ordering."); - Assert.True( - disposeWonCount > 0, - "Expected at least one iteration where Dispose() won the race against GetOrInitializeCache(); " + - "the coordinated race did not actually exercise this ordering."); } [Fact] From cddd78888fe9f12568a04e3f61049eb2b19f5e4a Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 18:34:07 -0700 Subject: [PATCH 09/10] Fix ExtendedSessionsCache shutdown/hand-off race with cached shim insertion Round 9 review found a genuine production race: ExtendedSessionsCache.Dispose() marked disposed under its lock but called Clear()/Dispose() on the underlying MemoryCache outside the lock. A GrpcOrchestrationRunner (or GrpcEntityRunner) holding the raw MemoryCache reference obtained before shutdown could call .Set(...) directly in the window after Clear() but before Dispose() completed, succeeding silently but never being evicted again -- while the runner recorded transferredShimToCache = true and skipped disposal, leaking the cached SHA1/shim on graceful shutdown racing with an in-flight extended-session orchestration. Fix: encapsulate all extended-session cache mutation inside ExtendedSessionsCache via three new internal synchronized methods -- TryGetCachedValue, RemoveCachedValue, TrySetCachedValue -- each guarded by the existing syncRoot lock and each checking disposed state. TrySetCachedValue returns false without mutating the cache if disposal has begun, so insertion and disposal are now atomic with respect to each other: either the entry is fully inserted before Dispose() acquires the lock (and will be evicted by the subsequent Clear()), or Dispose() wins first and insertion is rejected, in which case the caller retains and disposes the shim itself. - ExtendedSessionsCache.GetOrInitializeCache(double) and IsInitialized remain unchanged (pre-existing public API from main/PR #449). - GrpcOrchestrationRunner and GrpcEntityRunner now route all cache reads, removals, and insertions through the new synchronized methods instead of operating on the raw MemoryCache directly. - transferredShimToCache now reflects TrySetCachedValue's actual return value. - Added an end-to-end regression test (LoadAndRun_ExtendedSession_CacheDisposedDuringExecution_ShimIsDisposedImmediatelyNotLeaked) that disposes the cache mid-orchestration-execution and asserts the shim's cached SHA1 is disposed immediately rather than leaked. No public API changes, no GUID/replay behavior changes. Verified: full Worker.Grpc.Tests (147/147) and Worker.Tests (145/145) pass. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- src/Worker/Core/ExtendedSessionsCache.cs | 75 +++++++++++++++ src/Worker/Grpc/GrpcEntityRunner.cs | 12 ++- src/Worker/Grpc/GrpcOrchestrationRunner.cs | 39 +++++--- .../GrpcOrchestrationRunnerTests.cs | 95 +++++++++++++++++++ 4 files changed, 204 insertions(+), 17 deletions(-) diff --git a/src/Worker/Core/ExtendedSessionsCache.cs b/src/Worker/Core/ExtendedSessionsCache.cs index 3b9a6dd0..39635e31 100644 --- a/src/Worker/Core/ExtendedSessionsCache.cs +++ b/src/Worker/Core/ExtendedSessionsCache.cs @@ -96,4 +96,79 @@ public MemoryCache GetOrInitializeCache(double expirationScanFrequencyInSeconds) return this.extendedSessions; } } + + /// + /// Attempts to retrieve the cached value for the given key, if present and this cache has not been + /// disposed (nor is concurrently being disposed by another thread). Callers should use this instead + /// of calling directly on the + /// returned by , since this method is + /// synchronized with and therefore can never observe -- or throw from -- a + /// cache instance that is concurrently being torn down. + /// + /// The type of the cached value. + /// The cache key. + /// When this method returns, contains the cached value, if found. + /// true if a value was found; false if not found, or if this cache is disposed. + internal bool TryGetCachedValue(string key, out T? value) + { + lock (this.syncRoot) + { + if (this.disposed || this.extendedSessions is null) + { + value = default; + return false; + } + + return this.extendedSessions.TryGetValue(key, out value); + } + } + + /// + /// Removes the cached value for the given key, if present. This is a safe no-op if this cache has + /// already been disposed (or is concurrently being disposed). Synchronized with + /// for the same reason as . + /// + /// The cache key to remove. + internal void RemoveCachedValue(string key) + { + lock (this.syncRoot) + { + if (this.disposed || this.extendedSessions is null) + { + return; + } + + this.extendedSessions.Remove(key); + } + } + + /// + /// Attempts to insert or replace the cached value for the given key. Returns false without + /// modifying the cache if this has already been disposed, or is + /// concurrently being disposed by another thread -- in which case the caller retains ownership of + /// (and remains responsible for disposing it, if applicable) instead of + /// assuming the cache accepted it and will eventually evict and dispose it via a post-eviction + /// callback. Synchronized with so there is no window in which an entry can be + /// inserted after disposal has begun tearing the cache down (e.g. after Clear() has already + /// run but before the underlying itself has been disposed) -- an insertion + /// that would otherwise never be evicted or disposed again. + /// + /// The type of the value to cache. + /// The cache key. + /// The value to cache. + /// The cache entry options (e.g. sliding expiration, eviction callback). + /// true if the value was inserted; false if rejected because this cache is disposed. + internal bool TrySetCachedValue(string key, T value, MemoryCacheEntryOptions options) + { + lock (this.syncRoot) + { + if (this.disposed || this.extendedSessions is null) + { + return false; + } + + this.extendedSessions.Set(key, value, options); + return true; + } + } } diff --git a/src/Worker/Grpc/GrpcEntityRunner.cs b/src/Worker/Grpc/GrpcEntityRunner.cs index 28ebc67e..2c5b01f6 100644 --- a/src/Worker/Grpc/GrpcEntityRunner.cs +++ b/src/Worker/Grpc/GrpcEntityRunner.cs @@ -113,7 +113,7 @@ public static async Task LoadAndRunAsync( addToExtendedSessions = true; // If an entity state was provided, even if we already have one stored, we always want to use the provided state. - if (!entityStateIncluded && extendedSessions.TryGetValue(request.InstanceId, out string? entityState)) + if (!entityStateIncluded && extendedSessionsCache!.TryGetCachedValue(request.InstanceId, out string? entityState)) { batch.EntityState = entityState; stateCached = true; @@ -135,15 +135,19 @@ public static async Task LoadAndRunAsync( if (addToExtendedSessions) { - // addToExtendedSessions can only be set to true if extendedSessions is not null - extendedSessions!.Set( + // addToExtendedSessions can only be set to true if extendedSessionsCache is not null. + // TrySetCachedValue is synchronized with a concurrent ExtendedSessionsCache.Dispose() (see + // GrpcOrchestrationRunner for the full rationale); the entity's cached state is a plain + // string with nothing to dispose, but routing through the same encapsulated API keeps cache + // mutation consistently guarded against a racing shutdown. + extendedSessionsCache!.TrySetCachedValue( request.InstanceId, result.EntityState, new MemoryCacheEntryOptions { SlidingExpiration = TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds) }); } else { - extendedSessions?.Remove(request.InstanceId); + extendedSessionsCache?.RemoveCachedValue(request.InstanceId); } P.EntityBatchResult response = result.ToEntityBatchResult(); diff --git a/src/Worker/Grpc/GrpcOrchestrationRunner.cs b/src/Worker/Grpc/GrpcOrchestrationRunner.cs index 3d085da3..cc7a6a96 100644 --- a/src/Worker/Grpc/GrpcOrchestrationRunner.cs +++ b/src/Worker/Grpc/GrpcOrchestrationRunner.cs @@ -144,9 +144,14 @@ public static string LoadAndRun( if (isExtendedSession && extendedSessions != null) { + // extendedSessions is only non-null when extendedSessionsCache is also non-null. All reads, + // removals, and (later) insertions are routed through extendedSessionsCache's synchronized + // wrapper methods rather than operating on the raw MemoryCache directly, so every operation + // is atomic with respect to a concurrent Dispose() of the cache (see round-9 fix below). + // // If a history was provided, even if we already have an extended session stored, we always want to evict whatever state is in the cache and replace it with a new extended // session based on the provided history - if (!pastEventsIncluded && extendedSessions.TryGetValue(request.InstanceId, out ExtendedSessionState? extendedSessionState) && extendedSessionState is not null) + if (!pastEventsIncluded && extendedSessionsCache!.TryGetCachedValue(request.InstanceId, out ExtendedSessionState? extendedSessionState) && extendedSessionState is not null) { OrchestrationRuntimeState runtimeState = extendedSessionState!.RuntimeState; runtimeState.NewEvents.Clear(); @@ -158,12 +163,12 @@ public static string LoadAndRun( result = extendedSessionState.OrchestrationExecutor.ExecuteNewEvents(); if (extendedSessionState.OrchestrationExecutor.IsCompleted) { - extendedSessions.Remove(request.InstanceId); + extendedSessionsCache.RemoveCachedValue(request.InstanceId); } } else { - extendedSessions.Remove(request.InstanceId); + extendedSessionsCache!.RemoveCachedValue(request.InstanceId); addToExtendedSessions = true; } } @@ -214,26 +219,34 @@ public static string LoadAndRun( if (addToExtendedSessions && !executor.IsCompleted) { - // addToExtendedSessions can only be set to true if extendedSessions is not null. - // The shim is now owned by the cache; it must not be disposed here since the - // orchestration may resume via ExecuteNewEvents() on a future call. Instead, register - // an eviction callback so it's disposed exactly once, whenever this entry is removed - // for any reason (explicit Remove, sliding-expiration timeout, capacity eviction, or - // cache disposal). + // addToExtendedSessions can only be set to true if extendedSessionsCache is not + // null. The shim is now (attempted to be) owned by the cache; it must not be + // disposed here since the orchestration may resume via ExecuteNewEvents() on a + // future call. Register an eviction callback so it's disposed exactly once, + // whenever this entry is removed for any reason (explicit Remove, + // sliding-expiration timeout, capacity eviction, or cache disposal). MemoryCacheEntryOptions cacheEntryOptions = new() { SlidingExpiration = TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds), }; cacheEntryOptions.RegisterPostEvictionCallback(DisposeEvictedExtendedSession); - extendedSessions!.Set( + + // TrySetCachedValue is synchronized with a concurrent ExtendedSessionsCache.Dispose() + // (e.g. during a graceful worker shutdown that races with this in-flight execution). + // It returns false -- without inserting anything -- if the cache has already been (or + // is concurrently being) disposed, so there is no window in which an entry can be + // silently added after the cache has begun tearing down and would then never be + // evicted or disposed again. transferredShimToCache reflects the actual outcome, so + // the finally block below correctly retains and disposes the shim itself when the + // hand-off is rejected. + transferredShimToCache = extendedSessionsCache!.TrySetCachedValue( request.InstanceId, - new(runtimeState, shim, executor), + new ExtendedSessionState(runtimeState, shim, executor), cacheEntryOptions); - transferredShimToCache = true; } else { - extendedSessions?.Remove(request.InstanceId); + extendedSessionsCache?.RemoveCachedValue(request.InstanceId); } } finally diff --git a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs index 9c04b3bb..e4aa3b84 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -839,6 +839,56 @@ public async Task ExtendedSessionsCache_Dispose_DisposesCachedEntryExactlyOnce() Assert.Equal(1, spy.DisposeCount); } + [Fact] + public void LoadAndRun_ExtendedSession_CacheDisposedDuringExecution_ShimIsDisposedImmediatelyNotLeaked() + { + // Regression test for the round-9 shutdown/hand-off race, exercised through the full + // GrpcOrchestrationRunner.LoadAndRun pipeline. The orchestrator disposes the extended-sessions + // cache itself partway through its own execution -- simulating a graceful + // worker shutdown completing while this orchestration is still in flight, holding a MemoryCache + // reference that GrpcOrchestrationRunner obtained before the shutdown began. Because the + // orchestration does not complete on this execution (it awaits a sub-orchestration call), + // GrpcOrchestrationRunner attempts to hand its shim off to the cache afterward; with the round-9 + // fix, that hand-off is rejected (the cache is disposed), so the shim's wrapper is disposed + // immediately and synchronously in the `finally` block, rather than being silently leaked in a + // cache that will never evict or dispose it again. + var extendedSessions = new ExtendedSessionsCache(); + + // Obtain the cache reference before "shutdown" -- exactly as GrpcOrchestrationRunner does at the + // very start of LoadAndRun, well before the orchestrator's Dispose() call (below) runs. + extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds); + + var historyEvent = new Protobuf.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + ExecutionStarted = new Protobuf.ExecutionStartedEvent() + { + OrchestrationInstance = new Protobuf.OrchestrationInstance + { + InstanceId = TestInstanceId, + ExecutionId = TestExecutionId, + }, + } + }; + Protobuf.OrchestratorRequest orchestratorRequest = CreateOrchestratorRequest([historyEvent]); + orchestratorRequest.Properties.Add(new MapField() { + { "IncludeState", Value.ForBool(true) }, + { "IsExtendedSession", Value.ForBool(true) }, + { "ExtendedSessionIdleTimeoutInSeconds", Value.ForNumber(DefaultExtendedSessionIdleTimeoutInSeconds) } }); + byte[] requestBytes = orchestratorRequest.ToByteArray(); + string requestString = Convert.ToBase64String(requestBytes); + + var orchestrator = new DisposeCacheDuringExecutionOrchestrator(extendedSessions); + string responseString = GrpcOrchestrationRunner.LoadAndRun(requestString, orchestrator, extendedSessions); + + Assert.NotNull(orchestrator.CapturedHashAlgorithm); + Assert.Throws(() => orchestrator.CapturedHashAlgorithm!.ComputeHash([1, 2, 3])); + + Protobuf.OrchestratorResponse response = Protobuf.OrchestratorResponse.Parser.ParseFrom(Convert.FromBase64String(responseString)); + Assert.False(response.RequiresHistory); + } + [Fact] public void Null_ExtendedSessionsCache_IsOkay() { @@ -908,6 +958,21 @@ static SHA1 GetCachedHashAlgorithm(object extendedSessionState) ?? throw new InvalidOperationException("cachedHashAlgorithm was null; NewGuid() may not have run.")); } + // Like GetCachedHashAlgorithm above, but reaches directly into the TaskOrchestrationContext + // instance passed to an orchestrator's RunAsync -- which, per TaskOrchestrationShim, is exactly + // the shim's wrapperContext instance -- instead of going through a cached ExtendedSessionState. + // Used by DisposeCacheDuringExecutionOrchestrator, whose cache hand-off is rejected (round-9 fix), + // so its shim is never cached and thus unreachable via ExtendedSessionState afterward. + static SHA1 GetCachedHashAlgorithmFromContext(TaskOrchestrationContext context) + { + FieldInfo cachedHashAlgorithmField = context.GetType() + .GetField("cachedHashAlgorithm", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException( + "TaskOrchestrationContextWrapper.cachedHashAlgorithm was not found."); + return (SHA1)(cachedHashAlgorithmField.GetValue(context) + ?? throw new InvalidOperationException("cachedHashAlgorithm was null; NewGuid() may not have run.")); + } + // Eviction callbacks on the extended-sessions MemoryCache are dispatched via // Task.Factory.StartNew (i.e. asynchronously, on a background thread pool task) rather than // synchronously on the calling thread -- see Microsoft.Extensions.Caching.Memory's @@ -1001,6 +1066,36 @@ public override async Task RunAsync(TaskOrchestrationContext context, st } } + // Regression orchestrator for the round-9 shutdown/hand-off race: disposes the extended-sessions + // cache passed to its constructor partway through its own execution -- after calling NewGuid() so + // there is a live cached SHA1 to observe -- simulating a graceful worker shutdown completing while + // this orchestration is still in flight and holds a MemoryCache reference obtained before the + // shutdown began. It then awaits a sub-orchestration call (like CallSubOrchestrationOrchestrator) + // so the orchestration does not complete on this execution, forcing GrpcOrchestrationRunner to + // attempt a hand-off of its shim to the now-disposed cache afterward. + class DisposeCacheDuringExecutionOrchestrator : TaskOrchestrator + { + readonly ExtendedSessionsCache cacheToDisposeDuringExecution; + + public DisposeCacheDuringExecutionOrchestrator(ExtendedSessionsCache cacheToDisposeDuringExecution) + { + this.cacheToDisposeDuringExecution = cacheToDisposeDuringExecution; + } + + public SHA1? CapturedHashAlgorithm { get; private set; } + + public override async Task RunAsync(TaskOrchestrationContext context, string input) + { + context.NewGuid(); + this.CapturedHashAlgorithm = GetCachedHashAlgorithmFromContext(context); + + this.cacheToDisposeDuringExecution.Dispose(); + + await context.CallSubOrchestratorAsync(nameof(SimpleOrchestrator)); + return input; + } + } + // Minimal disposable spy used by the Dispose()/GetOrInitializeCache() race test to verify // exact-once disposal semantics precisely -- via a real, observable Dispose() call count -- rather // than only inferring disposal indirectly through the owning MemoryCache object becoming unusable. From 5a001cf98a830ec5c67f2e227f679410967ba404 Mon Sep 17 00:00:00 2001 From: Bernd Verst Date: Fri, 24 Jul 2026 18:37:43 -0700 Subject: [PATCH 10/10] Transplant #779 CI stability fix for WorkItemStreamConsumerTests This branch predates the deterministic CI repair applied on berndverst-reduce-shim-polling-load (commits c8b66aa and 8779f98), so PR #780's "Validate Build" job was failing on a flaky, unrelated baseline test: WorkItemStreamConsumerTests.PerItem_HeartbeatReset_KeepsTimerAlive intermittently reported SilentDisconnect instead of GracefulDrain under CI scheduling pressure. Root cause (from c8b66aa/8779f98): the test raced a real per-item wall-clock delay against the real 500ms silent-disconnect timer. Any scheduler delay in a continuation between the "item processed" signal and the next write could inflate an intended-short gap past the timeout, spuriously tripping SilentDisconnect even though production behavior (per-item timer reset) was correct. Fix (test-only, plus one purely-additive optional-parameter observability seam -- no shim polling or other production behavior changes): - WorkItemStreamConsumer.ConsumeAsync gains an optional trailing onSilentDisconnectTimerArmed callback, invoked synchronously every time the silent-disconnect timer is (re-)armed (once before the read loop, once per item). Defaults to null; the sole production call site is unaffected. - PerItem_HeartbeatReset_KeepsTimerAlive now uses this seam to assert the exact "armed"/"item" event interleaving directly, instead of inferring the per-item reset from elapsed real time. Runs in ~20ms with zero timing dependency. Cherry-picked directly from c8b66aa and 8779f98 (verified byte-identical result via `git diff 8779f98 -- ` producing no output); this branch's copies of both files were unmodified prior to this transplant, so the cherry-pick applied cleanly with no conflicts. Verified: full Worker.Grpc.Tests suite 147/147 passing; the previously-flaky test run 5x standalone, consistently ~20ms with no flakiness. No changes to NewGuid()/TaskOrchestrationContextWrapper or the round-9 cache fix in this commit. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: f9b8e1d4-b96d-41ec-adeb-a691aac6fee1 --- src/Worker/Grpc/WorkItemStreamConsumer.cs | 12 ++- .../Grpc.Tests/WorkItemStreamConsumerTests.cs | 76 +++++++++++-------- 2 files changed, 56 insertions(+), 32 deletions(-) diff --git a/src/Worker/Grpc/WorkItemStreamConsumer.cs b/src/Worker/Grpc/WorkItemStreamConsumer.cs index cce0b23a..9c39ef7b 100644 --- a/src/Worker/Grpc/WorkItemStreamConsumer.cs +++ b/src/Worker/Grpc/WorkItemStreamConsumer.cs @@ -60,13 +60,22 @@ internal static class WorkItemStreamConsumer /// reset retry counters that should only count consecutive transport failures. /// /// Outer worker cancellation token. + /// + /// Test-only observability hook invoked synchronously every time the silent-disconnect timer is + /// (re-)armed -- once before the read loop starts, and once per item, immediately before that + /// item is dispatched. Always in production; lets tests prove the + /// per-item reset actually happens (and in what order relative to dispatch) without depending on + /// real elapsed time. Never invoked when disables + /// detection. + /// /// The classified outcome plus whether any message was observed. public static async Task ConsumeAsync( Func> openStream, TimeSpan silentDisconnectTimeout, Action onItem, Action? onFirstMessage, - CancellationToken cancellation) + CancellationToken cancellation, + Action? onSilentDisconnectTimerArmed = null) { bool silentDisconnectEnabled = silentDisconnectTimeout > TimeSpan.Zero; @@ -79,6 +88,7 @@ void ArmSilentDisconnectTimer() if (silentDisconnectEnabled) { timeoutSource.CancelAfter(effectiveTimeout); + onSilentDisconnectTimerArmed?.Invoke(); } } diff --git a/test/Worker/Grpc.Tests/WorkItemStreamConsumerTests.cs b/test/Worker/Grpc.Tests/WorkItemStreamConsumerTests.cs index 2464c0c7..ba9104c1 100644 --- a/test/Worker/Grpc.Tests/WorkItemStreamConsumerTests.cs +++ b/test/Worker/Grpc.Tests/WorkItemStreamConsumerTests.cs @@ -2,7 +2,6 @@ // Licensed under the MIT License. using System.Runtime.CompilerServices; -using System.Threading.Channels; using Grpc.Core; using Microsoft.DurableTask.Worker.Grpc; using P = Microsoft.DurableTask.Protobuf; @@ -139,40 +138,55 @@ public async Task OuterCancellation_WithRpcCancelledFromStream_PropagatesExcepti [Fact] public async Task PerItem_HeartbeatReset_KeepsTimerAlive() { - // Feed one item, wait long enough that the original timer would have expired, then complete. - // Synchronize on the first item actually being processed so the second delay is measured from - // the consumer's timer reset instead of from the test thread's write timing. - Channel channel = Channel.CreateUnbounded(); - TimeSpan timeout = TimeSpan.FromMilliseconds(500); - TaskCompletionSource firstItemProcessed = new(TaskCreationOptions.RunContinuationsAsynchronously); - int itemCount = 0; - - Task consumeTask = WorkItemStreamConsumer.ConsumeAsync( - openStream: ct => channel.Reader.ReadAllAsync(ct), - silentDisconnectTimeout: timeout, - onItem: _ => - { - if (Interlocked.Increment(ref itemCount) == 1) - { - firstItemProcessed.TrySetResult(); - } - }, - onFirstMessage: null, - cancellation: CancellationToken.None); - - await Task.Delay(TimeSpan.FromMilliseconds(150)); - await channel.Writer.WriteAsync(new P.WorkItem { HealthPing = new P.HealthPing() }); - await firstItemProcessed.Task.WaitAsync(TimeSpan.FromSeconds(5)); - - // Without the per-item reset, the original timer would fire before this second item arrives. - await Task.Delay(TimeSpan.FromMilliseconds(400)); - await channel.Writer.WriteAsync(new P.WorkItem { HealthPing = new P.HealthPing() }); - channel.Writer.Complete(); + // Proves the per-item timer reset -- not just a single arm at loop start -- is what keeps the + // stream alive. Earlier versions of this test tried to prove the reset by racing real per-item + // delays (each comfortably under the timeout) against the real silent-disconnect timeout (so + // their sum comfortably exceeded it). That was still flaky under CI scheduling pressure: any + // continuation between the "item processed" signal and the next write could be delayed by the + // thread pool/scheduler, silently inflating an intended-short gap past the timeout even though + // production was correct. + // + // This version removes wall-clock timing from the assertion entirely. ConsumeAsync exposes a + // test-only observability hook that fires every time the silent-disconnect timer is (re-)armed: + // once before the read loop starts, and once per item, immediately before that item is + // dispatched to onItem. By recording the exact interleaving of "armed" and "item" events, the + // test proves the structural guarantee directly -- an arm precedes every item, and the total arm + // count is itemCount + 1 -- instead of inferring it from elapsed real time. A regression that + // only arms the timer once at loop start (and never re-arms it per item) fails this assertion + // deterministically, with no dependency on scheduler timing. + const int itemCount = 5; + List events = new(); + int itemIndex = 0; + + P.WorkItem[] items = new P.WorkItem[itemCount]; + for (int i = 0; i < itemCount; i++) + { + items[i] = new P.WorkItem { HealthPing = new P.HealthPing() }; + } - WorkItemStreamResult result = await consumeTask; + WorkItemStreamResult result = await WorkItemStreamConsumer.ConsumeAsync( + openStream: _ => StreamOf(items), + silentDisconnectTimeout: TimeSpan.FromMilliseconds(500), + onItem: _ => events.Add($"item{itemIndex++}"), + onFirstMessage: null, + cancellation: CancellationToken.None, + onSilentDisconnectTimerArmed: () => events.Add("armed")); result.Outcome.Should().Be(WorkItemStreamOutcome.GracefulDrain); result.FirstMessageObserved.Should().BeTrue(); + + // 1 initial arm (before the loop starts) + 1 re-arm per item. + events.Count(e => e == "armed").Should().Be(itemCount + 1); + + // Every item must be immediately preceded by its own re-arm, and the very first event overall + // is the initial pre-loop arm. + events[0].Should().Be("armed"); + for (int i = 0; i < itemCount; i++) + { + int armedIndex = 1 + (i * 2); + events[armedIndex].Should().Be("armed", "item {0} must be preceded by a timer re-arm", i); + events[armedIndex + 1].Should().Be($"item{i}"); + } } [Fact]