Add simple poison message handling for Azure Storage

src/DurableTask.AzureStorage/AzureStorageOrchestrationServiceSettings.cs

@@ -117,6 +117,11 @@ public class AzureStorageOrchestrationServiceSettings
         /// </summary>
         public int MaxConcurrentTaskEntityWorkItems { get; set; } = 100;
 
+        /// <summary>
+        /// Gets or sets the maximum dequeue count of any message before it is flagged as a "poison message".
+        /// The default value is 20.
+        /// </summary>
+        public int PoisonMessageDeuqueCountThreshold { get; set; } = 20;
         /// <summary>
         /// Gets or sets the maximum number of concurrent storage operations that can be executed in the context
         /// of a single orchestration instance.

src/DurableTask.AzureStorage/EntityTrackingStoreQueries.cs

@@ -235,9 +235,10 @@ bool OrchestrationIsRunning(OrchestrationStatus? status)
             {
                 // first, retrieve the entity scheduler state (= input of the orchestration state), possibly from blob storage.
                 string serializedSchedulerState;
-                if (MessageManager.TryGetLargeMessageReference(state.Input, out Uri blobUrl))
+                if (MessageManager.TryGetLargeMessageReference(state.Input, out Uri? blobUrl))
                 {
-                    serializedSchedulerState = await this.messageManager.DownloadAndDecompressAsBytesAsync(blobUrl);
+                    // we know blobUrl is not null because TryGetLargeMessageReference returned true
+                    serializedSchedulerState = await this.messageManager.DownloadAndDecompressAsBytesAsync(blobUrl!);
                 }
                 else
                 {

src/DurableTask.AzureStorage/MessageManager.cs

@@ -10,7 +10,7 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 //  ----------------------------------------------------------------------------------
-
+#nullable enable
 namespace DurableTask.AzureStorage
 {
     using System;
@@ -118,17 +118,18 @@ public async Task<string> SerializeMessageDataAsync(MessageData messageData)
         /// <returns>Actual string representation of message.</returns>
         public async Task<string> FetchLargeMessageIfNecessary(string message)
         {
-            if (TryGetLargeMessageReference(message, out Uri blobUrl))
+            if (TryGetLargeMessageReference(message, out Uri? blobUrl))
             {
-                return await this.DownloadAndDecompressAsBytesAsync(blobUrl);
+                // we know blobUrl is not null because TryGetLargeMessageReference returned true
+                return await this.DownloadAndDecompressAsBytesAsync(blobUrl!);
             }
             else
             {
                 return message;
             }
         }
 
-        internal static bool TryGetLargeMessageReference(string messagePayload, out Uri blobUrl)
+        internal static bool TryGetLargeMessageReference(string messagePayload, out Uri? blobUrl)
         {
             if (Uri.IsWellFormedUriString(messagePayload, UriKind.Absolute))
             {
@@ -314,6 +315,7 @@ public async Task<int> DeleteLargeMessageBlobs(string sanitizedInstanceId)
             return storageOperationCount;
         }
     }
+#nullable disable
 
 #if NETSTANDARD2_0
     class TypeNameSerializationBinder : ISerializationBinder

src/DurableTask.AzureStorage/Messaging/ControlQueue.cs

@@ -105,27 +105,27 @@ await batch.ParallelForEachAsync(async delegate (QueueMessage queueMessage)
                             MessageData messageData;
                             try
                             {
+                                // try to de-serialize message
                                 messageData = await this.messageManager.DeserializeQueueMessageAsync(
                                     queueMessage,
                                     this.storageQueue.Name);
+
+                                // if successful, check if it's a poison message. If so, we handle it
+                                // and log metadata about it as the de-serialization succeeded.
+                                await this.HandleIfPoisonMessageAsync(messageData);
                             }
-                            catch (Exception e)
+                            catch (Exception exception)
                             {
-                                // We have limited information about the details of the message
-                                // since we failed to deserialize it.
-                                this.settings.Logger.MessageFailure(
-                                    this.storageAccountName,
-                                    this.settings.TaskHubName,
-                                    queueMessage.Id /* MessageId */,
-                                    string.Empty /* InstanceId */,
-                                    string.Empty /* ExecutionId */,
-                                    this.storageQueue.Name,
-                                    string.Empty /* EventType */,
-                                    0 /* TaskEventId */,
-                                    e.ToString());
-
-                                // Abandon the message so we can try it again later.
-                                await this.AbandonMessageAsync(queueMessage);
+                                // Deserialization errors can be persistent, so we check if this is a poison message.
+                                bool isPoisonMessage = await this.TryHandlingDeserializationPoisonMessage(queueMessage, exception);
+                                if (isPoisonMessage)
+                                {
+                                    // we have already handled the poison message, so we move on.
+                                    return;
+                                }
+
+                                // This is not a poison message (at least not yet), so we abandon it to retry later.
+                                await this.AbandonMessageAsync(queueMessage, exception);
                                 return;
                             }
 
@@ -190,8 +190,22 @@ await batch.ParallelForEachAsync(async delegate (QueueMessage queueMessage)
         }
 
         // This overload is intended for cases where we aren't able to deserialize an instance of MessageData.
-        public Task AbandonMessageAsync(QueueMessage queueMessage)
+        public Task AbandonMessageAsync(QueueMessage queueMessage, Exception exception)
         {
+
+            // We have limited information about the details of the message
+            // since we failed to deserialize it.
+            this.settings.Logger.MessageFailure(
+                this.storageAccountName,
+                this.settings.TaskHubName,
+                queueMessage.Id /* MessageId */,
+                string.Empty /* InstanceId */,
+                string.Empty /* ExecutionId */,
+                this.storageQueue.Name,
+                string.Empty /* EventType */,
+                0 /* TaskEventId */,
+                exception.ToString());
+
             this.stats.PendingOrchestratorMessages.TryRemove(queueMessage.Id, out _);
             return base.AbandonMessageAsync(
                 queueMessage,

src/DurableTask.AzureStorage/Messaging/TaskHubQueue.cs

@@ -22,6 +22,7 @@ namespace DurableTask.AzureStorage.Messaging
     using DurableTask.AzureStorage.Storage;
     using DurableTask.Core;
     using DurableTask.Core.History;
+    using Microsoft.WindowsAzure.Storage.Table;
 
     abstract class TaskHubQueue
     {
@@ -57,6 +58,98 @@ public TaskHubQueue(
             this.backoffHelper = new BackoffPollingHelper(minPollingDelay, maxPollingDelay);
         }
 
+        public async Task HandleIfPoisonMessageAsync(MessageData messageData)
+        {
+            var queueMessage = messageData.OriginalQueueMessage;
+            var maxThreshold = this.settings.PoisonMessageDeuqueCountThreshold;
+
+            if (queueMessage.DequeueCount > maxThreshold)
+            {
+                string poisonMessageTableName = this.settings.TaskHubName.ToLowerInvariant() + "Poison";
+                Table poisonMessagesTable = this.azureStorageClient.GetTableReference(poisonMessageTableName);
+                await poisonMessagesTable.CreateIfNotExistsAsync();
+
+                // provide guidance, which is backend-specific
+                string guidance = $"Queue message ID '{queueMessage.Id}' was dequeued {queueMessage.DequeueCount} times," +
+                    $" which is greater than the threshold poison message threshold ({maxThreshold}). " +
+                    $"The message has been moved to the '{poisonMessageTableName}' table for manual review. " +
+                    $"This will fail the consuming orchestrator, activity, or entity";
+                messageData.TaskMessage.Event.PoisonGuidance = guidance;
+
+                // add to poison table
+                var poisonMessage = new DynamicTableEntity(queueMessage.Id, this.Name)
+                {
+                    Properties =
+                    {
+                        ["RawMessage"] = new EntityProperty(queueMessage.Message),
+                        ["Reason"] = new EntityProperty(guidance)
+                    }
+                };
+
+                await poisonMessagesTable.InsertAsync(poisonMessage);
+
+                // delete from queue so it doesn't get processed again.
+                await this.storageQueue.DeleteMessageAsync(queueMessage);
+
+                // since isPoison is `true`, we'll override the deserialized message
+                messageData.TaskMessage.Event.IsPoison = true;
+
+                this.settings.Logger.PoisonMessageDetected(
+                    this.storageAccountName,
+                    this.settings.TaskHubName,
+                    messageData.TaskMessage.Event.EventType.ToString(),
+                    messageData.TaskMessage.Event.EventId,
+                    messageData.OriginalQueueMessage.Id,
+                    messageData.TaskMessage.OrchestrationInstance.InstanceId,
+                    messageData.TaskMessage.OrchestrationInstance.ExecutionId,
+                    this.Name,
+                    messageData.OriginalQueueMessage.DequeueCount);
+            }
+        }
+
+        public async Task<bool> TryHandlingDeserializationPoisonMessage(QueueMessage queueMessage, Exception deserializationException)
+        {
+
+            var maxThreshold = this.settings.PoisonMessageDeuqueCountThreshold;
+            bool isPoisonMessage = queueMessage.DequeueCount > maxThreshold;
+            if (isPoisonMessage)
+            {
+                isPoisonMessage = true;
+                string guidance = $"Queue message ID '{queueMessage.Id}' was dequeued {queueMessage.DequeueCount} times," +
+                    $" which is greater than the threshold poison message threshold ({maxThreshold}). " +
+                    $"A de-serialization error ocurred: \n {deserializationException}";
+                var poisonMessage = new DynamicTableEntity(queueMessage.Id, this.Name)
+                {
+                    Properties =
+                    {
+                        ["RawMessage"] = new EntityProperty(queueMessage.Message),
+                        ["Reason"] = new EntityProperty(guidance)
+                    }
+                };
+
+                // add to poison table
+                string poisonMessageTableName = this.settings.TaskHubName.ToLowerInvariant() + "Poison";
+                Table poisonMessagesTable = this.azureStorageClient.GetTableReference(poisonMessageTableName);
+                await poisonMessagesTable.CreateIfNotExistsAsync();
+                await poisonMessagesTable.InsertAsync(poisonMessage);
+
+                // delete from queue so it doesn't get processed again.
+                await this.storageQueue.DeleteMessageAsync(queueMessage);
+
+                this.settings.Logger.PoisonMessageDetected(
+                    this.storageAccountName,
+                    this.settings.TaskHubName,
+                    string.Empty,
+                    0,
+                    string.Empty,
+                    string.Empty,
+                    string.Empty,
+                    this.Name,
+                    queueMessage.DequeueCount);
+            }
+            return isPoisonMessage;
+        }
+
         public string Name => this.storageQueue.Name;
 
         public Uri Uri => this.storageQueue.Uri;

src/DurableTask.AzureStorage/Messaging/WorkItemQueue.cs

@@ -10,7 +10,7 @@
 //  See the License for the specific language governing permissions and
 //  limitations under the License.
 //  ----------------------------------------------------------------------------------
-
+#nullable enable
 namespace DurableTask.AzureStorage.Messaging
 {
     using System;
@@ -30,26 +30,44 @@ public WorkItemQueue(
 
         protected override TimeSpan MessageVisibilityTimeout => this.settings.WorkItemQueueVisibilityTimeout;
 
-        public async Task<MessageData> GetMessageAsync(CancellationToken cancellationToken)
+        public async Task<MessageData?> GetMessageAsync(CancellationToken cancellationToken)
         {
             while (!cancellationToken.IsCancellationRequested)
             {
                 try
                 {
-                    QueueMessage queueMessage = await  this.storageQueue.GetMessageAsync(this.settings.WorkItemQueueVisibilityTimeout, cancellationToken);
+                    QueueMessage? queueMessage = await this.storageQueue.GetMessageAsync(this.settings.WorkItemQueueVisibilityTimeout, cancellationToken);
 
                     if (queueMessage == null)
                     {
                         await this.backoffHelper.WaitAsync(cancellationToken);
                         continue;
                     }
 
-                    MessageData data = await this.messageManager.DeserializeQueueMessageAsync(
-                        queueMessage,
-                        this.storageQueue.Name);
+                    try
+                    {
+                        MessageData data = await this.messageManager.DeserializeQueueMessageAsync(
+                            queueMessage,
+                            this.storageQueue.Name);
+
+                        // if successful, check if it's a poison message. If so, we handle it
+                        // and log metadata about it as the de-serialization succeeded.
+                        await this.HandleIfPoisonMessageAsync(data);
+                        this.backoffHelper.Reset();
+                        return data;
+                    }
+                    catch (Exception exception)
+                    {
+                        // Deserialization errors can be persistent, so we check if this is a poison message.
+                        bool isPoisonMessage = await this.TryHandlingDeserializationPoisonMessage(queueMessage, exception);
+                        if (isPoisonMessage)
+                        {
+                            // we have already handled the poison message, so we move on.
+                            continue;
+                        }
+                    }
+
 
-                    this.backoffHelper.Reset();
-                    return data;
                 }
                 catch (Exception e)
                 {

src/DurableTask.Core/History/HistoryEvent.cs

@@ -93,5 +93,16 @@ protected HistoryEvent(int eventId)
         /// Implementation for <see cref="IExtensibleDataObject.ExtensionData"/>.
         /// </summary>
         public ExtensionDataObject? ExtensionData { get; set; }
+
+        /// <summary>
+        /// Gets or sets whether this is a poison message.
+        /// </summary>
+        public bool IsPoison { get; set; } = false;
+
+        /// <summary>
+        /// Gets or sets user-facing details for why a message was labeled as poison.
+        /// This is to be set by each storage provider.
+        /// </summary>
+        public string PoisonGuidance { get; set; } = "";
     }
 }

src/DurableTask.Core/TaskActivityDispatcher.cs

@@ -22,6 +22,7 @@ namespace DurableTask.Core
     using DurableTask.Core.History;
     using DurableTask.Core.Logging;
     using DurableTask.Core.Middleware;
+    using DurableTask.Core.Serializing;
     using DurableTask.Core.Tracing;
 
     /// <summary>
@@ -192,6 +193,21 @@ await this.dispatchPipeline.RunAsync(dispatchContext, async _ =>
 
                         try
                         {
+                            if (scheduledEvent.IsPoison)
+                            {
+                                // if the activity is "poison", then we should not executed again. Instead, we'll manually fail the activity
+                                // by throwing an exception on behalf of the user-code. In the exception, we provide the storage-provider's guidance
+                                // on how to deal with the poison message.
+
+                                // We need to account for all possible deserialization modes, so we construct an exception valid in all modes.
+                                // TODO: revise - this is clunky
+                                var exception = new Exception(scheduledEvent.PoisonGuidance);
+                                var failureDetails = new FailureDetails(exception);
+                                var details = Utils.SerializeCause(exception, JsonDataConverter.Default);
+                                var taskFailure = new TaskFailureException(details, exception, details).WithFailureDetails(failureDetails);
+                                throw taskFailure;
+                            }
+
                             string? output = await taskActivity.RunAsync(context, scheduledEvent.Input);
                             responseEvent = new TaskCompletedEvent(-1, scheduledEvent.EventId, output);
                         }

src/DurableTask.Core/TaskOrchestrationExecutor.cs

@@ -197,6 +197,22 @@ void ProcessEvents(IEnumerable<HistoryEvent> events)
 
         void ProcessEvent(HistoryEvent historyEvent)
         {
+            if (historyEvent.IsPoison)
+            {
+                // If the message is labeled as "poison", then we should avoid processing it again.
+                // Therefore, we replace the event "in place" with an "ExecutionTerminatedEvent", so the
+                // orchestrator stops immediately.
+
+                var terminationEvent = new ExecutionTerminatedEvent(-1, historyEvent.PoisonGuidance);
+                historyEvent = terminationEvent;
+
+                // since replay is not guaranteed, we need to populate `this.result`
+                // with a completed task
+                var taskCompletionSource = new TaskCompletionSource<string>();
+                taskCompletionSource.SetResult("");
+                this.result = taskCompletionSource.Task;
+            }
+
             bool overrideSuspension = historyEvent.EventType == EventType.ExecutionResumed || historyEvent.EventType == EventType.ExecutionTerminated;
             if (this.context.IsSuspended && !overrideSuspension)
             {