Add simple poison message handling for Azure Storage

src/DurableTask.AzureStorage/DurableTask.AzureStorage.csproj

@@ -22,9 +22,10 @@
     <MajorVersion>1</MajorVersion>
     <MinorVersion>17</MinorVersion>
     <PatchVersion>1</PatchVersion>
-    <VersionPrefix>$(MajorVersion).$(MinorVersion).$(PatchVersion)</VersionPrefix>
+	<VersionPrefix>$(MajorVersion).$(MinorVersion).$(PatchVersion)</VersionPrefix>
     <FileVersion>$(VersionPrefix).0</FileVersion>
-    <!-- FileVersionRevision is expected to be set by the CI. This is useful for distinguishing between multiple builds of the same version. -->
+	<VersionSuffix>pmh.2</VersionSuffix>
+	  <!-- FileVersionRevision is expected to be set by the CI. This is useful for distinguishing between multiple builds of the same version. -->
     <FileVersion Condition="'$(FileVersionRevision)' != ''">$(VersionPrefix).$(FileVersionRevision)</FileVersion>
     <!-- The assembly version is only the major/minor pair, making it easier to do in-place upgrades -->
     <AssemblyVersion>$(MajorVersion).$(MinorVersion).0.0</AssemblyVersion>

src/DurableTask.AzureStorage/Messaging/ControlQueue.cs

@@ -19,9 +19,13 @@ namespace DurableTask.AzureStorage.Messaging
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
+    using Azure;
+    using Azure.Data.Tables;
     using DurableTask.AzureStorage.Monitoring;
     using DurableTask.AzureStorage.Partitioning;
     using DurableTask.AzureStorage.Storage;
+    using DurableTask.Core;
+    using Microsoft.WindowsAzure.Storage.Table;
 
     class ControlQueue : TaskHubQueue, IDisposable
     {
@@ -46,6 +50,36 @@ public ControlQueue(
 
         protected override TimeSpan MessageVisibilityTimeout => this.settings.ControlQueueVisibilityTimeout;
 
+        private async Task HandleIfPoisonMessageAsync(MessageData messageData)
+        {
+            var isPoison = false;
+            var queueMessage = messageData.OriginalQueueMessage;
+
+            // if deuque count is large, just flag it as poison. Don't even deserialize it!
+            if (queueMessage.DequeueCount > 5) // TODO: make configurable
+            {
+                var poisonMessage = new DynamicTableEntity(queueMessage.Id, this.Name)
+                {
+                    Properties =
+                    {
+                        ["TheFullMessage"] = new EntityProperty(queueMessage.Message)
+                    }
+                };
+
+                // add to poison table
+                var poisonMessagesTable = this.azureStorageClient.GetTableReference("PoisonMessagesTable");
+                await poisonMessagesTable.CreateIfNotExistsAsync();
+                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 w/ a suspend event
+                isPoison = true;
+            }
+            messageData.TaskMessage.Event.IsPoison = isPoison;
+        }
+
         public async Task<IReadOnlyList<MessageData>> GetMessagesAsync(CancellationToken cancellationToken)
         {
             using (var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(this.releaseCancellationToken, cancellationToken))
@@ -108,6 +142,9 @@ await batch.ParallelForEachAsync(async delegate (QueueMessage queueMessage)
                                 messageData = await this.messageManager.DeserializeQueueMessageAsync(
                                     queueMessage,
                                     this.storageQueue.Name);
+
+                                await this.HandleIfPoisonMessageAsync(messageData);
+
                             }
                             catch (Exception e)
                             {

src/DurableTask.AzureStorage/Messaging/WorkItemQueue.cs

@@ -17,6 +17,7 @@ namespace DurableTask.AzureStorage.Messaging
     using System.Threading;
     using System.Threading.Tasks;
     using DurableTask.AzureStorage.Storage;
+    using Microsoft.WindowsAzure.Storage.Table;
 
     class WorkItemQueue : TaskHubQueue
     {
@@ -30,6 +31,37 @@ public WorkItemQueue(
 
         protected override TimeSpan MessageVisibilityTimeout => this.settings.WorkItemQueueVisibilityTimeout;
 
+        private async Task HandleIfPoisonMessageAsync(MessageData messageData)
+        {
+            // TODO: put in superclass?
+            var isPoison = false;
+            var queueMessage = messageData.OriginalQueueMessage;
+
+            // if deuque count is large, just flag it as poison. Don't even deserialize it!
+            if (queueMessage.DequeueCount > 5) // TODO: make configurable
+            {
+                var poisonMessage = new DynamicTableEntity(queueMessage.Id, this.Name)
+                {
+                    Properties =
+                    {
+                        ["TheFullMessage"] = new EntityProperty(queueMessage.Message)
+                    }
+                };
+
+                // add to poison table
+                var poisonMessagesTable = this.azureStorageClient.GetTableReference("PoisonMessagesTable");
+                await poisonMessagesTable.CreateIfNotExistsAsync();
+                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 w/ a suspend event
+                isPoison = true;
+            }
+            messageData.TaskMessage.Event.IsPoison = isPoison;
+        }
+
         public async Task<MessageData> GetMessageAsync(CancellationToken cancellationToken)
         {
             while (!cancellationToken.IsCancellationRequested)
@@ -44,9 +76,11 @@ public async Task<MessageData> GetMessageAsync(CancellationToken cancellationTok
                         continue;
                     }
 
+                    // TODO: maybe the message manager should handle the poison?
                     MessageData data = await this.messageManager.DeserializeQueueMessageAsync(
                         queueMessage,
                         this.storageQueue.Name);
+                    await this.HandleIfPoisonMessageAsync(data);
 
                     this.backoffHelper.Reset();
                     return data;

src/DurableTask.Core/DurableTask.Core.csproj

@@ -20,6 +20,7 @@
     <MinorVersion>16</MinorVersion>
     <PatchVersion>2</PatchVersion>
     <VersionPrefix>$(MajorVersion).$(MinorVersion).$(PatchVersion)</VersionPrefix>
+	<VersionSuffix>pmh.2</VersionSuffix>
     <FileVersion>$(VersionPrefix).0</FileVersion>
     <!-- FileVersionRevision is expected to be set by the CI. This is useful for distinguishing between multiple builds of the same version. -->
     <FileVersion Condition="'$(FileVersionRevision)' != ''">$(VersionPrefix).$(FileVersionRevision)</FileVersion>

src/DurableTask.Core/History/HistoryEvent.cs

@@ -93,5 +93,10 @@ 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;
     }
 }

src/DurableTask.Core/TaskActivityDispatcher.cs

@@ -192,6 +192,12 @@ await this.dispatchPipeline.RunAsync(dispatchContext, async _ =>
 
                         try
                         {
+                            if (scheduledEvent.IsPoison)
+                            {
+                                var exception = new TaskFailureException("poison activity message detected!", details: "poison activity message detected!");
+                                throw exception;
+                            }
+
                             string? output = await taskActivity.RunAsync(context, scheduledEvent.Input);
                             responseEvent = new TaskCompletedEvent(-1, scheduledEvent.EventId, output);
                         }

src/DurableTask.Core/TaskOrchestrationExecutor.cs

@@ -197,6 +197,17 @@ void ProcessEvents(IEnumerable<HistoryEvent> events)
 
         void ProcessEvent(HistoryEvent historyEvent)
         {
+            if (historyEvent.IsPoison)
+            {
+                var terminationEvent = new ExecutionTerminatedEvent(-1, "detected poison!");
+                historyEvent = terminationEvent;
+
+                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)
             {