From 8bcbc7f0958b46766f4e7b5f3baf1291caf69ee8 Mon Sep 17 00:00:00 2001 From: Andy Staples Date: Fri, 6 Mar 2026 16:05:42 -0700 Subject: [PATCH 1/3] Fix extended sessions timeout issue --- .../OrchestrationSessionTests.cs | 260 ++++++++++++++++++ .../AzureStorageOrchestrationService.cs | 9 + .../Messaging/OrchestrationSession.cs | 15 +- .../OrchestrationSessionManager.cs | 18 ++ 4 files changed, 300 insertions(+), 2 deletions(-) create mode 100644 Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs diff --git a/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs b/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs new file mode 100644 index 000000000..a304246ec --- /dev/null +++ b/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs @@ -0,0 +1,260 @@ +// ---------------------------------------------------------------------------------- +// Copyright Microsoft Corporation +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// http://www.apache.org/licenses/LICENSE-2.0 +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ---------------------------------------------------------------------------------- + +namespace DurableTask.AzureStorage.Tests +{ + using System; + using System.Collections.Generic; + using System.Diagnostics; + using System.Reflection; + using System.Threading; + using System.Threading.Tasks; + using DurableTask.AzureStorage.Messaging; + using DurableTask.AzureStorage.Monitoring; + using DurableTask.AzureStorage.Tracking; + using Microsoft.VisualStudio.TestTools.UnitTesting; + using Moq; + + /// + /// Tests that validate the shutdown cancellation behavior for extended sessions. + /// These tests verify fixes for the issue where host shutdown was blocked by the + /// extended session idle timeout (default 30s) because no cancellation token was + /// propagated to the session wait, and isForced was ignored during StopAsync. + /// + [TestClass] + public class OrchestrationSessionTests + { + /// + /// Verifies that + /// exits immediately when the cancellation token is cancelled, rather than waiting for + /// the full timeout. This is the core mechanism behind Bug 1 fix. + /// + [TestMethod] + public async Task WaitAsync_CancellationToken_ExitsImmediately() + { + var resetEvent = new AsyncAutoResetEvent(signaled: false); + using var cts = new CancellationTokenSource(); + + // Start a wait with a long timeout (simulating ExtendedSessionIdleTimeoutInSeconds = 30s) + TimeSpan longTimeout = TimeSpan.FromSeconds(30); + Task waitTask = resetEvent.WaitAsync(longTimeout, cts.Token); + + // The wait should not have completed yet + Assert.IsFalse(waitTask.IsCompleted, "Wait should not complete immediately"); + + // Cancel the token (simulating shutdown) + var stopwatch = Stopwatch.StartNew(); + cts.Cancel(); + + // The wait should return false almost immediately (same as timeout) + bool result = await waitTask; + stopwatch.Stop(); + + Assert.IsFalse(result, "Cancellation should return false (no signal received)"); + + // Verify it completed quickly (well under the 30s timeout) + Assert.IsTrue( + stopwatch.ElapsedMilliseconds < 5000, + $"Cancellation should complete in under 5s, but took {stopwatch.ElapsedMilliseconds}ms"); + } + + /// + /// Verifies that the wait still works normally (returns true) when signaled, + /// even when a cancellation token is provided. + /// + [TestMethod] + public async Task WaitAsync_WithCancellationToken_SignalStillWorks() + { + var resetEvent = new AsyncAutoResetEvent(signaled: false); + using var cts = new CancellationTokenSource(); + + Task waitTask = resetEvent.WaitAsync(TimeSpan.FromSeconds(30), cts.Token); + Assert.IsFalse(waitTask.IsCompleted); + + // Signal the event (simulating new messages arriving) + resetEvent.Set(); + + // Wait for the task with a reasonable timeout + Task winner = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(5))); + Assert.IsTrue(winner == waitTask, "Signal should wake the waiter"); + Assert.IsTrue(waitTask.Result, "Wait result should be true when signaled"); + } + + /// + /// Verifies that the wait returns false on timeout even when a cancellation + /// token is provided but not cancelled (normal idle timeout behavior). + /// + [TestMethod] + public async Task WaitAsync_WithCancellationToken_TimeoutStillWorks() + { + var resetEvent = new AsyncAutoResetEvent(signaled: false); + using var cts = new CancellationTokenSource(); + + // Use a short timeout to keep the test fast + bool result = await resetEvent.WaitAsync(TimeSpan.FromMilliseconds(100), cts.Token); + + Assert.IsFalse(result, "Wait should return false on timeout"); + } + + /// + /// Verifies that cancellation works when multiple waiters are queued. + /// All waiters should return false when the token fires. + /// + [TestMethod] + public async Task WaitAsync_CancellationToken_MultipleWaiters() + { + var resetEvent = new AsyncAutoResetEvent(signaled: false); + using var cts = new CancellationTokenSource(); + + // Queue multiple waiters (simulating multiple sessions waiting during shutdown) + var waiters = new List>(); + for (int i = 0; i < 5; i++) + { + waiters.Add(resetEvent.WaitAsync(TimeSpan.FromSeconds(30), cts.Token)); + } + + // None should be completed yet + foreach (var waiter in waiters) + { + Assert.IsFalse(waiter.IsCompleted); + } + + // Cancel simulating shutdown + var stopwatch = Stopwatch.StartNew(); + cts.Cancel(); + + // All waiters should return false (cancelled = not signaled) + foreach (var waiter in waiters) + { + bool result = await waiter; + Assert.IsFalse(result, "Cancelled waiter should return false"); + } + + stopwatch.Stop(); + + // All should complete quickly + Assert.IsTrue( + stopwatch.ElapsedMilliseconds < 5000, + $"All waiters should complete in under 5s, but took {stopwatch.ElapsedMilliseconds}ms"); + } + + /// + /// Verifies that pre-cancelled tokens cause WaitAsync to return false immediately, + /// matching the behavior during shutdown if the token is cancelled before + /// FetchNewOrchestrationMessagesAsync is called. + /// + [TestMethod] + public async Task WaitAsync_AlreadyCancelledToken_ReturnsFalseImmediately() + { + var resetEvent = new AsyncAutoResetEvent(signaled: false); + using var cts = new CancellationTokenSource(); + cts.Cancel(); // Pre-cancel + + var stopwatch = Stopwatch.StartNew(); + bool result = await resetEvent.WaitAsync(TimeSpan.FromSeconds(30), cts.Token); + stopwatch.Stop(); + + Assert.IsFalse(result, "Pre-cancelled token should cause immediate false return"); + Assert.IsTrue( + stopwatch.ElapsedMilliseconds < 5000, + $"Should complete immediately, but took {stopwatch.ElapsedMilliseconds}ms"); + } + + /// + /// Verifies that a pre-cancelled token still returns true if the event + /// is already signaled (messages were already queued before shutdown). + /// + [TestMethod] + public async Task WaitAsync_AlreadySignaledAndCancelled_ReturnsTrue() + { + var resetEvent = new AsyncAutoResetEvent(signaled: true); + using var cts = new CancellationTokenSource(); + cts.Cancel(); + + // Even with a cancelled token, if the event is already signaled, + // WaitAsync should return true immediately without blocking + bool result = await resetEvent.WaitAsync(TimeSpan.FromSeconds(30), cts.Token); + Assert.IsTrue(result, "Already signaled event should return true even with cancelled token"); + } + + /// + /// Verifies that clears + /// all active sessions, causing + /// to return false. This is the core mechanism behind Bug 2 fix (isForced support). + /// + [TestMethod] + public void AbortAllSessions_ClearsActiveSessions() + { + var settings = new AzureStorageOrchestrationServiceSettings(); + var stats = new AzureStorageOrchestrationServiceStats(); + var trackingStore = new Mock(); + + using var manager = new OrchestrationSessionManager( + "testaccount", + settings, + stats, + trackingStore.Object); + + // Use reflection to add fake sessions to the internal dictionary, since constructing + // a real OrchestrationSession requires Azure infrastructure (ControlQueue) + var sessionsField = typeof(OrchestrationSessionManager) + .GetField("activeOrchestrationSessions", BindingFlags.NonPublic | BindingFlags.Instance); + var sessions = (Dictionary)sessionsField.GetValue(manager); + + // Add null entries to simulate active sessions (we only need the dictionary to be non-empty + // to test that AbortAllSessions clears it and IsControlQueueProcessingMessages returns false) + // Use the internal dictionary directly with a partition check + manager.GetStats(out _, out _, out int initialCount); + Assert.AreEqual(0, initialCount, "Should start with no active sessions"); + + // Manually insert entries to simulate active sessions + sessions["instance1"] = null; + sessions["instance2"] = null; + sessions["instance3"] = null; + + manager.GetStats(out _, out _, out int activeCount); + Assert.AreEqual(3, activeCount, "Should have 3 active sessions"); + + // Act: simulate forced shutdown by aborting all sessions + manager.AbortAllSessions(); + + // Assert: all sessions should be cleared + manager.GetStats(out _, out _, out int afterAbortCount); + Assert.AreEqual(0, afterAbortCount, "AbortAllSessions should clear all active sessions"); + } + + /// + /// Verifies that is safe to call + /// when there are no active sessions (idempotent). + /// + [TestMethod] + public void AbortAllSessions_NoSessions_DoesNotThrow() + { + var settings = new AzureStorageOrchestrationServiceSettings(); + var stats = new AzureStorageOrchestrationServiceStats(); + var trackingStore = new Mock(); + + using var manager = new OrchestrationSessionManager( + "testaccount", + settings, + stats, + trackingStore.Object); + + // Should not throw when called with no active sessions + manager.AbortAllSessions(); + + manager.GetStats(out _, out _, out int count); + Assert.AreEqual(0, count, "Should still have no active sessions"); + } + } +} diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs index 38cb7fac3..9c3471978 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs @@ -478,6 +478,15 @@ public async Task StopAsync(bool isForced) { this.shutdownSource.Cancel(); await this.statsLoop; + + if (isForced) + { + // When forced, immediately remove all active sessions so that + // partition draining completes without waiting for sessions to + // finish their idle timeout or in-flight work. + this.orchestrationSessionManager.AbortAllSessions(); + } + await this.appLeaseManager.StopAsync(); this.isStarted = false; } diff --git a/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs b/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs index 1b2e4a20e..b5a112741 100644 --- a/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs +++ b/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs @@ -17,6 +17,7 @@ namespace DurableTask.AzureStorage.Messaging using System.Collections.Generic; using System.IO; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Azure; using DurableTask.Core; @@ -26,6 +27,7 @@ namespace DurableTask.AzureStorage.Messaging sealed class OrchestrationSession : SessionBase, IOrchestrationSession { readonly TimeSpan idleTimeout; + readonly CancellationToken shutdownToken; readonly AsyncAutoResetEvent messagesAvailableEvent; readonly MessageCollection nextMessageBatch; @@ -41,10 +43,12 @@ public OrchestrationSession( DateTime lastCheckpointTime, object trackingStoreContext, TimeSpan idleTimeout, + CancellationToken shutdownToken, Guid traceActivityId) : base(settings, storageAccountName, orchestrationInstance, traceActivityId) { this.idleTimeout = idleTimeout; + this.shutdownToken = shutdownToken; this.ControlQueue = controlQueue ?? throw new ArgumentNullException(nameof(controlQueue)); this.CurrentMessageBatch = initialMessageBatch ?? throw new ArgumentNullException(nameof(initialMessageBatch)); this.RuntimeState = runtimeState ?? throw new ArgumentNullException(nameof(runtimeState)); @@ -98,9 +102,16 @@ public void AddOrReplaceMessages(IEnumerable messages) public async Task> FetchNewOrchestrationMessagesAsync( TaskOrchestrationWorkItem workItem) { - if (!await this.messagesAvailableEvent.WaitAsync(this.idleTimeout)) + try + { + if (!await this.messagesAvailableEvent.WaitAsync(this.idleTimeout, this.shutdownToken)) + { + return null; // timed-out + } + } + catch (OperationCanceledException) { - return null; // timed-out + return null; // shutting down } this.StartNewLogicalTraceScope(); diff --git a/src/DurableTask.AzureStorage/OrchestrationSessionManager.cs b/src/DurableTask.AzureStorage/OrchestrationSessionManager.cs index 7ecae232d..abf7a58b2 100644 --- a/src/DurableTask.AzureStorage/OrchestrationSessionManager.cs +++ b/src/DurableTask.AzureStorage/OrchestrationSessionManager.cs @@ -43,6 +43,8 @@ class OrchestrationSessionManager : IDisposable readonly ITrackingStore trackingStore; readonly DispatchQueue fetchRuntimeStateQueue; + CancellationToken shutdownToken; + public OrchestrationSessionManager( string queueAccountName, AzureStorageOrchestrationServiceSettings settings, @@ -61,6 +63,8 @@ public OrchestrationSessionManager( public void AddQueue(string partitionId, ControlQueue controlQueue, CancellationToken cancellationToken) { + this.shutdownToken = cancellationToken; + if (this.ownedControlQueues.TryAdd(partitionId, controlQueue)) { _ = Task.Run(() => this.DequeueLoop(partitionId, controlQueue, cancellationToken)); @@ -613,6 +617,7 @@ async Task ScheduleOrchestrationStatePrefetch( nextBatch.LastCheckpointTime, nextBatch.TrackingStoreContext, this.settings.ExtendedSessionIdleTimeout, + this.shutdownToken, traceActivityId); this.activeOrchestrationSessions.Add(instance.InstanceId, session); @@ -656,6 +661,19 @@ async Task ScheduleOrchestrationStatePrefetch( return null; } + /// + /// Immediately removes all active sessions, causing + /// to return false for all partitions. This unblocks so that + /// a forced shutdown can complete without waiting for sessions to drain naturally. + /// + public void AbortAllSessions() + { + lock (this.messageAndSessionLock) + { + this.activeOrchestrationSessions.Clear(); + } + } + public bool TryGetExistingSession(string instanceId, out OrchestrationSession session) { lock (this.messageAndSessionLock) From 4070c0ce42f6acc3a14a822d20cd573f79ed6cd6 Mon Sep 17 00:00:00 2001 From: andystaples <77818326+andystaples@users.noreply.github.com> Date: Fri, 6 Mar 2026 16:34:52 -0700 Subject: [PATCH 2/3] Potential fix for pull request finding 'Missed opportunity to use Select' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../OrchestrationSessionTests.cs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs b/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs index a304246ec..5a78b7cbb 100644 --- a/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs +++ b/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs @@ -16,6 +16,7 @@ namespace DurableTask.AzureStorage.Tests using System; using System.Collections.Generic; using System.Diagnostics; + using System.Linq; using System.Reflection; using System.Threading; using System.Threading.Tasks; @@ -134,11 +135,13 @@ public async Task WaitAsync_CancellationToken_MultipleWaiters() cts.Cancel(); // All waiters should return false (cancelled = not signaled) - foreach (var waiter in waiters) - { - bool result = await waiter; - Assert.IsFalse(result, "Cancelled waiter should return false"); - } + await Task.WhenAll( + waiters.Select( + async waiter => + { + bool result = await waiter; + Assert.IsFalse(result, "Cancelled waiter should return false"); + })); stopwatch.Stop(); From 5c38149727021ab4e7319c20c74d448aebdc462c Mon Sep 17 00:00:00 2001 From: Andy Staples Date: Fri, 6 Mar 2026 20:09:27 -0700 Subject: [PATCH 3/3] PR feedback --- .../OrchestrationSessionTests.cs | 57 ++++--------------- .../Messaging/OrchestrationSession.cs | 11 +--- 2 files changed, 12 insertions(+), 56 deletions(-) diff --git a/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs b/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs index a304246ec..1664ba66b 100644 --- a/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs +++ b/Test/DurableTask.AzureStorage.Tests/OrchestrationSessionTests.cs @@ -26,18 +26,14 @@ namespace DurableTask.AzureStorage.Tests using Moq; /// - /// Tests that validate the shutdown cancellation behavior for extended sessions. - /// These tests verify fixes for the issue where host shutdown was blocked by the - /// extended session idle timeout (default 30s) because no cancellation token was - /// propagated to the session wait, and isForced was ignored during StopAsync. + /// Tests for shutdown cancellation behavior with extended sessions. /// [TestClass] public class OrchestrationSessionTests { /// /// Verifies that - /// exits immediately when the cancellation token is cancelled, rather than waiting for - /// the full timeout. This is the core mechanism behind Bug 1 fix. + /// exits immediately when the cancellation token is cancelled. /// [TestMethod] public async Task WaitAsync_CancellationToken_ExitsImmediately() @@ -45,32 +41,25 @@ public async Task WaitAsync_CancellationToken_ExitsImmediately() var resetEvent = new AsyncAutoResetEvent(signaled: false); using var cts = new CancellationTokenSource(); - // Start a wait with a long timeout (simulating ExtendedSessionIdleTimeoutInSeconds = 30s) TimeSpan longTimeout = TimeSpan.FromSeconds(30); Task waitTask = resetEvent.WaitAsync(longTimeout, cts.Token); - // The wait should not have completed yet Assert.IsFalse(waitTask.IsCompleted, "Wait should not complete immediately"); - // Cancel the token (simulating shutdown) var stopwatch = Stopwatch.StartNew(); cts.Cancel(); - // The wait should return false almost immediately (same as timeout) bool result = await waitTask; stopwatch.Stop(); Assert.IsFalse(result, "Cancellation should return false (no signal received)"); - - // Verify it completed quickly (well under the 30s timeout) Assert.IsTrue( stopwatch.ElapsedMilliseconds < 5000, $"Cancellation should complete in under 5s, but took {stopwatch.ElapsedMilliseconds}ms"); } /// - /// Verifies that the wait still works normally (returns true) when signaled, - /// even when a cancellation token is provided. + /// Verifies that signaling still returns true when a cancellation token is provided. /// [TestMethod] public async Task WaitAsync_WithCancellationToken_SignalStillWorks() @@ -81,18 +70,15 @@ public async Task WaitAsync_WithCancellationToken_SignalStillWorks() Task waitTask = resetEvent.WaitAsync(TimeSpan.FromSeconds(30), cts.Token); Assert.IsFalse(waitTask.IsCompleted); - // Signal the event (simulating new messages arriving) resetEvent.Set(); - // Wait for the task with a reasonable timeout Task winner = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(5))); Assert.IsTrue(winner == waitTask, "Signal should wake the waiter"); Assert.IsTrue(waitTask.Result, "Wait result should be true when signaled"); } /// - /// Verifies that the wait returns false on timeout even when a cancellation - /// token is provided but not cancelled (normal idle timeout behavior). + /// Verifies that the wait returns false on timeout when a cancellation token is provided but not cancelled. /// [TestMethod] public async Task WaitAsync_WithCancellationToken_TimeoutStillWorks() @@ -100,15 +86,13 @@ public async Task WaitAsync_WithCancellationToken_TimeoutStillWorks() var resetEvent = new AsyncAutoResetEvent(signaled: false); using var cts = new CancellationTokenSource(); - // Use a short timeout to keep the test fast bool result = await resetEvent.WaitAsync(TimeSpan.FromMilliseconds(100), cts.Token); Assert.IsFalse(result, "Wait should return false on timeout"); } /// - /// Verifies that cancellation works when multiple waiters are queued. - /// All waiters should return false when the token fires. + /// Verifies that all queued waiters return false when the token is cancelled. /// [TestMethod] public async Task WaitAsync_CancellationToken_MultipleWaiters() @@ -116,24 +100,20 @@ public async Task WaitAsync_CancellationToken_MultipleWaiters() var resetEvent = new AsyncAutoResetEvent(signaled: false); using var cts = new CancellationTokenSource(); - // Queue multiple waiters (simulating multiple sessions waiting during shutdown) var waiters = new List>(); for (int i = 0; i < 5; i++) { waiters.Add(resetEvent.WaitAsync(TimeSpan.FromSeconds(30), cts.Token)); } - // None should be completed yet foreach (var waiter in waiters) { Assert.IsFalse(waiter.IsCompleted); } - // Cancel simulating shutdown var stopwatch = Stopwatch.StartNew(); cts.Cancel(); - // All waiters should return false (cancelled = not signaled) foreach (var waiter in waiters) { bool result = await waiter; @@ -142,16 +122,13 @@ public async Task WaitAsync_CancellationToken_MultipleWaiters() stopwatch.Stop(); - // All should complete quickly Assert.IsTrue( stopwatch.ElapsedMilliseconds < 5000, $"All waiters should complete in under 5s, but took {stopwatch.ElapsedMilliseconds}ms"); } /// - /// Verifies that pre-cancelled tokens cause WaitAsync to return false immediately, - /// matching the behavior during shutdown if the token is cancelled before - /// FetchNewOrchestrationMessagesAsync is called. + /// Verifies that a pre-cancelled token causes WaitAsync to return false immediately. /// [TestMethod] public async Task WaitAsync_AlreadyCancelledToken_ReturnsFalseImmediately() @@ -171,8 +148,7 @@ public async Task WaitAsync_AlreadyCancelledToken_ReturnsFalseImmediately() } /// - /// Verifies that a pre-cancelled token still returns true if the event - /// is already signaled (messages were already queued before shutdown). + /// Verifies that a pre-cancelled token still returns true if the event is already signaled. /// [TestMethod] public async Task WaitAsync_AlreadySignaledAndCancelled_ReturnsTrue() @@ -181,16 +157,12 @@ public async Task WaitAsync_AlreadySignaledAndCancelled_ReturnsTrue() using var cts = new CancellationTokenSource(); cts.Cancel(); - // Even with a cancelled token, if the event is already signaled, - // WaitAsync should return true immediately without blocking bool result = await resetEvent.WaitAsync(TimeSpan.FromSeconds(30), cts.Token); Assert.IsTrue(result, "Already signaled event should return true even with cancelled token"); } /// - /// Verifies that clears - /// all active sessions, causing - /// to return false. This is the core mechanism behind Bug 2 fix (isForced support). + /// Verifies that clears all active sessions. /// [TestMethod] public void AbortAllSessions_ClearsActiveSessions() @@ -205,19 +177,14 @@ public void AbortAllSessions_ClearsActiveSessions() stats, trackingStore.Object); - // Use reflection to add fake sessions to the internal dictionary, since constructing - // a real OrchestrationSession requires Azure infrastructure (ControlQueue) + // Use reflection to access the internal sessions dictionary. var sessionsField = typeof(OrchestrationSessionManager) .GetField("activeOrchestrationSessions", BindingFlags.NonPublic | BindingFlags.Instance); var sessions = (Dictionary)sessionsField.GetValue(manager); - // Add null entries to simulate active sessions (we only need the dictionary to be non-empty - // to test that AbortAllSessions clears it and IsControlQueueProcessingMessages returns false) - // Use the internal dictionary directly with a partition check manager.GetStats(out _, out _, out int initialCount); Assert.AreEqual(0, initialCount, "Should start with no active sessions"); - // Manually insert entries to simulate active sessions sessions["instance1"] = null; sessions["instance2"] = null; sessions["instance3"] = null; @@ -225,17 +192,14 @@ public void AbortAllSessions_ClearsActiveSessions() manager.GetStats(out _, out _, out int activeCount); Assert.AreEqual(3, activeCount, "Should have 3 active sessions"); - // Act: simulate forced shutdown by aborting all sessions manager.AbortAllSessions(); - // Assert: all sessions should be cleared manager.GetStats(out _, out _, out int afterAbortCount); Assert.AreEqual(0, afterAbortCount, "AbortAllSessions should clear all active sessions"); } /// - /// Verifies that is safe to call - /// when there are no active sessions (idempotent). + /// Verifies that is safe to call with no active sessions. /// [TestMethod] public void AbortAllSessions_NoSessions_DoesNotThrow() @@ -250,7 +214,6 @@ public void AbortAllSessions_NoSessions_DoesNotThrow() stats, trackingStore.Object); - // Should not throw when called with no active sessions manager.AbortAllSessions(); manager.GetStats(out _, out _, out int count); diff --git a/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs b/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs index b5a112741..d30a1b40f 100644 --- a/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs +++ b/src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs @@ -102,16 +102,9 @@ public void AddOrReplaceMessages(IEnumerable messages) public async Task> FetchNewOrchestrationMessagesAsync( TaskOrchestrationWorkItem workItem) { - try - { - if (!await this.messagesAvailableEvent.WaitAsync(this.idleTimeout, this.shutdownToken)) - { - return null; // timed-out - } - } - catch (OperationCanceledException) + if (!await this.messagesAvailableEvent.WaitAsync(this.idleTimeout, this.shutdownToken)) { - return null; // shutting down + return null; // timed-out or shutting down } this.StartNewLogicalTraceScope();