Skip to content
125 changes: 72 additions & 53 deletions src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs
Original file line number Diff line number Diff line change
Expand Up @@ -1098,80 +1098,86 @@ 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 we never need to call CalculateSize() on the same action more than once.
Comment thread
berndverst marked this conversation as resolved.
Outdated
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();
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();
if (totalSize <= maxChunkBytes)
{
Expand All @@ -1183,9 +1189,22 @@ 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.
List<P.OrchestratorAction> allActions = response.Actions.ToList();
if (actionSizes == null)
{
actionSizes = new int[allActions.Count];
for (int i = 0; i < allActions.Count; i++)
{
actionSizes[i] = allActions[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;

Expand All @@ -1204,7 +1223,7 @@ await this.ExecuteWithRetryAsync(

// Fill the chunk with actions until we reach the size limit
while (actionsCompletedSoFar < allActions.Count &&
TryAddAction(chunkedResponse.Actions, allActions[actionsCompletedSoFar], ref chunkPayloadSize, maxChunkBytes))
TryAddAction(chunkedResponse.Actions, allActions[actionsCompletedSoFar], actionSizes[actionsCompletedSoFar], ref chunkPayloadSize, maxChunkBytes))
{
actionsCompletedSoFar++;
}
Expand Down
Loading
Loading