Gate activity polling on app lease ownership

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -1544,6 +1544,15 @@ public async Task<TaskActivityWorkItem> LockNextTaskActivityWorkItem(
 
             using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, this.shutdownSource.Token))
             {
+                try
+                {
+                    await this.appLeaseManager.WaitForActivityOwnershipAsync(linkedCts.Token);
+                }
+                catch (OperationCanceledException)
+                {
+                    return null;
+                }
+
                 MessageData message = await this.workItemQueue.GetMessageAsync(linkedCts.Token);
 
                 if (message == null)

src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs

@@ -152,7 +152,13 @@ public class AzureStorageOrchestrationServiceSettings
         public TimeSpan MaxQueuePollingInterval { get; set; } = DefaultMaxQueuePollingInterval;
 
         /// <summary>
-        /// 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, workers wait for their <see cref="AppName"/> 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 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.
         /// </summary>
         public bool UseAppLease { get; set; } = true;
 

src/DurableTask.AzureStorage/Partitioning/AppLeaseManager.cs

@@ -42,7 +42,10 @@ sealed class AppLeaseManager
         readonly Blob appLeaseInfoBlob;
         readonly string appLeaseId;
         readonly AsyncManualResetEvent shutdownCompletedEvent;
+        readonly object activityOwnershipLock = new object();
 
+        TaskCompletionSource<object> activityOwnershipAvailable;
+        bool hasActivityOwnership;
         bool isLeaseOwner;
         int appLeaseIsStarted;
         Task renewTask;
@@ -78,12 +81,71 @@ public AppLeaseManager(
 
             this.isLeaseOwner = false;
             this.shutdownCompletedEvent = new AsyncManualResetEvent();
+            this.activityOwnershipAvailable = CreateActivityOwnershipSignal();
+        }
+
+        public async Task WaitForActivityOwnershipAsync(
+            CancellationToken cancellationToken)
+        {
+            while (true)
+            {
+                Task ownershipAvailableTask;
+                lock (this.activityOwnershipLock)
+                {
+                    if (this.hasActivityOwnership)
+                    {
+                        return;
+                    }
+
+                    ownershipAvailableTask = this.activityOwnershipAvailable.Task;
+                }
+
+                var canceled = new TaskCompletionSource<object>(
+                    TaskCreationOptions.RunContinuationsAsynchronously);
+                using (cancellationToken.Register(() => canceled.TrySetResult(null)))
+                {
+                    await Task.WhenAny(ownershipAvailableTask, canceled.Task);
+                }
+
+                cancellationToken.ThrowIfCancellationRequested();
+            }
+        }
+
+        internal void SetActivityOwnership(bool ownsLease)
+        {
+            TaskCompletionSource<object> 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);
+        }
+
+        static TaskCompletionSource<object> CreateActivityOwnershipSignal()
+        {
+            return new TaskCompletionSource<object>(
+                TaskCreationOptions.RunContinuationsAsynchronously);
         }
 
         public async Task StartAsync()
         {
             if (!this.appLeaseIsEnabled)
             {
+                this.SetActivityOwnership(ownsLease: true);
                 this.starterTokenSource = new CancellationTokenSource();
 
                 await Task.Factory.StartNew(() => this.PartitionManagerStarter(this.starterTokenSource.Token));
@@ -165,6 +227,8 @@ async Task AppLeaseManagerStarter(CancellationToken cancellationToken)
 
         public async Task StopAsync()
         {
+            this.SetActivityOwnership(ownsLease: false);
+
             if (this.starterTokenSource != null)
             {
                 this.starterTokenSource.Cancel();
@@ -253,6 +317,7 @@ async Task StartAppLeaseAsync()
             this.leaseRenewerCancellationTokenSource = new CancellationTokenSource();
 
             await this.partitionManager.StartAsync();
+            this.SetActivityOwnership(ownsLease: true);
 
             this.shutdownCompletedEvent.Reset();
 
@@ -267,6 +332,8 @@ async Task StopAppLeaseAsync()
                 return;
             }
 
+            this.SetActivityOwnership(ownsLease: false);
+
             await this.partitionManager.StopAsync();
 
             if (this.renewTask != null)
@@ -491,6 +558,7 @@ async Task<bool> RenewLeaseAsync()
                 {
                     renewed = false;
                     this.isLeaseOwner = false;
+                    this.SetActivityOwnership(ownsLease: false);
 
                     this.settings.Logger.LeaseRenewalFailed(
                         this.storageAccountName,