Fix renew/delete race in DT.AzureStorage (#1197)

Author: cgillum

src/DurableTask.AzureStorage/Messaging/TaskHubQueue.cs

@@ -332,8 +332,14 @@ await this.storageQueue.UpdateMessageAsync(
             }
             catch (Exception e)
             {
+                string details = $"Caller: {nameof(RenewMessageAsync)}";
+                if (e is DurableTaskStorageException storageException && storageException.ErrorCode != null)
+                {
+                    details += $", ErrorCode: {storageException.ErrorCode}";
+                }
+
                 // Message may have been processed and deleted already.
-                this.HandleMessagingExceptions(e, message, $"Caller: {nameof(RenewMessageAsync)}");
+                this.HandleMessagingExceptions(e, message, details);
             }
         }
 
@@ -342,46 +348,64 @@ public virtual async Task DeleteMessageAsync(MessageData message, SessionBase? s
             QueueMessage queueMessage = message.OriginalQueueMessage;
             TaskMessage taskMessage = message.TaskMessage;
 
-            this.settings.Logger.DeletingMessage(
-                this.storageAccountName,
-                this.settings.TaskHubName,
-                taskMessage.Event.EventType.ToString(),
-                Utils.GetTaskEventId(taskMessage.Event),
-                queueMessage.MessageId,
-                taskMessage.OrchestrationInstance.InstanceId,
-                taskMessage.OrchestrationInstance.ExecutionId,
-                this.storageQueue.Name,
-                message.SequenceNumber,
-                queueMessage.PopReceipt);
-
             bool haveRetried = false;
             while (true)
             {
+                this.settings.Logger.DeletingMessage(
+                    this.storageAccountName,
+                    this.settings.TaskHubName,
+                    taskMessage.Event.EventType.ToString(),
+                    Utils.GetTaskEventId(taskMessage.Event),
+                    queueMessage.MessageId,
+                    taskMessage.OrchestrationInstance.InstanceId,
+                    taskMessage.OrchestrationInstance.ExecutionId,
+                    this.storageQueue.Name,
+                    message.SequenceNumber,
+                    queueMessage.PopReceipt);
+
                 try
                 {
                     await this.storageQueue.DeleteMessageAsync(queueMessage, session?.TraceActivityId);
                 }
                 catch (Exception e)
                 {
-                    if (!haveRetried && this.IsMessageGoneException(e))
+                    // Delete operations can transiently fail if a delete operation races with a
+                    // message update operation. In this case, we retry the delete operation.
+                    if (!haveRetried && (IsMessageGoneException(e) || IsPopReceiptMismatch(e)))
                     {
                         haveRetried = true;
                         continue;
                     }
 
-                    this.HandleMessagingExceptions(e, message, $"Caller: {nameof(DeleteMessageAsync)}");
+                    string details = $"Caller: {nameof(DeleteMessageAsync)}";
+                    if (e is DurableTaskStorageException storageException && storageException.ErrorCode != null)
+                    {
+                        details += $", ErrorCode: {storageException.ErrorCode}";
+                    }
+
+                    this.HandleMessagingExceptions(e, message, details);
                 }
 
                 break;
             }
         }
 
-        private bool IsMessageGoneException(Exception e)
+        static bool IsMessageGoneException(Exception e)
         {
             DurableTaskStorageException? storageException = e as DurableTaskStorageException;
             return storageException?.HttpStatusCode == 404;
         }
 
+        static bool IsPopReceiptMismatch(Exception e)
+        {
+            if (e is DurableTaskStorageException storageException)
+            {
+                return storageException.IsPopReceiptMismatch;
+            }
+
+            return false;
+        }
+
         void HandleMessagingExceptions(Exception e, MessageData message, string details)
         {
             string messageId = message.OriginalQueueMessage.MessageId;
@@ -403,7 +427,7 @@ void HandleMessagingExceptions(
             string details,
             string popReceipt)
         {
-            if (this.IsMessageGoneException(e))
+            if (IsMessageGoneException(e))
             {
                 // Message may have been processed and deleted already.
                 this.settings.Logger.MessageGone(

src/DurableTask.AzureStorage/Storage/DurableTaskStorageException.cs

@@ -16,6 +16,7 @@ namespace DurableTask.AzureStorage.Storage
     using System;
     using Azure;
     using Azure.Storage.Blobs.Models;
+    using Azure.Storage.Queues.Models;
 
     [Serializable]
     class DurableTaskStorageException : Exception
@@ -40,12 +41,18 @@ public DurableTaskStorageException(RequestFailedException? requestFailedExceptio
             if (requestFailedException != null)
             {
                 this.HttpStatusCode = requestFailedException.Status;
-                this.LeaseLost = requestFailedException.ErrorCode != null && requestFailedException.ErrorCode == BlobErrorCode.LeaseLost;
+                this.ErrorCode = requestFailedException.ErrorCode;
+                this.LeaseLost = requestFailedException?.ErrorCode == BlobErrorCode.LeaseLost;
+                this.IsPopReceiptMismatch = requestFailedException?.ErrorCode == QueueErrorCode.PopReceiptMismatch;
             }
         }
 
         public int HttpStatusCode { get; }
 
+        public string? ErrorCode { get; }
+
         public bool LeaseLost { get; }
+
+        public bool IsPopReceiptMismatch { get; }
     }
 }

src/DurableTask.AzureStorage/Storage/Queue.cs

@@ -83,6 +83,8 @@ await this.queueClient
                     queueMessage.PopReceipt,
                     cancellationToken)
                 .DecorateFailure();
+
+            this.stats.MessagesUpdated.Increment();
         }
 
         public async Task<QueueMessage?> GetMessageAsync(TimeSpan visibilityTimeout, CancellationToken cancellationToken = default)