diff --git a/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs b/src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs index 74798cb45..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}"); @@ -1544,55 +1547,83 @@ 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); + Func onActivityMessageDequeued = this.OnActivityMessageDequeued; + if (onActivityMessageDequeued != null) + { + await onActivityMessageDequeued(); + } - TraceContextBase parentTraceContextBase = TraceContextBase.Restore(session.MessageData.SerializableTraceContext); - requestTraceContext.SetParentAndStart(parentTraceContextBase); - }); + if (!ownership.TryBeginDispatch()) + { + await this.workItemQueue.AbandonMessageAsync(message); + return null; + } - TraceMessageReceived(this.settings, session.MessageData, this.azureStorageClient.QueueAccountName); - session.TraceProcessingMessage(message, isExtendedSession: false, this.workItemQueue.Name); + Guid traceActivityId = Guid.NewGuid(); + var session = new ActivitySession(this.settings, this.azureStorageClient.QueueAccountName, message, traceActivityId); + session.StartNewLogicalTraceScope(); - 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 = null; + CorrelationTraceClient.Propagate( + () => + { + string name = $"{TraceConstants.Activity} {Utils.GetTargetClassName(((TaskScheduledEvent)session.MessageData.TaskMessage.Event)?.Name)}"; + requestTraceContext = TraceContextFactory.Create(name); - this.stats.ActiveActivityExecutions.Increment(); + TraceContextBase parentTraceContextBase = TraceContextBase.Restore(session.MessageData.SerializableTraceContext); + requestTraceContext.SetParentAndStart(parentTraceContextBase); + }); - return new TaskActivityWorkItem - { - Id = message.Id, - TaskMessage = session.MessageData.TaskMessage, - LockedUntilUtc = message.OriginalQueueMessage.NextVisibleOn.Value.UtcDateTime, + TraceMessageReceived(this.settings, session.MessageData, this.azureStorageClient.QueueAccountName); + session.TraceProcessingMessage(message, isExtendedSession: false, this.workItemQueue.Name); - TraceContextBase = requestTraceContext - }; + 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(); + + 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..5327df1bc 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,111 @@ 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) + { + if (renewerCancellation == null) + { + throw new InvalidOperationException( + "The app lease renewer has no cancellation source."); + } - this.shutdownCompletedEvent.Set(); + renewerCancellation.Cancel(); + await renewer; + } - this.leaseRenewerCancellationTokenSource?.Dispose(); + await this.partitionManager.StopAsync(); + } + finally + { + this.isLeaseOwner = false; + + renewerCancellation?.Dispose(); + this.leaseRenewerCancellationTokenSource = null; + this.renewTask = null; + + 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) + if (appLeaseInfo.DesiredSwapId == this.appLeaseId + && !string.IsNullOrEmpty(appLeaseInfo.OwnerId)) { - 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 +387,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 +407,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 +490,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 +521,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 +568,7 @@ async Task RenewLeaseAsync() { renewed = false; this.isLeaseOwner = false; + this.ownershipSignal.Reset(); this.settings.Logger.LeaseRenewalFailed( this.storageAccountName, @@ -577,6 +655,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(); + } + } + } +} diff --git a/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs b/src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs index 3fad8c98d..6fee697ef 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(); + } } @@ -145,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) { @@ -197,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 @@ -229,36 +272,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; + } + + 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); + } - // 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); + // Surface any unhandled exceptions + await this.partitionManagerTask; - if (timeoutTask.IsCompleted) + 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() @@ -424,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. @@ -498,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; @@ -895,13 +988,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(); } } } diff --git a/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs new file mode 100644 index 000000000..3b8619cfb --- /dev/null +++ b/test/DurableTask.AzureStorage.Tests/AppLeaseActivityTests.cs @@ -0,0 +1,1025 @@ +// ---------------------------------------------------------------------------------- +// 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.Concurrent; + 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.Extensions.Logging; + 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()); + try + { + await oldManager.CreateContainerIfNotExistsAsync(); + await oldManager.StartAsync(); + using (AppLeaseOwnershipSignal.AppLeaseOwnership oldOwnership = + await GetOwnershipAsync(oldManager)) + { + await newManager.StartAsync(); + + await newManager.ForceChangeAppLeaseAsync(); + using (AppLeaseOwnershipSignal.AppLeaseOwnership newOwnership = + await GetOwnershipAsync(newManager)) + { + Assert.IsTrue(oldOwnership.LostToken.IsCancellationRequested); + Assert.IsFalse(oldOwnership.TryBeginDispatch()); + } + } + } + finally + { + 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 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 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 ?? ""}."); + 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); + } + } + } + + [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 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() + { + 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 OwnershipLossAfterDequeueDoesNotStartProcessingTrace() + { + 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); + + 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); + }; + + TaskActivityWorkItem rejectedWorkItem = await LockActivityAsync(oldOwner); + + Assert.IsNull(rejectedWorkItem); + Assert.IsFalse(loggerFactory.HasEvent("ProcessingMessage")); + + TaskActivityWorkItem newOwnerWorkItem = await LockActivityAsync(newOwner); + Assert.IsNotNull(newOwnerWorkItem); + await newOwner.AbandonTaskActivityWorkItemAsync(newOwnerWorkItem); + } + finally + { + await StopAsync(newOwner); + await StopAsync(oldOwner); + } + } + + [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 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) + { + 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 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; + } + } + + 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; + + 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 7fdb42e04..f688bec58 100644 --- a/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs +++ b/test/DurableTask.AzureStorage.Tests/TestTablePartitionManager.cs @@ -726,6 +726,127 @@ 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 {