From adfed59f58e6128af4f5ab72816facf7dbcddeb7 Mon Sep 17 00:00:00 2001 From: wangbill Date: Thu, 3 Sep 2026 14:32:45 -0700 Subject: [PATCH 01/10] Gate activity polling on app lease ownership Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../AppLeaseActivityTests.cs | 539 ++++++++++++++++++ .../AzureStorageOrchestrationService.cs | 96 ++-- ...zureStorageOrchestrationServiceSettings.cs | 4 +- .../Partitioning/AppLeaseManager.cs | 132 ++++- .../Partitioning/AppLeaseOwnershipSignal.cs | 182 ++++++ 5 files changed, 885 insertions(+), 68 deletions(-) create mode 100644 Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs create mode 100644 src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs diff --git a/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs new file mode 100644 index 000000000..72f1c124f --- /dev/null +++ b/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -0,0 +1,539 @@ +// ---------------------------------------------------------------------------------- +// 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.Linq; + using System.Runtime.CompilerServices; + using System.Threading; + using System.Threading.Tasks; + using DurableTask.AzureStorage.Partitioning; + using DurableTask.AzureStorage.Storage; + using DurableTask.Core; + using DurableTask.Core.History; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class AppLeaseActivityTests + { + static readonly TimeSpan TestTimeout = TimeSpan.FromSeconds(15); + static readonly TimeSpan AutomaticFailoverTimeout = TimeSpan.FromSeconds(30); + + [TestMethod] + public async Task DifferentAppCannotDequeueActivityWhilePassive() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService owner = CreateService(taskHubName, "PrimaryApp", useAppLease: true); + AzureStorageOrchestrationService passive = CreateService(taskHubName, "SecondaryApp", useAppLease: true); + + try + { + await owner.CreateAsync(); + await owner.StartAsync(); + await WaitForOwnerAsync(owner); + await passive.StartAsync(); + await EnqueueActivityAsync(owner, "activity"); + + using (var cancellation = new CancellationTokenSource(TimeSpan.FromSeconds(1))) + { + TaskActivityWorkItem passiveWorkItem = await passive.LockNextTaskActivityWorkItem( + TestTimeout, + cancellation.Token); + Assert.IsNull(passiveWorkItem); + } + + TaskActivityWorkItem ownerWorkItem = await LockActivityAsync(owner); + Assert.IsNotNull(ownerWorkItem); + await owner.AbandonTaskActivityWorkItemAsync(ownerWorkItem); + } + finally + { + await StopAsync(passive); + await StopAsync(owner); + } + } + + [TestMethod] + public async Task SameAppWorkersCanDequeueActivitiesConcurrently() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService worker1 = CreateService(taskHubName, "SharedApp", useAppLease: true); + AzureStorageOrchestrationService worker2 = CreateService(taskHubName, "SharedApp", useAppLease: true); + + try + { + await worker1.CreateAsync(); + await worker1.StartAsync(); + await WaitForOwnerAsync(worker1); + await worker2.StartAsync(); + await EnqueueActivityAsync(worker1, "activity-1"); + await EnqueueActivityAsync(worker1, "activity-2"); + + Task lock1 = LockActivityAsync(worker1); + Task lock2 = LockActivityAsync(worker2); + TaskActivityWorkItem[] workItems = await Task.WhenAll(lock1, lock2); + + Assert.IsNotNull(workItems[0]); + Assert.IsNotNull(workItems[1]); + Assert.AreNotEqual(workItems[0].Id, workItems[1].Id); + + await worker1.AbandonTaskActivityWorkItemAsync(workItems[0]); + await worker2.AbandonTaskActivityWorkItemAsync(workItems[1]); + } + finally + { + await StopAsync(worker2); + await StopAsync(worker1); + } + } + + [TestMethod] + public async Task AppLeaseDisabledAllowsDifferentAppsToDequeueActivities() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService worker1 = CreateService(taskHubName, "App1", useAppLease: false); + AzureStorageOrchestrationService worker2 = CreateService(taskHubName, "App2", useAppLease: false); + + try + { + await worker1.CreateAsync(); + await worker1.StartAsync(); + await worker2.StartAsync(); + await EnqueueActivityAsync(worker1, "activity-1"); + await EnqueueActivityAsync(worker1, "activity-2"); + + TaskActivityWorkItem[] workItems = await Task.WhenAll( + LockActivityAsync(worker1), + LockActivityAsync(worker2)); + + Assert.IsNotNull(workItems[0]); + Assert.IsNotNull(workItems[1]); + + await worker1.AbandonTaskActivityWorkItemAsync(workItems[0]); + await worker2.AbandonTaskActivityWorkItemAsync(workItems[1]); + } + finally + { + await StopAsync(worker2); + await StopAsync(worker1); + } + } + + [TestMethod] + public async Task ForcedFailoverCancelsOldOwnerPollAndEnablesNewOwner() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService oldOwner = CreateService(taskHubName, "OldApp", useAppLease: true); + AzureStorageOrchestrationService newOwner = CreateService(taskHubName, "NewApp", useAppLease: true); + + try + { + await oldOwner.CreateAsync(); + await oldOwner.StartAsync(); + await WaitForOwnerAsync(oldOwner); + await newOwner.StartAsync(); + + Task oldOwnerPoll = oldOwner.LockNextTaskActivityWorkItem( + TestTimeout, + CancellationToken.None); + + await newOwner.ForceChangeAppLeaseAsync(); + await WaitForOwnerAsync(newOwner); + await EnqueueActivityAsync(newOwner, "activity"); + + TaskActivityWorkItem newOwnerWorkItem = await LockActivityAsync(newOwner); + Assert.IsNotNull(newOwnerWorkItem); + Assert.IsNull(await WithTimeoutAsync(oldOwnerPoll)); + + await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); + } + finally + { + await StopAsync(newOwner); + await StopAsync(oldOwner); + } + } + + [TestMethod] + public async Task ForcedFailoverQuiescesOldAppBeforeNewOwnerStarts() + { + string taskHubName = GetTaskHubName(); + TimeSpan renewInterval = TimeSpan.FromSeconds(2); + AzureStorageOrchestrationServiceSettings oldSettings = CreateSettings( + taskHubName, + "OldApp", + useAppLease: true, + renewInterval: renewInterval); + AzureStorageOrchestrationServiceSettings newSettings = CreateSettings( + taskHubName, + "NewApp", + useAppLease: true, + renewInterval: renewInterval); + AppLeaseManager oldManager = CreateAppLeaseManager( + oldSettings, + new TestPartitionManager()); + AppLeaseManager newManager = CreateAppLeaseManager( + newSettings, + new TestPartitionManager()); + AppLeaseOwnershipSignal.AppLeaseOwnership oldOwnership = null; + AppLeaseOwnershipSignal.AppLeaseOwnership newOwnership = null; + + try + { + await oldManager.CreateContainerIfNotExistsAsync(); + await oldManager.StartAsync(); + oldOwnership = await GetOwnershipAsync(oldManager); + await newManager.StartAsync(); + + await newManager.ForceChangeAppLeaseAsync(); + newOwnership = await GetOwnershipAsync(newManager); + + Assert.IsTrue(oldOwnership.LostToken.IsCancellationRequested); + Assert.IsFalse(oldOwnership.TryBeginDispatch()); + } + finally + { + newOwnership?.Dispose(); + oldOwnership?.Dispose(); + await newManager.StopAsync(); + await oldManager.StopAsync(); + } + } + + [TestMethod] + public async Task LeaseExpirationEnablesNewOwner() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService oldOwner = CreateService(taskHubName, "OldApp", useAppLease: true); + AzureStorageOrchestrationService newOwner = CreateService(taskHubName, "NewApp", useAppLease: true); + + try + { + await oldOwner.CreateAsync(); + await oldOwner.StartAsync(); + await WaitForOwnerAsync(oldOwner); + await newOwner.StartAsync(); + await oldOwner.StopAsync(isForced: true); + oldOwner = null; + await EnqueueActivityAsync(newOwner, "activity"); + + TaskActivityWorkItem newOwnerWorkItem = + await LockActivityAsync(newOwner, AutomaticFailoverTimeout); + + Assert.IsNotNull(newOwnerWorkItem); + await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); + } + finally + { + await StopAsync(newOwner); + await StopAsync(oldOwner); + } + } + + [TestMethod] + public async Task AppLeaseManagerCanRestartAfterPartitionManagerStopFails() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationServiceSettings settings = + CreateSettings(taskHubName, "PrimaryApp", useAppLease: true); + var partitionManager = new StopFailingPartitionManager(); + AppLeaseManager manager = CreateAppLeaseManager(settings, partitionManager); + + await manager.CreateContainerIfNotExistsAsync(); + await manager.StartAsync(); + await WaitForOwnershipAsync(manager); + + await Assert.ThrowsExceptionAsync(() => manager.StopAsync()); + + await manager.StartAsync(); + await WaitForOwnershipAsync(manager); + + await WithTimeoutAsync(manager.StopAsync()); + Assert.AreEqual(2, partitionManager.StartCount); + Assert.AreEqual(2, partitionManager.StopCount); + } + + [TestMethod] + public async Task LeaseLossDoesNotCancelDispatchedActivity() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService oldOwner = CreateService(taskHubName, "OldApp", useAppLease: true); + AzureStorageOrchestrationService newOwner = CreateService(taskHubName, "NewApp", useAppLease: true); + + try + { + await oldOwner.CreateAsync(); + await oldOwner.StartAsync(); + await WaitForOwnerAsync(oldOwner); + await newOwner.StartAsync(); + await EnqueueActivityAsync(oldOwner, "in-flight"); + + TaskActivityWorkItem inFlightWorkItem = await LockActivityAsync(oldOwner); + await newOwner.ForceChangeAppLeaseAsync(); + await EnqueueActivityAsync(newOwner, "after-failover"); + TaskActivityWorkItem newOwnerWorkItem = await LockActivityAsync(newOwner); + + TaskActivityWorkItem renewedWorkItem = + await oldOwner.RenewTaskActivityWorkItemLockAsync(inFlightWorkItem); + Assert.IsTrue(renewedWorkItem.LockedUntilUtc > DateTime.UtcNow); + + await oldOwner.AbandonTaskActivityWorkItemAsync(inFlightWorkItem); + await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); + } + finally + { + await StopAsync(newOwner); + await StopAsync(oldOwner); + } + } + + [TestMethod] + public async Task OwnershipLossAfterDequeuePreventsDispatch() + { + var ownershipSignal = new AppLeaseOwnershipSignal(); + ownershipSignal.Set(); + + using (AppLeaseOwnershipSignal.AppLeaseOwnership ownership = + await ownershipSignal.WaitAsync(CancellationToken.None)) + { + ownershipSignal.Reset(); + + Assert.IsTrue(ownership.LostToken.IsCancellationRequested); + Assert.IsFalse(ownership.TryBeginDispatch()); + } + } + + [TestMethod] + public async Task PassiveActivityPollStopsOnCancellationAndShutdown() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService owner = CreateService(taskHubName, "PrimaryApp", useAppLease: true); + AzureStorageOrchestrationService passive = CreateService(taskHubName, "SecondaryApp", useAppLease: true); + + try + { + await owner.CreateAsync(); + await owner.StartAsync(); + await WaitForOwnerAsync(owner); + await passive.StartAsync(); + + using (var cancellation = new CancellationTokenSource()) + { + Task canceledPoll = passive.LockNextTaskActivityWorkItem( + TestTimeout, + cancellation.Token); + cancellation.Cancel(); + Assert.IsNull(await WithTimeoutAsync(canceledPoll)); + } + + Task shutdownPoll = passive.LockNextTaskActivityWorkItem( + TestTimeout, + CancellationToken.None); + await passive.StopAsync(isForced: true); + Assert.IsNull(await WithTimeoutAsync(shutdownPoll)); + passive = null; + } + finally + { + await StopAsync(passive); + await StopAsync(owner); + } + } + + static AzureStorageOrchestrationService CreateService( + string taskHubName, + string appName, + bool useAppLease, + TimeSpan? renewInterval = null) + { + return new AzureStorageOrchestrationService( + CreateSettings(taskHubName, appName, useAppLease, renewInterval)); + } + + static AzureStorageOrchestrationServiceSettings CreateSettings( + string taskHubName, + string appName, + bool useAppLease, + TimeSpan? renewInterval = null) + { + return new AzureStorageOrchestrationServiceSettings + { + AppName = appName, + AppLeaseOptions = new AppLeaseOptions + { + AcquireInterval = TimeSpan.FromMilliseconds(200), + LeaseInterval = TimeSpan.FromSeconds(15), + RenewInterval = renewInterval ?? TimeSpan.FromMilliseconds(200), + }, + MaxQueuePollingInterval = TimeSpan.FromMilliseconds(50), + PartitionCount = 1, + StorageAccountClientProvider = new StorageAccountClientProvider( + TestHelpers.GetTestStorageAccountConnectionString()), + TaskHubName = taskHubName, + UseAppLease = useAppLease, + WorkerId = Guid.NewGuid().ToString("N"), + }; + } + + static string GetTaskHubName() + { + return "applease" + Guid.NewGuid().ToString("N").Substring(0, 16); + } + + static async Task EnqueueActivityAsync( + AzureStorageOrchestrationService service, + string activityName) + { + var instance = new OrchestrationInstance + { + ExecutionId = Guid.NewGuid().ToString("N"), + InstanceId = Guid.NewGuid().ToString("N"), + }; + + await service.WorkItemQueue.AddMessageAsync( + new TaskMessage + { + Event = new TaskScheduledEvent(0, activityName), + OrchestrationInstance = instance, + }, + instance); + } + + static async Task WaitForOwnerAsync(AzureStorageOrchestrationService service) + { + await TestHelpers.WaitFor( + () => service.OwnedControlQueues.Any(), + TestTimeout); + } + + static async Task WaitForOwnershipAsync(AppLeaseManager manager) + { + using (await GetOwnershipAsync(manager)) + { + } + } + + static async Task GetOwnershipAsync( + AppLeaseManager manager) + { + using (var cancellation = new CancellationTokenSource(TestTimeout)) + { + return await manager.WaitForOwnershipAsync(cancellation.Token); + } + } + + static AppLeaseManager CreateAppLeaseManager( + AzureStorageOrchestrationServiceSettings settings, + IPartitionManager partitionManager) + { + string taskHubName = settings.TaskHubName.ToLowerInvariant(); + return new AppLeaseManager( + new AzureStorageClient(settings), + partitionManager, + taskHubName + "-applease", + taskHubName + "-appleaseinfo", + settings.AppLeaseOptions); + } + + static async Task LockActivityAsync( + AzureStorageOrchestrationService service, + TimeSpan? timeout = null) + { + TimeSpan effectiveTimeout = timeout ?? TestTimeout; + using (var cancellation = new CancellationTokenSource(effectiveTimeout)) + { + return await service.LockNextTaskActivityWorkItem(effectiveTimeout, cancellation.Token); + } + } + + static async Task WithTimeoutAsync(Task task) + { + Task completedTask = await Task.WhenAny(task, Task.Delay(TestTimeout)); + Assert.AreSame(task, completedTask, "The operation did not complete before the test timeout."); + return await task; + } + + static async Task WithTimeoutAsync(Task task) + { + Task completedTask = await Task.WhenAny(task, Task.Delay(TestTimeout)); + Assert.AreSame(task, completedTask, "The operation did not complete before the test timeout."); + await task; + } + + static async Task StopAsync(AzureStorageOrchestrationService service) + { + if (service != null) + { + await service.StopAsync(isForced: true); + } + } + + sealed class StopFailingPartitionManager : IPartitionManager + { + public int StartCount { get; private set; } + + public int StopCount { get; private set; } + + public Task StartAsync() + { + this.StartCount++; + return Task.CompletedTask; + } + + public Task StopAsync() + { + this.StopCount++; + return this.StopCount == 1 + ? Task.FromException(new InvalidOperationException("Simulated stop failure.")) + : Task.CompletedTask; + } + + public Task CreateLeaseStore() => Task.CompletedTask; + + public Task CreateLease(string leaseName) => Task.CompletedTask; + + public Task DeleteLeases() => Task.CompletedTask; + + public async IAsyncEnumerable GetOwnershipBlobLeasesAsync( + [EnumeratorCancellation] + CancellationToken cancellationToken = default) + { + await Task.CompletedTask; + yield break; + } + } + + sealed class TestPartitionManager : IPartitionManager + { + public Task StartAsync() => Task.CompletedTask; + + public Task StopAsync() => Task.CompletedTask; + + public Task CreateLeaseStore() => Task.CompletedTask; + + public Task CreateLease(string leaseName) => Task.CompletedTask; + + public Task DeleteLeases() => Task.CompletedTask; + + public async IAsyncEnumerable GetOwnershipBlobLeasesAsync( + [EnumeratorCancellation] + CancellationToken cancellationToken = default) + { + await Task.CompletedTask; + yield break; + } + } + } +} diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs index 74798cb45..99f07aadb 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs @@ -1544,55 +1544,77 @@ public async Task LockNextTaskActivityWorkItem( using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, this.shutdownSource.Token)) { - MessageData message = await this.workItemQueue.GetMessageAsync(linkedCts.Token); - - if (message == null) + AppLeaseOwnershipSignal.AppLeaseOwnership ownership; + try + { + ownership = await this.appLeaseManager.WaitForOwnershipAsync(linkedCts.Token); + } + catch (OperationCanceledException) { - // shutting down return null; } + using (ownership) + { + using (var receiveCts = CancellationTokenSource.CreateLinkedTokenSource( + linkedCts.Token, + ownership.LostToken)) + { + MessageData message = await this.workItemQueue.GetMessageAsync(receiveCts.Token); - Guid traceActivityId = Guid.NewGuid(); - var session = new ActivitySession(this.settings, this.azureStorageClient.QueueAccountName, message, traceActivityId); - session.StartNewLogicalTraceScope(); + if (message == null) + { + // shutting down, canceled, or app lease ownership was lost + return null; + } - // correlation - TraceContextBase requestTraceContext = null; - CorrelationTraceClient.Propagate( - () => - { - string name = $"{TraceConstants.Activity} {Utils.GetTargetClassName(((TaskScheduledEvent)session.MessageData.TaskMessage.Event)?.Name)}"; - requestTraceContext = TraceContextFactory.Create(name); + Guid traceActivityId = Guid.NewGuid(); + var session = new ActivitySession(this.settings, this.azureStorageClient.QueueAccountName, message, traceActivityId); + session.StartNewLogicalTraceScope(); - TraceContextBase parentTraceContextBase = TraceContextBase.Restore(session.MessageData.SerializableTraceContext); - requestTraceContext.SetParentAndStart(parentTraceContextBase); - }); + TraceContextBase requestTraceContext = null; + CorrelationTraceClient.Propagate( + () => + { + string name = $"{TraceConstants.Activity} {Utils.GetTargetClassName(((TaskScheduledEvent)session.MessageData.TaskMessage.Event)?.Name)}"; + requestTraceContext = TraceContextFactory.Create(name); - TraceMessageReceived(this.settings, session.MessageData, this.azureStorageClient.QueueAccountName); - session.TraceProcessingMessage(message, isExtendedSession: false, this.workItemQueue.Name); + TraceContextBase parentTraceContextBase = TraceContextBase.Restore(session.MessageData.SerializableTraceContext); + requestTraceContext.SetParentAndStart(parentTraceContextBase); + }); - if (!this.activeActivitySessions.TryAdd(message.Id, session)) - { - // This means we're already processing this message. This is never expected since the message - // should be kept invisible via background calls to RenewTaskActivityWorkItemLockAsync. - this.settings.Logger.AssertFailure( - this.azureStorageClient.QueueAccountName, - this.settings.TaskHubName, - $"Work item queue message with ID = {message.Id} is being processed multiple times concurrently."); - return null; - } + TraceMessageReceived(this.settings, session.MessageData, this.azureStorageClient.QueueAccountName); + session.TraceProcessingMessage(message, isExtendedSession: false, this.workItemQueue.Name); - this.stats.ActiveActivityExecutions.Increment(); + if (!ownership.TryBeginDispatch()) + { + await this.workItemQueue.AbandonMessageAsync(message); + return null; + } - return new TaskActivityWorkItem - { - Id = message.Id, - TaskMessage = session.MessageData.TaskMessage, - LockedUntilUtc = message.OriginalQueueMessage.NextVisibleOn.Value.UtcDateTime, + if (!this.activeActivitySessions.TryAdd(message.Id, session)) + { + // This means we're already processing this message. This is never expected since the message + // should be kept invisible via background calls to RenewTaskActivityWorkItemLockAsync. + this.settings.Logger.AssertFailure( + this.azureStorageClient.QueueAccountName, + this.settings.TaskHubName, + $"Work item queue message with ID = {message.Id} is being processed multiple times concurrently."); + return null; + } - TraceContextBase = requestTraceContext - }; + this.stats.ActiveActivityExecutions.Increment(); + + return new TaskActivityWorkItem + { + Id = message.Id, + TaskMessage = session.MessageData.TaskMessage, + LockedUntilUtc = message.OriginalQueueMessage.NextVisibleOn.Value.UtcDateTime, + + TraceContextBase = requestTraceContext + }; + } + } } } diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs index e614da92f..58df21089 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs @@ -152,7 +152,9 @@ public class AzureStorageOrchestrationServiceSettings public TimeSpan MaxQueuePollingInterval { get; set; } = DefaultMaxQueuePollingInterval; /// - /// If true, takes a lease on the task hub container, allowing for only one app to process messages in a task hub at a time. + /// If true, takes a lease on the task hub container so that only workers with the lease-owning + /// process orchestration, entity, or activity messages. + /// Workers that share the same may process messages concurrently. /// public bool UseAppLease { get; set; } = true; diff --git a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs index 9be4dc366..d3c1afab0 100644 --- a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs @@ -42,6 +42,7 @@ sealed class AppLeaseManager readonly Blob appLeaseInfoBlob; readonly string appLeaseId; readonly AsyncManualResetEvent shutdownCompletedEvent; + readonly AppLeaseOwnershipSignal ownershipSignal; bool isLeaseOwner; int appLeaseIsStarted; @@ -78,12 +79,20 @@ public AppLeaseManager( this.isLeaseOwner = false; this.shutdownCompletedEvent = new AsyncManualResetEvent(); + this.ownershipSignal = new AppLeaseOwnershipSignal(); + } + + public Task WaitForOwnershipAsync( + CancellationToken cancellationToken) + { + return this.ownershipSignal.WaitAsync(cancellationToken); } public async Task StartAsync() { if (!this.appLeaseIsEnabled) { + this.ownershipSignal.Set(); this.starterTokenSource = new CancellationTokenSource(); await Task.Factory.StartNew(() => this.PartitionManagerStarter(this.starterTokenSource.Token)); @@ -138,7 +147,7 @@ async Task AppLeaseManagerStarter(CancellationToken cancellationToken) { try { - while (!await this.TryAcquireAppLeaseAsync()) + while (!await this.TryAcquireAppLeaseAsync(cancellationToken)) { await Task.Delay(this.settings.AppLeaseOptions.AcquireInterval, cancellationToken); } @@ -165,6 +174,8 @@ async Task AppLeaseManagerStarter(CancellationToken cancellationToken) public async Task StopAsync() { + this.ownershipSignal.Reset(); + if (this.starterTokenSource != null) { this.starterTokenSource.Cancel(); @@ -250,58 +261,104 @@ async Task StartAppLeaseAsync() throw new InvalidOperationException("AppLeaseManager has already started"); } - this.leaseRenewerCancellationTokenSource = new CancellationTokenSource(); - - await this.partitionManager.StartAsync(); - this.shutdownCompletedEvent.Reset(); + this.leaseRenewerCancellationTokenSource = new CancellationTokenSource(); - this.renewTask = await Task.Factory.StartNew(() => this.LeaseRenewer(leaseRenewerCancellationTokenSource.Token)); + try + { + await this.partitionManager.StartAsync(); + this.ownershipSignal.Set(); + this.renewTask = await Task.Factory.StartNew( + () => this.LeaseRenewer(this.leaseRenewerCancellationTokenSource.Token)); + } + catch + { + this.ownershipSignal.Reset(); + this.isLeaseOwner = false; + Interlocked.Exchange(ref this.appLeaseIsStarted, 0); + this.leaseRenewerCancellationTokenSource.Dispose(); + this.leaseRenewerCancellationTokenSource = null; + this.shutdownCompletedEvent.Set(); + throw; + } } - async Task StopAppLeaseAsync() + async Task StopAppLeaseAsync(bool calledFromLeaseRenewer = false) { - if (Interlocked.CompareExchange(ref this.appLeaseIsStarted, 0, 1) != 1) + int previousState = Interlocked.CompareExchange(ref this.appLeaseIsStarted, 2, 1); + if (previousState != 1) { - //idempotent + if (previousState == 2 && !calledFromLeaseRenewer) + { + await this.shutdownCompletedEvent.WaitAsync( + Timeout.InfiniteTimeSpan, + CancellationToken.None); + } + return; } - await this.partitionManager.StopAsync(); + CancellationTokenSource renewerCancellation = this.leaseRenewerCancellationTokenSource; + Task renewer = this.renewTask; - if (this.renewTask != null) + try { - this.leaseRenewerCancellationTokenSource.Cancel(); - await this.renewTask; - } + this.ownershipSignal.Reset(); - this.isLeaseOwner = false; + if (!calledFromLeaseRenewer && renewer != null) + { + renewerCancellation.Cancel(); + await renewer; + } + + await this.partitionManager.StopAsync(); + } + finally + { + this.isLeaseOwner = false; - this.shutdownCompletedEvent.Set(); + renewerCancellation?.Dispose(); + this.leaseRenewerCancellationTokenSource = null; + this.renewTask = null; - this.leaseRenewerCancellationTokenSource?.Dispose(); + Interlocked.Exchange(ref this.appLeaseIsStarted, 0); + this.shutdownCompletedEvent.Set(); + } } - async Task TryAcquireAppLeaseAsync() + async Task TryAcquireAppLeaseAsync(CancellationToken cancellationToken) { AppLeaseInfo appLeaseInfo = await this.GetAppLeaseInfoAsync(); bool leaseAcquired; if (appLeaseInfo.DesiredSwapId == this.appLeaseId) { - leaseAcquired = await this.ChangeLeaseAsync(appLeaseInfo.OwnerId); + leaseAcquired = await this.ChangeLeaseAsync(appLeaseInfo.OwnerId, cancellationToken); } else { leaseAcquired = await this.TryAcquireLeaseAsync(); } + if (leaseAcquired) + { + AppLeaseInfo currentAppLeaseInfo = await this.GetAppLeaseInfoAsync(); + TimeSpan transitionDelay = + currentAppLeaseInfo.TransitionUntilUtc.GetValueOrDefault() - DateTime.UtcNow; + if (transitionDelay > TimeSpan.Zero) + { + await Task.Delay(transitionDelay, cancellationToken); + } + } + this.isLeaseOwner = leaseAcquired; return leaseAcquired; } - async Task ChangeLeaseAsync(string currentLeaseId) + async Task ChangeLeaseAsync( + string currentLeaseId, + CancellationToken cancellationToken) { this.settings.Logger.PartitionManagerInfo( this.storageAccountName, @@ -323,12 +380,18 @@ async Task ChangeLeaseAsync(string currentLeaseId) await this.appLeaseContainer.ChangeLeaseAsync(this.appLeaseId, currentLeaseId); + DateTime transitionUntilUtc = DateTime.UtcNow.Add(this.options.RenewInterval); var appLeaseInfo = new AppLeaseInfo() { OwnerId = this.appLeaseId, + TransitionUntilUtc = transitionUntilUtc, }; await this.UpdateAppLeaseInfoBlob(appLeaseInfo); + + // Give the previous app's renewers time to observe the changed lease and stop + // dequeueing work before this app begins processing. + await Task.Delay(this.options.RenewInterval, cancellationToken); leaseAcquired = true; this.settings.Logger.LeaseAcquisitionSucceeded( @@ -337,14 +400,6 @@ async Task ChangeLeaseAsync(string currentLeaseId) this.workerName, this.appLeaseContainerName, LeaseType); - - // When changing the lease over to another app, the paritions will still be listened to on the first app until the AppLeaseManager - // renew task fails to renew the lease. To avoid potential split brain we must delay before the new lease holder can start - // listening to the partitions. - if (this.settings.UseLegacyPartitionManagement == true) - { - await Task.Delay(this.settings.AppLeaseOptions.RenewInterval); - } } catch (DurableTaskStorageException e) { @@ -428,7 +483,7 @@ async Task LeaseRenewer(CancellationToken cancellationToken) break; } - await Task.Delay(this.options.RenewInterval, this.leaseRenewerCancellationTokenSource.Token); + await Task.Delay(this.options.RenewInterval, cancellationToken); } catch (OperationCanceledException) { @@ -459,7 +514,22 @@ async Task LeaseRenewer(CancellationToken cancellationToken) this.appLeaseContainerName, "Lease renewer task completing. Stopping AppLeaseManager."); - await this.StopAppLeaseAsync(); + if (!cancellationToken.IsCancellationRequested) + { + try + { + await this.StopAppLeaseAsync(calledFromLeaseRenewer: true); + } + catch (Exception ex) + { + this.settings.Logger.PartitionManagerError( + this.storageAccountName, + this.taskHub, + this.workerName, + this.appLeaseContainerName, + $"Failed to stop AppLeaseManager after losing the app lease. AppLeaseId: {this.appLeaseId} Exception: {ex}"); + } + } } async Task RenewLeaseAsync() @@ -491,6 +561,7 @@ async Task RenewLeaseAsync() { renewed = false; this.isLeaseOwner = false; + this.ownershipSignal.Reset(); this.settings.Logger.LeaseRenewalFailed( this.storageAccountName, @@ -577,6 +648,7 @@ private class AppLeaseInfo { public string OwnerId { get; set; } public string DesiredSwapId { get; set; } + public DateTime? TransitionUntilUtc { get; set; } } } } diff --git a/src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs b/src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs new file mode 100644 index 000000000..507aac107 --- /dev/null +++ b/src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs @@ -0,0 +1,182 @@ +// ---------------------------------------------------------------------------------- +// 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.Partitioning +{ + using System; + using System.Threading; + using System.Threading.Tasks; + + sealed class AppLeaseOwnershipSignal + { + readonly object syncLock = new object(); + + TaskCompletionSource ownershipAvailable = CreateCompletionSource(); + OwnershipEpoch currentEpoch; + + public void Set() + { + TaskCompletionSource available; + lock (this.syncLock) + { + if (this.currentEpoch != null) + { + return; + } + + this.currentEpoch = new OwnershipEpoch(); + available = this.ownershipAvailable; + } + + available.TrySetResult(null); + } + + public void Reset() + { + OwnershipEpoch epoch; + lock (this.syncLock) + { + epoch = this.currentEpoch; + if (epoch == null) + { + return; + } + + this.currentEpoch = null; + this.ownershipAvailable = CreateCompletionSource(); + } + + epoch.Deactivate(); + } + + public async Task WaitAsync(CancellationToken cancellationToken) + { + while (true) + { + Task availableTask; + lock (this.syncLock) + { + AppLeaseOwnership ownership = this.currentEpoch?.TryAcquire(); + if (ownership != null) + { + return ownership; + } + + availableTask = this.ownershipAvailable.Task; + } + + var canceled = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + using (cancellationToken.Register(() => canceled.TrySetCanceled())) + { + await Task.WhenAny(availableTask, canceled.Task); + } + + cancellationToken.ThrowIfCancellationRequested(); + } + } + + static TaskCompletionSource CreateCompletionSource() + { + return new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + } + + internal sealed class OwnershipEpoch + { + readonly object syncLock = new object(); + readonly CancellationTokenSource lostSource = new CancellationTokenSource(); + + bool active = true; + int referenceCount = 1; + + public AppLeaseOwnership TryAcquire() + { + lock (this.syncLock) + { + if (!this.active) + { + return null; + } + + this.referenceCount++; + return new AppLeaseOwnership(this); + } + } + + public bool TryBeginDispatch() + { + lock (this.syncLock) + { + return this.active; + } + } + + public void Deactivate() + { + lock (this.syncLock) + { + if (!this.active) + { + return; + } + + this.active = false; + } + + this.lostSource.Cancel(); + this.Release(); + } + + public void Release() + { + bool dispose; + lock (this.syncLock) + { + this.referenceCount--; + dispose = this.referenceCount == 0; + } + + if (dispose) + { + this.lostSource.Dispose(); + } + } + + public CancellationToken LostToken => this.lostSource.Token; + } + + public sealed class AppLeaseOwnership : IDisposable + { + OwnershipEpoch epoch; + + internal AppLeaseOwnership(OwnershipEpoch epoch) + { + this.epoch = epoch; + } + + public CancellationToken LostToken => this.epoch.LostToken; + + public bool TryBeginDispatch() + { + return this.epoch.TryBeginDispatch(); + } + + public void Dispose() + { + OwnershipEpoch epoch = Interlocked.Exchange(ref this.epoch, null); + epoch?.Release(); + } + } + } +} From c9a49ef6e725a876aea00ea7f4a4392b2480df71 Mon Sep 17 00:00:00 2001 From: wangbill Date: Tue, 8 Sep 2026 09:19:40 -0700 Subject: [PATCH 02/10] Make table partition manager restartable Create shutdown token sources per run and serialize lifecycle transitions so a previous app-lease owner can resume partition processing after reacquisition. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../AppLeaseActivityTests.cs | 186 ++++++++++++++++++ .../Partitioning/TablePartitionManager.cs | 181 ++++++++++++----- 2 files changed, 323 insertions(+), 44 deletions(-) diff --git a/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs index 72f1c124f..14334487f 100644 --- a/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs +++ b/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -242,6 +242,80 @@ public async Task LeaseExpirationEnablesNewOwner() } } + [TestMethod] + public async Task PreviousOwnerProcessesOrchestrationAfterReacquiringAppLease() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationServiceSettings firstSettings = + CreateSettings(taskHubName, "FirstApp", useAppLease: true); + AzureStorageOrchestrationServiceSettings secondSettings = + CreateSettings(taskHubName, "SecondApp", useAppLease: true); + firstSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); + secondSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); + + var firstService = new AzureStorageOrchestrationService(firstSettings); + var secondService = new AzureStorageOrchestrationService(secondSettings); + var firstWorker = new TaskHubWorker(firstService); + var secondWorker = new TaskHubWorker(secondService); + firstWorker.AddTaskOrchestrations(typeof(WaitForSignalOrchestration)); + secondWorker.AddTaskOrchestrations(typeof(WaitForSignalOrchestration)); + + bool firstWorkerStarted = false; + bool secondWorkerStarted = false; + + try + { + await firstService.CreateAsync(); + await firstWorker.StartAsync(); + firstWorkerStarted = true; + await WaitForOwnerAsync(firstService); + + var client = new TaskHubClient(firstService); + OrchestrationInstance instance = + await client.CreateOrchestrationInstanceAsync( + typeof(WaitForSignalOrchestration), + input: null); + await WaitForOrchestrationStatusAsync( + client, + instance, + OrchestrationStatus.Running); + + await secondWorker.StartAsync(); + secondWorkerStarted = true; + await secondService.ForceChangeAppLeaseAsync(); + await WaitForOwnerAsync(secondService, secondSettings, instance.InstanceId); + + await secondWorker.StopAsync(isForced: true); + secondWorkerStarted = false; + await WaitForAppLeaseOwnerAsync(firstSettings, AutomaticFailoverTimeout); + await WaitForOwnerAsync( + firstService, + firstSettings, + instance.InstanceId, + AutomaticFailoverTimeout); + await client.RaiseEventAsync(instance, "complete", "done"); + + OrchestrationState state = + await client.WaitForOrchestrationAsync(instance, TestTimeout); + + Assert.IsNotNull(state); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("\"done\"", state.Output); + } + finally + { + if (secondWorkerStarted) + { + await secondWorker.StopAsync(isForced: true); + } + + if (firstWorkerStarted) + { + await firstWorker.StopAsync(isForced: true); + } + } + } + [TestMethod] public async Task AppLeaseManagerCanRestartAfterPartitionManagerStopFails() { @@ -418,6 +492,89 @@ await TestHelpers.WaitFor( TestTimeout); } + static async Task WaitForOwnerAsync( + AzureStorageOrchestrationService service, + AzureStorageOrchestrationServiceSettings settings, + string instanceId, + TimeSpan? timeout = null) + { + uint partitionIndex = + Fnv1aHashHelper.ComputeHash(instanceId) % (uint)settings.PartitionCount; + string queueName = + AzureStorageOrchestrationService.GetControlQueueName( + settings.TaskHubName, + (int)partitionIndex); + + await TestHelpers.WaitFor( + () => service.OwnedControlQueues.Any(queue => queue.Name == queueName), + timeout ?? TestTimeout); + } + + static async Task WaitForAppLeaseOwnerAsync( + AzureStorageOrchestrationServiceSettings settings, + TimeSpan? timeout = null) + { + string taskHubName = settings.TaskHubName.ToLowerInvariant(); + Blob appLeaseInfoBlob = new AzureStorageClient(settings) + .GetBlobContainerReference(taskHubName + "-applease") + .GetBlobReference(taskHubName + "-appleaseinfo"); + byte[] appLeaseIdBytes = + BitConverter.GetBytes(Fnv1aHashHelper.ComputeHash(settings.AppName)); + Array.Resize(ref appLeaseIdBytes, 16); + string expectedOwnerId = new Guid(appLeaseIdBytes).ToString(); + + using (var cancellation = new CancellationTokenSource(timeout ?? TestTimeout)) + { + try + { + while (true) + { + string json = await appLeaseInfoBlob.DownloadTextAsync(); + var info = Utils.DeserializeFromJson(json); + if (info.OwnerId == expectedOwnerId) + { + return; + } + + await Task.Delay(TimeSpan.FromMilliseconds(100), cancellation.Token); + } + } + catch (OperationCanceledException) + { + Assert.Fail("The previous owner did not reacquire the app lease."); + } + } + } + + static async Task WaitForOrchestrationStatusAsync( + TaskHubClient client, + OrchestrationInstance instance, + OrchestrationStatus expectedStatus) + { + using (var cancellation = new CancellationTokenSource(TestTimeout)) + { + try + { + while (true) + { + OrchestrationState state = + await client.GetOrchestrationStateAsync(instance); + if (state?.OrchestrationStatus == expectedStatus) + { + return; + } + + await Task.Delay(TimeSpan.FromMilliseconds(100), cancellation.Token); + } + } + catch (OperationCanceledException) + { + Assert.Fail( + $"The orchestration did not reach {expectedStatus} before the test timeout."); + } + } + } + static async Task WaitForOwnershipAsync(AppLeaseManager manager) { using (await GetOwnershipAsync(manager)) @@ -535,5 +692,34 @@ public async IAsyncEnumerable GetOwnershipBlobLeasesAsync( yield break; } } + + sealed class WaitForSignalOrchestration : TaskOrchestration + { + TaskCompletionSource signal; + + public override async Task RunTask( + OrchestrationContext context, + string input) + { + this.signal = new TaskCompletionSource(); + return await this.signal.Task; + } + + public override void OnEvent( + OrchestrationContext context, + string name, + string input) + { + if (name == "complete") + { + this.signal?.TrySetResult(input); + } + } + } + + sealed class TestAppLeaseInfo + { + public string OwnerId { get; set; } + } } } diff --git a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs index 3fad8c98d..58fa2b815 100644 --- a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs @@ -40,14 +40,17 @@ sealed class TablePartitionManager : IPartitionManager, IDisposable readonly AzureStorageClient azureStorageClient; readonly AzureStorageOrchestrationService service; readonly AzureStorageOrchestrationServiceSettings settings; - readonly CancellationTokenSource gracefulShutdownTokenSource; - readonly CancellationTokenSource forcefulShutdownTokenSource; + readonly SemaphoreSlim lifecycleLock; readonly string storageAccountName; readonly Table partitionTable; readonly TableLeaseManager tableLeaseManager; readonly LeaseCollectionBalancerOptions options; + CancellationTokenSource? gracefulShutdownTokenSource; + CancellationTokenSource? forcefulShutdownTokenSource; Task partitionManagerTask; + bool isStarted; + bool isStopping; /// /// Constructor to initiate new instances of TablePartitionManager. @@ -68,8 +71,7 @@ public TablePartitionManager( LeaseInterval = this.settings.LeaseInterval, ShouldStealLeases = true }; - this.gracefulShutdownTokenSource = new CancellationTokenSource(); - this.forcefulShutdownTokenSource = new CancellationTokenSource(); + this.lifecycleLock = new SemaphoreSlim(1, 1); this.partitionTable = azureStorageClient.GetTableReference(this.settings.PartitionTableName); this.tableLeaseManager = new TableLeaseManager(this.partitionTable, this.service, this.settings, this.storageAccountName, this.options); this.partitionManagerTask = Task.CompletedTask; @@ -79,19 +81,52 @@ public TablePartitionManager( /// /// Starts the partition management loop for the current worker. /// - Task IPartitionManager.StartAsync() + async Task IPartitionManager.StartAsync() { - // Run the partition manager loop in the background - this.partitionManagerTask = this.PartitionManagerLoop( - this.gracefulShutdownTokenSource.Token, - this.forcefulShutdownTokenSource.Token); - this.settings.Logger.PartitionManagerInfo( - this.storageAccountName, - this.settings.TaskHubName, - this.settings.WorkerId, - partitionId: NotApplicable, - details: $"Started the background partition manager loop to acquire and balance partitions."); - return Task.CompletedTask; + await this.lifecycleLock.WaitAsync(); + try + { + if (this.isStarted) + { + throw new InvalidOperationException( + $"{nameof(TablePartitionManager)} has already started"); + } + + if (this.isStopping) + { + try + { + await this.partitionManagerTask; + } + finally + { + if (this.partitionManagerTask.IsCompleted) + { + this.DisposeRun(); + } + } + } + + var gracefulShutdown = new CancellationTokenSource(); + var forcefulShutdown = new CancellationTokenSource(); + this.gracefulShutdownTokenSource = gracefulShutdown; + this.forcefulShutdownTokenSource = forcefulShutdown; + this.partitionManagerTask = this.PartitionManagerLoop( + gracefulShutdown.Token, + forcefulShutdown.Token); + this.isStarted = true; + + this.settings.Logger.PartitionManagerInfo( + this.storageAccountName, + this.settings.TaskHubName, + this.settings.WorkerId, + partitionId: NotApplicable, + details: $"Started the background partition manager loop to acquire and balance partitions."); + } + finally + { + this.lifecycleLock.Release(); + } } @@ -229,36 +264,93 @@ async Task PartitionManagerLoop(CancellationToken gracefulShutdownToken, Cancell /// async Task IPartitionManager.StopAsync() { - this.gracefulShutdownTokenSource.Cancel(); - this.settings.Logger.PartitionManagerInfo( - this.storageAccountName, - this.settings.TaskHubName, - this.settings.WorkerId, - partitionId: NotApplicable, - "Started draining the in-memory messages of all owned control queues for shutdown."); + await this.lifecycleLock.WaitAsync(); + try + { + if (!this.isStarted) + { + if (this.isStopping) + { + try + { + await this.partitionManagerTask; + } + finally + { + if (this.partitionManagerTask.IsCompleted) + { + this.DisposeRun(); + } + } + } + + return; + } - // Wait 10 minutes for the partition manager to shutdown gracefully. Otherwise force a shutdown. - var timeout = TimeSpan.FromMinutes(10); - var timeoutTask = Task.Delay(Timeout.Infinite, this.forcefulShutdownTokenSource.Token); - this.forcefulShutdownTokenSource.CancelAfter(timeout); - await Task.WhenAny(this.partitionManagerTask, timeoutTask); + this.isStarted = false; + this.isStopping = true; + CancellationTokenSource gracefulShutdown = + this.gracefulShutdownTokenSource + ?? throw new InvalidOperationException("The graceful shutdown token source is missing."); + CancellationTokenSource forcefulShutdown = + this.forcefulShutdownTokenSource + ?? throw new InvalidOperationException("The forceful shutdown token source is missing."); - if (timeoutTask.IsCompleted) + gracefulShutdown.Cancel(); + this.settings.Logger.PartitionManagerInfo( + this.storageAccountName, + this.settings.TaskHubName, + this.settings.WorkerId, + partitionId: NotApplicable, + "Started draining the in-memory messages of all owned control queues for shutdown."); + + // Wait 10 minutes for the partition manager to shutdown gracefully. Otherwise force a shutdown. + var timeout = TimeSpan.FromMinutes(10); + var timeoutTask = Task.Delay(Timeout.Infinite, forcefulShutdown.Token); + forcefulShutdown.CancelAfter(timeout); + + try + { + Task completedTask = await Task.WhenAny(this.partitionManagerTask, timeoutTask); + if (completedTask == timeoutTask) + { + throw new TimeoutException( + $"Timed-out waiting for the partition manager to shut down. Timeout duration: {timeout}", + timeoutTask.Exception?.InnerException); + } + + // Surface any unhandled exceptions + await this.partitionManagerTask; + + this.settings.Logger.PartitionManagerInfo( + this.storageAccountName, + this.settings.TaskHubName, + this.settings.WorkerId, + partitionId: NotApplicable, + "Table partition manager stopped successfully."); + } + finally + { + if (this.partitionManagerTask.IsCompleted) + { + this.DisposeRun(); + } + } + } + finally { - throw new TimeoutException( - $"Timed-out waiting for the partition manager to shut down. Timeout duration: {timeout}", - timeoutTask.Exception?.InnerException); + this.lifecycleLock.Release(); } + } - // Surface any unhandled exceptions - await this.partitionManagerTask; - - this.settings.Logger.PartitionManagerInfo( - this.storageAccountName, - this.settings.TaskHubName, - this.settings.WorkerId, - partitionId: NotApplicable, - "Table partition manager stopped successfully."); + void DisposeRun() + { + this.gracefulShutdownTokenSource?.Dispose(); + this.forcefulShutdownTokenSource?.Dispose(); + this.gracefulShutdownTokenSource = null; + this.forcefulShutdownTokenSource = null; + this.partitionManagerTask = Task.CompletedTask; + this.isStopping = false; } async Task IPartitionManager.CreateLeaseStore() @@ -895,13 +987,14 @@ internal void SimulateUnhealthyWorker(CancellationToken testToken) // used for internal testing internal void KillLoop() { - this.forcefulShutdownTokenSource.Cancel(); + this.forcefulShutdownTokenSource?.Cancel(); } public void Dispose() { - this.gracefulShutdownTokenSource.Dispose(); - this.forcefulShutdownTokenSource.Dispose(); + this.gracefulShutdownTokenSource?.Dispose(); + this.forcefulShutdownTokenSource?.Dispose(); + this.lifecycleLock.Dispose(); } } } From d75cdaea5e99b64f7523252843856722026f7b51 Mon Sep 17 00:00:00 2001 From: wangbill Date: Tue, 8 Sep 2026 09:51:30 -0700 Subject: [PATCH 03/10] Stop table partition loop on forced cancellation Handle forceful cancellation before generic and graceful cancellation paths so timed-out shutdowns terminate the old loop and allow a safe restart. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../Partitioning/TablePartitionManager.cs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs index 58fa2b815..85090d413 100644 --- a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs @@ -180,6 +180,10 @@ async Task PartitionManagerLoop(CancellationToken gracefulShutdownToken, Cancell consecutiveFailureCount = 0; } + catch (OperationCanceledException) when (forcefulShutdownToken.IsCancellationRequested) + { + break; + } // Exception Status 412 represents an out of date ETag. We already logged this. catch (DurableTaskStorageException ex) when (ex.HttpStatusCode == (int)HttpStatusCode.PreconditionFailed) { @@ -232,6 +236,10 @@ async Task PartitionManagerLoop(CancellationToken gracefulShutdownToken, Cancell await Task.Delay(timeToSleep, gracefulShutdownToken); } } + catch (OperationCanceledException) when (forcefulShutdownToken.IsCancellationRequested) + { + break; + } catch (OperationCanceledException) when (gracefulShutdownToken.IsCancellationRequested) { // Shutdown requested, but we still need to release all leases From 9d20ca01fcc485ceb6983c04977c0a14485eea03 Mon Sep 17 00:00:00 2001 From: wangbill Date: Tue, 8 Sep 2026 09:51:43 -0700 Subject: [PATCH 04/10] Test forced table partition cancellation restart Exercise graceful drain, forced loop cancellation, repeated stop, and real orchestration progress after restarting the same service. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../TestTablePartitionManager.cs | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs b/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs index 7fdb42e04..c98452e30 100644 --- a/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs +++ b/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs @@ -726,6 +726,95 @@ await WaitForConditionAsync( Assert.AreEqual(1, service.OwnedControlQueues.Count()); } + [TestMethod] + public async Task ForcedCancellationCompletesStopAndAllowsWorkAfterRestart() + { + string taskHubName = + "forcecancel" + Guid.NewGuid().ToString("N").Substring(0, 12); + var settings = new AzureStorageOrchestrationServiceSettings + { + StorageAccountClientProvider = new StorageAccountClientProvider(this.connection), + TaskHubName = taskHubName, + PartitionCount = 1, + LeaseAcquireInterval = TimeSpan.FromMilliseconds(200), + MaxQueuePollingInterval = TimeSpan.FromMilliseconds(50), + WorkerId = "0", + UseAppLease = false, + UseTablePartitionManagement = true, + }; + var service = new AzureStorageOrchestrationService(settings); + var client = new TaskHubClient(service); + + await service.StartAsync(); + await WaitForConditionAsync( + TimeSpan.FromSeconds(5), + t => new ValueTask(service.OwnedControlQueues.Any())); + + OrchestrationInstance instance = + await client.CreateOrchestrationInstanceAsync( + typeof(HelloOrchestrator), + input: null); + using var lockCancellation = new CancellationTokenSource(TimeSpan.FromSeconds(10)); + TaskOrchestrationWorkItem workItem = + await service.LockNextTaskOrchestrationWorkItemAsync( + TimeSpan.FromSeconds(10), + lockCancellation.Token); + Assert.IsNotNull(workItem); + + Task stopTask = service.StopAsync(isForced: false); + await WaitForConditionAsync( + TimeSpan.FromSeconds(5), + t => new ValueTask( + service.OwnedControlQueues.Single().IsReleased)); + Assert.IsFalse( + stopTask.IsCompleted, + "Graceful stop should wait for the locked orchestration work item to drain."); + + service.KillPartitionManagerLoop(); + + Task completedTask = + await Task.WhenAny(stopTask, Task.Delay(TimeSpan.FromSeconds(5))); + Assert.AreSame( + stopTask, + completedTask, + "Forceful cancellation did not complete the timed-out stop."); + await Assert.ThrowsExceptionAsync(() => stopTask); + + await service.AbandonTaskOrchestrationWorkItemAsync(workItem); + await service.ReleaseTaskOrchestrationWorkItemAsync(workItem); + await WaitForConditionAsync( + TimeSpan.FromSeconds(5), + t => new ValueTask(!service.OwnedControlQueues.Any())); + + Task secondStopTask = service.StopAsync(isForced: false); + completedTask = + await Task.WhenAny( + secondStopTask, + Task.Delay(TimeSpan.FromSeconds(5))); + Assert.AreSame( + secondStopTask, + completedTask, + "A repeated stop did not observe termination of the old partition manager loop."); + await secondStopTask; + + var worker = new TaskHubWorker(service); + worker.AddTaskOrchestrations(typeof(HelloOrchestrator)); + worker.AddTaskActivities(typeof(Hello)); + await worker.StartAsync(); + + OrchestrationState state = + await client.WaitForOrchestrationAsync( + instance, + TimeSpan.FromSeconds(30)); + + Assert.IsNotNull(state); + Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); + Assert.AreEqual("\"Hello, world!\"", state.Output); + + await worker.StopAsync(isForced: true); + await service.DeleteAsync(); + } + [KnownType(typeof(Hello))] internal class HelloOrchestrator : TaskOrchestration { From 03b57f7c6f0f53963c61437683374e60273173b2 Mon Sep 17 00:00:00 2001 From: wangbill Date: Tue, 8 Sep 2026 10:11:40 -0700 Subject: [PATCH 05/10] Preserve forced partition stop semantics Restore timeout reporting when the force signal and partition loop complete together, and ensure app-lease tests are tracked under the case-correct test project path. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../Partitioning/TablePartitionManager.cs | 4 +-- .../AppLeaseActivityTests.cs | 0 .../TestTablePartitionManager.cs | 32 +++++++++++++++++++ 3 files changed, 34 insertions(+), 2 deletions(-) rename {Test => test}/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs (100%) diff --git a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs index 85090d413..2c61b8943 100644 --- a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs @@ -319,8 +319,8 @@ async Task IPartitionManager.StopAsync() try { - Task completedTask = await Task.WhenAny(this.partitionManagerTask, timeoutTask); - if (completedTask == timeoutTask) + await Task.WhenAny(this.partitionManagerTask, timeoutTask); + if (timeoutTask.IsCompleted) { throw new TimeoutException( $"Timed-out waiting for the partition manager to shut down. Timeout duration: {timeout}", diff --git a/Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs similarity index 100% rename from Test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs rename to test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs diff --git a/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs b/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs index c98452e30..f688bec58 100644 --- a/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs +++ b/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs @@ -815,6 +815,38 @@ await client.WaitForOrchestrationAsync( await service.DeleteAsync(); } + [TestMethod] + public async Task ForceSignalStillReportsTimeoutAfterLoopStops() + { + string taskHubName = + "forcetimeout" + Guid.NewGuid().ToString("N").Substring(0, 12); + var settings = new AzureStorageOrchestrationServiceSettings + { + StorageAccountClientProvider = new StorageAccountClientProvider(this.connection), + TaskHubName = taskHubName, + PartitionCount = 1, + LeaseAcquireInterval = TimeSpan.FromMilliseconds(200), + WorkerId = "0", + UseAppLease = false, + UseTablePartitionManagement = true, + }; + var service = new AzureStorageOrchestrationService(settings); + + await service.StartAsync(); + await WaitForConditionAsync( + TimeSpan.FromSeconds(5), + t => new ValueTask(service.OwnedControlQueues.Any())); + + service.KillPartitionManagerLoop(); + await Task.Delay(TimeSpan.FromSeconds(1)); + + await Assert.ThrowsExceptionAsync( + () => service.StopAsync(isForced: false)); + + await service.StopAsync(isForced: false); + await service.DeleteAsync(); + } + [KnownType(typeof(Hello))] internal class HelloOrchestrator : TaskOrchestration { From d84f4a83fe5d7789aac74b78e2013582224526e2 Mon Sep 17 00:00:00 2001 From: wangbill Date: Tue, 8 Sep 2026 12:46:35 -0700 Subject: [PATCH 06/10] Prevent stale table partition readers Only resume control-queue polling for partitions currently owned by the worker and not draining. Recover stale drain state on same-worker renewal and cover four-partition A-B-A dispatch behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../Partitioning/TablePartitionManager.cs | 13 +- .../AppLeaseActivityTests.cs | 186 ++++++++++++++++++ 2 files changed, 189 insertions(+), 10 deletions(-) diff --git a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs index 2c61b8943..6fee697ef 100644 --- a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs @@ -524,16 +524,8 @@ public async Task ReadAndWriteTableAsync(bool isShuttingDown, throw; } - // Ensure worker is listening to the control queue iff either: - // 1) worker just claimed the lease, - // 2) worker was already the owner in the partitions table and is not actively draining the queue. - // Note that during draining, we renew the lease but do not want to listen to new messages. - // Otherwise, we'll never finish draining our in-memory messages. - // When draining completes, and the worker may decide to release the lease. In that moment, - // IsDrainingPartition can still be true but renewedLease can be false — without checking - // !releasedLease, the worker could incorrectly resume listening just before releasing the lease. - bool isRenewingToDrainQueue = renewedLease && response.IsDrainingPartition && !releasedLease; - if (claimedLease || !isRenewingToDrainQueue) + // Only an owner of this non-draining partition should listen to its control queue. + if (partition.CurrentOwner == this.workerName && !partition.IsDraining) { // Notify the orchestration session manager that we acquired a lease for one of the partitions. // This will cause it to start reading control queue messages for that partition. @@ -598,6 +590,7 @@ void RenewOrReleaseMyLease( if (partition.NextOwner == null) { // We still own the lease and nobody is trying to steal it. + partition.IsDraining = false; ownershipLeaseCount++; this.RenewLease(partition); renewedLease = true; diff --git a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs index 14334487f..924c42b9a 100644 --- a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs +++ b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -316,6 +316,113 @@ await WaitForOwnerAsync( } } + [TestMethod] + public async Task ReleasedPartitionsDoNotResumeListeningBeforeFailback() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationServiceSettings firstSettings = + CreateSettings(taskHubName, "FirstApp", useAppLease: true); + AzureStorageOrchestrationServiceSettings secondSettings = + CreateSettings(taskHubName, "SecondApp", useAppLease: true); + firstSettings.PartitionCount = 4; + firstSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); + secondSettings.PartitionCount = 4; + secondSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); + secondSettings.ControlQueueBufferThreshold = 0; + + var firstService = new AzureStorageOrchestrationService(firstSettings); + var secondService = new AzureStorageOrchestrationService(secondSettings); + var firstWorker = new TaskHubWorker(firstService); + var secondWorker = new TaskHubWorker(secondService); + firstWorker.AddTaskOrchestrations( + new WorkerIdentityOrchestrationCreator("FirstApp")); + secondWorker.AddTaskOrchestrations( + new WorkerIdentityOrchestrationCreator("SecondApp")); + + bool firstWorkerStarted = false; + bool secondWorkerStarted = false; + + try + { + await firstService.CreateAsync(); + await firstWorker.StartAsync(); + firstWorkerStarted = true; + await WaitForAppLeaseOwnerAsync(firstSettings); + await WaitForTableOwnershipAsync(firstService, firstSettings); + await WaitForReceivingQueueCountAsync(firstService, firstSettings.PartitionCount); + + await secondWorker.StartAsync(); + secondWorkerStarted = true; + await secondService.ForceChangeAppLeaseAsync(); + await WaitForAppLeaseOwnerAsync(secondSettings); + await WaitForTableOwnershipAsync(secondService, secondSettings); + await WaitForReceivingQueueCountAsync(secondService, secondSettings.PartitionCount); + + string[] oldReceivingQueues = firstService.OwnedControlQueues + .Where(queue => !queue.IsReleased) + .Select(queue => queue.Name) + .ToArray(); + if (oldReceivingQueues.Length == 0) + { + secondSettings.ControlQueueBufferThreshold = 1000; + } + + var client = new TaskHubClient(firstService); + OrchestrationInstance secondOwnerInstance = + await client.CreateOrchestrationInstanceAsync( + WorkerIdentityOrchestrationCreator.OrchestrationName, + version: string.Empty, + input: null); + OrchestrationState secondOwnerState = + await client.WaitForOrchestrationAsync(secondOwnerInstance, TestTimeout); + + Assert.AreEqual( + 0, + oldReceivingQueues.Length, + $"The previous app resumed receiving from released partitions: {string.Join(", ", oldReceivingQueues)}. " + + $"New orchestration output: {secondOwnerState?.Output ?? ""}."); + Assert.IsNotNull(secondOwnerState); + Assert.AreEqual(OrchestrationStatus.Completed, secondOwnerState.OrchestrationStatus); + Assert.AreEqual("\"SecondApp\"", secondOwnerState.Output); + + await secondWorker.StopAsync(isForced: true); + secondWorkerStarted = false; + await WaitForAppLeaseOwnerAsync(firstSettings, AutomaticFailoverTimeout); + await WaitForTableOwnershipAsync( + firstService, + firstSettings, + AutomaticFailoverTimeout); + await WaitForReceivingQueueCountAsync( + firstService, + firstSettings.PartitionCount, + AutomaticFailoverTimeout); + + OrchestrationInstance firstOwnerInstance = + await client.CreateOrchestrationInstanceAsync( + WorkerIdentityOrchestrationCreator.OrchestrationName, + version: string.Empty, + input: null); + OrchestrationState firstOwnerState = + await client.WaitForOrchestrationAsync(firstOwnerInstance, TestTimeout); + + Assert.IsNotNull(firstOwnerState); + Assert.AreEqual(OrchestrationStatus.Completed, firstOwnerState.OrchestrationStatus); + Assert.AreEqual("\"FirstApp\"", firstOwnerState.Output); + } + finally + { + if (secondWorkerStarted) + { + await secondWorker.StopAsync(isForced: true); + } + + if (firstWorkerStarted) + { + await firstWorker.StopAsync(isForced: true); + } + } + } + [TestMethod] public async Task AppLeaseManagerCanRestartAfterPartitionManagerStopFails() { @@ -546,6 +653,49 @@ static async Task WaitForAppLeaseOwnerAsync( } } + static async Task WaitForTableOwnershipAsync( + AzureStorageOrchestrationService service, + AzureStorageOrchestrationServiceSettings settings, + TimeSpan? timeout = null) + { + using (var cancellation = new CancellationTokenSource(timeout ?? TestTimeout)) + { + try + { + while (true) + { + List leases = + await service.ListTableLeasesAsync().ToListAsync(); + if (leases.Count == settings.PartitionCount && + leases.All( + lease => + lease.CurrentOwner == settings.WorkerId && + !lease.IsDraining)) + { + return; + } + + await Task.Delay(TimeSpan.FromMilliseconds(100), cancellation.Token); + } + } + catch (OperationCanceledException) + { + Assert.Fail( + $"Worker {settings.WorkerId} did not own all non-draining table partitions."); + } + } + } + + static Task WaitForReceivingQueueCountAsync( + AzureStorageOrchestrationService service, + int expectedCount, + TimeSpan? timeout = null) + { + return TestHelpers.WaitFor( + () => service.OwnedControlQueues.Count(queue => !queue.IsReleased) == expectedCount, + timeout ?? TestTimeout); + } + static async Task WaitForOrchestrationStatusAsync( TaskHubClient client, OrchestrationInstance instance, @@ -672,6 +822,42 @@ public async IAsyncEnumerable GetOwnershipBlobLeasesAsync( } } + sealed class WorkerIdentityOrchestrationCreator : ObjectCreator + { + public const string OrchestrationName = "AppLeaseWorkerIdentity"; + + readonly string workerIdentity; + + public WorkerIdentityOrchestrationCreator(string workerIdentity) + { + this.workerIdentity = workerIdentity; + this.Name = OrchestrationName; + this.Version = string.Empty; + } + + public override TaskOrchestration Create() + { + return new WorkerIdentityOrchestration(this.workerIdentity); + } + } + + sealed class WorkerIdentityOrchestration : TaskOrchestration + { + readonly string workerIdentity; + + public WorkerIdentityOrchestration(string workerIdentity) + { + this.workerIdentity = workerIdentity; + } + + public override Task RunTask( + OrchestrationContext context, + string input) + { + return Task.FromResult(this.workerIdentity); + } + } + sealed class TestPartitionManager : IPartitionManager { public Task StartAsync() => Task.CompletedTask; From c913f6bf7788bcb87e0628341eb2fd505a51e04e Mon Sep 17 00:00:00 2001 From: wangbill Date: Tue, 8 Sep 2026 16:49:50 -0700 Subject: [PATCH 07/10] Address app lease review feedback Move the activity ownership fence ahead of tracing, handle forced acquisition before an owner is recorded, and tighten lifecycle/test invariants. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../AzureStorageOrchestrationService.cs | 21 ++- .../Partitioning/AppLeaseManager.cs | 9 +- .../AppLeaseActivityTests.cs | 158 +++++++++++++++--- 3 files changed, 159 insertions(+), 29 deletions(-) diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs index 99f07aadb..3981828b5 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs @@ -198,6 +198,9 @@ public AzureStorageOrchestrationService(AzureStorageOrchestrationServiceSettings internal ITrackingStore TrackingStore => this.trackingStore; + // Intended only for use by tests that need to coordinate the post-dequeue ownership race. + internal Func OnActivityMessageDequeued { get; set; } + internal static string GetControlQueueName(string taskHub, int partitionIndex) { return GetQueueName(taskHub, $"control-{partitionIndex:00}"); @@ -1568,6 +1571,18 @@ public async Task LockNextTaskActivityWorkItem( return null; } + Func onActivityMessageDequeued = this.OnActivityMessageDequeued; + if (onActivityMessageDequeued != null) + { + await onActivityMessageDequeued(); + } + + if (!ownership.TryBeginDispatch()) + { + await this.workItemQueue.AbandonMessageAsync(message); + return null; + } + Guid traceActivityId = Guid.NewGuid(); var session = new ActivitySession(this.settings, this.azureStorageClient.QueueAccountName, message, traceActivityId); session.StartNewLogicalTraceScope(); @@ -1586,12 +1601,6 @@ public async Task LockNextTaskActivityWorkItem( TraceMessageReceived(this.settings, session.MessageData, this.azureStorageClient.QueueAccountName); session.TraceProcessingMessage(message, isExtendedSession: false, this.workItemQueue.Name); - if (!ownership.TryBeginDispatch()) - { - await this.workItemQueue.AbandonMessageAsync(message); - return null; - } - if (!this.activeActivitySessions.TryAdd(message.Id, session)) { // This means we're already processing this message. This is never expected since the message diff --git a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs index d3c1afab0..5327df1bc 100644 --- a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs @@ -307,6 +307,12 @@ await this.shutdownCompletedEvent.WaitAsync( if (!calledFromLeaseRenewer && renewer != null) { + if (renewerCancellation == null) + { + throw new InvalidOperationException( + "The app lease renewer has no cancellation source."); + } + renewerCancellation.Cancel(); await renewer; } @@ -331,7 +337,8 @@ async Task TryAcquireAppLeaseAsync(CancellationToken cancellationToken) AppLeaseInfo appLeaseInfo = await this.GetAppLeaseInfoAsync(); bool leaseAcquired; - if (appLeaseInfo.DesiredSwapId == this.appLeaseId) + if (appLeaseInfo.DesiredSwapId == this.appLeaseId + && !string.IsNullOrEmpty(appLeaseInfo.OwnerId)) { leaseAcquired = await this.ChangeLeaseAsync(appLeaseInfo.OwnerId, cancellationToken); } diff --git a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs index 924c42b9a..3b8619cfb 100644 --- a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs +++ b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -14,6 +14,7 @@ namespace DurableTask.AzureStorage.Tests { using System; + using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; using System.Runtime.CompilerServices; @@ -23,6 +24,7 @@ namespace DurableTask.AzureStorage.Tests using DurableTask.AzureStorage.Storage; using DurableTask.Core; using DurableTask.Core.History; + using Microsoft.Extensions.Logging; using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] @@ -187,26 +189,26 @@ public async Task ForcedFailoverQuiescesOldAppBeforeNewOwnerStarts() AppLeaseManager newManager = CreateAppLeaseManager( newSettings, new TestPartitionManager()); - AppLeaseOwnershipSignal.AppLeaseOwnership oldOwnership = null; - AppLeaseOwnershipSignal.AppLeaseOwnership newOwnership = null; - try { await oldManager.CreateContainerIfNotExistsAsync(); await oldManager.StartAsync(); - oldOwnership = await GetOwnershipAsync(oldManager); - await newManager.StartAsync(); - - await newManager.ForceChangeAppLeaseAsync(); - newOwnership = await GetOwnershipAsync(newManager); + using (AppLeaseOwnershipSignal.AppLeaseOwnership oldOwnership = + await GetOwnershipAsync(oldManager)) + { + await newManager.StartAsync(); - Assert.IsTrue(oldOwnership.LostToken.IsCancellationRequested); - Assert.IsFalse(oldOwnership.TryBeginDispatch()); + await newManager.ForceChangeAppLeaseAsync(); + using (AppLeaseOwnershipSignal.AppLeaseOwnership newOwnership = + await GetOwnershipAsync(newManager)) + { + Assert.IsTrue(oldOwnership.LostToken.IsCancellationRequested); + Assert.IsFalse(oldOwnership.TryBeginDispatch()); + } + } } finally { - newOwnership?.Dispose(); - oldOwnership?.Dispose(); await newManager.StopAsync(); await oldManager.StopAsync(); } @@ -381,7 +383,11 @@ await client.CreateOrchestrationInstanceAsync( oldReceivingQueues.Length, $"The previous app resumed receiving from released partitions: {string.Join(", ", oldReceivingQueues)}. " + $"New orchestration output: {secondOwnerState?.Output ?? ""}."); - Assert.IsNotNull(secondOwnerState); + if (secondOwnerState == null) + { + Assert.Fail("The new owner did not complete the orchestration."); + } + Assert.AreEqual(OrchestrationStatus.Completed, secondOwnerState.OrchestrationStatus); Assert.AreEqual("\"SecondApp\"", secondOwnerState.Output); @@ -446,6 +452,33 @@ public async Task AppLeaseManagerCanRestartAfterPartitionManagerStopFails() Assert.AreEqual(2, partitionManager.StopCount); } + [TestMethod] + public async Task ForceChangeBeforeFirstOwnerAcquiresLease() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationServiceSettings settings = + CreateSettings(taskHubName, "PrimaryApp", useAppLease: true); + var partitionManager = new TestPartitionManager(); + AppLeaseManager manager = CreateAppLeaseManager(settings, partitionManager); + + try + { + await manager.CreateContainerIfNotExistsAsync(); + + await manager.ForceChangeAppLeaseAsync(); + + using (var cancellation = new CancellationTokenSource(TimeSpan.FromSeconds(3))) + using (await manager.WaitForOwnershipAsync(cancellation.Token)) + { + Assert.AreEqual(1, partitionManager.StartCount); + } + } + finally + { + await manager.StopAsync(); + } + } + [TestMethod] public async Task LeaseLossDoesNotCancelDispatchedActivity() { @@ -481,18 +514,46 @@ public async Task LeaseLossDoesNotCancelDispatchedActivity() } [TestMethod] - public async Task OwnershipLossAfterDequeuePreventsDispatch() + public async Task OwnershipLossAfterDequeueDoesNotStartProcessingTrace() { - var ownershipSignal = new AppLeaseOwnershipSignal(); - ownershipSignal.Set(); + string taskHubName = GetTaskHubName(); + var loggerFactory = new RecordingLoggerFactory(); + AzureStorageOrchestrationServiceSettings oldSettings = + CreateSettings(taskHubName, "OldApp", useAppLease: true); + oldSettings.LoggerFactory = loggerFactory; + AzureStorageOrchestrationService oldOwner = + new AzureStorageOrchestrationService(oldSettings); + AzureStorageOrchestrationService newOwner = + CreateService(taskHubName, "NewApp", useAppLease: true); - using (AppLeaseOwnershipSignal.AppLeaseOwnership ownership = - await ownershipSignal.WaitAsync(CancellationToken.None)) + try { - ownershipSignal.Reset(); + await oldOwner.CreateAsync(); + await oldOwner.StartAsync(); + await WaitForOwnerAsync(oldOwner); + await newOwner.StartAsync(); + await EnqueueActivityAsync(oldOwner, "ownership-race"); + + oldOwner.OnActivityMessageDequeued = async () => + { + oldOwner.OnActivityMessageDequeued = null; + await newOwner.ForceChangeAppLeaseAsync(); + await WaitForOwnerAsync(newOwner); + }; + + TaskActivityWorkItem rejectedWorkItem = await LockActivityAsync(oldOwner); + + Assert.IsNull(rejectedWorkItem); + Assert.IsFalse(loggerFactory.HasEvent("ProcessingMessage")); - Assert.IsTrue(ownership.LostToken.IsCancellationRequested); - Assert.IsFalse(ownership.TryBeginDispatch()); + TaskActivityWorkItem newOwnerWorkItem = await LockActivityAsync(newOwner); + Assert.IsNotNull(newOwnerWorkItem); + await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); + } + finally + { + await StopAsync(newOwner); + await StopAsync(oldOwner); } } @@ -860,7 +921,13 @@ public override Task RunTask( sealed class TestPartitionManager : IPartitionManager { - public Task StartAsync() => Task.CompletedTask; + public int StartCount { get; private set; } + + public Task StartAsync() + { + this.StartCount++; + return Task.CompletedTask; + } public Task StopAsync() => Task.CompletedTask; @@ -879,6 +946,53 @@ public async IAsyncEnumerable GetOwnershipBlobLeasesAsync( } } + sealed class RecordingLoggerFactory : ILoggerFactory + { + readonly ConcurrentQueue events = new ConcurrentQueue(); + + public void AddProvider(ILoggerProvider provider) + { + } + + public ILogger CreateLogger(string categoryName) + { + return new RecordingLogger(this.events); + } + + public void Dispose() + { + } + + public bool HasEvent(string name) + { + return this.events.Any(e => e.Name == name); + } + + sealed class RecordingLogger : ILogger + { + readonly ConcurrentQueue events; + + public RecordingLogger(ConcurrentQueue events) + { + this.events = events; + } + + public IDisposable BeginScope(TState state) => null; + + public bool IsEnabled(LogLevel logLevel) => true; + + public void Log( + LogLevel logLevel, + EventId eventId, + TState state, + Exception exception, + Func formatter) + { + this.events.Enqueue(eventId); + } + } + } + sealed class WaitForSignalOrchestration : TaskOrchestration { TaskCompletionSource signal; From 9b04d91a67efd261045e7ce824623efdfaefad64 Mon Sep 17 00:00:00 2001 From: wangbill Date: Wed, 9 Sep 2026 08:47:13 -0700 Subject: [PATCH 08/10] Narrow app lease changes to activity gating Restore table partition and app lease lifecycle behavior to main while retaining ownership-aware activity receive admission and focused regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- ...zureStorageOrchestrationServiceSettings.cs | 8 +- .../Partitioning/AppLeaseManager.cs | 132 +-- .../Partitioning/TablePartitionManager.cs | 202 ++--- .../AppLeaseActivityTests.cs | 773 +++--------------- .../TestTablePartitionManager.cs | 121 --- 5 files changed, 194 insertions(+), 1042 deletions(-) diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs index 58df21089..f04e039c6 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs @@ -152,9 +152,11 @@ public class AzureStorageOrchestrationServiceSettings public TimeSpan MaxQueuePollingInterval { get; set; } = DefaultMaxQueuePollingInterval; /// - /// If true, takes a lease on the task hub container so that only workers with the lease-owning - /// process orchestration, entity, or activity messages. - /// Workers that share the same may process messages concurrently. + /// If true, workers wait for their to own the app lease before receiving + /// new activity messages. Workers sharing that app name may receive activities concurrently. + /// Observed lease loss stops pending activity receives and rejects dequeued activities before + /// dispatch; already dispatched activities are not canceled. Orchestration and entity message + /// processing retain their existing behavior. /// public bool UseAppLease { get; set; } = true; diff --git a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs index 5327df1bc..4fcda8917 100644 --- a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs @@ -147,7 +147,7 @@ async Task AppLeaseManagerStarter(CancellationToken cancellationToken) { try { - while (!await this.TryAcquireAppLeaseAsync(cancellationToken)) + while (!await this.TryAcquireAppLeaseAsync()) { await Task.Delay(this.settings.AppLeaseOptions.AcquireInterval, cancellationToken); } @@ -261,111 +261,61 @@ async Task StartAppLeaseAsync() throw new InvalidOperationException("AppLeaseManager has already started"); } - this.shutdownCompletedEvent.Reset(); this.leaseRenewerCancellationTokenSource = new CancellationTokenSource(); - try - { - await this.partitionManager.StartAsync(); - this.ownershipSignal.Set(); - this.renewTask = await Task.Factory.StartNew( - () => this.LeaseRenewer(this.leaseRenewerCancellationTokenSource.Token)); - } - catch - { - this.ownershipSignal.Reset(); - this.isLeaseOwner = false; - Interlocked.Exchange(ref this.appLeaseIsStarted, 0); - this.leaseRenewerCancellationTokenSource.Dispose(); - this.leaseRenewerCancellationTokenSource = null; - this.shutdownCompletedEvent.Set(); - throw; - } + await this.partitionManager.StartAsync(); + this.ownershipSignal.Set(); + + this.shutdownCompletedEvent.Reset(); + + this.renewTask = await Task.Factory.StartNew(() => this.LeaseRenewer(leaseRenewerCancellationTokenSource.Token)); } - async Task StopAppLeaseAsync(bool calledFromLeaseRenewer = false) + async Task StopAppLeaseAsync() { - int previousState = Interlocked.CompareExchange(ref this.appLeaseIsStarted, 2, 1); - if (previousState != 1) + if (Interlocked.CompareExchange(ref this.appLeaseIsStarted, 0, 1) != 1) { - if (previousState == 2 && !calledFromLeaseRenewer) - { - await this.shutdownCompletedEvent.WaitAsync( - Timeout.InfiniteTimeSpan, - CancellationToken.None); - } - + //idempotent return; } - CancellationTokenSource renewerCancellation = this.leaseRenewerCancellationTokenSource; - Task renewer = this.renewTask; - - try - { - this.ownershipSignal.Reset(); - - if (!calledFromLeaseRenewer && renewer != null) - { - if (renewerCancellation == null) - { - throw new InvalidOperationException( - "The app lease renewer has no cancellation source."); - } + this.ownershipSignal.Reset(); - renewerCancellation.Cancel(); - await renewer; - } + await this.partitionManager.StopAsync(); - await this.partitionManager.StopAsync(); - } - finally + if (this.renewTask != null) { - this.isLeaseOwner = false; + this.leaseRenewerCancellationTokenSource.Cancel(); + await this.renewTask; + } - renewerCancellation?.Dispose(); - this.leaseRenewerCancellationTokenSource = null; - this.renewTask = null; + this.isLeaseOwner = false; - Interlocked.Exchange(ref this.appLeaseIsStarted, 0); - this.shutdownCompletedEvent.Set(); - } + this.shutdownCompletedEvent.Set(); + + this.leaseRenewerCancellationTokenSource?.Dispose(); } - async Task TryAcquireAppLeaseAsync(CancellationToken cancellationToken) + async Task TryAcquireAppLeaseAsync() { AppLeaseInfo appLeaseInfo = await this.GetAppLeaseInfoAsync(); bool leaseAcquired; - if (appLeaseInfo.DesiredSwapId == this.appLeaseId - && !string.IsNullOrEmpty(appLeaseInfo.OwnerId)) + if (appLeaseInfo.DesiredSwapId == this.appLeaseId) { - leaseAcquired = await this.ChangeLeaseAsync(appLeaseInfo.OwnerId, cancellationToken); + leaseAcquired = await this.ChangeLeaseAsync(appLeaseInfo.OwnerId); } else { leaseAcquired = await this.TryAcquireLeaseAsync(); } - if (leaseAcquired) - { - AppLeaseInfo currentAppLeaseInfo = await this.GetAppLeaseInfoAsync(); - TimeSpan transitionDelay = - currentAppLeaseInfo.TransitionUntilUtc.GetValueOrDefault() - DateTime.UtcNow; - if (transitionDelay > TimeSpan.Zero) - { - await Task.Delay(transitionDelay, cancellationToken); - } - } - this.isLeaseOwner = leaseAcquired; return leaseAcquired; } - async Task ChangeLeaseAsync( - string currentLeaseId, - CancellationToken cancellationToken) + async Task ChangeLeaseAsync(string currentLeaseId) { this.settings.Logger.PartitionManagerInfo( this.storageAccountName, @@ -387,18 +337,12 @@ async Task ChangeLeaseAsync( await this.appLeaseContainer.ChangeLeaseAsync(this.appLeaseId, currentLeaseId); - DateTime transitionUntilUtc = DateTime.UtcNow.Add(this.options.RenewInterval); var appLeaseInfo = new AppLeaseInfo() { OwnerId = this.appLeaseId, - TransitionUntilUtc = transitionUntilUtc, }; await this.UpdateAppLeaseInfoBlob(appLeaseInfo); - - // Give the previous app's renewers time to observe the changed lease and stop - // dequeueing work before this app begins processing. - await Task.Delay(this.options.RenewInterval, cancellationToken); leaseAcquired = true; this.settings.Logger.LeaseAcquisitionSucceeded( @@ -407,6 +351,14 @@ async Task ChangeLeaseAsync( this.workerName, this.appLeaseContainerName, LeaseType); + + // When changing the lease over to another app, the paritions will still be listened to on the first app until the AppLeaseManager + // renew task fails to renew the lease. To avoid potential split brain we must delay before the new lease holder can start + // listening to the partitions. + if (this.settings.UseLegacyPartitionManagement == true) + { + await Task.Delay(this.settings.AppLeaseOptions.RenewInterval); + } } catch (DurableTaskStorageException e) { @@ -490,7 +442,7 @@ async Task LeaseRenewer(CancellationToken cancellationToken) break; } - await Task.Delay(this.options.RenewInterval, cancellationToken); + await Task.Delay(this.options.RenewInterval, this.leaseRenewerCancellationTokenSource.Token); } catch (OperationCanceledException) { @@ -521,22 +473,7 @@ async Task LeaseRenewer(CancellationToken cancellationToken) this.appLeaseContainerName, "Lease renewer task completing. Stopping AppLeaseManager."); - if (!cancellationToken.IsCancellationRequested) - { - try - { - await this.StopAppLeaseAsync(calledFromLeaseRenewer: true); - } - catch (Exception ex) - { - this.settings.Logger.PartitionManagerError( - this.storageAccountName, - this.taskHub, - this.workerName, - this.appLeaseContainerName, - $"Failed to stop AppLeaseManager after losing the app lease. AppLeaseId: {this.appLeaseId} Exception: {ex}"); - } - } + await this.StopAppLeaseAsync(); } async Task RenewLeaseAsync() @@ -655,7 +592,6 @@ private class AppLeaseInfo { public string OwnerId { get; set; } public string DesiredSwapId { get; set; } - public DateTime? TransitionUntilUtc { get; set; } } } } diff --git a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs index 6fee697ef..3fad8c98d 100644 --- a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs @@ -40,17 +40,14 @@ sealed class TablePartitionManager : IPartitionManager, IDisposable readonly AzureStorageClient azureStorageClient; readonly AzureStorageOrchestrationService service; readonly AzureStorageOrchestrationServiceSettings settings; - readonly SemaphoreSlim lifecycleLock; + readonly CancellationTokenSource gracefulShutdownTokenSource; + readonly CancellationTokenSource forcefulShutdownTokenSource; readonly string storageAccountName; readonly Table partitionTable; readonly TableLeaseManager tableLeaseManager; readonly LeaseCollectionBalancerOptions options; - CancellationTokenSource? gracefulShutdownTokenSource; - CancellationTokenSource? forcefulShutdownTokenSource; Task partitionManagerTask; - bool isStarted; - bool isStopping; /// /// Constructor to initiate new instances of TablePartitionManager. @@ -71,7 +68,8 @@ public TablePartitionManager( LeaseInterval = this.settings.LeaseInterval, ShouldStealLeases = true }; - this.lifecycleLock = new SemaphoreSlim(1, 1); + this.gracefulShutdownTokenSource = new CancellationTokenSource(); + this.forcefulShutdownTokenSource = new CancellationTokenSource(); this.partitionTable = azureStorageClient.GetTableReference(this.settings.PartitionTableName); this.tableLeaseManager = new TableLeaseManager(this.partitionTable, this.service, this.settings, this.storageAccountName, this.options); this.partitionManagerTask = Task.CompletedTask; @@ -81,52 +79,19 @@ public TablePartitionManager( /// /// Starts the partition management loop for the current worker. /// - async Task IPartitionManager.StartAsync() + Task IPartitionManager.StartAsync() { - await this.lifecycleLock.WaitAsync(); - try - { - if (this.isStarted) - { - throw new InvalidOperationException( - $"{nameof(TablePartitionManager)} has already started"); - } - - if (this.isStopping) - { - try - { - await this.partitionManagerTask; - } - finally - { - if (this.partitionManagerTask.IsCompleted) - { - this.DisposeRun(); - } - } - } - - var gracefulShutdown = new CancellationTokenSource(); - var forcefulShutdown = new CancellationTokenSource(); - this.gracefulShutdownTokenSource = gracefulShutdown; - this.forcefulShutdownTokenSource = forcefulShutdown; - this.partitionManagerTask = this.PartitionManagerLoop( - gracefulShutdown.Token, - forcefulShutdown.Token); - this.isStarted = true; - - this.settings.Logger.PartitionManagerInfo( - this.storageAccountName, - this.settings.TaskHubName, - this.settings.WorkerId, - partitionId: NotApplicable, - details: $"Started the background partition manager loop to acquire and balance partitions."); - } - finally - { - this.lifecycleLock.Release(); - } + // Run the partition manager loop in the background + this.partitionManagerTask = this.PartitionManagerLoop( + this.gracefulShutdownTokenSource.Token, + this.forcefulShutdownTokenSource.Token); + this.settings.Logger.PartitionManagerInfo( + this.storageAccountName, + this.settings.TaskHubName, + this.settings.WorkerId, + partitionId: NotApplicable, + details: $"Started the background partition manager loop to acquire and balance partitions."); + return Task.CompletedTask; } @@ -180,10 +145,6 @@ async Task PartitionManagerLoop(CancellationToken gracefulShutdownToken, Cancell consecutiveFailureCount = 0; } - catch (OperationCanceledException) when (forcefulShutdownToken.IsCancellationRequested) - { - break; - } // Exception Status 412 represents an out of date ETag. We already logged this. catch (DurableTaskStorageException ex) when (ex.HttpStatusCode == (int)HttpStatusCode.PreconditionFailed) { @@ -236,10 +197,6 @@ async Task PartitionManagerLoop(CancellationToken gracefulShutdownToken, Cancell await Task.Delay(timeToSleep, gracefulShutdownToken); } } - catch (OperationCanceledException) when (forcefulShutdownToken.IsCancellationRequested) - { - break; - } catch (OperationCanceledException) when (gracefulShutdownToken.IsCancellationRequested) { // Shutdown requested, but we still need to release all leases @@ -272,93 +229,36 @@ async Task PartitionManagerLoop(CancellationToken gracefulShutdownToken, Cancell /// async Task IPartitionManager.StopAsync() { - await this.lifecycleLock.WaitAsync(); - try - { - if (!this.isStarted) - { - if (this.isStopping) - { - try - { - await this.partitionManagerTask; - } - finally - { - if (this.partitionManagerTask.IsCompleted) - { - this.DisposeRun(); - } - } - } - - return; - } - - this.isStarted = false; - this.isStopping = true; - CancellationTokenSource gracefulShutdown = - this.gracefulShutdownTokenSource - ?? throw new InvalidOperationException("The graceful shutdown token source is missing."); - CancellationTokenSource forcefulShutdown = - this.forcefulShutdownTokenSource - ?? throw new InvalidOperationException("The forceful shutdown token source is missing."); - - gracefulShutdown.Cancel(); - this.settings.Logger.PartitionManagerInfo( - this.storageAccountName, - this.settings.TaskHubName, - this.settings.WorkerId, - partitionId: NotApplicable, - "Started draining the in-memory messages of all owned control queues for shutdown."); - - // Wait 10 minutes for the partition manager to shutdown gracefully. Otherwise force a shutdown. - var timeout = TimeSpan.FromMinutes(10); - var timeoutTask = Task.Delay(Timeout.Infinite, forcefulShutdown.Token); - forcefulShutdown.CancelAfter(timeout); - - try - { - await Task.WhenAny(this.partitionManagerTask, timeoutTask); - if (timeoutTask.IsCompleted) - { - throw new TimeoutException( - $"Timed-out waiting for the partition manager to shut down. Timeout duration: {timeout}", - timeoutTask.Exception?.InnerException); - } + this.gracefulShutdownTokenSource.Cancel(); + this.settings.Logger.PartitionManagerInfo( + this.storageAccountName, + this.settings.TaskHubName, + this.settings.WorkerId, + partitionId: NotApplicable, + "Started draining the in-memory messages of all owned control queues for shutdown."); - // Surface any unhandled exceptions - await this.partitionManagerTask; + // Wait 10 minutes for the partition manager to shutdown gracefully. Otherwise force a shutdown. + var timeout = TimeSpan.FromMinutes(10); + var timeoutTask = Task.Delay(Timeout.Infinite, this.forcefulShutdownTokenSource.Token); + this.forcefulShutdownTokenSource.CancelAfter(timeout); + await Task.WhenAny(this.partitionManagerTask, timeoutTask); - this.settings.Logger.PartitionManagerInfo( - this.storageAccountName, - this.settings.TaskHubName, - this.settings.WorkerId, - partitionId: NotApplicable, - "Table partition manager stopped successfully."); - } - finally - { - if (this.partitionManagerTask.IsCompleted) - { - this.DisposeRun(); - } - } - } - finally + if (timeoutTask.IsCompleted) { - this.lifecycleLock.Release(); + throw new TimeoutException( + $"Timed-out waiting for the partition manager to shut down. Timeout duration: {timeout}", + timeoutTask.Exception?.InnerException); } - } - void DisposeRun() - { - this.gracefulShutdownTokenSource?.Dispose(); - this.forcefulShutdownTokenSource?.Dispose(); - this.gracefulShutdownTokenSource = null; - this.forcefulShutdownTokenSource = null; - this.partitionManagerTask = Task.CompletedTask; - this.isStopping = false; + // Surface any unhandled exceptions + await this.partitionManagerTask; + + this.settings.Logger.PartitionManagerInfo( + this.storageAccountName, + this.settings.TaskHubName, + this.settings.WorkerId, + partitionId: NotApplicable, + "Table partition manager stopped successfully."); } async Task IPartitionManager.CreateLeaseStore() @@ -524,8 +424,16 @@ public async Task ReadAndWriteTableAsync(bool isShuttingDown, throw; } - // Only an owner of this non-draining partition should listen to its control queue. - if (partition.CurrentOwner == this.workerName && !partition.IsDraining) + // Ensure worker is listening to the control queue iff either: + // 1) worker just claimed the lease, + // 2) worker was already the owner in the partitions table and is not actively draining the queue. + // Note that during draining, we renew the lease but do not want to listen to new messages. + // Otherwise, we'll never finish draining our in-memory messages. + // When draining completes, and the worker may decide to release the lease. In that moment, + // IsDrainingPartition can still be true but renewedLease can be false — without checking + // !releasedLease, the worker could incorrectly resume listening just before releasing the lease. + bool isRenewingToDrainQueue = renewedLease && response.IsDrainingPartition && !releasedLease; + if (claimedLease || !isRenewingToDrainQueue) { // Notify the orchestration session manager that we acquired a lease for one of the partitions. // This will cause it to start reading control queue messages for that partition. @@ -590,7 +498,6 @@ void RenewOrReleaseMyLease( if (partition.NextOwner == null) { // We still own the lease and nobody is trying to steal it. - partition.IsDraining = false; ownershipLeaseCount++; this.RenewLease(partition); renewedLease = true; @@ -988,14 +895,13 @@ internal void SimulateUnhealthyWorker(CancellationToken testToken) // used for internal testing internal void KillLoop() { - this.forcefulShutdownTokenSource?.Cancel(); + this.forcefulShutdownTokenSource.Cancel(); } public void Dispose() { - this.gracefulShutdownTokenSource?.Dispose(); - this.forcefulShutdownTokenSource?.Dispose(); - this.lifecycleLock.Dispose(); + this.gracefulShutdownTokenSource.Dispose(); + this.forcefulShutdownTokenSource.Dispose(); } } } diff --git a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs index 3b8619cfb..5d04af120 100644 --- a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs +++ b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -15,23 +15,22 @@ namespace DurableTask.AzureStorage.Tests { using System; using System.Collections.Concurrent; - using System.Collections.Generic; using System.Linq; - using System.Runtime.CompilerServices; using System.Threading; using System.Threading.Tasks; + using DurableTask.AzureStorage.Messaging; using DurableTask.AzureStorage.Partitioning; - using DurableTask.AzureStorage.Storage; using DurableTask.Core; using DurableTask.Core.History; + using DurableTask.Core.Settings; using Microsoft.Extensions.Logging; using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] + [DoNotParallelize] public class AppLeaseActivityTests { static readonly TimeSpan TestTimeout = TimeSpan.FromSeconds(15); - static readonly TimeSpan AutomaticFailoverTimeout = TimeSpan.FromSeconds(30); [TestMethod] public async Task DifferentAppCannotDequeueActivityWhilePassive() @@ -83,9 +82,9 @@ public async Task SameAppWorkersCanDequeueActivitiesConcurrently() await EnqueueActivityAsync(worker1, "activity-1"); await EnqueueActivityAsync(worker1, "activity-2"); - Task lock1 = LockActivityAsync(worker1); - Task lock2 = LockActivityAsync(worker2); - TaskActivityWorkItem[] workItems = await Task.WhenAll(lock1, lock2); + TaskActivityWorkItem[] workItems = await Task.WhenAll( + LockActivityAsync(worker1), + LockActivityAsync(worker2)); Assert.IsNotNull(workItems[0]); Assert.IsNotNull(workItems[1]); @@ -134,431 +133,149 @@ public async Task AppLeaseDisabledAllowsDifferentAppsToDequeueActivities() } [TestMethod] - public async Task ForcedFailoverCancelsOldOwnerPollAndEnablesNewOwner() + public async Task OwnershipResetCancelsOldEpochAndBlocksNewWaiters() { - string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationService oldOwner = CreateService(taskHubName, "OldApp", useAppLease: true); - AzureStorageOrchestrationService newOwner = CreateService(taskHubName, "NewApp", useAppLease: true); + var ownershipSignal = new AppLeaseOwnershipSignal(); + ownershipSignal.Set(); - try + using (AppLeaseOwnershipSignal.AppLeaseOwnership ownership = + await ownershipSignal.WaitAsync(CancellationToken.None)) { - await oldOwner.CreateAsync(); - await oldOwner.StartAsync(); - await WaitForOwnerAsync(oldOwner); - await newOwner.StartAsync(); - - Task oldOwnerPoll = oldOwner.LockNextTaskActivityWorkItem( - TestTimeout, - CancellationToken.None); + ownershipSignal.Reset(); - await newOwner.ForceChangeAppLeaseAsync(); - await WaitForOwnerAsync(newOwner); - await EnqueueActivityAsync(newOwner, "activity"); + Assert.IsTrue(ownership.LostToken.IsCancellationRequested); + Assert.IsFalse(ownership.TryBeginDispatch()); - TaskActivityWorkItem newOwnerWorkItem = await LockActivityAsync(newOwner); - Assert.IsNotNull(newOwnerWorkItem); - Assert.IsNull(await WithTimeoutAsync(oldOwnerPoll)); - - await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); - } - finally - { - await StopAsync(newOwner); - await StopAsync(oldOwner); - } - } - - [TestMethod] - public async Task ForcedFailoverQuiescesOldAppBeforeNewOwnerStarts() - { - string taskHubName = GetTaskHubName(); - TimeSpan renewInterval = TimeSpan.FromSeconds(2); - AzureStorageOrchestrationServiceSettings oldSettings = CreateSettings( - taskHubName, - "OldApp", - useAppLease: true, - renewInterval: renewInterval); - AzureStorageOrchestrationServiceSettings newSettings = CreateSettings( - taskHubName, - "NewApp", - useAppLease: true, - renewInterval: renewInterval); - AppLeaseManager oldManager = CreateAppLeaseManager( - oldSettings, - new TestPartitionManager()); - AppLeaseManager newManager = CreateAppLeaseManager( - newSettings, - new TestPartitionManager()); - try - { - await oldManager.CreateContainerIfNotExistsAsync(); - await oldManager.StartAsync(); - using (AppLeaseOwnershipSignal.AppLeaseOwnership oldOwnership = - await GetOwnershipAsync(oldManager)) + using (var cancellation = new CancellationTokenSource(TimeSpan.FromMilliseconds(100))) { - await newManager.StartAsync(); - - await newManager.ForceChangeAppLeaseAsync(); - using (AppLeaseOwnershipSignal.AppLeaseOwnership newOwnership = - await GetOwnershipAsync(newManager)) - { - Assert.IsTrue(oldOwnership.LostToken.IsCancellationRequested); - Assert.IsFalse(oldOwnership.TryBeginDispatch()); - } + await Assert.ThrowsExceptionAsync( + () => ownershipSignal.WaitAsync(cancellation.Token)); } } - finally - { - await newManager.StopAsync(); - await oldManager.StopAsync(); - } } [TestMethod] - public async Task LeaseExpirationEnablesNewOwner() + public async Task OwnershipLossCancelsPendingQueueReceive() { string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationService oldOwner = CreateService(taskHubName, "OldApp", useAppLease: true); - AzureStorageOrchestrationService newOwner = CreateService(taskHubName, "NewApp", useAppLease: true); + AzureStorageOrchestrationService service = + CreateService(taskHubName, "PrimaryApp", useAppLease: true); + var ownershipSignal = new AppLeaseOwnershipSignal(); + ownershipSignal.Set(); try { - await oldOwner.CreateAsync(); - await oldOwner.StartAsync(); - await WaitForOwnerAsync(oldOwner); - await newOwner.StartAsync(); - await oldOwner.StopAsync(isForced: true); - oldOwner = null; - await EnqueueActivityAsync(newOwner, "activity"); + await service.CreateAsync(); + await service.StartAsync(); + await WaitForOwnerAsync(service); - TaskActivityWorkItem newOwnerWorkItem = - await LockActivityAsync(newOwner, AutomaticFailoverTimeout); - - Assert.IsNotNull(newOwnerWorkItem); - await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); - } - finally - { - await StopAsync(newOwner); - await StopAsync(oldOwner); - } - } - - [TestMethod] - public async Task PreviousOwnerProcessesOrchestrationAfterReacquiringAppLease() - { - string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationServiceSettings firstSettings = - CreateSettings(taskHubName, "FirstApp", useAppLease: true); - AzureStorageOrchestrationServiceSettings secondSettings = - CreateSettings(taskHubName, "SecondApp", useAppLease: true); - firstSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); - secondSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); - - var firstService = new AzureStorageOrchestrationService(firstSettings); - var secondService = new AzureStorageOrchestrationService(secondSettings); - var firstWorker = new TaskHubWorker(firstService); - var secondWorker = new TaskHubWorker(secondService); - firstWorker.AddTaskOrchestrations(typeof(WaitForSignalOrchestration)); - secondWorker.AddTaskOrchestrations(typeof(WaitForSignalOrchestration)); - - bool firstWorkerStarted = false; - bool secondWorkerStarted = false; - - try - { - await firstService.CreateAsync(); - await firstWorker.StartAsync(); - firstWorkerStarted = true; - await WaitForOwnerAsync(firstService); - - var client = new TaskHubClient(firstService); - OrchestrationInstance instance = - await client.CreateOrchestrationInstanceAsync( - typeof(WaitForSignalOrchestration), - input: null); - await WaitForOrchestrationStatusAsync( - client, - instance, - OrchestrationStatus.Running); - - await secondWorker.StartAsync(); - secondWorkerStarted = true; - await secondService.ForceChangeAppLeaseAsync(); - await WaitForOwnerAsync(secondService, secondSettings, instance.InstanceId); - - await secondWorker.StopAsync(isForced: true); - secondWorkerStarted = false; - await WaitForAppLeaseOwnerAsync(firstSettings, AutomaticFailoverTimeout); - await WaitForOwnerAsync( - firstService, - firstSettings, - instance.InstanceId, - AutomaticFailoverTimeout); - await client.RaiseEventAsync(instance, "complete", "done"); - - OrchestrationState state = - await client.WaitForOrchestrationAsync(instance, TestTimeout); - - Assert.IsNotNull(state); - Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); - Assert.AreEqual("\"done\"", state.Output); - } - finally - { - if (secondWorkerStarted) + using (AppLeaseOwnershipSignal.AppLeaseOwnership ownership = + await ownershipSignal.WaitAsync(CancellationToken.None)) + using (var receiveCancellation = + CancellationTokenSource.CreateLinkedTokenSource(ownership.LostToken)) { - await secondWorker.StopAsync(isForced: true); - } + Task pendingReceive = + service.WorkItemQueue.GetMessageAsync(receiveCancellation.Token); - if (firstWorkerStarted) - { - await firstWorker.StopAsync(isForced: true); - } - } - } - - [TestMethod] - public async Task ReleasedPartitionsDoNotResumeListeningBeforeFailback() - { - string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationServiceSettings firstSettings = - CreateSettings(taskHubName, "FirstApp", useAppLease: true); - AzureStorageOrchestrationServiceSettings secondSettings = - CreateSettings(taskHubName, "SecondApp", useAppLease: true); - firstSettings.PartitionCount = 4; - firstSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); - secondSettings.PartitionCount = 4; - secondSettings.LeaseAcquireInterval = TimeSpan.FromMilliseconds(200); - secondSettings.ControlQueueBufferThreshold = 0; + ownershipSignal.Reset(); - var firstService = new AzureStorageOrchestrationService(firstSettings); - var secondService = new AzureStorageOrchestrationService(secondSettings); - var firstWorker = new TaskHubWorker(firstService); - var secondWorker = new TaskHubWorker(secondService); - firstWorker.AddTaskOrchestrations( - new WorkerIdentityOrchestrationCreator("FirstApp")); - secondWorker.AddTaskOrchestrations( - new WorkerIdentityOrchestrationCreator("SecondApp")); - - bool firstWorkerStarted = false; - bool secondWorkerStarted = false; - - try - { - await firstService.CreateAsync(); - await firstWorker.StartAsync(); - firstWorkerStarted = true; - await WaitForAppLeaseOwnerAsync(firstSettings); - await WaitForTableOwnershipAsync(firstService, firstSettings); - await WaitForReceivingQueueCountAsync(firstService, firstSettings.PartitionCount); - - await secondWorker.StartAsync(); - secondWorkerStarted = true; - await secondService.ForceChangeAppLeaseAsync(); - await WaitForAppLeaseOwnerAsync(secondSettings); - await WaitForTableOwnershipAsync(secondService, secondSettings); - await WaitForReceivingQueueCountAsync(secondService, secondSettings.PartitionCount); - - string[] oldReceivingQueues = firstService.OwnedControlQueues - .Where(queue => !queue.IsReleased) - .Select(queue => queue.Name) - .ToArray(); - if (oldReceivingQueues.Length == 0) - { - secondSettings.ControlQueueBufferThreshold = 1000; + Assert.IsNull(await WithTimeoutAsync(pendingReceive)); } - - var client = new TaskHubClient(firstService); - OrchestrationInstance secondOwnerInstance = - await client.CreateOrchestrationInstanceAsync( - WorkerIdentityOrchestrationCreator.OrchestrationName, - version: string.Empty, - input: null); - OrchestrationState secondOwnerState = - await client.WaitForOrchestrationAsync(secondOwnerInstance, TestTimeout); - - Assert.AreEqual( - 0, - oldReceivingQueues.Length, - $"The previous app resumed receiving from released partitions: {string.Join(", ", oldReceivingQueues)}. " + - $"New orchestration output: {secondOwnerState?.Output ?? ""}."); - if (secondOwnerState == null) - { - Assert.Fail("The new owner did not complete the orchestration."); - } - - Assert.AreEqual(OrchestrationStatus.Completed, secondOwnerState.OrchestrationStatus); - Assert.AreEqual("\"SecondApp\"", secondOwnerState.Output); - - await secondWorker.StopAsync(isForced: true); - secondWorkerStarted = false; - await WaitForAppLeaseOwnerAsync(firstSettings, AutomaticFailoverTimeout); - await WaitForTableOwnershipAsync( - firstService, - firstSettings, - AutomaticFailoverTimeout); - await WaitForReceivingQueueCountAsync( - firstService, - firstSettings.PartitionCount, - AutomaticFailoverTimeout); - - OrchestrationInstance firstOwnerInstance = - await client.CreateOrchestrationInstanceAsync( - WorkerIdentityOrchestrationCreator.OrchestrationName, - version: string.Empty, - input: null); - OrchestrationState firstOwnerState = - await client.WaitForOrchestrationAsync(firstOwnerInstance, TestTimeout); - - Assert.IsNotNull(firstOwnerState); - Assert.AreEqual(OrchestrationStatus.Completed, firstOwnerState.OrchestrationStatus); - Assert.AreEqual("\"FirstApp\"", firstOwnerState.Output); } finally { - if (secondWorkerStarted) - { - await secondWorker.StopAsync(isForced: true); - } - - if (firstWorkerStarted) - { - await firstWorker.StopAsync(isForced: true); - } + await StopAsync(service); } } [TestMethod] - public async Task AppLeaseManagerCanRestartAfterPartitionManagerStopFails() + public async Task OwnershipLossAfterDequeueAbandonsWithoutStartingTrace() { string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationServiceSettings settings = - CreateSettings(taskHubName, "PrimaryApp", useAppLease: true); - var partitionManager = new StopFailingPartitionManager(); - AppLeaseManager manager = CreateAppLeaseManager(settings, partitionManager); - - await manager.CreateContainerIfNotExistsAsync(); - await manager.StartAsync(); - await WaitForOwnershipAsync(manager); - - await Assert.ThrowsExceptionAsync(() => manager.StopAsync()); - - await manager.StartAsync(); - await WaitForOwnershipAsync(manager); - - await WithTimeoutAsync(manager.StopAsync()); - Assert.AreEqual(2, partitionManager.StartCount); - Assert.AreEqual(2, partitionManager.StopCount); - } - - [TestMethod] - public async Task ForceChangeBeforeFirstOwnerAcquiresLease() - { - string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationServiceSettings settings = + var loggerFactory = new RecordingLoggerFactory(); + AzureStorageOrchestrationServiceSettings ownerSettings = CreateSettings(taskHubName, "PrimaryApp", useAppLease: true); - var partitionManager = new TestPartitionManager(); - AppLeaseManager manager = CreateAppLeaseManager(settings, partitionManager); + ownerSettings.LoggerFactory = loggerFactory; + AzureStorageOrchestrationService owner = + new AzureStorageOrchestrationService(ownerSettings); + AzureStorageOrchestrationService recoveryReader = + CreateService(taskHubName, "RecoveryReader", useAppLease: false); + CorrelationSettings previousCorrelationSettings = CorrelationSettings.Current; try { - await manager.CreateContainerIfNotExistsAsync(); - - await manager.ForceChangeAppLeaseAsync(); - - using (var cancellation = new CancellationTokenSource(TimeSpan.FromSeconds(3))) - using (await manager.WaitForOwnershipAsync(cancellation.Token)) + CorrelationSettings.Current = new CorrelationSettings { - Assert.AreEqual(1, partitionManager.StartCount); - } - } - finally - { - await manager.StopAsync(); - } - } + EnableDistributedTracing = true, + Protocol = Protocol.W3CTraceContext, + }; + CorrelationTraceContext.Current = null; - [TestMethod] - public async Task LeaseLossDoesNotCancelDispatchedActivity() - { - string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationService oldOwner = CreateService(taskHubName, "OldApp", useAppLease: true); - AzureStorageOrchestrationService newOwner = CreateService(taskHubName, "NewApp", useAppLease: true); + await owner.CreateAsync(); + await owner.StartAsync(); + await WaitForOwnerAsync(owner); + await EnqueueActivityAsync(owner, "ownership-race"); - try - { - await oldOwner.CreateAsync(); - await oldOwner.StartAsync(); - await WaitForOwnerAsync(oldOwner); - await newOwner.StartAsync(); - await EnqueueActivityAsync(oldOwner, "in-flight"); + owner.OnActivityMessageDequeued = async () => + { + owner.OnActivityMessageDequeued = null; + await StopAsync(owner); + }; - TaskActivityWorkItem inFlightWorkItem = await LockActivityAsync(oldOwner); - await newOwner.ForceChangeAppLeaseAsync(); - await EnqueueActivityAsync(newOwner, "after-failover"); - TaskActivityWorkItem newOwnerWorkItem = await LockActivityAsync(newOwner); + TaskActivityWorkItem rejectedWorkItem = await LockActivityAsync(owner); + owner = null; - TaskActivityWorkItem renewedWorkItem = - await oldOwner.RenewTaskActivityWorkItemLockAsync(inFlightWorkItem); - Assert.IsTrue(renewedWorkItem.LockedUntilUtc > DateTime.UtcNow); + Assert.IsNull(rejectedWorkItem); + Assert.IsNull(CorrelationTraceContext.Current); + Assert.IsFalse(loggerFactory.HasEvent("ReceivedMessage")); + Assert.IsFalse(loggerFactory.HasEvent("ProcessingMessage")); - await oldOwner.AbandonTaskActivityWorkItemAsync(inFlightWorkItem); - await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); + await recoveryReader.StartAsync(); + TaskActivityWorkItem recoveredWorkItem = await LockActivityAsync(recoveryReader); + Assert.IsNotNull(recoveredWorkItem); + await recoveryReader.AbandonTaskActivityWorkItemAsync(recoveredWorkItem); } finally { - await StopAsync(newOwner); - await StopAsync(oldOwner); + CorrelationTraceContext.Current = null; + CorrelationSettings.Current = previousCorrelationSettings; + await StopAsync(recoveryReader); + await StopAsync(owner); } } [TestMethod] - public async Task OwnershipLossAfterDequeueDoesNotStartProcessingTrace() + public async Task ClosingGateDoesNotCancelDispatchedActivity() { string taskHubName = GetTaskHubName(); - var loggerFactory = new RecordingLoggerFactory(); - AzureStorageOrchestrationServiceSettings oldSettings = - CreateSettings(taskHubName, "OldApp", useAppLease: true); - oldSettings.LoggerFactory = loggerFactory; - AzureStorageOrchestrationService oldOwner = - new AzureStorageOrchestrationService(oldSettings); - AzureStorageOrchestrationService newOwner = - CreateService(taskHubName, "NewApp", useAppLease: true); + AzureStorageOrchestrationService owner = + CreateService(taskHubName, "PrimaryApp", useAppLease: true); try { - await oldOwner.CreateAsync(); - await oldOwner.StartAsync(); - await WaitForOwnerAsync(oldOwner); - await newOwner.StartAsync(); - await EnqueueActivityAsync(oldOwner, "ownership-race"); - - oldOwner.OnActivityMessageDequeued = async () => - { - oldOwner.OnActivityMessageDequeued = null; - await newOwner.ForceChangeAppLeaseAsync(); - await WaitForOwnerAsync(newOwner); - }; + await owner.CreateAsync(); + await owner.StartAsync(); + await WaitForOwnerAsync(owner); + await EnqueueActivityAsync(owner, "in-flight"); - TaskActivityWorkItem rejectedWorkItem = await LockActivityAsync(oldOwner); + TaskActivityWorkItem inFlightWorkItem = await LockActivityAsync(owner); + await owner.StopAsync(isForced: true); - Assert.IsNull(rejectedWorkItem); - Assert.IsFalse(loggerFactory.HasEvent("ProcessingMessage")); + TaskActivityWorkItem renewedWorkItem = + await owner.RenewTaskActivityWorkItemLockAsync(inFlightWorkItem); + Assert.IsTrue(renewedWorkItem.LockedUntilUtc > DateTime.UtcNow); - TaskActivityWorkItem newOwnerWorkItem = await LockActivityAsync(newOwner); - Assert.IsNotNull(newOwnerWorkItem); - await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); + await owner.AbandonTaskActivityWorkItemAsync(inFlightWorkItem); + owner = null; } finally { - await StopAsync(newOwner); - await StopAsync(oldOwner); + await StopAsync(owner); } } [TestMethod] - public async Task PassiveActivityPollStopsOnCancellationAndShutdown() + public async Task PassiveWaitHonorsCallerCancellationAndServiceShutdown() { string taskHubName = GetTaskHubName(); AzureStorageOrchestrationService owner = CreateService(taskHubName, "PrimaryApp", useAppLease: true); @@ -597,18 +314,16 @@ public async Task PassiveActivityPollStopsOnCancellationAndShutdown() static AzureStorageOrchestrationService CreateService( string taskHubName, string appName, - bool useAppLease, - TimeSpan? renewInterval = null) + bool useAppLease) { return new AzureStorageOrchestrationService( - CreateSettings(taskHubName, appName, useAppLease, renewInterval)); + CreateSettings(taskHubName, appName, useAppLease)); } static AzureStorageOrchestrationServiceSettings CreateSettings( string taskHubName, string appName, - bool useAppLease, - TimeSpan? renewInterval = null) + bool useAppLease) { return new AzureStorageOrchestrationServiceSettings { @@ -617,7 +332,7 @@ static AzureStorageOrchestrationServiceSettings CreateSettings( { AcquireInterval = TimeSpan.FromMilliseconds(200), LeaseInterval = TimeSpan.FromSeconds(15), - RenewInterval = renewInterval ?? TimeSpan.FromMilliseconds(200), + RenewInterval = TimeSpan.FromMilliseconds(200), }, MaxQueuePollingInterval = TimeSpan.FromMilliseconds(50), PartitionCount = 1, @@ -631,7 +346,7 @@ static AzureStorageOrchestrationServiceSettings CreateSettings( static string GetTaskHubName() { - return "applease" + Guid.NewGuid().ToString("N").Substring(0, 16); + return "activitygate" + Guid.NewGuid().ToString("N").Substring(0, 12); } static async Task EnqueueActivityAsync( @@ -660,169 +375,14 @@ await TestHelpers.WaitFor( TestTimeout); } - static async Task WaitForOwnerAsync( - AzureStorageOrchestrationService service, - AzureStorageOrchestrationServiceSettings settings, - string instanceId, - TimeSpan? timeout = null) - { - uint partitionIndex = - Fnv1aHashHelper.ComputeHash(instanceId) % (uint)settings.PartitionCount; - string queueName = - AzureStorageOrchestrationService.GetControlQueueName( - settings.TaskHubName, - (int)partitionIndex); - - await TestHelpers.WaitFor( - () => service.OwnedControlQueues.Any(queue => queue.Name == queueName), - timeout ?? TestTimeout); - } - - static async Task WaitForAppLeaseOwnerAsync( - AzureStorageOrchestrationServiceSettings settings, - TimeSpan? timeout = null) - { - string taskHubName = settings.TaskHubName.ToLowerInvariant(); - Blob appLeaseInfoBlob = new AzureStorageClient(settings) - .GetBlobContainerReference(taskHubName + "-applease") - .GetBlobReference(taskHubName + "-appleaseinfo"); - byte[] appLeaseIdBytes = - BitConverter.GetBytes(Fnv1aHashHelper.ComputeHash(settings.AppName)); - Array.Resize(ref appLeaseIdBytes, 16); - string expectedOwnerId = new Guid(appLeaseIdBytes).ToString(); - - using (var cancellation = new CancellationTokenSource(timeout ?? TestTimeout)) - { - try - { - while (true) - { - string json = await appLeaseInfoBlob.DownloadTextAsync(); - var info = Utils.DeserializeFromJson(json); - if (info.OwnerId == expectedOwnerId) - { - return; - } - - await Task.Delay(TimeSpan.FromMilliseconds(100), cancellation.Token); - } - } - catch (OperationCanceledException) - { - Assert.Fail("The previous owner did not reacquire the app lease."); - } - } - } - - static async Task WaitForTableOwnershipAsync( - AzureStorageOrchestrationService service, - AzureStorageOrchestrationServiceSettings settings, - TimeSpan? timeout = null) - { - using (var cancellation = new CancellationTokenSource(timeout ?? TestTimeout)) - { - try - { - while (true) - { - List leases = - await service.ListTableLeasesAsync().ToListAsync(); - if (leases.Count == settings.PartitionCount && - leases.All( - lease => - lease.CurrentOwner == settings.WorkerId && - !lease.IsDraining)) - { - return; - } - - await Task.Delay(TimeSpan.FromMilliseconds(100), cancellation.Token); - } - } - catch (OperationCanceledException) - { - Assert.Fail( - $"Worker {settings.WorkerId} did not own all non-draining table partitions."); - } - } - } - - static Task WaitForReceivingQueueCountAsync( - AzureStorageOrchestrationService service, - int expectedCount, - TimeSpan? timeout = null) - { - return TestHelpers.WaitFor( - () => service.OwnedControlQueues.Count(queue => !queue.IsReleased) == expectedCount, - timeout ?? TestTimeout); - } - - static async Task WaitForOrchestrationStatusAsync( - TaskHubClient client, - OrchestrationInstance instance, - OrchestrationStatus expectedStatus) - { - using (var cancellation = new CancellationTokenSource(TestTimeout)) - { - try - { - while (true) - { - OrchestrationState state = - await client.GetOrchestrationStateAsync(instance); - if (state?.OrchestrationStatus == expectedStatus) - { - return; - } - - await Task.Delay(TimeSpan.FromMilliseconds(100), cancellation.Token); - } - } - catch (OperationCanceledException) - { - Assert.Fail( - $"The orchestration did not reach {expectedStatus} before the test timeout."); - } - } - } - - static async Task WaitForOwnershipAsync(AppLeaseManager manager) - { - using (await GetOwnershipAsync(manager)) - { - } - } - - static async Task GetOwnershipAsync( - AppLeaseManager manager) - { - using (var cancellation = new CancellationTokenSource(TestTimeout)) - { - return await manager.WaitForOwnershipAsync(cancellation.Token); - } - } - - static AppLeaseManager CreateAppLeaseManager( - AzureStorageOrchestrationServiceSettings settings, - IPartitionManager partitionManager) - { - string taskHubName = settings.TaskHubName.ToLowerInvariant(); - return new AppLeaseManager( - new AzureStorageClient(settings), - partitionManager, - taskHubName + "-applease", - taskHubName + "-appleaseinfo", - settings.AppLeaseOptions); - } - static async Task LockActivityAsync( - AzureStorageOrchestrationService service, - TimeSpan? timeout = null) + AzureStorageOrchestrationService service) { - TimeSpan effectiveTimeout = timeout ?? TestTimeout; - using (var cancellation = new CancellationTokenSource(effectiveTimeout)) + using (var cancellation = new CancellationTokenSource(TestTimeout)) { - return await service.LockNextTaskActivityWorkItem(effectiveTimeout, cancellation.Token); + return await service.LockNextTaskActivityWorkItem( + TestTimeout, + cancellation.Token); } } @@ -833,116 +393,14 @@ static async Task WithTimeoutAsync(Task task) return await task; } - static async Task WithTimeoutAsync(Task task) - { - Task completedTask = await Task.WhenAny(task, Task.Delay(TestTimeout)); - Assert.AreSame(task, completedTask, "The operation did not complete before the test timeout."); - await task; - } - static async Task StopAsync(AzureStorageOrchestrationService service) { if (service != null) { - await service.StopAsync(isForced: true); - } - } - - sealed class StopFailingPartitionManager : IPartitionManager - { - public int StartCount { get; private set; } - - public int StopCount { get; private set; } - - public Task StartAsync() - { - this.StartCount++; - return Task.CompletedTask; - } - - public Task StopAsync() - { - this.StopCount++; - return this.StopCount == 1 - ? Task.FromException(new InvalidOperationException("Simulated stop failure.")) - : Task.CompletedTask; - } - - public Task CreateLeaseStore() => Task.CompletedTask; - - public Task CreateLease(string leaseName) => Task.CompletedTask; - - public Task DeleteLeases() => Task.CompletedTask; - - public async IAsyncEnumerable GetOwnershipBlobLeasesAsync( - [EnumeratorCancellation] - CancellationToken cancellationToken = default) - { - await Task.CompletedTask; - yield break; - } - } - - sealed class WorkerIdentityOrchestrationCreator : ObjectCreator - { - public const string OrchestrationName = "AppLeaseWorkerIdentity"; - - readonly string workerIdentity; - - public WorkerIdentityOrchestrationCreator(string workerIdentity) - { - this.workerIdentity = workerIdentity; - this.Name = OrchestrationName; - this.Version = string.Empty; - } - - public override TaskOrchestration Create() - { - return new WorkerIdentityOrchestration(this.workerIdentity); - } - } - - sealed class WorkerIdentityOrchestration : TaskOrchestration - { - readonly string workerIdentity; - - public WorkerIdentityOrchestration(string workerIdentity) - { - this.workerIdentity = workerIdentity; - } - - public override Task RunTask( - OrchestrationContext context, - string input) - { - return Task.FromResult(this.workerIdentity); - } - } - - sealed class TestPartitionManager : IPartitionManager - { - public int StartCount { get; private set; } - - public Task StartAsync() - { - this.StartCount++; - return Task.CompletedTask; - } - - public Task StopAsync() => Task.CompletedTask; - - public Task CreateLeaseStore() => Task.CompletedTask; - - public Task CreateLease(string leaseName) => Task.CompletedTask; - - public Task DeleteLeases() => Task.CompletedTask; - - public async IAsyncEnumerable GetOwnershipBlobLeasesAsync( - [EnumeratorCancellation] - CancellationToken cancellationToken = default) - { - await Task.CompletedTask; - yield break; + Task stopTask = service.StopAsync(isForced: true); + Task completedTask = await Task.WhenAny(stopTask, Task.Delay(TestTimeout)); + Assert.AreSame(stopTask, completedTask, "Service shutdown did not complete before the test timeout."); + await stopTask; } } @@ -992,34 +450,5 @@ public void Log( } } } - - sealed class WaitForSignalOrchestration : TaskOrchestration - { - TaskCompletionSource signal; - - public override async Task RunTask( - OrchestrationContext context, - string input) - { - this.signal = new TaskCompletionSource(); - return await this.signal.Task; - } - - public override void OnEvent( - OrchestrationContext context, - string name, - string input) - { - if (name == "complete") - { - this.signal?.TrySetResult(input); - } - } - } - - sealed class TestAppLeaseInfo - { - public string OwnerId { get; set; } - } } } diff --git a/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs b/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs index f688bec58..7fdb42e04 100644 --- a/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs +++ b/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs @@ -726,127 +726,6 @@ await WaitForConditionAsync( Assert.AreEqual(1, service.OwnedControlQueues.Count()); } - [TestMethod] - public async Task ForcedCancellationCompletesStopAndAllowsWorkAfterRestart() - { - string taskHubName = - "forcecancel" + Guid.NewGuid().ToString("N").Substring(0, 12); - var settings = new AzureStorageOrchestrationServiceSettings - { - StorageAccountClientProvider = new StorageAccountClientProvider(this.connection), - TaskHubName = taskHubName, - PartitionCount = 1, - LeaseAcquireInterval = TimeSpan.FromMilliseconds(200), - MaxQueuePollingInterval = TimeSpan.FromMilliseconds(50), - WorkerId = "0", - UseAppLease = false, - UseTablePartitionManagement = true, - }; - var service = new AzureStorageOrchestrationService(settings); - var client = new TaskHubClient(service); - - await service.StartAsync(); - await WaitForConditionAsync( - TimeSpan.FromSeconds(5), - t => new ValueTask(service.OwnedControlQueues.Any())); - - OrchestrationInstance instance = - await client.CreateOrchestrationInstanceAsync( - typeof(HelloOrchestrator), - input: null); - using var lockCancellation = new CancellationTokenSource(TimeSpan.FromSeconds(10)); - TaskOrchestrationWorkItem workItem = - await service.LockNextTaskOrchestrationWorkItemAsync( - TimeSpan.FromSeconds(10), - lockCancellation.Token); - Assert.IsNotNull(workItem); - - Task stopTask = service.StopAsync(isForced: false); - await WaitForConditionAsync( - TimeSpan.FromSeconds(5), - t => new ValueTask( - service.OwnedControlQueues.Single().IsReleased)); - Assert.IsFalse( - stopTask.IsCompleted, - "Graceful stop should wait for the locked orchestration work item to drain."); - - service.KillPartitionManagerLoop(); - - Task completedTask = - await Task.WhenAny(stopTask, Task.Delay(TimeSpan.FromSeconds(5))); - Assert.AreSame( - stopTask, - completedTask, - "Forceful cancellation did not complete the timed-out stop."); - await Assert.ThrowsExceptionAsync(() => stopTask); - - await service.AbandonTaskOrchestrationWorkItemAsync(workItem); - await service.ReleaseTaskOrchestrationWorkItemAsync(workItem); - await WaitForConditionAsync( - TimeSpan.FromSeconds(5), - t => new ValueTask(!service.OwnedControlQueues.Any())); - - Task secondStopTask = service.StopAsync(isForced: false); - completedTask = - await Task.WhenAny( - secondStopTask, - Task.Delay(TimeSpan.FromSeconds(5))); - Assert.AreSame( - secondStopTask, - completedTask, - "A repeated stop did not observe termination of the old partition manager loop."); - await secondStopTask; - - var worker = new TaskHubWorker(service); - worker.AddTaskOrchestrations(typeof(HelloOrchestrator)); - worker.AddTaskActivities(typeof(Hello)); - await worker.StartAsync(); - - OrchestrationState state = - await client.WaitForOrchestrationAsync( - instance, - TimeSpan.FromSeconds(30)); - - Assert.IsNotNull(state); - Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus); - Assert.AreEqual("\"Hello, world!\"", state.Output); - - await worker.StopAsync(isForced: true); - await service.DeleteAsync(); - } - - [TestMethod] - public async Task ForceSignalStillReportsTimeoutAfterLoopStops() - { - string taskHubName = - "forcetimeout" + Guid.NewGuid().ToString("N").Substring(0, 12); - var settings = new AzureStorageOrchestrationServiceSettings - { - StorageAccountClientProvider = new StorageAccountClientProvider(this.connection), - TaskHubName = taskHubName, - PartitionCount = 1, - LeaseAcquireInterval = TimeSpan.FromMilliseconds(200), - WorkerId = "0", - UseAppLease = false, - UseTablePartitionManagement = true, - }; - var service = new AzureStorageOrchestrationService(settings); - - await service.StartAsync(); - await WaitForConditionAsync( - TimeSpan.FromSeconds(5), - t => new ValueTask(service.OwnedControlQueues.Any())); - - service.KillPartitionManagerLoop(); - await Task.Delay(TimeSpan.FromSeconds(1)); - - await Assert.ThrowsExceptionAsync( - () => service.StopAsync(isForced: false)); - - await service.StopAsync(isForced: false); - await service.DeleteAsync(); - } - [KnownType(typeof(Hello))] internal class HelloOrchestrator : TaskOrchestration { From d60976e7827c1db6864641353567b7e3d9114d67 Mon Sep 17 00:00:00 2001 From: wangbill Date: Wed, 9 Sep 2026 11:51:38 -0700 Subject: [PATCH 09/10] Simplify activity app lease gating Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../AzureStorageOrchestrationService.cs | 107 +++++----- ...zureStorageOrchestrationServiceSettings.cs | 6 +- .../Partitioning/AppLeaseManager.cs | 82 +++++++- .../Partitioning/AppLeaseOwnershipSignal.cs | 182 ------------------ .../AppLeaseActivityTests.cs | 113 ++++++++--- 5 files changed, 214 insertions(+), 276 deletions(-) delete mode 100644 src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs index 3981828b5..a4b5143a6 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs @@ -1547,83 +1547,74 @@ public async Task LockNextTaskActivityWorkItem( using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, this.shutdownSource.Token)) { - AppLeaseOwnershipSignal.AppLeaseOwnership ownership; try { - ownership = await this.appLeaseManager.WaitForOwnershipAsync(linkedCts.Token); + await this.appLeaseManager.WaitForActivityOwnershipAsync(linkedCts.Token); } catch (OperationCanceledException) { return null; } - using (ownership) - { - using (var receiveCts = CancellationTokenSource.CreateLinkedTokenSource( - linkedCts.Token, - ownership.LostToken)) - { - MessageData message = await this.workItemQueue.GetMessageAsync(receiveCts.Token); + MessageData message = await this.workItemQueue.GetMessageAsync(linkedCts.Token); - if (message == null) - { - // shutting down, canceled, or app lease ownership was lost - return null; - } + if (message == null) + { + // shutting down or canceled + return null; + } - Func onActivityMessageDequeued = this.OnActivityMessageDequeued; - if (onActivityMessageDequeued != null) - { - await onActivityMessageDequeued(); - } + Func onActivityMessageDequeued = this.OnActivityMessageDequeued; + if (onActivityMessageDequeued != null) + { + await onActivityMessageDequeued(); + } - if (!ownership.TryBeginDispatch()) - { - await this.workItemQueue.AbandonMessageAsync(message); - return null; - } + if (!this.appLeaseManager.HasActivityOwnership) + { + await this.workItemQueue.AbandonMessageAsync(message); + return null; + } - Guid traceActivityId = Guid.NewGuid(); - var session = new ActivitySession(this.settings, this.azureStorageClient.QueueAccountName, message, traceActivityId); - session.StartNewLogicalTraceScope(); + Guid traceActivityId = Guid.NewGuid(); + var session = new ActivitySession(this.settings, this.azureStorageClient.QueueAccountName, message, traceActivityId); + session.StartNewLogicalTraceScope(); - TraceContextBase requestTraceContext = null; - CorrelationTraceClient.Propagate( - () => - { - string name = $"{TraceConstants.Activity} {Utils.GetTargetClassName(((TaskScheduledEvent)session.MessageData.TaskMessage.Event)?.Name)}"; - requestTraceContext = TraceContextFactory.Create(name); + TraceContextBase requestTraceContext = null; + CorrelationTraceClient.Propagate( + () => + { + string name = $"{TraceConstants.Activity} {Utils.GetTargetClassName(((TaskScheduledEvent)session.MessageData.TaskMessage.Event)?.Name)}"; + requestTraceContext = TraceContextFactory.Create(name); - TraceContextBase parentTraceContextBase = TraceContextBase.Restore(session.MessageData.SerializableTraceContext); - requestTraceContext.SetParentAndStart(parentTraceContextBase); - }); + TraceContextBase parentTraceContextBase = TraceContextBase.Restore(session.MessageData.SerializableTraceContext); + requestTraceContext.SetParentAndStart(parentTraceContextBase); + }); - TraceMessageReceived(this.settings, session.MessageData, this.azureStorageClient.QueueAccountName); - session.TraceProcessingMessage(message, isExtendedSession: false, this.workItemQueue.Name); + TraceMessageReceived(this.settings, session.MessageData, this.azureStorageClient.QueueAccountName); + session.TraceProcessingMessage(message, isExtendedSession: false, this.workItemQueue.Name); - if (!this.activeActivitySessions.TryAdd(message.Id, session)) - { - // This means we're already processing this message. This is never expected since the message - // should be kept invisible via background calls to RenewTaskActivityWorkItemLockAsync. - this.settings.Logger.AssertFailure( - this.azureStorageClient.QueueAccountName, - this.settings.TaskHubName, - $"Work item queue message with ID = {message.Id} is being processed multiple times concurrently."); - return null; - } + if (!this.activeActivitySessions.TryAdd(message.Id, session)) + { + // This means we're already processing this message. This is never expected since the message + // should be kept invisible via background calls to RenewTaskActivityWorkItemLockAsync. + this.settings.Logger.AssertFailure( + this.azureStorageClient.QueueAccountName, + this.settings.TaskHubName, + $"Work item queue message with ID = {message.Id} is being processed multiple times concurrently."); + return null; + } - this.stats.ActiveActivityExecutions.Increment(); + this.stats.ActiveActivityExecutions.Increment(); - return new TaskActivityWorkItem - { - Id = message.Id, - TaskMessage = session.MessageData.TaskMessage, - LockedUntilUtc = message.OriginalQueueMessage.NextVisibleOn.Value.UtcDateTime, + return new TaskActivityWorkItem + { + Id = message.Id, + TaskMessage = session.MessageData.TaskMessage, + LockedUntilUtc = message.OriginalQueueMessage.NextVisibleOn.Value.UtcDateTime, - TraceContextBase = requestTraceContext - }; - } - } + TraceContextBase = requestTraceContext + }; } } diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs index f04e039c6..4a24f6dca 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs @@ -154,8 +154,10 @@ public class AzureStorageOrchestrationServiceSettings /// /// If true, workers wait for their to own the app lease before receiving /// new activity messages. Workers sharing that app name may receive activities concurrently. - /// Observed lease loss stops pending activity receives and rejects dequeued activities before - /// dispatch; already dispatched activities are not canceled. Orchestration and entity message + /// Ownership loss does not cancel an activity receive that already started, which may continue + /// polling until it gets a message or caller or service-shutdown cancellation occurs. Current + /// ownership is checked again before dispatch; rejected messages are abandoned before tracing + /// starts. Already dispatched activities are not canceled. Orchestration and entity message /// processing retain their existing behavior. /// public bool UseAppLease { get; set; } = true; diff --git a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs index 4fcda8917..0113e52c2 100644 --- a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs @@ -42,8 +42,10 @@ sealed class AppLeaseManager readonly Blob appLeaseInfoBlob; readonly string appLeaseId; readonly AsyncManualResetEvent shutdownCompletedEvent; - readonly AppLeaseOwnershipSignal ownershipSignal; + readonly object activityOwnershipLock = new object(); + TaskCompletionSource activityOwnershipAvailable; + bool hasActivityOwnership; bool isLeaseOwner; int appLeaseIsStarted; Task renewTask; @@ -79,20 +81,82 @@ public AppLeaseManager( this.isLeaseOwner = false; this.shutdownCompletedEvent = new AsyncManualResetEvent(); - this.ownershipSignal = new AppLeaseOwnershipSignal(); + this.activityOwnershipAvailable = CreateActivityOwnershipSignal(); } - public Task WaitForOwnershipAsync( + public async Task WaitForActivityOwnershipAsync( CancellationToken cancellationToken) { - return this.ownershipSignal.WaitAsync(cancellationToken); + while (true) + { + Task ownershipAvailableTask; + lock (this.activityOwnershipLock) + { + if (this.hasActivityOwnership) + { + return; + } + + ownershipAvailableTask = this.activityOwnershipAvailable.Task; + } + + var canceled = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + using (cancellationToken.Register(() => canceled.TrySetResult(null))) + { + await Task.WhenAny(ownershipAvailableTask, canceled.Task); + } + + cancellationToken.ThrowIfCancellationRequested(); + } + } + + internal void SetActivityOwnership(bool ownsLease) + { + TaskCompletionSource ownershipAvailable = null; + lock (this.activityOwnershipLock) + { + if (this.hasActivityOwnership == ownsLease) + { + return; + } + + this.hasActivityOwnership = ownsLease; + if (ownsLease) + { + ownershipAvailable = this.activityOwnershipAvailable; + } + else + { + this.activityOwnershipAvailable = CreateActivityOwnershipSignal(); + } + } + + ownershipAvailable?.TrySetResult(null); + } + + internal bool HasActivityOwnership + { + get + { + lock (this.activityOwnershipLock) + { + return this.hasActivityOwnership; + } + } + } + + static TaskCompletionSource CreateActivityOwnershipSignal() + { + return new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); } public async Task StartAsync() { if (!this.appLeaseIsEnabled) { - this.ownershipSignal.Set(); + this.SetActivityOwnership(ownsLease: true); this.starterTokenSource = new CancellationTokenSource(); await Task.Factory.StartNew(() => this.PartitionManagerStarter(this.starterTokenSource.Token)); @@ -174,7 +238,7 @@ async Task AppLeaseManagerStarter(CancellationToken cancellationToken) public async Task StopAsync() { - this.ownershipSignal.Reset(); + this.SetActivityOwnership(ownsLease: false); if (this.starterTokenSource != null) { @@ -264,7 +328,7 @@ async Task StartAppLeaseAsync() this.leaseRenewerCancellationTokenSource = new CancellationTokenSource(); await this.partitionManager.StartAsync(); - this.ownershipSignal.Set(); + this.SetActivityOwnership(ownsLease: true); this.shutdownCompletedEvent.Reset(); @@ -279,7 +343,7 @@ async Task StopAppLeaseAsync() return; } - this.ownershipSignal.Reset(); + this.SetActivityOwnership(ownsLease: false); await this.partitionManager.StopAsync(); @@ -505,7 +569,7 @@ async Task RenewLeaseAsync() { renewed = false; this.isLeaseOwner = false; - this.ownershipSignal.Reset(); + this.SetActivityOwnership(ownsLease: false); this.settings.Logger.LeaseRenewalFailed( this.storageAccountName, diff --git a/src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs b/src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs deleted file mode 100644 index 507aac107..000000000 --- a/src/DurableTask.AzureStorage/Partitioning/AppLeaseOwnershipSignal.cs +++ /dev/null @@ -1,182 +0,0 @@ -// ---------------------------------------------------------------------------------- -// 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.Partitioning -{ - using System; - using System.Threading; - using System.Threading.Tasks; - - sealed class AppLeaseOwnershipSignal - { - readonly object syncLock = new object(); - - TaskCompletionSource ownershipAvailable = CreateCompletionSource(); - OwnershipEpoch currentEpoch; - - public void Set() - { - TaskCompletionSource available; - lock (this.syncLock) - { - if (this.currentEpoch != null) - { - return; - } - - this.currentEpoch = new OwnershipEpoch(); - available = this.ownershipAvailable; - } - - available.TrySetResult(null); - } - - public void Reset() - { - OwnershipEpoch epoch; - lock (this.syncLock) - { - epoch = this.currentEpoch; - if (epoch == null) - { - return; - } - - this.currentEpoch = null; - this.ownershipAvailable = CreateCompletionSource(); - } - - epoch.Deactivate(); - } - - public async Task WaitAsync(CancellationToken cancellationToken) - { - while (true) - { - Task availableTask; - lock (this.syncLock) - { - AppLeaseOwnership ownership = this.currentEpoch?.TryAcquire(); - if (ownership != null) - { - return ownership; - } - - availableTask = this.ownershipAvailable.Task; - } - - var canceled = new TaskCompletionSource( - TaskCreationOptions.RunContinuationsAsynchronously); - using (cancellationToken.Register(() => canceled.TrySetCanceled())) - { - await Task.WhenAny(availableTask, canceled.Task); - } - - cancellationToken.ThrowIfCancellationRequested(); - } - } - - static TaskCompletionSource CreateCompletionSource() - { - return new TaskCompletionSource( - TaskCreationOptions.RunContinuationsAsynchronously); - } - - internal sealed class OwnershipEpoch - { - readonly object syncLock = new object(); - readonly CancellationTokenSource lostSource = new CancellationTokenSource(); - - bool active = true; - int referenceCount = 1; - - public AppLeaseOwnership TryAcquire() - { - lock (this.syncLock) - { - if (!this.active) - { - return null; - } - - this.referenceCount++; - return new AppLeaseOwnership(this); - } - } - - public bool TryBeginDispatch() - { - lock (this.syncLock) - { - return this.active; - } - } - - public void Deactivate() - { - lock (this.syncLock) - { - if (!this.active) - { - return; - } - - this.active = false; - } - - this.lostSource.Cancel(); - this.Release(); - } - - public void Release() - { - bool dispose; - lock (this.syncLock) - { - this.referenceCount--; - dispose = this.referenceCount == 0; - } - - if (dispose) - { - this.lostSource.Dispose(); - } - } - - public CancellationToken LostToken => this.lostSource.Token; - } - - public sealed class AppLeaseOwnership : IDisposable - { - OwnershipEpoch epoch; - - internal AppLeaseOwnership(OwnershipEpoch epoch) - { - this.epoch = epoch; - } - - public CancellationToken LostToken => this.epoch.LostToken; - - public bool TryBeginDispatch() - { - return this.epoch.TryBeginDispatch(); - } - - public void Dispose() - { - OwnershipEpoch epoch = Interlocked.Exchange(ref this.epoch, null); - epoch?.Release(); - } - } - } -} diff --git a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs index 5d04af120..180e7183d 100644 --- a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs +++ b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -16,6 +16,7 @@ namespace DurableTask.AzureStorage.Tests using System; using System.Collections.Concurrent; using System.Linq; + using System.Reflection; using System.Threading; using System.Threading.Tasks; using DurableTask.AzureStorage.Messaging; @@ -133,35 +134,44 @@ public async Task AppLeaseDisabledAllowsDifferentAppsToDequeueActivities() } [TestMethod] - public async Task OwnershipResetCancelsOldEpochAndBlocksNewWaiters() + public async Task ClosedGatePreventsNewActivityReceive() { - var ownershipSignal = new AppLeaseOwnershipSignal(); - ownershipSignal.Set(); + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService service = + CreateService(taskHubName, "PrimaryApp", useAppLease: true); - using (AppLeaseOwnershipSignal.AppLeaseOwnership ownership = - await ownershipSignal.WaitAsync(CancellationToken.None)) + try { - ownershipSignal.Reset(); - - Assert.IsTrue(ownership.LostToken.IsCancellationRequested); - Assert.IsFalse(ownership.TryBeginDispatch()); + await service.CreateAsync(); + await service.StartAsync(); + await WaitForOwnerAsync(service); + await EnqueueActivityAsync(service, "blocked"); + GetAppLeaseManager(service).SetActivityOwnership(ownsLease: false); - using (var cancellation = new CancellationTokenSource(TimeSpan.FromMilliseconds(100))) + using (var cancellation = new CancellationTokenSource(TimeSpan.FromMilliseconds(250))) { - await Assert.ThrowsExceptionAsync( - () => ownershipSignal.WaitAsync(cancellation.Token)); + TaskActivityWorkItem blockedWorkItem = + await service.LockNextTaskActivityWorkItem(TestTimeout, cancellation.Token); + Assert.IsNull(blockedWorkItem); } + + GetAppLeaseManager(service).SetActivityOwnership(ownsLease: true); + TaskActivityWorkItem recoveredWorkItem = await LockActivityAsync(service); + Assert.IsNotNull(recoveredWorkItem); + await service.AbandonTaskActivityWorkItemAsync(recoveredWorkItem); + } + finally + { + await StopAsync(service); } } [TestMethod] - public async Task OwnershipLossCancelsPendingQueueReceive() + public async Task OwnershipLossDoesNotCancelPendingReceive() { string taskHubName = GetTaskHubName(); AzureStorageOrchestrationService service = CreateService(taskHubName, "PrimaryApp", useAppLease: true); - var ownershipSignal = new AppLeaseOwnershipSignal(); - ownershipSignal.Set(); try { @@ -169,18 +179,59 @@ public async Task OwnershipLossCancelsPendingQueueReceive() await service.StartAsync(); await WaitForOwnerAsync(service); - using (AppLeaseOwnershipSignal.AppLeaseOwnership ownership = - await ownershipSignal.WaitAsync(CancellationToken.None)) - using (var receiveCancellation = - CancellationTokenSource.CreateLinkedTokenSource(ownership.LostToken)) - { - Task pendingReceive = - service.WorkItemQueue.GetMessageAsync(receiveCancellation.Token); + Task pendingReceive = + service.LockNextTaskActivityWorkItem(TestTimeout, CancellationToken.None); + Assert.IsFalse(pendingReceive.IsCompleted); - ownershipSignal.Reset(); + GetAppLeaseManager(service).SetActivityOwnership(ownsLease: false); - Assert.IsNull(await WithTimeoutAsync(pendingReceive)); - } + Task completionAfterLoss = await Task.WhenAny( + pendingReceive, + Task.Delay(TimeSpan.FromMilliseconds(500))); + Assert.AreNotSame( + pendingReceive, + completionAfterLoss, + "Ownership loss alone must not cancel a receive that already started."); + + await EnqueueActivityAsync(service, "rejected-while-closed"); + Assert.IsNull(await WithTimeoutAsync(pendingReceive)); + + GetAppLeaseManager(service).SetActivityOwnership(ownsLease: true); + TaskActivityWorkItem recoveredWorkItem = await LockActivityAsync(service); + Assert.IsNotNull(recoveredWorkItem); + await service.AbandonTaskActivityWorkItemAsync(recoveredWorkItem); + } + finally + { + await StopAsync(service); + } + } + + [TestMethod] + public async Task OwnershipRegainAllowsPendingReceiveToDispatch() + { + string taskHubName = GetTaskHubName(); + AzureStorageOrchestrationService service = + CreateService(taskHubName, "PrimaryApp", useAppLease: true); + + try + { + await service.CreateAsync(); + await service.StartAsync(); + await WaitForOwnerAsync(service); + + Task pendingReceive = + service.LockNextTaskActivityWorkItem(TestTimeout, CancellationToken.None); + Assert.IsFalse(pendingReceive.IsCompleted); + + AppLeaseManager appLeaseManager = GetAppLeaseManager(service); + appLeaseManager.SetActivityOwnership(ownsLease: false); + appLeaseManager.SetActivityOwnership(ownsLease: true); + + await EnqueueActivityAsync(service, "accepted-after-regain"); + TaskActivityWorkItem workItem = await WithTimeoutAsync(pendingReceive); + Assert.IsNotNull(workItem); + await service.AbandonTaskActivityWorkItemAsync(workItem); } finally { @@ -375,6 +426,18 @@ await TestHelpers.WaitFor( TestTimeout); } + static AppLeaseManager GetAppLeaseManager(AzureStorageOrchestrationService service) + { + FieldInfo field = typeof(AzureStorageOrchestrationService).GetField( + "appLeaseManager", + BindingFlags.Instance | BindingFlags.NonPublic); + Assert.IsNotNull(field); + + var appLeaseManager = field.GetValue(service) as AppLeaseManager; + Assert.IsNotNull(appLeaseManager); + return appLeaseManager; + } + static async Task LockActivityAsync( AzureStorageOrchestrationService service) { From bdcec1d36c392e0dcbd4d241a308098155413019 Mon Sep 17 00:00:00 2001 From: wangbill Date: Wed, 9 Sep 2026 15:59:33 -0700 Subject: [PATCH 10/10] Gate activities only before receive Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a197517-3aa5-4206-bc63-dfecd954bd05 --- .../AzureStorageOrchestrationService.cs | 17 +- ...zureStorageOrchestrationServiceSettings.cs | 12 +- .../Partitioning/AppLeaseManager.cs | 11 -- .../AppLeaseActivityTests.cs | 186 ++---------------- 4 files changed, 22 insertions(+), 204 deletions(-) diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs index a4b5143a6..f794f0c2a 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs @@ -198,9 +198,6 @@ public AzureStorageOrchestrationService(AzureStorageOrchestrationServiceSettings internal ITrackingStore TrackingStore => this.trackingStore; - // Intended only for use by tests that need to coordinate the post-dequeue ownership race. - internal Func OnActivityMessageDequeued { get; set; } - internal static string GetControlQueueName(string taskHub, int partitionIndex) { return GetQueueName(taskHub, $"control-{partitionIndex:00}"); @@ -1560,26 +1557,16 @@ public async Task LockNextTaskActivityWorkItem( if (message == null) { - // shutting down or canceled + // shutting down return null; } - Func onActivityMessageDequeued = this.OnActivityMessageDequeued; - if (onActivityMessageDequeued != null) - { - await onActivityMessageDequeued(); - } - - if (!this.appLeaseManager.HasActivityOwnership) - { - await this.workItemQueue.AbandonMessageAsync(message); - return null; - } Guid traceActivityId = Guid.NewGuid(); var session = new ActivitySession(this.settings, this.azureStorageClient.QueueAccountName, message, traceActivityId); session.StartNewLogicalTraceScope(); + // correlation TraceContextBase requestTraceContext = null; CorrelationTraceClient.Propagate( () => diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs index 4a24f6dca..058fa2462 100644 --- a/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs +++ b/src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs @@ -152,13 +152,13 @@ public class AzureStorageOrchestrationServiceSettings public TimeSpan MaxQueuePollingInterval { get; set; } = DefaultMaxQueuePollingInterval; /// - /// If true, workers wait for their to own the app lease before receiving - /// new activity messages. Workers sharing that app name may receive activities concurrently. + /// If true, workers wait for their to own the app lease before starting + /// each new activity receive. Workers sharing that app name may receive activities concurrently. /// Ownership loss does not cancel an activity receive that already started, which may continue - /// polling until it gets a message or caller or service-shutdown cancellation occurs. Current - /// ownership is checked again before dispatch; rejected messages are abandoned before tracing - /// starts. Already dispatched activities are not canceled. Orchestration and entity message - /// processing retain their existing behavior. + /// polling and execute a returned activity after ownership is lost. The next receive waits for + /// ownership. This local gate is cooperative, not an atomic or exactly-once ownership boundary. + /// Already dispatched activities are not canceled. Orchestration and entity message processing + /// retain their existing behavior. /// public bool UseAppLease { get; set; } = true; diff --git a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs index 0113e52c2..b99a3fc27 100644 --- a/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs +++ b/src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs @@ -135,17 +135,6 @@ internal void SetActivityOwnership(bool ownsLease) ownershipAvailable?.TrySetResult(null); } - internal bool HasActivityOwnership - { - get - { - lock (this.activityOwnershipLock) - { - return this.hasActivityOwnership; - } - } - } - static TaskCompletionSource CreateActivityOwnershipSignal() { return new TaskCompletionSource( diff --git a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs index 180e7183d..84f3adf60 100644 --- a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs +++ b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -14,7 +14,6 @@ namespace DurableTask.AzureStorage.Tests { using System; - using System.Collections.Concurrent; using System.Linq; using System.Reflection; using System.Threading; @@ -23,8 +22,6 @@ namespace DurableTask.AzureStorage.Tests using DurableTask.AzureStorage.Partitioning; using DurableTask.Core; using DurableTask.Core.History; - using DurableTask.Core.Settings; - using Microsoft.Extensions.Logging; using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] @@ -133,39 +130,6 @@ public async Task AppLeaseDisabledAllowsDifferentAppsToDequeueActivities() } } - [TestMethod] - public async Task ClosedGatePreventsNewActivityReceive() - { - string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationService service = - CreateService(taskHubName, "PrimaryApp", useAppLease: true); - - try - { - await service.CreateAsync(); - await service.StartAsync(); - await WaitForOwnerAsync(service); - await EnqueueActivityAsync(service, "blocked"); - GetAppLeaseManager(service).SetActivityOwnership(ownsLease: false); - - using (var cancellation = new CancellationTokenSource(TimeSpan.FromMilliseconds(250))) - { - TaskActivityWorkItem blockedWorkItem = - await service.LockNextTaskActivityWorkItem(TestTimeout, cancellation.Token); - Assert.IsNull(blockedWorkItem); - } - - GetAppLeaseManager(service).SetActivityOwnership(ownsLease: true); - TaskActivityWorkItem recoveredWorkItem = await LockActivityAsync(service); - Assert.IsNotNull(recoveredWorkItem); - await service.AbandonTaskActivityWorkItemAsync(recoveredWorkItem); - } - finally - { - await StopAsync(service); - } - } - [TestMethod] public async Task OwnershipLossDoesNotCancelPendingReceive() { @@ -193,105 +157,29 @@ public async Task OwnershipLossDoesNotCancelPendingReceive() completionAfterLoss, "Ownership loss alone must not cancel a receive that already started."); - await EnqueueActivityAsync(service, "rejected-while-closed"); - Assert.IsNull(await WithTimeoutAsync(pendingReceive)); - - GetAppLeaseManager(service).SetActivityOwnership(ownsLease: true); - TaskActivityWorkItem recoveredWorkItem = await LockActivityAsync(service); - Assert.IsNotNull(recoveredWorkItem); - await service.AbandonTaskActivityWorkItemAsync(recoveredWorkItem); - } - finally - { - await StopAsync(service); - } - } - - [TestMethod] - public async Task OwnershipRegainAllowsPendingReceiveToDispatch() - { - string taskHubName = GetTaskHubName(); - AzureStorageOrchestrationService service = - CreateService(taskHubName, "PrimaryApp", useAppLease: true); - - try - { - await service.CreateAsync(); - await service.StartAsync(); - await WaitForOwnerAsync(service); - - Task pendingReceive = - service.LockNextTaskActivityWorkItem(TestTimeout, CancellationToken.None); - Assert.IsFalse(pendingReceive.IsCompleted); - - AppLeaseManager appLeaseManager = GetAppLeaseManager(service); - appLeaseManager.SetActivityOwnership(ownsLease: false); - appLeaseManager.SetActivityOwnership(ownsLease: true); - - await EnqueueActivityAsync(service, "accepted-after-regain"); + await EnqueueActivityAsync(service, "received-after-loss"); TaskActivityWorkItem workItem = await WithTimeoutAsync(pendingReceive); - Assert.IsNotNull(workItem); + Assert.IsNotNull( + workItem, + "A receive that started while owned may admit a message after ownership is lost."); await service.AbandonTaskActivityWorkItemAsync(workItem); - } - finally - { - await StopAsync(service); - } - } - [TestMethod] - public async Task OwnershipLossAfterDequeueAbandonsWithoutStartingTrace() - { - string taskHubName = GetTaskHubName(); - var loggerFactory = new RecordingLoggerFactory(); - AzureStorageOrchestrationServiceSettings ownerSettings = - CreateSettings(taskHubName, "PrimaryApp", useAppLease: true); - ownerSettings.LoggerFactory = loggerFactory; - AzureStorageOrchestrationService owner = - new AzureStorageOrchestrationService(ownerSettings); - AzureStorageOrchestrationService recoveryReader = - CreateService(taskHubName, "RecoveryReader", useAppLease: false); - CorrelationSettings previousCorrelationSettings = CorrelationSettings.Current; - - try - { - CorrelationSettings.Current = new CorrelationSettings - { - EnableDistributedTracing = true, - Protocol = Protocol.W3CTraceContext, - }; - CorrelationTraceContext.Current = null; - - await owner.CreateAsync(); - await owner.StartAsync(); - await WaitForOwnerAsync(owner); - await EnqueueActivityAsync(owner, "ownership-race"); - - owner.OnActivityMessageDequeued = async () => + await EnqueueActivityAsync(service, "blocked-next-receive"); + using (var cancellation = new CancellationTokenSource(TimeSpan.FromMilliseconds(250))) { - owner.OnActivityMessageDequeued = null; - await StopAsync(owner); - }; - - TaskActivityWorkItem rejectedWorkItem = await LockActivityAsync(owner); - owner = null; - - Assert.IsNull(rejectedWorkItem); - Assert.IsNull(CorrelationTraceContext.Current); - Assert.IsFalse(loggerFactory.HasEvent("ReceivedMessage")); - Assert.IsFalse(loggerFactory.HasEvent("ProcessingMessage")); + TaskActivityWorkItem blockedWorkItem = + await service.LockNextTaskActivityWorkItem(TestTimeout, cancellation.Token); + Assert.IsNull(blockedWorkItem, "The next receive must wait while ownership is closed."); + } - await recoveryReader.StartAsync(); - TaskActivityWorkItem recoveredWorkItem = await LockActivityAsync(recoveryReader); + GetAppLeaseManager(service).SetActivityOwnership(ownsLease: true); + TaskActivityWorkItem recoveredWorkItem = await LockActivityAsync(service); Assert.IsNotNull(recoveredWorkItem); - await recoveryReader.AbandonTaskActivityWorkItemAsync(recoveredWorkItem); + await service.AbandonTaskActivityWorkItemAsync(recoveredWorkItem); } finally { - CorrelationTraceContext.Current = null; - CorrelationSettings.Current = previousCorrelationSettings; - await StopAsync(recoveryReader); - await StopAsync(owner); + await StopAsync(service); } } @@ -467,51 +355,5 @@ static async Task StopAsync(AzureStorageOrchestrationService service) } } - sealed class RecordingLoggerFactory : ILoggerFactory - { - readonly ConcurrentQueue events = new ConcurrentQueue(); - - public void AddProvider(ILoggerProvider provider) - { - } - - public ILogger CreateLogger(string categoryName) - { - return new RecordingLogger(this.events); - } - - public void Dispose() - { - } - - public bool HasEvent(string name) - { - return this.events.Any(e => e.Name == name); - } - - sealed class RecordingLogger : ILogger - { - readonly ConcurrentQueue events; - - public RecordingLogger(ConcurrentQueue events) - { - this.events = events; - } - - public IDisposable BeginScope(TState state) => null; - - public bool IsEnabled(LogLevel logLevel) => true; - - public void Log( - LogLevel logLevel, - EventId eventId, - TState state, - Exception exception, - Func formatter) - { - this.events.Enqueue(eventId); - } - } - } } }