Perf: avoid repeated protobuf sizing during orchestration response chunking

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -32,6 +32,29 @@ class Processor
     {
         static readonly Google.Protobuf.WellKnownTypes.Empty EmptyMessage = new();
 
+        /// <summary>
+        /// Test-only hook invoked immediately after an individual action's size is computed during the
+        /// fail-fast validation pass in <see cref="CompleteOrchestratorTaskWithChunkingAsync"/>. Always
+        /// <see langword="null"/> in production (adds no overhead beyond a null check); used exclusively
+        /// by unit tests (accessed via reflection) to deterministically verify that validation stops at
+        /// the first oversized action instead of computing the size of every action.
+        /// </summary>
+#pragma warning disable CS0649 // Field is never assigned to in production code - it is set via reflection by tests.
+        static Action<int>? testActionSizedHook;
+
+        /// <summary>
+        /// Test-only hook invoked immediately after the *whole response* is sized (via
+        /// <c>response.CalculateSize()</c>) in <see cref="CompleteOrchestratorTaskWithChunkingAsync"/>.
+        /// Always <see langword="null"/> in production (adds no overhead beyond a null check). Protobuf's
+        /// whole-response sizing recursively sizes every action internally without going through
+        /// <see cref="testActionSizedHook"/>, so this separate hook exists to let unit tests prove that
+        /// whole-response sizing is never reached when fail-fast validation rejects an oversized action -
+        /// a regression that reintroduces whole-response sizing *before* validation would otherwise go
+        /// undetected by <see cref="testActionSizedHook"/> alone.
+        /// </summary>
+        static Action? testResponseSizedHook;
+#pragma warning restore CS0649
+
         readonly GrpcDurableTaskWorker worker;
         readonly TaskHubSidecarServiceClient client;
         readonly DurableTaskShimFactory shimFactory;
@@ -1098,81 +1121,89 @@ async Task CompleteOrchestratorTaskWithChunkingAsync(
             int maxChunkBytes,
             CancellationToken cancellationToken)
         {
-            // Validate that no single action exceeds the maximum chunk size
-            static P.TaskFailureDetails? ValidateActionsSize(IEnumerable<P.OrchestratorAction> actions, int maxChunkBytes)
+            // Helper to add an action to the current chunk if it fits, using a precomputed action size
+            // so validation and chunk packing do not need to recalculate its serialized size.
+            static bool TryAddAction(
+                Google.Protobuf.Collections.RepeatedField<P.OrchestratorAction> dest,
+                P.OrchestratorAction action,
+                int actionSize,
+                ref int currentSize,
+                int maxChunkBytes)
+            {
+                if (currentSize + actionSize > maxChunkBytes && currentSize > 0)
+                {
+                    return false;
+                }
+
+                dest.Add(action);
+                currentSize += actionSize;
+                return true;
+            }
+
+            bool largePayloads = this.worker.grpcOptions.Capabilities.Contains(P.WorkerCapability.LargePayloads);
+
+            // When the LargePayloads capability is not present, validate that no single action
+            // exceeds the maximum chunk size *before* doing any whole-response sizing. This must
+            // stay fail-fast: as soon as an oversized action is found, we fail the orchestration
+            // immediately without computing the size of any later action. Only when every action
+            // is confirmed to individually fit do we keep the sizes we already computed here, so
+            // they can be reused below instead of being recalculated.
+            int[]? actionSizes = null;
+            if (!largePayloads)
             {
-                foreach (P.OrchestratorAction action in actions)
+                actionSizes = new int[response.Actions.Count];
+                for (int i = 0; i < response.Actions.Count; i++)
                 {
+                    P.OrchestratorAction action = response.Actions[i];
                     int actionSize = action.CalculateSize();
+                    testActionSizedHook?.Invoke(action.Id);
                     if (actionSize > maxChunkBytes)
                     {
                         // TODO: large payload doc is not available yet on aka.ms, add doc link to below error message
                         string errorMessage = $"A single orchestrator action of type {action.OrchestratorActionTypeCase} with id {action.Id} " +
                             $"exceeds the {maxChunkBytes / 1024.0 / 1024.0:F2}MB limit: {actionSize / 1024.0 / 1024.0:F2}MB. " +
                             "Enable large-payload externalization to Azure Blob Storage to support oversized actions.";
-                        return new P.TaskFailureDetails
+                        P.TaskFailureDetails validationFailure = new()
                         {
                             ErrorType = typeof(InvalidOperationException).FullName,
                             ErrorMessage = errorMessage,
                             IsNonRetriable = true,
                         };
-                    }
-                }
 
-                return null;
-            }
-
-            P.TaskFailureDetails? validationFailure = this.worker.grpcOptions.Capabilities.Contains(P.WorkerCapability.LargePayloads)
-                ? null
-                : ValidateActionsSize(response.Actions, maxChunkBytes);
-            if (validationFailure != null)
-            {
-                // Complete the orchestration with a failed status and failure details
-                P.OrchestratorResponse failureResponse = new()
-                {
-                    InstanceId = response.InstanceId,
-                    CompletionToken = response.CompletionToken,
-                    OrchestrationTraceContext = response.OrchestrationTraceContext,
-                    Actions =
-                    {
-                        new P.OrchestratorAction
+                        // Complete the orchestration with a failed status and failure details
+                        P.OrchestratorResponse failureResponse = new()
                         {
-                            CompleteOrchestration = new P.CompleteOrchestrationAction
+                            InstanceId = response.InstanceId,
+                            CompletionToken = response.CompletionToken,
+                            OrchestrationTraceContext = response.OrchestrationTraceContext,
+                            Actions =
                             {
-                                OrchestrationStatus = P.OrchestrationStatus.Failed,
-                                FailureDetails = validationFailure,
+                                new P.OrchestratorAction
+                                {
+                                    CompleteOrchestration = new P.CompleteOrchestrationAction
+                                    {
+                                        OrchestrationStatus = P.OrchestrationStatus.Failed,
+                                        FailureDetails = validationFailure,
+                                    },
+                                },
                             },
-                        },
-                    },
-                };
+                        };
 
-                await this.ExecuteWithRetryAsync(
-                    async () => await this.client.CompleteOrchestratorTaskAsync(failureResponse, cancellationToken: cancellationToken),
-                    nameof(this.client.CompleteOrchestratorTaskAsync),
-                    cancellationToken);
-                return;
-            }
+                        await this.ExecuteWithRetryAsync(
+                            async () => await this.client.CompleteOrchestratorTaskAsync(failureResponse, cancellationToken: cancellationToken),
+                            nameof(this.client.CompleteOrchestratorTaskAsync),
+                            cancellationToken);
+                        return;
+                    }
 
