Azure Storage lease management correctness fixes (#724)

- Fix misleading "lease lost" log message during shutdown - Make lease release serial instead of parallel for better determinism - Cleanup subscriptions during shutdown that were causing unexpected lease management operations
Azure · May 5, 2022 · d215370 · d215370
1 parent 665119b
commit d215370
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 20 deletions.
diff --git a/src/DurableTask.AzureStorage/Messaging/ControlQueue.cs b/src/DurableTask.AzureStorage/Messaging/ControlQueue.cs
@@ -214,16 +214,19 @@ public override Task DeleteMessageAsync(MessageData message, SessionBase? sessio
 
         public void Release(CloseReason? reason, string caller)
         {
-            this.releaseTokenSource.Cancel();
+            if (!this.IsReleased)
+            {
+                this.releaseTokenSource.Cancel();
 
-            this.IsReleased = true;
+                this.IsReleased = true;
 
-            this.settings.Logger.PartitionManagerInfo(
-                this.storageAccountName,
-                this.settings.TaskHubName,
-                this.settings.WorkerId,
-                this.Name,
-                $"{caller} is releasing partition {this.Name} for reason: {reason}");
+                this.settings.Logger.PartitionManagerInfo(
+                    this.storageAccountName,
+                    this.settings.TaskHubName,
+                    this.settings.WorkerId,
+                    this.Name,
+                    $"{caller} is releasing partition {this.Name} for reason: {reason}");
+            }
         }
 
         public virtual void Dispose()

diff --git a/src/DurableTask.AzureStorage/Partitioning/LeaseCollectionBalancer.cs b/src/DurableTask.AzureStorage/Partitioning/LeaseCollectionBalancer.cs
@@ -340,7 +340,7 @@ async Task<IDictionary<string, T>> TakeLeasesAsync()
                         this.accountName,
                         this.taskHub,
                         this.workerName,
-                        string.Empty /* partitionId */,
+                        lease.PartitionId,
                         $"Skipping {this.leaseType} lease aquiring for {lease.PartitionId}");
                     continue;
                 }

diff --git a/src/DurableTask.AzureStorage/Partitioning/SafePartitionManager.cs b/src/DurableTask.AzureStorage/Partitioning/SafePartitionManager.cs
@@ -33,6 +33,9 @@ class SafePartitionManager : IPartitionManager
         readonly BlobLeaseManager ownershipLeaseManager;
         readonly LeaseCollectionBalancer<BlobLease> ownershipLeaseCollectionManager;
 
+        IDisposable intentLeaseSubscription;
+        IDisposable ownershipLeaseSubscription;
+
         public SafePartitionManager(
             AzureStorageOrchestrationService service,
             AzureStorageClient azureStorageClient,
@@ -122,24 +125,25 @@ Task IPartitionManager.DeleteLeases()
         async Task IPartitionManager.StartAsync()
         {
             await this.intentLeaseCollectionManager.InitializeAsync();
-            await this.intentLeaseCollectionManager.SubscribeAsync(
-                    this.service.OnIntentLeaseAquiredAsync,
-                    this.service.OnIntentLeaseReleasedAsync);
+            this.intentLeaseSubscription = await this.intentLeaseCollectionManager.SubscribeAsync(
+                this.service.OnIntentLeaseAquiredAsync,
+                this.service.OnIntentLeaseReleasedAsync);
             await this.intentLeaseCollectionManager.StartAsync();
 
             await this.ownershipLeaseCollectionManager.InitializeAsync();
-            await this.ownershipLeaseCollectionManager.SubscribeAsync(
-                    this.service.OnOwnershipLeaseAquiredAsync,
-                    this.service.OnOwnershipLeaseReleasedAsync);
+            this.ownershipLeaseSubscription = await this.ownershipLeaseCollectionManager.SubscribeAsync(
+                this.service.OnOwnershipLeaseAquiredAsync,
+                this.service.OnOwnershipLeaseReleasedAsync);
             await this.ownershipLeaseCollectionManager.StartAsync();
         }
 
-        Task IPartitionManager.StopAsync()
+        async Task IPartitionManager.StopAsync()
         {
-            return Task.WhenAll(
-                this.intentLeaseCollectionManager.StopAsync(),
-                this.ownershipLeaseCollectionManager.StopAsync()
-            );
+            await this.intentLeaseCollectionManager.StopAsync();
+            this.intentLeaseSubscription?.Dispose();
+
+            await this.ownershipLeaseCollectionManager.StopAsync();
+            this.ownershipLeaseSubscription?.Dispose();
         }
 
         Task IPartitionManager.CreateLease(string leaseName)