Delete persisted output during Azure Storage rewind

src/DurableTask.AzureStorage/Logging/LogHelper.cs

@@ -759,14 +759,15 @@ internal void GeneralWarning(
             string account,
             string taskHub,
             string details,
-            string instanceId = null)
+            string instanceId = null,
+            Exception exception = null)
         {
             var logEvent = new LogEvents.GeneralWarning(
                 account,
                 taskHub,
                 details,
                 instanceId ?? string.Empty);
-            this.WriteStructuredLog(logEvent);
+            this.WriteStructuredLog(logEvent, exception);
         }
 
         internal void SplitBrainDetected(

src/DurableTask.AzureStorage/MessageManager.cs

@@ -253,10 +253,21 @@ public Task<string> DownloadAndDecompressAsBytesAsync(Uri blobUri, CancellationT
             return DownloadAndDecompressAsBytesAsync(blob, cancellationToken);
         }
 
-        public Task<bool> DeleteBlobAsync(string blobName, CancellationToken cancellationToken = default)
+        public async Task<bool> DeleteBlobAsync(string blobName, CancellationToken cancellationToken = default)
         {
             Blob blob = this.blobContainer.GetBlobReference(blobName);
-            return blob.DeleteIfExistsAsync(cancellationToken);
+            try
+            {
+                return await blob.DeleteIfExistsAsync(cancellationToken);
+            }
+            catch (AggregateException ex) when (!cancellationToken.IsCancellationRequested)
+            {
+                throw new DurableTaskStorageException("Azure Storage retries failed while deleting a blob.", ex);
+            }
+            catch (OperationCanceledException ex) when (!cancellationToken.IsCancellationRequested)
+            {
+                throw new DurableTaskStorageException("Azure Storage timed out while deleting a blob.", ex);
+            }
         }
 
         private async Task<string> DownloadAndDecompressAsBytesAsync(Blob blob, CancellationToken cancellationToken = default)

src/DurableTask.AzureStorage/Tracking/AzureTableTrackingStore.cs

@@ -283,6 +283,7 @@ public override async IAsyncEnumerable<string> RewindHistoryAsync(string instanc
             ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
             bool hasFailedSubOrchestrations = false;
+            var blobsToDelete = new List<string>();
             string partitionFilter = AzureTableQueryFilter.PartitionKeyEquals(instanceId);
 
             string orchestratorStartedFilter = $"{partitionFilter} and {nameof(HistoryEvent.EventType)} eq '{nameof(EventType.OrchestratorStarted)}'";
@@ -361,6 +362,13 @@ public override async IAsyncEnumerable<string> RewindHistoryAsync(string instanc
                         break;
                 }
 
+                if (entity.GetString(nameof(HistoryEvent.EventType)) == nameof(EventType.ExecutionCompleted))
+                {
+                    // GenericEvent replay ignores the terminal payload, so remove its blob references in the same ETag-guarded replace.
+                    RemovePropertyAndTrackBlob(entity, nameof(ExecutionCompletedEvent.Result), blobsToDelete);
+                    RemovePropertyAndTrackBlob(entity, nameof(ExecutionCompletedEvent.FailureDetails), blobsToDelete);
+                }
+
                 // "clear" failure event by making RewindEvent: replay ignores row while dummy event preserves rowKey
                 entity[nameof(TaskFailedEvent.Reason)] = "Rewound: " + entity.GetString(nameof(HistoryEvent.EventType));
                 entity[nameof(TaskFailedEvent.EventType)] = nameof(EventType.GenericEvent);
@@ -371,12 +379,40 @@ public override async IAsyncEnumerable<string> RewindHistoryAsync(string instanc
             // reset orchestration status in instance store table
             await this.UpdateStatusForRewindAsync(instanceId, cancellationToken);
 
+            // Delete only after both the history pointers and the Instances-table Output reference are gone.
+            await this.DeleteRewindBlobsAsync(instanceId, blobsToDelete, cancellationToken);
+
             if (!hasFailedSubOrchestrations)
             {
                 yield return instanceId;
             }
         }
 
+        async Task DeleteRewindBlobsAsync(string instanceId, IEnumerable<string> blobNames, CancellationToken cancellationToken)
+        {
+            foreach (string blobName in blobNames)
+            {
+                cancellationToken.ThrowIfCancellationRequested();
+
+                try
+                {
+                    await this.messageManager.DeleteBlobAsync(blobName, cancellationToken);
+                }
+                catch (DurableTaskStorageException ex)
+                {
+                    this.settings.Logger.GeneralWarning(
+                        this.azureStorageClient.BlobAccountName,
+                        this.settings.TaskHubName,
+                        $"Failed to delete unreferenced rewind blob '{blobName}'. The blob will remain until the orchestration is purged. " +
+                        $"Storage status code: {ex.HttpStatusCode}; error code: '{ex.ErrorCode}'.",
+                        instanceId,
+                        ex);
+                }
+            }
+
+            cancellationToken.ThrowIfCancellationRequested();
+        }
+
         /// <inheritdoc />
         public override async IAsyncEnumerable<OrchestrationState> GetStateAsync(string instanceId, bool allExecutions, bool fetchInput, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
@@ -856,15 +892,23 @@ public override async Task<bool> SetNewExecutionAsync(
         /// <inheritdoc />
         public override async Task UpdateStatusForRewindAsync(string instanceId, CancellationToken cancellationToken = default)
         {
-            string sanitizedInstanceId = KeySanitation.EscapePartitionKey(instanceId);
-            TableEntity entity = new TableEntity(sanitizedInstanceId, "")
+            string filter = $"{AzureTableQueryFilter.PartitionKeyEquals(instanceId)} and {AzureTableQueryFilter.ColumnEquals(RowKeyProperty, string.Empty)}";
+            TableEntity entity = await this.InstancesTable
+                .ExecuteQueryAsync<TableEntity>(filter, 1, cancellationToken: cancellationToken)
+                .FirstOrDefaultAsync();
+
+            if (entity == null)
             {
-                ["RuntimeStatus"] = OrchestrationStatus.Pending.ToString("G"),
-                ["LastUpdatedTime"] = DateTime.UtcNow,
-            };
+                throw new DurableTaskStorageException($"The orchestration instance '{instanceId}' does not exist.");
+            }
+
+            // Merge cannot remove a table property, so replace the complete row using its current ETag.
+            entity.Remove(OutputProperty);
+            entity["RuntimeStatus"] = OrchestrationStatus.Pending.ToString("G");
+            entity["LastUpdatedTime"] = DateTime.UtcNow;
 
             Stopwatch stopwatch = Stopwatch.StartNew();
-            await this.InstancesTable.MergeEntityAsync(entity, ETag.All, cancellationToken);
+            await this.InstancesTable.ReplaceEntityAsync(entity, entity.ETag, cancellationToken);
 
             // We don't have enough information to get the episode number.
             // It's also not important to have for this particular trace.
@@ -1378,6 +1422,20 @@ static string GetBlobPropertyName(string originalPropertyName)
             return originalPropertyName + "BlobName";
         }
 
+        static void RemovePropertyAndTrackBlob(TableEntity entity, string propertyName, List<string> blobsToDelete)
+        {
+            string blobPropertyName = GetBlobPropertyName(propertyName);
+            if (entity.TryGetValue(blobPropertyName, out object value) &&
+                value is string blobName &&
+                !string.IsNullOrEmpty(blobName))
+            {
+                blobsToDelete.Add(blobName);
+            }
+
+            entity.Remove(propertyName);
+            entity.Remove(blobPropertyName);
+        }
+
         static string GetBlobName(TableEntity entity, string property)
         {
             string sanitizedInstanceId = entity.PartitionKey;

test/DurableTask.AzureStorage.Tests/AzureStorageScenarioTests.cs

@@ -16,13 +16,16 @@ namespace DurableTask.AzureStorage.Tests
     using Azure.Data.Tables;
     using Azure.Storage.Blobs;
     using Azure.Storage.Blobs.Models;
+    using Azure.Storage.Blobs.Specialized;
+    using DurableTask.AzureStorage.Logging;
     using DurableTask.AzureStorage.Storage;
     using DurableTask.AzureStorage.Tracking;
     using DurableTask.Core;
     using DurableTask.Core.Exceptions;
     using DurableTask.Core.History;
     using DurableTask.Core.Settings;
     using Microsoft.Practices.EnterpriseLibrary.SemanticLogging.Utility;
+    using Microsoft.Extensions.Logging;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
     using Moq;
     using Newtonsoft.Json;
@@ -1510,6 +1513,178 @@ public async Task RewindActivityFail()
             }
         }
 
+        [DataTestMethod]
+        [DataRow(false)]
+        [DataRow(true)]
+        public async Task RewindLargeFailure_HandlesBlobDeletionFailure(bool failBlobDeletion)
+        {
+            var logger = new Mock<ILogger>();
+            var loggerFactory = new Mock<ILoggerFactory>();
+            loggerFactory
+                .Setup(factory => factory.CreateLogger(It.IsAny<string>()))
+                .Returns(logger.Object);
+
+            using (TestOrchestrationHost host = TestHelpers.GetTestOrchestrationHost(
+                enableExtendedSessions: false,
+                modifySettingsAction: settings => settings.LoggerFactory = loggerFactory.Object))
+            {
+                Orchestrations.RewindLargeFailure.ShouldFail = true;
+                host.ErrorPropagationMode = ErrorPropagationMode.UseFailureDetails;
+                await host.StartAsync();
+
+                string failureMessage = this.GenerateMediumRandomStringPayload().ToString();
+                var client = await host.StartOrchestrationAsync(
+                    typeof(Orchestrations.RewindLargeFailure),
+                    input: failureMessage);
+                OrchestrationState failed = await client.WaitForCompletionAsync(TimeSpan.FromSeconds(30));
+                Assert.IsNotNull(failed);
+                Assert.AreEqual(OrchestrationStatus.Failed, failed.OrchestrationStatus);
+
+                var trackingStore = (AzureTableTrackingStore)host.service.TrackingStore;
+                string instanceFilter = $"{AzureTableQueryFilter.PartitionKeyEquals(client.InstanceId)} and " +
+                    $"{AzureTableQueryFilter.ColumnEquals(nameof(ITableEntity.RowKey), string.Empty)}";
+                TableEntity failedInstance = await trackingStore.InstancesTable
+                    .ExecuteQueryAsync<TableEntity>(instanceFilter, 1)
+                    .FirstOrDefaultAsync();
+                string outputBlobUrl = failedInstance.GetString("Output");
+                Assert.IsTrue(Uri.IsWellFormedUriString(outputBlobUrl, UriKind.Absolute));
+
+                string failedCompletionFilter = $"{AzureTableQueryFilter.PartitionKeyEquals(client.InstanceId)} and " +
+                    $"{AzureTableQueryFilter.ColumnEquals(nameof(OrchestrationInstance.ExecutionId), failed.OrchestrationInstance.ExecutionId)} and " +
+                    $"{AzureTableQueryFilter.ColumnEquals(nameof(HistoryEvent.EventType), nameof(EventType.ExecutionCompleted))}";
+                TableEntity failedCompletion = (await trackingStore.HistoryTable
+                    .ExecuteQueryAsync<TableEntity>(failedCompletionFilter)
+                    .ToListAsync())
+                    .Single();
+                string resultBlobName = failedCompletion.GetString("ResultBlobName");
+                Assert.IsNotNull(resultBlobName);
+                Assert.IsTrue(new Uri(outputBlobUrl).AbsolutePath.EndsWith(resultBlobName, StringComparison.Ordinal));
+                string failureDetailsBlobName = failedCompletion.GetString("FailureDetailsBlobName");
+                Assert.IsNotNull(failureDetailsBlobName);
+
+                string[] outputBlobNames = new[]
+                {
+                    resultBlobName,
+                    failureDetailsBlobName,
+                };
+                var blobServiceClient = new BlobServiceClient(TestHelpers.GetTestStorageAccountConnectionString());
+                BlobContainerClient container = blobServiceClient.GetBlobContainerClient($"{host.TaskHub.ToLowerInvariant()}-largemessages");
+                foreach (string blobName in outputBlobNames)
+                {
+                    Assert.IsTrue((await container.GetBlobClient(blobName).ExistsAsync()).Value);
+                }
+
+                if (!failBlobDeletion)
+                {
+                    using var canceled = new CancellationTokenSource();
+                    canceled.Cancel();
+                    await Assert.ThrowsExceptionAsync<TaskCanceledException>(
+                        async () => await trackingStore.RewindHistoryAsync(client.InstanceId, canceled.Token).ToListAsync());
+
+                    TableEntity instanceAfterCancellation = await trackingStore.InstancesTable
+                        .ExecuteQueryAsync<TableEntity>(instanceFilter, 1)
+                        .FirstOrDefaultAsync();
+                    Assert.AreEqual(OrchestrationStatus.Failed.ToString(), instanceAfterCancellation.GetString("RuntimeStatus"));
+                    Assert.AreEqual(outputBlobUrl, instanceAfterCancellation.GetString("Output"));
+                }
+
+                BlobClient resultBlob = container.GetBlobClient(resultBlobName);
+                BlobLeaseClient resultBlobLease = null;
+                if (failBlobDeletion)
+                {
+                    resultBlobLease = resultBlob.GetBlobLeaseClient();
+                    await resultBlobLease.AcquireAsync(TimeSpan.FromSeconds(15));
+                }
+
+                Orchestrations.RewindLargeFailure.ShouldFail = false;
+                try
+                {
+                    if (failBlobDeletion)
+                    {
+                        await client.RewindAsync("Retry despite the persisted-output cleanup failure.");
+                    }
+                    else
+                    {
+                        CollectionAssert.AreEqual(
+                            new[] { client.InstanceId },
+                            await trackingStore.RewindHistoryAsync(client.InstanceId).ToListAsync());
+                    }
+
+                    TableEntity rewoundInstance = await trackingStore.InstancesTable
+                        .ExecuteQueryAsync<TableEntity>(instanceFilter, 1)
+                        .FirstOrDefaultAsync();
+                    if (failBlobDeletion)
+                    {
+                        Assert.AreNotEqual(outputBlobUrl, rewoundInstance.GetString("Output"));
+                    }
+                    else
+                    {
+                        Assert.AreEqual(OrchestrationStatus.Pending.ToString(), rewoundInstance.GetString("RuntimeStatus"));
+                        Assert.IsFalse(rewoundInstance.ContainsKey("Output"));
+                    }
+
+                    string historyFilter = $"{AzureTableQueryFilter.PartitionKeyEquals(client.InstanceId)} and " +
+                        $"{AzureTableQueryFilter.ColumnEquals(nameof(ITableEntity.RowKey), failedCompletion.RowKey)}";
+                    TableEntity rewoundCompletion = (await trackingStore.HistoryTable
+                        .ExecuteQueryAsync<TableEntity>(historyFilter)
+                        .ToListAsync())
+                        .Single();
+
+                    Assert.AreEqual(nameof(EventType.GenericEvent), rewoundCompletion.GetString(nameof(HistoryEvent.EventType)));
+                    Assert.IsFalse(rewoundCompletion.ContainsKey("Result"));
+                    Assert.IsFalse(rewoundCompletion.ContainsKey("ResultBlobName"));
+                    Assert.IsFalse(rewoundCompletion.ContainsKey("FailureDetails"));
+                    Assert.IsFalse(rewoundCompletion.ContainsKey("FailureDetailsBlobName"));
+
+                    Assert.AreEqual(failBlobDeletion, (await resultBlob.ExistsAsync()).Value);
+                    Assert.IsFalse((await container.GetBlobClient(failureDetailsBlobName).ExistsAsync()).Value);
+
+                    var cleanupWarnings = logger.Invocations
+                        .Where(invocation => invocation.Arguments.Count > 2)
+                        .Where(invocation => invocation.Arguments[2] is LogEvents.GeneralWarning warning &&
+                            warning.Details.Contains(resultBlobName))
+                        .ToList();
+                    if (failBlobDeletion)
+                    {
+                        Assert.AreEqual(1, cleanupWarnings.Count);
+                        var cleanupWarning = (LogEvents.GeneralWarning)cleanupWarnings[0].Arguments[2];
+                        Assert.AreEqual(client.InstanceId, cleanupWarning.InstanceId);
+                        StringAssert.Contains(cleanupWarning.Details, "LeaseIdMissing");
+                        Assert.IsInstanceOfType(cleanupWarnings[0].Arguments[3], typeof(DurableTaskStorageException));
+                    }
+                    else
+                    {
+                        Assert.AreEqual(0, cleanupWarnings.Count);
+                    }
+                }
+                finally
+                {
+                    if (resultBlobLease != null)
+                    {
+                        await resultBlobLease.ReleaseAsync();
+                        await resultBlob.DeleteIfExistsAsync();
+                    }
+                }
+
+                if (!failBlobDeletion)
+                {
+                    await client.RewindAsync("Retry the persisted-output cleanup.");
+                }
+
+                OrchestrationState completed = await client.WaitForCompletionAsync(TimeSpan.FromSeconds(30));
+                Assert.AreEqual(OrchestrationStatus.Completed, completed?.OrchestrationStatus);
+                Assert.AreEqual("\"Done\"", completed?.Output);
+
+                TableEntity completedInstance = await trackingStore.InstancesTable
+                    .ExecuteQueryAsync<TableEntity>(instanceFilter, 1)
+                    .FirstOrDefaultAsync();
+                Assert.AreEqual(OrchestrationStatus.Completed.ToString(), completedInstance.GetString("RuntimeStatus"));
+                Assert.AreEqual("\"Done\"", completedInstance.GetString("Output"));
+
+                await host.StopAsync();
+            }
+        }
+
         [TestMethod]
         public async Task RewindMultipleActivityFail()
         {
@@ -4857,6 +5032,21 @@ public override async Task<string> RunTask(OrchestrationContext context, string
                 }
             }
 
+            internal class RewindLargeFailure : TaskOrchestration<string, string>
+            {
+                public static bool ShouldFail = true;
+
+                public override Task<string> RunTask(OrchestrationContext context, string message)
+                {
+                    if (ShouldFail)
+                    {
+                        throw new Exception(message);
+                    }
+
+                    return Task.FromResult("Done");
+                }
+            }
+
             [KnownType(typeof(Activities.Throw))]
             internal class Throw : TaskOrchestration<string, string>
             {