reset pending orchestrations when worker restart

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -553,7 +553,7 @@ internal async Task OnIntentLeaseAquiredAsync(BlobPartitionLease lease)
         {
             var controlQueue = new ControlQueue(this.azureStorageClient, lease.PartitionId, this.messageManager);
             await controlQueue.CreateIfNotExistsAsync();
-            this.orchestrationSessionManager.ResumeListeningIfOwnQueue(lease.PartitionId, controlQueue, this.shutdownSource.Token);
+            await this.orchestrationSessionManager.ResumeListeningIfOwnQueue(lease.PartitionId, controlQueue, this.shutdownSource.Token);
         }
 
         internal Task OnIntentLeaseReleasedAsync(BlobPartitionLease lease, CloseReason reason)
@@ -572,21 +572,20 @@ internal async Task OnOwnershipLeaseAquiredAsync(BlobPartitionLease lease)
             this.allControlQueues[lease.PartitionId] = controlQueue;
         }
 
-        internal void DropLostControlQueue(TablePartitionLease partition)
+        internal async Task DropLostControlQueue(TablePartitionLease partition)
         {
             // If lease is lost but we're still dequeuing messages, remove the queue
             if (this.allControlQueues.TryGetValue(partition.RowKey, out ControlQueue controlQueue) &&
                 this.OwnedControlQueues.Contains(controlQueue) &&
                 partition.CurrentOwner != this.settings.WorkerId)
             {
-                this.orchestrationSessionManager.RemoveQueue(partition.RowKey, CloseReason.LeaseLost, nameof(DropLostControlQueue));
+                await this.orchestrationSessionManager.RemoveQueue(partition.RowKey, CloseReason.LeaseLost, nameof(DropLostControlQueue));
             }
         }
 
         internal Task OnOwnershipLeaseReleasedAsync(BlobPartitionLease lease, CloseReason reason)
         {
-            this.orchestrationSessionManager.RemoveQueue(lease.PartitionId, reason, "Ownership LeaseCollectionBalancer");
-            return Utils.CompletedTask;
+            return this.orchestrationSessionManager.RemoveQueue(lease.PartitionId, reason, "Ownership LeaseCollectionBalancer");
         }
 
         internal async Task OnTableLeaseAcquiredAsync(TablePartitionLease lease)
@@ -1249,7 +1248,13 @@ await this.CommitOutboundQueueMessages(
                 {
                     var messages = session.DeferredMessages.ToList();
                     session.DeferredMessages.Clear();
-                    this.orchestrationSessionManager.AddMessageToPendingOrchestration(session.ControlQueue, messages, session.TraceActivityId, CancellationToken.None);
+                    IReadOnlyList<MessageData> messagesToAbandon = this.orchestrationSessionManager.AddMessageToPendingOrchestration(
+                        session.ControlQueue,
+                        messages,
+                        session.TraceActivityId,
+                        CancellationToken.None);
+
+                    await this.orchestrationSessionManager.AbandonMessagesForDrainAsync(session.ControlQueue, messagesToAbandon);
                 }
             }
             // Handle the case where the 'ETag' has changed, which implies another worker has taken over this work item while
@@ -1524,8 +1529,11 @@ async Task ReleaseSessionAsync(string instanceId)
             if (this.orchestrationSessionManager.TryReleaseSession(
                 instanceId,
                 this.shutdownSource.Token,
-                out OrchestrationSession session))
+                out OrchestrationSession session,
+                out IReadOnlyList<MessageData> messagesToAbandon))
             {
+                await this.orchestrationSessionManager.AbandonMessagesForDrainAsync(session.ControlQueue, messagesToAbandon);
+
                 // Some messages may need to be discarded
                 await session.DiscardedMessages.ParallelForEachAsync(
                     this.settings.MaxStorageOperationConcurrency,

src/DurableTask.AzureStorage/Messaging/ControlQueue.cs

@@ -23,6 +23,7 @@ namespace DurableTask.AzureStorage.Messaging
     using DurableTask.AzureStorage.Monitoring;
     using DurableTask.AzureStorage.Partitioning;
     using DurableTask.AzureStorage.Storage;
+    using DurableTask.Core;
 
     class ControlQueue : TaskHubQueue, IDisposable
     {
@@ -209,6 +210,50 @@ public override Task AbandonMessageAsync(MessageData message, SessionBase? sessi
             return base.AbandonMessageAsync(message, session);
         }
 
+        /// <summary>
+        /// Abandons a message with zero visibility timeout so it becomes immediately visible
+        /// for another partition owner to pick up. This is used during drain to avoid stranding
+        /// messages that were dequeued but not yet promoted to active sessions.
+        /// </summary>
+        public async Task AbandonMessageForDrainAsync(MessageData message)
+        {
+            this.stats.PendingOrchestratorMessages.TryRemove(message.OriginalQueueMessage.MessageId, out _);
+
+            QueueMessage queueMessage = message.OriginalQueueMessage;
+            TaskMessage taskMessage = message.TaskMessage;
+            OrchestrationInstance instance = taskMessage.OrchestrationInstance;
+
+            this.settings.Logger.AbandoningMessage(
+                this.storageAccountName,
+                this.settings.TaskHubName,
+                taskMessage.Event.EventType.ToString(),
+                Utils.GetTaskEventId(taskMessage.Event),
+                queueMessage.MessageId,
+                instance.InstanceId,
+                instance.ExecutionId,
+                this.storageQueue.Name,
+                message.SequenceNumber,
+                queueMessage.PopReceipt,
+                visibilityTimeoutSeconds: 0);
+
+            try
+            {
+                await this.storageQueue.UpdateMessageAsync(
+                    queueMessage,
+                    TimeSpan.Zero,
+                    clientRequestId: null);
+            }
+            catch (DurableTaskStorageException e)
+            {
+                this.settings.Logger.PartitionManagerWarning(
+                    this.storageAccountName,
+                    this.settings.TaskHubName,
+                    this.settings.WorkerId,
+                    this.Name,
+                    $"Failed to abandon message {queueMessage.MessageId} during drain: {e.Message}");
+            }
+        }
+
         public override Task DeleteMessageAsync(MessageData message, SessionBase? session = null)
         {
             this.stats.PendingOrchestratorMessages.TryRemove(message.OriginalQueueMessage.MessageId, out _);