Fix payload retries and entity cleanup after storage failures

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Keep externalized payload tokens stable across completion retries and release entity locks when payload storage permanently fails an orchestration.
+- Add the worker-owned `SetOrchestratorResponsePreparation` infrastructure API for response preparation before final worker entity cleanup and chunking.
 
 ## v1.27.1
 - Propagate tags from entity-started orchestrations ([#812](https://github.com/microsoft/durabletask-dotnet/pull/812))

src/Extensions/AzureBlobPayloads/DependencyInjection/DurableTaskWorkerBuilderExtensions.AzureBlobPayloads.cs

@@ -101,7 +101,12 @@ static IDurableTaskWorkerBuilder UseExternalizedPayloadsCore(IDurableTaskWorkerB
                 // Register an interceptor rather than moving Channel onto an intercepted CallInvoker.
                 // Clearing Channel would disable the worker's gRPC channel recreation, and requiring a
                 // pre-built Channel/CallInvoker would rule out the Address-only configuration.
-                opt.Interceptors.Add(new AzureBlobPayloadsSideCarInterceptor(store, opts));
+                AzureBlobPayloadsSideCarInterceptor interceptor = new(store, opts);
+                opt.Interceptors.Add(interceptor);
+
+                // Storage failures must become terminal responses before the worker releases entity locks.
+                // Unary interception still handles other RPCs; known tokens make the second pass a no-op.
+                opt.SetOrchestratorResponsePreparation(interceptor.PrepareOrchestratorResponseAsync);
 
                 opt.Capabilities.Add(P.WorkerCapability.LargePayloads);
                 opt.ConfigureVersioningExemptions([nameof(BlobPurgeJobOrchestrator)], AutoPurgeActivityNames);

src/Extensions/AzureBlobPayloads/Interceptors/AzureBlobPayloadsSideCarInterceptor.cs

@@ -27,6 +27,50 @@ public sealed class AzureBlobPayloadsSideCarInterceptor(PayloadStore payloadStor
 
     Action<object>? SharedMessageLockAcquiredForTest { get; set; }
 
+    /// <summary>
+    /// Externalizes a response before the worker finalizes terminal entity cleanup.
+    /// </summary>
+    /// <param name="response">The orchestration response.</param>
+    /// <param name="cancellation">Cancellation token.</param>
+    /// <returns>A task representing response preparation.</returns>
+    internal async Task PrepareOrchestratorResponseAsync(P.OrchestratorResponse response, CancellationToken cancellation)
+    {
+        try
+        {
+            await this.ExternalizeOrchestratorResponseAsync(response, cancellation);
+        }
+        catch (Exception ex) when (IsPermanentStorageFailure(ex))
+        {
+            // Retain entity releases when replacing actions with a Failed completion so
+            // the orchestration terminates without leaving its entities locked.
+            P.OrchestratorAction[] unlocks = response.Actions.Where(action =>
+                action.SendEntityMessage?.EntityUnlockSent is not null
+                || (action.SendEvent is { Name: "release", Instance: { } instance }
+                    && global::DurableTask.Core.Common.Entities.IsEntityInstance(instance.InstanceId))).ToArray();
+            response.Actions.Clear();
+            response.Actions.AddRange(unlocks);
+            response.CustomStatus = null;
+#pragma warning disable CS0612 // isPartial/chunkIndex are deprecated but still required for chunked response wire compatibility.
+            response.IsPartial = false;
+            response.ChunkIndex = null;
+#pragma warning restore CS0612
+            response.Actions.Add(new P.OrchestratorAction
+            {
+                CompleteOrchestration = new P.CompleteOrchestrationAction
+                {
+                    OrchestrationStatus = P.OrchestrationStatus.Failed,
+                    FailureDetails = new P.TaskFailureDetails
+                    {
+                        ErrorType = ex.GetType().FullName,
+                        ErrorMessage = ex.Message,
+                        StackTrace = ex.StackTrace,
+                        IsNonRetriable = true,
+                    },
+                },
+            });
+        }
+    }
+
     /// <inheritdoc/>
     protected override async Task ExternalizeRequestPayloadsAsync<TRequest>(TRequest request, CancellationToken cancellation)
     {
@@ -76,37 +120,7 @@ protected override async Task ExternalizeRequestPayloadsAsync<TRequest>(TRequest
 
                 break;
             case P.OrchestratorResponse r:
-                try
-                {
-                    await this.ExternalizeOrchestratorResponseAsync(r, cancellation);
-                }
-                catch (Exception ex) when (IsPermanentStorageFailure(ex))
-                {
-                    // Permanent failure during orchestration response externalization.
-                    // Replace all actions with a single Failed completion so the orchestration
-                    // terminates instead of being abandoned and redelivered indefinitely.
-                    r.Actions.Clear();
-                    r.CustomStatus = null;
-#pragma warning disable CS0612 // isPartial/chunkIndex are deprecated but still required for chunked response wire compatibility.
-                    r.IsPartial = false;
-                    r.ChunkIndex = null;
-#pragma warning restore CS0612
-                    r.Actions.Add(new P.OrchestratorAction
-                    {
-                        CompleteOrchestration = new P.CompleteOrchestrationAction
-                        {
-                            OrchestrationStatus = P.OrchestrationStatus.Failed,
-                            FailureDetails = new P.TaskFailureDetails
-                            {
-                                ErrorType = ex.GetType().FullName,
-                                ErrorMessage = ex.Message,
-                                StackTrace = ex.StackTrace,
-                                IsNonRetriable = true,
-                            },
-                        },
-                    });
-                }
-
+                await this.PrepareOrchestratorResponseAsync(r, cancellation);
                 break;
             case P.EntityBatchResult r:
                 await this.ExternalizeEntityBatchResultAsync(r, cancellation);

src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs

@@ -174,7 +174,7 @@ public override AsyncServerStreamingCall<TResponse> AsyncServerStreamingCall<TRe
     protected abstract Task ResolveResponsePayloadsAsync<TResponse>(TResponse response, CancellationToken cancellation);
 
     /// <summary>
-    /// Externalizes a payload if it exceeds the threshold.
+    /// Externalizes a payload if it exceeds the threshold and is not already a known payload token.
     /// </summary>
     /// <param name="value">The value to potentially externalize.</param>
     /// <param name="cancellation">Cancellation token.</param>
@@ -232,11 +232,11 @@ private protected async Task<string> ExternalizePayloadAsync(
     /// Calculates a value's UTF-8 size and determines whether it requires externalization.
     /// </summary>
     /// <param name="value">The value to inspect.</param>
-    /// <param name="size">The calculated UTF-8 size, or zero when <paramref name="value"/> is empty.</param>
-    /// <returns><c>true</c> when the value meets the configured externalization threshold.</returns>
+    /// <param name="size">The calculated UTF-8 size, or zero when <paramref name="value"/> is empty or a known payload token.</param>
+    /// <returns><c>true</c> when a non-token value meets the configured externalization threshold.</returns>
     private protected bool TryGetExternalizationSize(string? value, out int size)
     {
-        if (string.IsNullOrEmpty(value))
+        if (string.IsNullOrEmpty(value) || this.payloadStore.IsKnownPayloadToken(value!))
         {
             size = 0;
             return false;

src/Shared/Grpc/ProtoUtils.cs

@@ -442,29 +442,6 @@ internal static P.OrchestratorResponse ConstructOrchestratorResponse(
 
                     break;
                 case OrchestratorActionType.OrchestrationComplete:
-
-                    if (entityConversionState is not null)
-                    {
-                        // as a precaution, unlock any entities that were not unlocked for some reason, before
-                        // completing the orchestration.
-                        foreach ((string target, string criticalSectionId) in entityConversionState.ResetObligations())
-                        {
-                            response.Actions.Add(new P.OrchestratorAction
-                            {
-                                Id = action.Id,
-                                SendEntityMessage = new P.SendEntityMessageAction
-                                {
-                                    EntityUnlockSent = new P.EntityUnlockSentEvent
-                                    {
-                                        CriticalSectionId = criticalSectionId,
-                                        TargetInstanceId = target,
-                                        ParentInstanceId = entityConversionState.CurrentInstance?.InstanceId,
-                                    },
-                                },
-                            });
-                        }
-                    }
-
                     var completeAction = (OrchestrationCompleteOrchestratorAction)action;
                     protoAction.CompleteOrchestration = new P.CompleteOrchestrationAction
                     {
@@ -491,6 +468,10 @@ internal static P.OrchestratorResponse ConstructOrchestratorResponse(
             }
 
             response.Actions.Add(protoAction);
+            if (protoAction.CompleteOrchestration is not null)
+            {
+                entityConversionState?.InsertMissingUnlocks(response, protoAction);
+            }
         }
 
         return response;
@@ -1302,6 +1283,32 @@ public void RemoveUnlockObligation(string target)
             this.unlockObligations?.Remove(target);
         }
 
+        /// <summary>
+        /// Inserts any outstanding entity releases before a terminal completion.
+        /// </summary>
+        /// <param name="response">The response containing the completion.</param>
+        /// <param name="completionAction">The completion action already added to the response.</param>
+        public void InsertMissingUnlocks(P.OrchestratorResponse response, P.OrchestratorAction completionAction)
+        {
+            int index = response.Actions.IndexOf(completionAction);
+            foreach ((string target, string criticalSectionId) in this.ResetObligations())
+            {
+                response.Actions.Insert(index++, new P.OrchestratorAction
+                {
+                    Id = completionAction.Id,
+                    SendEntityMessage = new P.SendEntityMessageAction
+                    {
+                        EntityUnlockSent = new P.EntityUnlockSentEvent
+                        {
+                            CriticalSectionId = criticalSectionId,
+                            TargetInstanceId = target,
+                            ParentInstanceId = this.CurrentInstance?.InstanceId,
+                        },
+                    },
+                });
+            }
+        }
+
         /// <summary>
         /// Returns the remaining unlock obligations, and clears the list.
         /// </summary>

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -979,11 +979,19 @@ await this.ExecuteWithRetryAsync(
                 };
             }
 
+            if (this.internalOptions.PrepareOrchestratorResponseAsync is { } prepareResponse)
+            {
+                await prepareResponse(response, cancellationToken);
+            }
+
             var completeOrchestrationAction = response.Actions.FirstOrDefault(
                 a => a.CompleteOrchestration is not null);
 
             if (completeOrchestrationAction is not null)
             {
+                // Cover direct or newly terminal responses; converter cleanup already drains its obligations.
+                entityConversionState?.InsertMissingUnlocks(response, completeOrchestrationAction);
+
                 if (completeOrchestrationAction.CompleteOrchestration.OrchestrationStatus == P.OrchestrationStatus.Failed)
                 {
                     traceActivity?.SetStatus(

src/Worker/Grpc/GrpcDurableTaskWorkerOptions.cs

@@ -135,6 +135,12 @@ internal class InternalOptions
         /// </summary>
         public bool InsertEntityUnlocksOnCompletion { get; set; }
 
+        /// <summary>
+        /// Gets or sets optional payload preparation that runs before final worker entity cleanup and chunking.
+        /// Permanent payload failures can turn a pending response into a terminal completion.
+        /// </summary>
+        public Func<P.OrchestratorResponse, CancellationToken, Task>? PrepareOrchestratorResponseAsync { get; set; }
+
         /// <summary>
         /// Gets or sets the maximum amount of time to wait for the Hello handshake with the backend before
         /// treating the connect attempt as failed and retrying. A non-positive value disables the deadline.

src/Worker/Grpc/Internal/InternalOptionsExtensions.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using P = Microsoft.DurableTask.Protobuf;
 
 namespace Microsoft.DurableTask.Worker.Grpc.Internal;
 
@@ -167,4 +168,24 @@ public static void SetCallInvokerPublisher(
     {
         options.Internal.CallInvokerPublisher = publisher ?? throw new ArgumentNullException(nameof(publisher));
     }
+
+    /// <summary>
+    /// Sets an asynchronous callback that prepares a built orchestration response before the worker's final
+    /// entity-unlock cleanup and before the response is chunked or sent.
+    /// </summary>
+    /// <param name="options">The gRPC worker options.</param>
+    /// <param name="preparation">The callback receiving the mutable response and the worker cancellation token.</param>
+    /// <remarks>
+    /// This is an internal API that supports the DurableTask infrastructure and not subject to
+    /// the same compatibility standards as public APIs. It may be changed or removed without notice in
+    /// any release. The callback must finish all response mutations before its task completes.
+    /// </remarks>
+    /// <exception cref="ArgumentNullException"><paramref name="options"/> or <paramref name="preparation"/> is null.</exception>
+    public static void SetOrchestratorResponsePreparation(
+        this GrpcDurableTaskWorkerOptions options,
+        Func<P.OrchestratorResponse, CancellationToken, Task> preparation)
+    {
+        Check.NotNull(options);
+        options.Internal.PrepareOrchestratorResponseAsync = preparation ?? throw new ArgumentNullException(nameof(preparation));
+    }
 }

test/Client/Grpc.Tests/ProtoUtilsTests.cs

@@ -2,12 +2,81 @@
 // Licensed under the MIT License.
 
 using System.Collections.Immutable;
+using DurableTask.Core.Command;
+using DTCore = DurableTask.Core;
 using P = Microsoft.DurableTask.Protobuf;
+using SharedProtoUtils = Microsoft.DurableTask.ProtoUtils;
 
 namespace Microsoft.DurableTask.Client.Grpc.Tests;
 
 public class ProtoUtilsTests
 {
+    [Theory]
+    [InlineData(null)]
+    [InlineData(false)]
+    [InlineData(true)]
+    public void ConstructOrchestratorResponse_TerminalCleanup_IsCompleteAndIdempotent(bool? insertMissingUnlocks)
+    {
+        // Arrange
+        const string InstanceId = "instance";
+        const string CriticalSectionId = "11111111-1111-1111-1111-111111111111";
+        const int CompletionId = 42;
+        string[] entities = ["@counter@one", "@counter@two"];
+        SharedProtoUtils.EntityConversionState? state = insertMissingUnlocks.HasValue
+            ? new(insertMissingUnlocks.Value)
+            : null;
+        state?.SetOrchestrationInstance(new DTCore.OrchestrationInstance
+        {
+            InstanceId = InstanceId,
+            ExecutionId = "execution",
+        });
+        state?.AddUnlockObligations(new P.EntityLockRequestedEvent
+        {
+            CriticalSectionId = CriticalSectionId,
+            LockSet = { entities },
+        });
+        OrchestrationCompleteOrchestratorAction completion = new()
+        {
+            Id = CompletionId,
+            OrchestrationStatus = DTCore.OrchestrationStatus.Completed,
+            Result = "result",
+        };
+
+        // Act
+        P.OrchestratorResponse response = SharedProtoUtils.ConstructOrchestratorResponse(
+            InstanceId,
+            "execution",
+            customStatus: null,
+            actions: [completion],
+            completionToken: "completion-token",
+            entityConversionState: state,
+            orchestrationActivity: null);
+        P.OrchestratorResponse converted = response.Clone();
+        state?.InsertMissingUnlocks(response, response.Actions.Last());
+
+        // Assert
+        int expectedUnlocks = insertMissingUnlocks == true ? entities.Length : 0;
+        converted.Actions.Should().HaveCount(expectedUnlocks + 1);
+        converted.Actions.Should().OnlyContain(action => action.Id == CompletionId);
+        converted.Actions.Last().CompleteOrchestration.OrchestrationStatus.Should().Be(P.OrchestrationStatus.Completed);
+        converted.Actions.Last().CompleteOrchestration.Result.Should().Be("result");
+        P.EntityUnlockSentEvent[] unlocks = converted.Actions.Take(expectedUnlocks)
+            .Select(action => Assert.IsType<P.EntityUnlockSentEvent>(action.SendEntityMessage?.EntityUnlockSent))
+            .ToArray();
+        unlocks.Select(unlock => unlock.TargetInstanceId).Should().Equal(entities.Take(expectedUnlocks));
+        foreach (P.EntityUnlockSentEvent unlock in unlocks)
+        {
+            unlock.CriticalSectionId.Should().Be(CriticalSectionId);
+            unlock.ParentInstanceId.Should().Be(InstanceId);
+        }
+        converted.CompletionToken.Should().Be("completion-token");
+        response.Should().Be(converted);
+        if (state is not null)
+        {
+            state.ResetObligations().Should().BeEmpty();
+        }
+    }
+
     [Fact]
     public void GetAllStatuses_ReturnsExpectedStatuses()
     {