reset pending orchestrations when worker restart

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -1249,7 +1249,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 +1530,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

@@ -19,10 +19,12 @@ namespace DurableTask.AzureStorage.Messaging
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
+    using Azure;
     using Azure.Storage.Queues.Models;
     using DurableTask.AzureStorage.Monitoring;
     using DurableTask.AzureStorage.Partitioning;
     using DurableTask.AzureStorage.Storage;
+    using DurableTask.Core;
 
     class ControlQueue : TaskHubQueue, IDisposable
     {
@@ -209,6 +211,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 (RequestFailedException 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 _);