diff --git a/src/Worker/Core/ExtendedSessionsCache.cs b/src/Worker/Core/ExtendedSessionsCache.cs index 59df2536..39635e31 100644 --- a/src/Worker/Core/ExtendedSessionsCache.cs +++ b/src/Worker/Core/ExtendedSessionsCache.cs @@ -11,7 +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; + bool disposed; /// /// Gets a value indicating whether returns whether or not the cache has been initialized. @@ -24,7 +35,38 @@ public class ExtendedSessionsCache : IDisposable /// public void Dispose() { - this.extendedSessions?.Dispose(); + MemoryCache? cacheToDispose; + lock (this.syncRoot) + { + 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 + // 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 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. + cacheToDispose?.Clear(); + cacheToDispose?.Dispose(); GC.SuppressFinalize(this); } @@ -36,13 +78,97 @@ 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) + { + if (this.disposed) + { + throw new ObjectDisposedException(nameof(ExtendedSessionsCache)); + } + + this.extendedSessions ??= new MemoryCache(new MemoryCacheOptions + { + ExpirationScanFrequency = TimeSpan.FromSeconds(expirationScanFrequencyInSeconds / 5), + }); + + 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) { - ExpirationScanFrequency = TimeSpan.FromSeconds(expirationScanFrequencyInSeconds / 5), - }); + if (this.disposed || this.extendedSessions is null) + { + return false; + } - return this.extendedSessions; + this.extendedSessions.Set(key, value, options); + return true; + } } } diff --git a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs index c92ee5d6..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. @@ -33,6 +33,22 @@ 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. + // 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; + /// /// Initializes a new instance of the class. /// @@ -434,16 +450,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]; @@ -458,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..b7c8f1de 100644 --- a/src/Worker/Core/Shims/TaskOrchestrationShim.cs +++ b/src/Worker/Core/Shims/TaskOrchestrationShim.cs @@ -13,8 +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. +/// /// -partial class TaskOrchestrationShim : TaskOrchestration +partial class TaskOrchestrationShim : TaskOrchestration, IDisposable { readonly ITaskOrchestrator implementation; readonly OrchestrationInvocationContext invocationContext; @@ -64,6 +72,14 @@ public TaskOrchestrationShim( innerContext.ErrorDataConverter = converterShim; object? input = this.DataConverter.Deserialize(rawInput, this.implementation.InputType); + + // 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); string instanceId = innerContext.OrchestrationInstance.InstanceId; @@ -118,4 +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(); + } } 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/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 5fbe2228..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; } } @@ -197,25 +202,64 @@ 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 - extendedSessions!.Set( - request.InstanceId, - new(runtimeState, shim, executor), - new MemoryCacheEntryOptions { SlidingExpiration = TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds) }); + 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 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); + + // 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 ExtendedSessionState(runtimeState, shim, executor), + cacheEntryOptions); + } + else + { + extendedSessionsCache?.RemoveCachedValue(request.InstanceId); + } } - else + finally { - extendedSessions?.Remove(request.InstanceId); + 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(); + } } } } @@ -232,4 +276,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/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/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs index e8335e79..5544470c 100644 --- a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs @@ -2,7 +2,9 @@ // Licensed under the MIT License. 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; @@ -391,6 +393,203 @@ 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", CultureInfo.InvariantCulture, 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", CultureInfo.InvariantCulture, 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", CultureInfo.InvariantCulture, 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", CultureInfo.InvariantCulture, 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); + } + + [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", CultureInfo.InvariantCulture, 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")!; @@ -508,15 +707,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(); 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..e4aa3b84 100644 --- a/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs +++ b/test/Worker/Grpc.Tests/GrpcOrchestrationRunnerTests.cs @@ -1,6 +1,9 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Diagnostics; +using System.Reflection; +using System.Security.Cryptography; using Google.Protobuf; using Google.Protobuf.Collections; using Google.Protobuf.WellKnownTypes; @@ -455,6 +458,437 @@ public void PastEventIncluded_Means_ExtendedSession_Evicted() Assert.True(extendedSessions.GetOrInitializeCache(DefaultExtendedSessionIdleTimeoutInSeconds).TryGetValue(TestInstanceId, out extendedSession)); } + [Fact] + 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 + { + 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 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) }, + { "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 _)); + + await WaitUntilDisposedAsync(cachedHashAlgorithm, TimeSpan.FromSeconds(10)); + } + + [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. + // + // 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 + { + 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 _)); + + 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] + 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 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() + { + // 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. + // + // 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. 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 -- + // 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. + TimeSpan barrierTimeout = TimeSpan.FromSeconds(10); + + for (int iteration = 0; iteration < 50; iteration++) + { + var extendedSessions = new ExtendedSessionsCache(); + using var barrier = new Barrier(2); + + Task getOrInitTask = Task.Run(() => + { + if (!barrier.SignalAndWait(barrierTimeout)) + { + throw new TimeoutException( + "Barrier synchronization timed out waiting for both racing tasks to start."); + } + + 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(() => + { + if (!barrier.SignalAndWait(barrierTimeout)) + { + throw new TimeoutException( + "Barrier synchronization timed out waiting for both racing tasks to start."); + } + + 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 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] + 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() { @@ -498,6 +932,95 @@ 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.")); + } + + // 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 + // 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); + } + } + + // 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() @@ -529,4 +1052,59 @@ 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; + } + } + + // 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. + sealed class CountingDisposable : IDisposable + { + int disposeCount; + + public int DisposeCount => Volatile.Read(ref this.disposeCount); + + public void Dispose() => Interlocked.Increment(ref this.disposeCount); + } } 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]