Fix payload retries and entity cleanup after storage failures

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Keep externalized payload tokens stable across completion retries and release entity locks when payload storage permanently fails an orchestration.
 - Preserve the original orchestration version when restarting through the orchestration-service client shim ([#463](https://github.com/microsoft/durabletask-dotnet/issues/463)).
 - Support configurable scheduler token audiences and government defaults ([#806](https://github.com/microsoft/durabletask-dotnet/pull/806))
 

src/Extensions/AzureBlobPayloads/AzureBlobPayloads.csproj

@@ -30,4 +30,12 @@
   <ItemGroup>
     <InternalsVisibleTo Include="Microsoft.DurableTask.Grpc.IntegrationTests" Key="$(StrongNamePublicKey)" />
   </ItemGroup>
+
+  <!-- Reuse friend-visible validation helpers, but retain local compiler compatibility types. -->
+  <Target Name="RemoveDuplicateSharedValidation" AfterTargets="IncludeSharedItems">
+    <ItemGroup>
+      <Compile Remove="$([MSBuild]::NormalizePath('$(SharedRoot)Core\Validation\Check.cs'))" />
+      <Compile Remove="$([MSBuild]::NormalizePath('$(SharedRoot)Core\Validation\Verify.cs'))" />
+    </ItemGroup>
+  </Target>
 </Project>

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

@@ -100,7 +100,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.Internal.PrepareOrchestratorResponseAsync = 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;
@@ -1291,6 +1272,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,18 @@ 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)
             {
+                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 terminal 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/Worker.Grpc.csproj

@@ -22,6 +22,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <InternalsVisibleTo Include="Microsoft.DurableTask.Extensions.AzureBlobPayloads" Key="$(StrongNamePublicKey)" />
     <!-- The AzureBlobPayloads auto-purge activities share this worker's gRPC transport, so their tests drive a
          real worker through a channel recreate. Reaching that path quickly needs the internal connect-loop
          knobs (recreate threshold), which are deliberately not part of any supported surface. -->