-            // Helper to add an action to the current chunk if it fits
-            static bool TryAddAction(
-                Google.Protobuf.Collections.RepeatedField<P.OrchestratorAction> dest,
-                P.OrchestratorAction action,
-                ref int currentSize,
-                int maxChunkBytes)
-            {
-                int actionSize = action.CalculateSize();
-                if (currentSize + actionSize > maxChunkBytes && currentSize > 0)
-                {
-                    return false;
+                    actionSizes[i] = actionSize;
                 }
-
-                dest.Add(action);
-                currentSize += actionSize;
-                return true;
             }
 
-            // Check if the entire response fits in one chunk
+            // Calculate the whole response size exactly once. If it fits in a single chunk, send
+            // it directly without any further per-action work.
             int totalSize = response.CalculateSize();
+            testResponseSizedHook?.Invoke();
             if (totalSize <= maxChunkBytes)
             {
                 // Response fits in one chunk, send it directly (isPartial defaults to false)
@@ -1183,9 +1214,21 @@ await this.ExecuteWithRetryAsync(
                 return;
             }
 
+            // Response is too large to fit in a single chunk. Reuse the action sizes computed
+            // above during validation if we have them (LargePayloads not present); otherwise (no
+            // validation was needed) compute each action's serialized size exactly once here. Either
+            // way, the cached sizes are reused for chunk packing below instead of being recalculated.
+            if (actionSizes == null)
+            {
+                actionSizes = new int[response.Actions.Count];
+                for (int i = 0; i < response.Actions.Count; i++)
+                {
+                    actionSizes[i] = response.Actions[i].CalculateSize();
+                }
+            }
+
             // Response is too large, split into multiple chunks
             int actionsCompletedSoFar = 0, chunkIndex = 0;
-            List<P.OrchestratorAction> allActions = response.Actions.ToList();
             bool isPartial = true;
             bool isChunkedMode = false;
 
@@ -1203,15 +1246,15 @@ await this.ExecuteWithRetryAsync(
                 int chunkPayloadSize = 0;
 
                 // Fill the chunk with actions until we reach the size limit
-                while (actionsCompletedSoFar < allActions.Count &&
-                       TryAddAction(chunkedResponse.Actions, allActions[actionsCompletedSoFar], ref chunkPayloadSize, maxChunkBytes))
+                while (actionsCompletedSoFar < response.Actions.Count &&
+                       TryAddAction(chunkedResponse.Actions, response.Actions[actionsCompletedSoFar], actionSizes[actionsCompletedSoFar], ref chunkPayloadSize, maxChunkBytes))
                 {
                     actionsCompletedSoFar++;
                 }
 
                 // Determine if this is a partial chunk (more actions remaining)
 #pragma warning disable CS0612 // isPartial/chunkIndex are deprecated but still required for chunked response wire compatibility.
-                isPartial = actionsCompletedSoFar < allActions.Count;
+                isPartial = actionsCompletedSoFar < response.Actions.Count;
                 chunkedResponse.IsPartial = isPartial;
 
                 // Only activate chunked mode when we actually need multiple chunks.