diff --git a/src/Worker/Core/ExtendedSessionState.cs b/src/Worker/Core/ExtendedSessionState.cs index 4d67e047..380fcb4e 100644 --- a/src/Worker/Core/ExtendedSessionState.cs +++ b/src/Worker/Core/ExtendedSessionState.cs @@ -10,6 +10,12 @@ namespace Microsoft.DurableTask.Worker; /// public class ExtendedSessionState { + const long DisposedOwnership = -1; + const long RunnerOwnership = 0; + + long nextCacheGeneration; + long ownership = RunnerOwnership; + /// /// Initializes a new instance of the class. /// @@ -37,4 +43,90 @@ public ExtendedSessionState(OrchestrationRuntimeState state, TaskOrchestration t /// Gets or sets the saved TaskOrchestrationExecutor. /// public TaskOrchestrationExecutor OrchestrationExecutor { get; set; } + + /// + /// Attempts to transfer ownership from a runner to a new cache generation. + /// + /// The new cache generation. + /// true if ownership was transferred; otherwise, false. + internal bool TryTransferToCache(out long generation) + { + generation = Interlocked.Increment(ref this.nextCacheGeneration); + if (generation <= RunnerOwnership) + { + throw new InvalidOperationException("The extended-session cache generation overflowed."); + } + + return Interlocked.CompareExchange(ref this.ownership, generation, RunnerOwnership) + == RunnerOwnership; + } + + /// + /// Attempts to transfer ownership from the current cache generation to a runner. + /// + /// The cache generation that previously owned the session. + /// true if ownership was transferred; otherwise, false. + internal bool TryTakeFromCache(out long generation) + { + while (true) + { + generation = Volatile.Read(ref this.ownership); + if (generation == RunnerOwnership) + { + // Values inserted directly through the public MemoryCache API predate generation tracking. + return true; + } + + if (generation == DisposedOwnership) + { + return false; + } + + if (Interlocked.CompareExchange(ref this.ownership, RunnerOwnership, generation) == generation) + { + return true; + } + } + } + + /// + /// Attempts to reclaim a specific cache generation after a failed insertion. + /// + /// The cache generation to reclaim. + /// true if ownership returned to the runner; otherwise, false. + internal bool TryTakeCacheGeneration(long generation) + { + return generation > RunnerOwnership + && Interlocked.CompareExchange(ref this.ownership, RunnerOwnership, generation) == generation; + } + + /// + /// Disposes the session if it is currently owned by a runner. + /// + internal void DisposeRunnerOwned() + { + if (Interlocked.CompareExchange(ref this.ownership, DisposedOwnership, RunnerOwnership) + == RunnerOwnership) + { + this.DisposeTaskOrchestration(); + } + } + + /// + /// Disposes the session if it is currently owned by the specified cache generation. + /// + /// The cache generation requesting disposal. + internal void DisposeCacheGeneration(long generation) + { + if (generation > RunnerOwnership + && Interlocked.CompareExchange(ref this.ownership, DisposedOwnership, generation) == generation) + { + this.DisposeTaskOrchestration(); + } + } + + void DisposeTaskOrchestration() + { + (this.TaskOrchestration as IDisposable)?.Dispose(); + } } diff --git a/src/Worker/Core/ExtendedSessionsCache.cs b/src/Worker/Core/ExtendedSessionsCache.cs index 59df2536..c22cf9c0 100644 --- a/src/Worker/Core/ExtendedSessionsCache.cs +++ b/src/Worker/Core/ExtendedSessionsCache.cs @@ -11,12 +11,22 @@ 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. + /// Gets a value indicating whether the cache has been initialized. /// - /// True if the cache has been initialized, false otherwise. public bool IsInitialized => this.extendedSessions is not null; /// @@ -24,25 +34,257 @@ 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); } /// - /// Gets the cache for extended sessions if it has already been initialized, or otherwise initializes it with the given expiration scan frequency. + /// Gets the cache for extended sessions if it has already been initialized, or otherwise initializes it + /// with an expiration scan frequency derived from the supplied duration. /// /// - /// The expiration scan frequency of the cache, in seconds. - /// This specifies how often the cache checks for stale items, and evicts them. + /// The duration, in seconds, used to derive the expiration scan frequency. The cache checks for stale + /// items every one-fifth of this duration. /// /// 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) { - ExpirationScanFrequency = TimeSpan.FromSeconds(expirationScanFrequencyInSeconds / 5), - }); + if (this.disposed || this.extendedSessions is null) + { + value = default; + return false; + } - return this.extendedSessions; + return this.extendedSessions.TryGetValue(key, out value); + } + } + + /// + /// Atomically removes an extended session from the cache and transfers ownership to the caller. + /// + /// The cache key. + /// When this method returns, contains the runner-owned session, if found. + /// true if ownership was transferred; otherwise, false. + internal bool TryTakeExtendedSession(string key, out ExtendedSessionState? value) + { + lock (this.syncRoot) + { + if (this.disposed + || this.extendedSessions is null + || !this.extendedSessions.TryGetValue(key, out value) + || value is null) + { + value = null; + return false; + } + + if (!value.TryTakeFromCache(out _)) + { + this.extendedSessions.Remove(key); + value = null; + return false; + } + + try + { + // Transfer ownership before removing the entry. MemoryCache queues callbacks + // asynchronously, so the callback for the old generation will now be harmless. + this.extendedSessions.Remove(key); + return true; + } + catch + { + // The session is no longer safely cache-owned. Dispose the runner lease rather than + // risk either leaking it or returning a state that may still be reachable from the cache. + value.DisposeRunnerOwned(); + value = null; + throw; + } + } + } + + /// + /// 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; + } + } + + /// + /// Atomically transfers a runner-owned extended session into the cache using a fresh ownership generation. + /// + /// The cache key. + /// The runner-owned extended session. + /// The sliding expiration for the cache entry. + /// + /// true if ownership was transferred; false if the cache rejected the transfer. + /// + internal bool TryStoreExtendedSession( + string key, + ExtendedSessionState value, + TimeSpan slidingExpiration) + { + lock (this.syncRoot) + { + if (this.disposed || this.extendedSessions is null) + { + return false; + } + + // Validate all caller-controlled options before transferring ownership. In particular, + // a sub-tick positive timeout rounds to TimeSpan.Zero, which SlidingExpiration rejects. + MemoryCacheEntryOptions options = new() + { + SlidingExpiration = slidingExpiration, + }; + + if (!value.TryTransferToCache(out long generation)) + { + return false; + } + + try + { + options.RegisterPostEvictionCallback(DisposeEvictedExtendedSession, generation); + this.extendedSessions.Set(key, value, options); + return true; + } + catch + { + // Set normally either commits or throws before insertion. Defensively handle an + // implementation that commits and then throws by removing this exact value first. + // Its queued callback and the ownership transfer below race safely on the generation. + try + { + if (this.extendedSessions.TryGetValue(key, out object? cachedValue) + && ReferenceEquals(cachedValue, value)) + { + this.extendedSessions.Remove(key); + } + } + catch (ObjectDisposedException) + { + // The publicly exposed MemoryCache may have been disposed directly. It no longer + // owns a usable entry, so ownership can still be returned to the runner below. + } + + value.TryTakeCacheGeneration(generation); + throw; + } + } + } + + static void DisposeEvictedExtendedSession(object key, object? value, EvictionReason reason, object? state) + { + if (value is ExtendedSessionState sessionState && state is long generation) + { + sessionState.DisposeCacheGeneration(generation); + } } } diff --git a/src/Worker/Core/Shims/DurableTaskShimFactory.cs b/src/Worker/Core/Shims/DurableTaskShimFactory.cs index 3ae1cbc7..61903ffb 100644 --- a/src/Worker/Core/Shims/DurableTaskShimFactory.cs +++ b/src/Worker/Core/Shims/DurableTaskShimFactory.cs @@ -85,8 +85,12 @@ public TaskOrchestration CreateOrchestration( { Check.NotDefault(name); Check.NotNull(orchestrator); - OrchestrationInvocationContext context = new(name, this.options, this.loggerFactory, parent); - return new TaskOrchestrationShim(context, orchestrator); + return this.CreateOrchestrationCore( + name, + orchestrator, + properties: null, + parent, + reuseNewGuidHashAlgorithm: false); } /// @@ -108,8 +112,12 @@ public TaskOrchestration CreateOrchestration( Check.NotDefault(name); Check.NotNull(orchestrator); Check.NotNull(properties); - OrchestrationInvocationContext context = new(name, this.options, this.loggerFactory, parent); - return new TaskOrchestrationShim(context, orchestrator, properties); + return this.CreateOrchestrationCore( + name, + orchestrator, + properties, + parent, + reuseNewGuidHashAlgorithm: false); } /// @@ -153,4 +161,71 @@ public TaskEntity CreateEntity(TaskName name, ITaskEntity entity, EntityId entit ILogger logger = this.loggerFactory.CreateLogger(entity.GetType()); return new TaskEntityShim(this.options.DataConverter, entity, entityId, logger); } + + /// + /// Creates an orchestration shim whose deterministic-GUID hash algorithm may be reused. + /// The caller must dispose the returned shim or transfer it to an owner that will. + /// + /// The invoked orchestration name. + /// The orchestration to wrap. + /// The orchestration parent details, if any. + /// A new orchestration shim with managed lifetime. + internal TaskOrchestration CreateOrchestrationWithManagedLifetime( + TaskName name, + ITaskOrchestrator orchestrator, + ParentOrchestrationInstance? parent) + { + Check.NotDefault(name); + Check.NotNull(orchestrator); + return this.CreateOrchestrationCore( + name, + orchestrator, + properties: null, + parent, + reuseNewGuidHashAlgorithm: true); + } + + /// + /// Creates an orchestration shim whose deterministic-GUID hash algorithm may be reused. + /// The caller must dispose the returned shim or transfer it to an owner that will. + /// + /// The invoked orchestration name. + /// The orchestration to wrap. + /// Configuration for the orchestration. + /// The orchestration parent details, if any. + /// A new orchestration shim with managed lifetime. + internal TaskOrchestration CreateOrchestrationWithManagedLifetime( + TaskName name, + ITaskOrchestrator orchestrator, + IReadOnlyDictionary properties, + ParentOrchestrationInstance? parent) + { + Check.NotDefault(name); + Check.NotNull(orchestrator); + Check.NotNull(properties); + return this.CreateOrchestrationCore( + name, + orchestrator, + properties, + parent, + reuseNewGuidHashAlgorithm: true); + } + + TaskOrchestrationShim CreateOrchestrationCore( + TaskName name, + ITaskOrchestrator orchestrator, + IReadOnlyDictionary? properties, + ParentOrchestrationInstance? parent, + bool reuseNewGuidHashAlgorithm) + { + OrchestrationInvocationContext context = new( + name, + this.options, + this.loggerFactory, + parent, + reuseNewGuidHashAlgorithm); + return properties is null + ? new TaskOrchestrationShim(context, orchestrator) + : new TaskOrchestrationShim(context, orchestrator, properties); + } } diff --git a/src/Worker/Core/Shims/OrchestrationInvocationContext.cs b/src/Worker/Core/Shims/OrchestrationInvocationContext.cs index 26bc92e3..f2780fef 100644 --- a/src/Worker/Core/Shims/OrchestrationInvocationContext.cs +++ b/src/Worker/Core/Shims/OrchestrationInvocationContext.cs @@ -12,8 +12,13 @@ namespace Microsoft.DurableTask.Worker.Shims; /// The Durable Task worker options. /// The logger factory for this orchestration. /// The orchestration parent details. +/// +/// Whether this invocation may cache its deterministic-GUID hash algorithm. This is enabled only for +/// internal execution paths that own the complete shim lifetime. +/// record OrchestrationInvocationContext( TaskName Name, DurableTaskWorkerOptions Options, ILoggerFactory LoggerFactory, - ParentOrchestrationInstance? Parent = null); + ParentOrchestrationInstance? Parent = null, + bool ReuseNewGuidHashAlgorithm = false); diff --git a/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs b/src/Worker/Core/Shims/TaskOrchestrationContextWrapper.cs index c92ee5d6..50c27e07 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,18 @@ sealed partial class TaskOrchestrationContextWrapper : TaskOrchestrationContext bool preserveUnprocessedEventsOnContinueAsNew; TaskOrchestrationEntityContext? entityFeature; + // Internal worker paths with complete shim-lifetime ownership cache and reuse this instance across + // NewGuid() calls. Public factory paths use a fresh, per-call SHA1 instead because TaskOrchestration + // has no disposal hook through which those callers could release a cached instance. + // + // 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,15 +446,34 @@ 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, - and the goal is to generate a deterministic Guid. We cannot update to SHA2-based algorithms without breaking - customers' inflight orchestrations. */ + bool reuseHashAlgorithm = this.invocationContext.ReuseNewGuidHashAlgorithm; + SHA1 hashAlgorithm = reuseHashAlgorithm + ? this.cachedHashAlgorithm ??= SHA1.Create() + : 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. */ + + byte[] hashByteArray; + try { + // Resetting before every use makes the reusable and per-call modes byte-for-byte identical. + hashAlgorithm.Initialize(); hashAlgorithm.TransformBlock(namespaceValueByteArray, 0, namespaceValueByteArray.Length, null, 0); hashAlgorithm.TransformFinalBlock(nameByteArray, 0, nameByteArray.Length); - hashByteArray = hashAlgorithm.Hash; + + // HashAlgorithm.Hash is nullable in its API surface, but is guaranteed to be populated + // immediately after TransformFinalBlock completes successfully. + hashByteArray = hashAlgorithm.Hash + ?? throw new InvalidOperationException( + "SHA1.Hash was unexpectedly null after TransformFinalBlock."); + } + finally + { + if (!reuseHashAlgorithm) + { + hashAlgorithm.Dispose(); + } } #pragma warning restore CA5350 // Do Not Use Weak Cryptographic Algorithms -- not for cryptography @@ -458,6 +489,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..6a081252 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,16 @@ 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. + // Execute itself is async (it awaits orchestrator code across yield points), but callers never + // invoke it again -- concurrently or otherwise -- until a previous Execute call on this shim has + // fully completed (returned or thrown). That sequential lifecycle guarantee, not synchronous + // execution, is what ensures a previous wrapper is no longer 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 +136,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..df5613f7 100644 --- a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs +++ b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs @@ -773,15 +773,29 @@ await this.ExecuteWithRetryAsync( _ => null, }; - 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(); + TaskOrchestration shim = this.shimFactory.CreateOrchestrationWithManagedLifetime( + name, + orchestrator!, + parent); + 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/GrpcInstanceRunnerUtils.cs b/src/Worker/Grpc/GrpcInstanceRunnerUtils.cs index abc2e2e2..9953e35b 100644 --- a/src/Worker/Grpc/GrpcInstanceRunnerUtils.cs +++ b/src/Worker/Grpc/GrpcInstanceRunnerUtils.cs @@ -56,7 +56,22 @@ internal static void ParseRequestPropertiesAndInitializeCache( { extendedSessionIdleTimeoutInSeconds = extendedSessionIdleTimeout; isExtendedSession = extendedSession; - extendedSessions = extendedSessionsCache?.GetOrInitializeCache(extendedSessionIdleTimeoutInSeconds); + + try + { + extendedSessions = extendedSessionsCache?.GetOrInitializeCache(extendedSessionIdleTimeoutInSeconds); + } + catch (ObjectDisposedException) + { + // The extended sessions cache has already been disposed -- e.g. a request raced with the + // final stage of a graceful worker shutdown. This is not a caller error: fall back to the + // same "no cache available" behavior already used when extendedSessionsCache is null to + // begin with (extendedSessions stays null; isExtendedSession is left as parsed above), so + // the caller degrades to requesting full history instead of failing the call outright. + // Do not widen this catch beyond ObjectDisposedException -- any other failure here is + // unexpected and should propagate. + extendedSessions = null; + } } if (properties.TryGetValue("IncludeState", out object? includeStateObj) diff --git a/src/Worker/Grpc/GrpcOrchestrationRunner.cs b/src/Worker/Grpc/GrpcOrchestrationRunner.cs index 5fbe2228..78846fd6 100644 --- a/src/Worker/Grpc/GrpcOrchestrationRunner.cs +++ b/src/Worker/Grpc/GrpcOrchestrationRunner.cs @@ -144,26 +144,48 @@ public static string LoadAndRun( if (isExtendedSession && extendedSessions != null) { + // extendedSessions is only non-null when extendedSessionsCache is also non-null. All reads, + // removals, and insertions are routed through the synchronized cache wrapper so ownership + // transfers are atomic with respect to expiration, removal, and concurrent disposal. + // // 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!.TryTakeExtendedSession( + request.InstanceId, + out ExtendedSessionState? extendedSessionState) + && extendedSessionState is not null) { - OrchestrationRuntimeState runtimeState = extendedSessionState!.RuntimeState; - runtimeState.NewEvents.Clear(); - foreach (HistoryEvent newEvent in newEvents) + ExtendedSessionState? runnerOwnedSession = extendedSessionState; + try { - runtimeState.AddEvent(newEvent); - } + OrchestrationRuntimeState runtimeState = extendedSessionState.RuntimeState; + runtimeState.NewEvents.Clear(); + foreach (HistoryEvent newEvent in newEvents) + { + runtimeState.AddEvent(newEvent); + } - result = extendedSessionState.OrchestrationExecutor.ExecuteNewEvents(); - if (extendedSessionState.OrchestrationExecutor.IsCompleted) + result = extendedSessionState.OrchestrationExecutor.ExecuteNewEvents(); + if (!extendedSessionState.OrchestrationExecutor.IsCompleted + && extendedSessionsCache.TryStoreExtendedSession( + request.InstanceId, + extendedSessionState, + TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds))) + { + runnerOwnedSession = null; + } + } + finally { - extendedSessions.Remove(request.InstanceId); + // Completion, execution failure, or a rejected cache hand-off all leave the + // runner owning the session and therefore responsible for disposing its shim. + runnerOwnedSession?.DisposeRunnerOwned(); } } else { - extendedSessions.Remove(request.InstanceId); + extendedSessionsCache!.RemoveCachedValue(request.InstanceId); addToExtendedSessions = true; } } @@ -195,27 +217,52 @@ public static string LoadAndRun( DurableTaskShimFactory factory = services is null ? DurableTaskShimFactory.Default : 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(); + TaskOrchestration shim = factory.CreateOrchestrationWithManagedLifetime( + orchestratorName, + implementation, + properties, + parent); - if (addToExtendedSessions && !executor.IsCompleted) + ExtendedSessionState? runnerOwnedSession = null; + bool transferredSessionToCache = 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); + runnerOwnedSession = new ExtendedSessionState(runtimeState, shim, executor); + result = executor.Execute(); + + if (addToExtendedSessions && !executor.IsCompleted) + { + transferredSessionToCache = extendedSessionsCache!.TryStoreExtendedSession( + request.InstanceId, + runnerOwnedSession, + TimeSpan.FromSeconds(extendedSessionIdleTimeoutInSeconds)); + if (transferredSessionToCache) + { + runnerOwnedSession = null; + } + } + else + { + extendedSessionsCache?.RemoveCachedValue(request.InstanceId); + } } - else + finally { - extendedSessions?.Remove(request.InstanceId); + if (runnerOwnedSession is not null) + { + runnerOwnedSession.DisposeRunnerOwned(); + } + else if (!transferredSessionToCache) + { + // The executor failed before its ExtendedSessionState could assume ownership. + (shim as IDisposable)?.Dispose(); + } } } } diff --git a/test/Worker/Core.Tests/ExtendedSessionsCacheTests.cs b/test/Worker/Core.Tests/ExtendedSessionsCacheTests.cs new file mode 100644 index 00000000..a0b39a77 --- /dev/null +++ b/test/Worker/Core.Tests/ExtendedSessionsCacheTests.cs @@ -0,0 +1,115 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Reflection; +using DurableTask.Core; +using Microsoft.Extensions.Caching.Memory; + +namespace Microsoft.DurableTask.Worker; + +public class ExtendedSessionsCacheTests +{ + static readonly FieldInfo OwnershipField = typeof(ExtendedSessionState) + .GetField("ownership", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException($"{nameof(ExtendedSessionState)}.ownership was not found."); + + [Fact] + public void TakeAndReinsert_StaleGenerationCannotDisposeFreshOwner() + { + // Arrange + using var cache = new ExtendedSessionsCache(); + cache.GetOrInitializeCache(30); + var orchestration = new CountingTaskOrchestration(); + var session = new ExtendedSessionState(null!, orchestration, null!); + + // Act + cache.TryStoreExtendedSession("instance", session, TimeSpan.FromSeconds(30)).Should().BeTrue(); + long oldGeneration = GetOwnership(session); + cache.TryTakeExtendedSession("instance", out ExtendedSessionState? taken).Should().BeTrue(); + taken.Should().BeSameAs(session); + cache.TryStoreExtendedSession("instance", session, TimeSpan.FromSeconds(30)).Should().BeTrue(); + long freshGeneration = GetOwnership(session); + + session.DisposeCacheGeneration(oldGeneration); + + // Assert + freshGeneration.Should().BeGreaterThan(oldGeneration); + orchestration.DisposeCount.Should().Be(0); + + cache.RemoveCachedValue("instance"); + session.DisposeCacheGeneration(freshGeneration); + session.DisposeCacheGeneration(oldGeneration); + session.DisposeCacheGeneration(freshGeneration); + orchestration.DisposeCount.Should().Be(1); + } + + [Fact] + public void Store_WhenMemoryCacheThrows_ReturnsOwnershipToRunner() + { + // Arrange + var cache = new ExtendedSessionsCache(); + MemoryCache memoryCache = cache.GetOrInitializeCache(30); + memoryCache.Dispose(); + var orchestration = new CountingTaskOrchestration(); + var session = new ExtendedSessionState(null!, orchestration, null!); + + // Act + Action store = () => cache.TryStoreExtendedSession( + "instance", + session, + TimeSpan.FromSeconds(30)); + + // Assert + store.Should().Throw(); + session.DisposeRunnerOwned(); + session.DisposeRunnerOwned(); + orchestration.DisposeCount.Should().Be(1); + } + + [Fact] + public void Store_WhenSlidingExpirationIsRejected_RetainsRunnerOwnership() + { + // Arrange + using var cache = new ExtendedSessionsCache(); + cache.GetOrInitializeCache(30); + var orchestration = new CountingTaskOrchestration(); + var session = new ExtendedSessionState(null!, orchestration, null!); + + // Act + Action store = () => cache.TryStoreExtendedSession( + "instance", + session, + TimeSpan.Zero); + + // Assert + store.Should().Throw(); + session.DisposeRunnerOwned(); + session.DisposeRunnerOwned(); + orchestration.DisposeCount.Should().Be(1); + } + + static long GetOwnership(ExtendedSessionState session) + { + return OwnershipField.GetValue(session) is long ownership + ? ownership + : throw new InvalidOperationException( + $"{nameof(ExtendedSessionState)}.ownership was null or had an unexpected type."); + } + + sealed class CountingTaskOrchestration : TaskOrchestration, IDisposable + { + int disposeCount; + + public int DisposeCount => Volatile.Read(ref this.disposeCount); + + public void Dispose() => Interlocked.Increment(ref this.disposeCount); + + public override Task Execute(OrchestrationContext context, string input) + => throw new NotImplementedException(); + + public override string? GetStatus() => throw new NotImplementedException(); + + public override void RaiseEvent(OrchestrationContext context, string name, string input) + => throw new NotImplementedException(); + } +} diff --git a/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs b/test/Worker/Core.Tests/Shims/TaskOrchestrationContextWrapperTests.cs index e8335e79..768f1d68 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; @@ -15,6 +17,11 @@ public class TaskOrchestrationContextWrapperTests .GetMethod(nameof(TaskOrchestrationContextWrapper.CompleteExternalEvent), BindingFlags.Instance | BindingFlags.NonPublic) ?? throw new InvalidOperationException($"{nameof(TaskOrchestrationContextWrapper)}.{nameof(TaskOrchestrationContextWrapper.CompleteExternalEvent)} 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 Ctor_NullParent_Populates() { @@ -391,10 +398,263 @@ await wrapper.CallSubOrchestratorAsync( innerContext.LastSubOrchestrationVersion.Should().Be(string.Empty); } + [Theory] + [InlineData(false)] + [InlineData(true)] + public void NewGuid_FixedInputs_ProducesStableDeterministicValue(bool reuseNewGuidHashAlgorithm) + { + // 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, + reuseNewGuidHashAlgorithm); + 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, + ReuseNewGuidHashAlgorithm: true); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + // 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 NewGuid_DefaultMode_DoesNotCacheHashAlgorithmInstance() + { + // Arrange + TestOrchestrationContext innerContext = new(); + OrchestrationInvocationContext invocationContext = new( + "Test", + new(), + NullLoggerFactory.Instance, + null); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + // Act + wrapper.NewGuid(); + object? afterFirstCall = CachedHashAlgorithmField.GetValue(wrapper); + wrapper.NewGuid(); + object? afterSecondCall = CachedHashAlgorithmField.GetValue(wrapper); + + // Assert + afterFirstCall.Should().BeNull(); + afterSecondCall.Should().BeNull(); + } + + [Fact] + public void Dispose_ReleasesCachedHashAlgorithm() + { + // Arrange + TestOrchestrationContext innerContext = new(); + OrchestrationInvocationContext invocationContext = new( + "Test", + new(), + NullLoggerFactory.Instance, + null, + ReuseNewGuidHashAlgorithm: true); + TaskOrchestrationContextWrapper wrapper = new(innerContext, invocationContext, "input"); + + wrapper.NewGuid(); + SHA1 cachedInstance = (SHA1)(CachedHashAlgorithmField.GetValue(wrapper) + ?? throw new InvalidOperationException("NewGuid() did not populate cachedHashAlgorithm.")); + + // 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, + ReuseNewGuidHashAlgorithm: true); + 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, + ReuseNewGuidHashAlgorithm: true); + 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, + ReuseNewGuidHashAlgorithm: true); + 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")!; - return (IReadOnlyDictionary)tagsProperty.GetValue(innerContext.LastScheduledTaskOptions)!; + ScheduleTaskOptions options = innerContext.LastScheduledTaskOptions + ?? throw new InvalidOperationException("No scheduled-task options were captured."); + PropertyInfo tagsProperty = options.GetType().GetProperty("Tags") + ?? throw new InvalidOperationException($"{options.GetType().FullName}.Tags was not found."); + return tagsProperty.GetValue(options) as IReadOnlyDictionary + ?? throw new InvalidOperationException($"{options.GetType().FullName}.Tags was null or had an unexpected type."); } static void InvokeCompleteExternalEvent(TaskOrchestrationContextWrapper wrapper, string eventName, string rawEventPayload) @@ -506,17 +766,32 @@ public override void SendEvent(OrchestrationInstance orchestrationInstance, stri } } - class TestOrchestrationContext : OrchestrationContext + sealed 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..c43156ba --- /dev/null +++ b/test/Worker/Core.Tests/Shims/TaskOrchestrationShimTests.cs @@ -0,0 +1,206 @@ +// 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) + ?? throw new InvalidOperationException("NewGuid() did not populate cachedHashAlgorithm.")); + + // Act + shim.Dispose(); + + // Assert + CachedHashAlgorithmField.GetValue(wrapperContext).Should().BeNull(); + Action useAfterDispose = () => cachedInstance.ComputeHash(new byte[] { 1, 2, 3 }); + useAfterDispose.Should().Throw(); + } + + [Fact] + public async Task PublicFactory_CreateOrchestration_DoesNotCacheHashAlgorithm() + { + // Arrange + DurableTaskShimFactory factory = new(); + CapturingNewGuidOrchestrator orchestrator = new(); + TaskOrchestration shim = factory.CreateOrchestration("Test", orchestrator); + TestOrchestrationContext innerContext = new(); + + // Act + await shim.Execute(innerContext, "null"); + + // Assert + orchestrator.AfterFirstCall.Should().BeNull(); + orchestrator.AfterSecondCall.Should().BeNull(); + } + + [Fact] + public async Task InternalFactory_CreateOrchestration_ReusesHashAlgorithmUntilDisposed() + { + // Arrange + DurableTaskShimFactory factory = new(); + CapturingNewGuidOrchestrator orchestrator = new(); + TaskOrchestration shim = factory.CreateOrchestrationWithManagedLifetime( + "Test", + orchestrator, + parent: null); + TestOrchestrationContext innerContext = new(); + + // Act + await shim.Execute(innerContext, "null"); + + // Assert + SHA1 cachedHashAlgorithm = orchestrator.AfterFirstCall.Should().BeAssignableTo().Which; + orchestrator.AfterSecondCall.Should().BeSameAs(cachedHashAlgorithm); + + ((IDisposable)shim).Dispose(); + Action useAfterDispose = () => cachedHashAlgorithm.ComputeHash([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, + ReuseNewGuidHashAlgorithm: true); + TestOrchestrationContext innerContext = new(); + return new TaskOrchestrationContextWrapper(innerContext, invocationContext, deserializedInput: null); + } + + sealed class CapturingNewGuidOrchestrator : ITaskOrchestrator + { + public Type InputType => typeof(object); + + public Type OutputType => typeof(object); + + public object? AfterFirstCall { get; private set; } + + public object? AfterSecondCall { get; private set; } + + public Task RunAsync(TaskOrchestrationContext context, object? input) + { + context.NewGuid(); + this.AfterFirstCall = CachedHashAlgorithmField.GetValue(context); + context.NewGuid(); + this.AfterSecondCall = CachedHashAlgorithmField.GetValue(context); + return Task.FromResult(input); + } + } + + 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..97dd081c 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; @@ -14,6 +17,15 @@ public class GrpcOrchestrationRunnerTests const string TestExecutionId = "execution_id"; const int DefaultExtendedSessionIdleTimeoutInSeconds = 30; + static readonly FieldInfo ExtendedSessionOwnershipField = typeof(ExtendedSessionState) + .GetField("ownership", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException($"{nameof(ExtendedSessionState)}.ownership was not found."); + + static readonly MethodInfo DisposeCacheGenerationMethod = typeof(ExtendedSessionState) + .GetMethod("DisposeCacheGeneration", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException( + $"{nameof(ExtendedSessionState)}.DisposeCacheGeneration was not found."); + [Fact] public void EmptyOrNullParameters_Throw_Exceptions() { @@ -217,10 +229,10 @@ public void IsExtendedSessionFalse_Means_NoExtendedSessionStored() } /// - /// These tests verify that a malformed/nonexistent "IncludeState" parameter means that the worker will attempt to - /// fulfill the orchestration request and not request a history for it. However, it is of course very undesirable that a + /// These tests verify that a malformed/nonexistent "IncludeState" parameter means that the worker will attempt to + /// fulfill the orchestration request and not request a history for it. However, it is of course very undesirable that a /// history is not attached to the request, but no history is requested by the worker due to a malformed "IncludeState" parameter - /// even when it needs one to fulfill the request. This would need to be checked on whatever side is calling this SDK. + /// even when it needs one to fulfill the request. This would need to be checked on whatever side is calling this SDK. /// [Fact] public void MalformedIncludeStateParameter_Means_NoHistoryRequired() @@ -404,7 +416,7 @@ public async Task Stale_ExtendedSessions_Evicted_Async() Assert.True(extendedSessions.IsInitialized); Assert.True(extendedSessions.GetOrInitializeCache(extendedSessionIdleTimeout).TryGetValue(TestInstanceId, out object? extendedSession)); - // Wait for longer than the timeout to account for finite cache scan for stale items frequency + // 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 extendedSession)); @@ -455,6 +467,573 @@ 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 extended-sessions cache 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; per + // ExtendedSessionsCache.TryStoreExtendedSession's disposal invariant, 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 async Task LoadAndRun_ResumedExtendedSession_CacheDisposedWhileExecutionBlocked_DisposesAfterRejectedReinsert() + { + // Arrange + var extendedSessions = new ExtendedSessionsCache(); + var orchestrator = new BlockingResumedOrchestrator(); + Protobuf.HistoryEvent executionStarted = CreateExecutionStartedEvent(); + Protobuf.OrchestratorRequest initialRequest = CreateExtendedSessionRequest( + [executionStarted], + includeState: true); + + GrpcOrchestrationRunner.LoadAndRun( + Convert.ToBase64String(initialRequest.ToByteArray()), + orchestrator, + extendedSessions); + + MemoryCache cache = extendedSessions.GetOrInitializeCache( + DefaultExtendedSessionIdleTimeoutInSeconds); + Assert.True(cache.TryGetValue(TestInstanceId, out object? cachedValue)); + ExtendedSessionState cachedSession = Assert.IsType(cachedValue); + long oldGeneration = GetOwnershipGeneration(cachedSession); + SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(cachedSession); + + Protobuf.HistoryEvent resumeEvent = CreateEventRaisedEvent("Resume"); + Protobuf.OrchestratorRequest resumedRequest = CreateExtendedSessionRequest( + [resumeEvent], + includeState: false); + Task resumedExecution = Task.Run(() => GrpcOrchestrationRunner.LoadAndRun( + Convert.ToBase64String(resumedRequest.ToByteArray()), + orchestrator, + extendedSessions)); + + try + { + Assert.True( + orchestrator.WaitUntilResumedExecutionIsBlocked(TimeSpan.FromSeconds(10)), + "the resumed ExecuteNewEvents call did not reach its deterministic blocking point"); + + // Act + extendedSessions.Dispose(); + + // MemoryCache dispatches the old entry's callback asynchronously. Invoke the exact + // generation check directly as well so this assertion never depends on callback timing. + InvokeDisposeCacheGeneration(cachedSession, oldGeneration); + + // Assert + Action useWhileRunnerOwned = () => cachedHashAlgorithm.ComputeHash([1, 2, 3]); + useWhileRunnerOwned.Should().NotThrow(); + } + finally + { + orchestrator.ReleaseResumedExecution(); + } + + string responseString = await resumedExecution.WaitAsync(TimeSpan.FromSeconds(10)); + Protobuf.OrchestratorResponse response = Protobuf.OrchestratorResponse.Parser.ParseFrom( + Convert.FromBase64String(responseString)); + + response.RequiresHistory.Should().BeFalse(); + Action useAfterRejectedReinsert = () => cachedHashAlgorithm.ComputeHash([1, 2, 3]); + useAfterRejectedReinsert.Should().Throw(); + } + + [Fact] + public void LoadAndRun_ResumedExtendedSession_StaleGenerationCannotDisposeReinsertedSession() + { + // Arrange + using var extendedSessions = new ExtendedSessionsCache(); + var orchestrator = new BlockingResumedOrchestrator(); + Protobuf.OrchestratorRequest initialRequest = CreateExtendedSessionRequest( + [CreateExecutionStartedEvent()], + includeState: true); + GrpcOrchestrationRunner.LoadAndRun( + Convert.ToBase64String(initialRequest.ToByteArray()), + orchestrator, + extendedSessions); + + MemoryCache cache = extendedSessions.GetOrInitializeCache( + DefaultExtendedSessionIdleTimeoutInSeconds); + Assert.True(cache.TryGetValue(TestInstanceId, out object? cachedValue)); + ExtendedSessionState cachedSession = Assert.IsType(cachedValue); + long oldGeneration = GetOwnershipGeneration(cachedSession); + SHA1 cachedHashAlgorithm = GetCachedHashAlgorithm(cachedSession); + + orchestrator.ReleaseResumedExecution(); + Protobuf.OrchestratorRequest resumedRequest = CreateExtendedSessionRequest( + [CreateEventRaisedEvent("Resume")], + includeState: false); + + // Act + GrpcOrchestrationRunner.LoadAndRun( + Convert.ToBase64String(resumedRequest.ToByteArray()), + orchestrator, + extendedSessions); + + // Assert + Assert.True(cache.TryGetValue(TestInstanceId, out object? reinsertedValue)); + reinsertedValue.Should().BeSameAs(cachedSession); + long freshGeneration = GetOwnershipGeneration(cachedSession); + freshGeneration.Should().BeGreaterThan(oldGeneration); + + InvokeDisposeCacheGeneration(cachedSession, oldGeneration); + Action useAfterStaleCallback = () => cachedHashAlgorithm.ComputeHash([1, 2, 3]); + useAfterStaleCallback.Should().NotThrow(); + + cache.Remove(TestInstanceId); + InvokeDisposeCacheGeneration(cachedSession, freshGeneration); + Action useAfterCurrentCallback = () => cachedHashAlgorithm.ComputeHash([1, 2, 3]); + useAfterCurrentCallback.Should().Throw(); + } + + [Fact] + public void LoadAndRun_ExtendedSession_SubTickTimeoutRejected_DisposesRunnerOwnedShim() + { + // Arrange + using var extendedSessions = new ExtendedSessionsCache(); + var orchestrator = new CaptureHashThenCallSubOrchestrationOrchestrator(); + Protobuf.OrchestratorRequest request = CreateExtendedSessionRequest( + [CreateExecutionStartedEvent()], + includeState: true); + request.Properties["ExtendedSessionIdleTimeoutInSeconds"] = Value.ForNumber(double.Epsilon); + + // Act + Action run = () => GrpcOrchestrationRunner.LoadAndRun( + Convert.ToBase64String(request.ToByteArray()), + orchestrator, + extendedSessions); + + // Assert + run.Should().Throw(); + orchestrator.CapturedHashAlgorithm.Should().NotBeNull(); + Action useAfterRejectedStore = () => orchestrator.CapturedHashAlgorithm!.ComputeHash([1, 2, 3]); + useAfterRejectedStore.Should().Throw(); + } + [Fact] public void Null_ExtendedSessionsCache_IsOkay() { @@ -498,6 +1077,154 @@ public void Null_ExtendedSessionsCache_IsOkay() Assert.Equal(Protobuf.OrchestrationStatus.Completed, response.Actions[0].CompleteOrchestration.OrchestrationStatus); } + [Fact] + public void AlreadyDisposedExtendedSessionsCache_IsOkay() + { + // A non-null ExtendedSessionsCache that has already been disposed models a request that raced + // with the tail end of a graceful worker shutdown: GrpcOrchestrationRunner obtained a reference + // to the cache before shutdown began, but by the time it calls GetOrInitializeCache, the cache + // has already been disposed. This should degrade the same way a null cache does (see + // Null_ExtendedSessionsCache_IsOkay above) -- falling back to non-extended-session behavior -- + // rather than letting an unhandled ObjectDisposedException fail the call outright. + var extendedSessions = new ExtendedSessionsCache(); + extendedSessions.Dispose(); + + 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); + + string responseString = GrpcOrchestrationRunner.LoadAndRun(requestString, new SimpleOrchestrator(), extendedSessions); + + Protobuf.OrchestratorResponse response = Protobuf.OrchestratorResponse.Parser.ParseFrom(Convert.FromBase64String(responseString)); + Assert.Single(response.Actions); + Assert.NotNull(response.Actions[0].CompleteOrchestration); + 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.GetType().FullName}.TaskOrchestration was not found."); + object shim = taskOrchestrationProperty.GetValue(extendedSessionState) + ?? throw new InvalidOperationException( + $"{extendedSessionState.GetType().FullName}.TaskOrchestration was null."); + + FieldInfo wrapperContextField = shim.GetType() + .GetField("wrapperContext", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException($"{shim.GetType().FullName}.wrapperContext was not found."); + object wrapperContext = wrapperContextField.GetValue(shim) + ?? throw new InvalidOperationException($"{shim.GetType().FullName}.wrapperContext was null."); + + FieldInfo cachedHashAlgorithmField = wrapperContext.GetType() + .GetField("cachedHashAlgorithm", BindingFlags.Instance | BindingFlags.NonPublic) + ?? throw new InvalidOperationException( + $"{wrapperContext.GetType().FullName}.cachedHashAlgorithm was not found."); + return (SHA1)(cachedHashAlgorithmField.GetValue(wrapperContext) + ?? throw new InvalidOperationException( + $"{wrapperContext.GetType().FullName}.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 per + // ExtendedSessionsCache.TryStoreExtendedSession's disposal invariant, 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( + $"{context.GetType().FullName}.cachedHashAlgorithm was not found."); + return (SHA1)(cachedHashAlgorithmField.GetValue(context) + ?? throw new InvalidOperationException( + $"{context.GetType().FullName}.cachedHashAlgorithm was null; NewGuid() may not have run.")); + } + + static long GetOwnershipGeneration(ExtendedSessionState sessionState) + { + return ExtendedSessionOwnershipField.GetValue(sessionState) is long generation + ? generation + : throw new InvalidOperationException( + $"{nameof(ExtendedSessionState)}.ownership was null or had an unexpected type."); + } + + static void InvokeDisposeCacheGeneration(ExtendedSessionState sessionState, long generation) + { + DisposeCacheGenerationMethod.Invoke(sessionState, [generation]); + } + + // 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() @@ -513,7 +1240,55 @@ static Protobuf.OrchestratorRequest CreateOrchestratorRequest(IEnumerable + static Protobuf.OrchestratorRequest CreateExtendedSessionRequest( + IEnumerable newEvents, + bool includeState) + { + Protobuf.OrchestratorRequest request = CreateOrchestratorRequest(newEvents); + request.Properties.Add(new MapField() + { + { "IncludeState", Value.ForBool(includeState) }, + { "IsExtendedSession", Value.ForBool(true) }, + { + "ExtendedSessionIdleTimeoutInSeconds", + Value.ForNumber(DefaultExtendedSessionIdleTimeoutInSeconds) + }, + }); + return request; + } + + static Protobuf.HistoryEvent CreateExecutionStartedEvent() + { + return new Protobuf.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + ExecutionStarted = new Protobuf.ExecutionStartedEvent + { + OrchestrationInstance = new Protobuf.OrchestrationInstance + { + InstanceId = TestInstanceId, + ExecutionId = TestExecutionId, + }, + }, + }; + } + + static Protobuf.HistoryEvent CreateEventRaisedEvent(string name) + { + return new Protobuf.HistoryEvent + { + EventId = 0, + Timestamp = Timestamp.FromDateTime(DateTime.UtcNow), + EventRaised = new Protobuf.EventRaisedEvent + { + Name = name, + Input = "\"resume\"", + }, + }; + } + + sealed class SimpleOrchestrator : TaskOrchestrator { public override Task RunAsync(TaskOrchestrationContext context, string input) { @@ -521,7 +1296,7 @@ public override Task RunAsync(TaskOrchestrationContext context, string i } } - class CallSubOrchestrationOrchestrator : TaskOrchestrator + sealed class CallSubOrchestrationOrchestrator : TaskOrchestrator { public override async Task RunAsync(TaskOrchestrationContext context, string input) { @@ -529,4 +1304,99 @@ 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. + sealed class NewGuidThenCallSubOrchestrationOrchestrator : TaskOrchestrator + { + public override async Task RunAsync(TaskOrchestrationContext context, string input) + { + context.NewGuid(); + await context.CallSubOrchestratorAsync(nameof(SimpleOrchestrator)); + return input; + } + } + + sealed class CaptureHashThenCallSubOrchestrationOrchestrator : TaskOrchestrator + { + public SHA1? CapturedHashAlgorithm { get; private set; } + + public override async Task RunAsync(TaskOrchestrationContext context, string input) + { + context.NewGuid(); + this.CapturedHashAlgorithm = GetCachedHashAlgorithmFromContext(context); + await context.CallSubOrchestratorAsync(nameof(SimpleOrchestrator)); + return input; + } + } + + // Regression orchestrator for the extended-sessions cache 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. + sealed 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; + } + } + + sealed class BlockingResumedOrchestrator : TaskOrchestrator + { + readonly ManualResetEventSlim resumedExecutionBlocked = new(); + readonly ManualResetEventSlim releaseResumedExecution = new(); + + public override async Task RunAsync(TaskOrchestrationContext context, string input) + { + context.NewGuid(); + await context.WaitForExternalEvent("Resume"); + + this.resumedExecutionBlocked.Set(); + if (!this.releaseResumedExecution.Wait(TimeSpan.FromSeconds(10))) + { + throw new TimeoutException("Timed out waiting for the test to release resumed execution."); + } + + await context.CallSubOrchestratorAsync(nameof(SimpleOrchestrator)); + return input; + } + + public bool WaitUntilResumedExecutionIsBlocked(TimeSpan timeout) + => this.resumedExecutionBlocked.Wait(timeout); + + public void ReleaseResumedExecution() => this.releaseResumedExecution.Set(); + } + + // 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); + } }