diff --git a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs index 5dd18d52..3eef0b35 100644 --- a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs +++ b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs @@ -1130,11 +1130,40 @@ async Task CompleteOrchestratorTaskWithChunkingAsync( int maxChunkBytes, CancellationToken cancellationToken) { - // Validate that no single action exceeds the maximum chunk size - static P.TaskFailureDetails? ValidateActionsSize(IEnumerable 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 dest, + P.OrchestratorAction action, + int actionSize, + ref int currentSize, + int maxChunkBytes) { - foreach (P.OrchestratorAction action in actions) + 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) + { + 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) { @@ -1142,68 +1171,45 @@ async Task CompleteOrchestratorTaskWithChunkingAsync( 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 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) { @@ -1215,9 +1221,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 allActions = response.Actions.ToList(); bool isPartial = true; bool isChunkedMode = false; @@ -1235,15 +1253,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. diff --git a/test/Worker/Grpc.Tests/CompleteOrchestratorTaskWithChunkingTests.cs b/test/Worker/Grpc.Tests/CompleteOrchestratorTaskWithChunkingTests.cs new file mode 100644 index 00000000..92aeb39f --- /dev/null +++ b/test/Worker/Grpc.Tests/CompleteOrchestratorTaskWithChunkingTests.cs @@ -0,0 +1,509 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Reflection; +using FluentAssertions; +using Grpc.Core; +using Microsoft.DurableTask.Worker; +using Microsoft.DurableTask.Worker.Grpc; +using Microsoft.Extensions.Logging.Abstractions; +using Microsoft.Extensions.Options; +using Moq; +using P = Microsoft.DurableTask.Protobuf; + +namespace Microsoft.DurableTask.Worker.Grpc.Tests; + +/// +/// Focused unit tests for Processor.CompleteOrchestratorTaskWithChunkingAsync, covering the +/// size-boundary decisions, the capability combinations, +/// oversized-single-action failure behavior, chunk wire-compatibility, and retry behavior. These tests +/// guard against regressions in the optimization that avoids recalculating protobuf action sizes +/// (see https://github.com/microsoft/durabletask-dotnet/issues/773). +/// +public class CompleteOrchestratorTaskWithChunkingTests +{ + [Fact] + public async Task ResponseFitsExactlyAtLimit_SendsOriginalResponseDirectly_NoChunking() + { + // Arrange + P.OrchestratorResponse response = BuildResponse( + "instance-1", + BuildScheduleTaskAction(0, 16), + BuildScheduleTaskAction(1, 16)); + int maxChunkBytes = response.CalculateSize(); + + using Fixture fixture = Fixture.Create(); + + // Act + await fixture.InvokeAsync(response, maxChunkBytes); + + // Assert - the exact same response instance is sent, unmodified, in a single call. + fixture.Sent.Should().HaveCount(1); + fixture.Sent[0].Should().BeSameAs(response); +#pragma warning disable CS0612 // IsPartial/ChunkIndex are deprecated but still part of the wire contract. + fixture.Sent[0].IsPartial.Should().BeFalse(); + fixture.Sent[0].ChunkIndex.Should().BeNull(); +#pragma warning restore CS0612 + } + + [Fact] + public async Task ResponseOneByteOverLimit_ActionsStillFitIndividually_ProducesSingleNonPartialChunk() + { + // Arrange - the whole response is one byte too large, but rebuilding a chunk from the + // (already fitting) individual actions still yields just one, non-partial chunk. This + // exercises the chunking code path (not the direct-send fast path) while confirming the + // resulting chunk still preserves all wire-compatibility fields correctly. + P.OrchestratorResponse response = BuildResponse( + "instance-2", + BuildScheduleTaskAction(0, 16), + BuildScheduleTaskAction(1, 16)); + response.CustomStatus = "status"; + response.OrchestrationTraceContext = new P.OrchestrationTraceContext { SpanID = "span-1" }; + int exactSize = response.CalculateSize(); + + using Fixture fixture = Fixture.Create(); + + // Act + await fixture.InvokeAsync(response, exactSize - 1); + + // Assert + fixture.Sent.Should().HaveCount(1); + P.OrchestratorResponse sent = fixture.Sent[0]; + sent.Should().NotBeSameAs(response); // Went through the chunking construction path. + sent.InstanceId.Should().Be(response.InstanceId); + sent.CompletionToken.Should().Be(response.CompletionToken); + sent.CustomStatus.Should().Be(response.CustomStatus); + sent.OrchestrationTraceContext.SpanID.Should().Be(response.OrchestrationTraceContext.SpanID); + sent.NumEventsProcessed.Should().BeNull(); + sent.Actions.Should().HaveCount(2); +#pragma warning disable CS0612 + sent.IsPartial.Should().BeFalse(); + sent.ChunkIndex.Should().BeNull(); +#pragma warning restore CS0612 + } + + [Fact] + public async Task ActionExactlyFillsRemainingChunkSpace_IsIncludedInSameChunk() + { + // Arrange - two actions whose combined size exactly equals maxChunkBytes. The boundary + // condition in TryAddAction (currentSize + actionSize > maxChunkBytes) must treat "equal" + // as fitting, so both actions land in the same, single, non-partial chunk. + P.OrchestratorAction action0 = BuildScheduleTaskAction(0, 16); + P.OrchestratorAction action1 = BuildScheduleTaskAction(1, 32); + int size0 = action0.CalculateSize(); + int size1 = action1.CalculateSize(); + + P.OrchestratorResponse response = BuildResponse("instance-3", action0, action1); + using Fixture fixture = Fixture.Create(); + + // Act + await fixture.InvokeAsync(response, size0 + size1); + + // Assert + fixture.Sent.Should().HaveCount(1); + fixture.Sent[0].Actions.Should().HaveCount(2); + fixture.Sent[0].Actions[0].Id.Should().Be(0); + fixture.Sent[0].Actions[1].Id.Should().Be(1); +#pragma warning disable CS0612 + fixture.Sent[0].IsPartial.Should().BeFalse(); +#pragma warning restore CS0612 + } + + [Fact] + public async Task ActionExceedsRemainingChunkSpaceByOneByte_IsMovedToNextChunk() + { + // Arrange - same two actions, but maxChunkBytes is one byte less than their combined size. + // The second action must now overflow into its own, second chunk. + P.OrchestratorAction action0 = BuildScheduleTaskAction(0, 16); + P.OrchestratorAction action1 = BuildScheduleTaskAction(1, 32); + int size0 = action0.CalculateSize(); + int size1 = action1.CalculateSize(); + + P.OrchestratorResponse response = BuildResponse("instance-4", action0, action1); + using Fixture fixture = Fixture.Create(); + + // Act + await fixture.InvokeAsync(response, size0 + size1 - 1); + + // Assert + fixture.Sent.Should().HaveCount(2); + fixture.Sent[0].Actions.Should().ContainSingle(a => a.Id == 0); + fixture.Sent[1].Actions.Should().ContainSingle(a => a.Id == 1); +#pragma warning disable CS0612 + fixture.Sent[0].IsPartial.Should().BeTrue(); + fixture.Sent[0].ChunkIndex.Should().Be(0); + fixture.Sent[1].IsPartial.Should().BeFalse(); + fixture.Sent[1].ChunkIndex.Should().Be(1); +#pragma warning restore CS0612 + } + + [Fact] + public async Task NoLargePayloadsCapability_SingleOversizedAction_ReturnsFailedCompletionResponse() + { + // Arrange - one small action and one large action that alone exceeds maxChunkBytes. Without + // the LargePayloads capability, this must fail the orchestration instead of sending the + // oversized action. + P.OrchestratorAction small = BuildScheduleTaskAction(0, 16); + P.OrchestratorAction large = BuildScheduleTaskAction(1, 2048); + int largeSize = large.CalculateSize(); + int maxChunkBytes = largeSize - 1; + + P.OrchestratorResponse response = BuildResponse("instance-5", small, large); + response.OrchestrationTraceContext = new P.OrchestrationTraceContext { SpanID = "span-5" }; + using Fixture fixture = Fixture.Create(largePayloads: false); + + // Act + await fixture.InvokeAsync(response, maxChunkBytes); + + // Assert + fixture.Sent.Should().HaveCount(1); + P.OrchestratorResponse failure = fixture.Sent[0]; + failure.InstanceId.Should().Be(response.InstanceId); + failure.CompletionToken.Should().Be(response.CompletionToken); + failure.OrchestrationTraceContext.SpanID.Should().Be("span-5"); + failure.Actions.Should().HaveCount(1); + P.OrchestratorAction failureAction = failure.Actions[0]; + failureAction.CompleteOrchestration.Should().NotBeNull(); + failureAction.CompleteOrchestration.OrchestrationStatus.Should().Be(P.OrchestrationStatus.Failed); + failureAction.CompleteOrchestration.FailureDetails.IsNonRetriable.Should().BeTrue(); + failureAction.CompleteOrchestration.FailureDetails.ErrorType.Should().Be(typeof(InvalidOperationException).FullName); + string expectedMessage = $"A single orchestrator action of type ScheduleTask with id 1 " + + $"exceeds the {maxChunkBytes / 1024.0 / 1024.0:F2}MB limit: {largeSize / 1024.0 / 1024.0:F2}MB. " + + "Enable large-payload externalization to Azure Blob Storage to support oversized actions."; + failureAction.CompleteOrchestration.FailureDetails.ErrorMessage.Should().Be(expectedMessage); + } + + [Fact] + public async Task NoLargePayloadsCapability_FirstActionOversized_FailsOnFirstOffender_DoesNotEvaluateLaterActions() + { + // Arrange - regression coverage for a fail-fast ordering bug: without the LargePayloads + // capability, validation must stop at the *first* oversized action instead of first sizing + // the whole response and/or every action. Action id=0 is oversized, and several later + // actions (ids 1-3) are ALSO individually oversized with distinguishable ids. If the + // algorithm regresses back to "size everything, then scan for a failure", it could still + // produce *a* failure, but this asserts it fails on id=0 specifically - proving iteration + // stopped at the very first offender rather than continuing to size/scan later actions. + P.OrchestratorAction action0 = BuildScheduleTaskAction(0, 2048); // oversized - first, must win + P.OrchestratorAction action1 = BuildScheduleTaskAction(1, 2048); // also oversized + P.OrchestratorAction action2 = BuildScheduleTaskAction(2, 2048); // also oversized + P.OrchestratorAction action3 = BuildScheduleTaskAction(3, 2048); // also oversized + int maxChunkBytes = action0.CalculateSize() - 1; + + P.OrchestratorResponse response = BuildResponse("instance-9", action0, action1, action2, action3); + using Fixture fixture = Fixture.Create(largePayloads: false); + + // Act + await fixture.InvokeAsync(response, maxChunkBytes); + + // Assert - exactly one response was sent (no partial chunks for the other oversized + // actions), and it is a failure referencing id=0. + fixture.Sent.Should().HaveCount(1); + P.OrchestratorResponse failure = fixture.Sent[0]; + failure.Actions.Should().HaveCount(1); + P.OrchestratorAction failureAction = failure.Actions[0]; + failureAction.CompleteOrchestration.Should().NotBeNull(); + failureAction.CompleteOrchestration.OrchestrationStatus.Should().Be(P.OrchestrationStatus.Failed); + failureAction.CompleteOrchestration.FailureDetails.ErrorMessage.Should().Contain("with id 0 "); + failureAction.CompleteOrchestration.FailureDetails.ErrorMessage.Should().NotContain("with id 1 "); + failureAction.CompleteOrchestration.FailureDetails.ErrorMessage.Should().NotContain("with id 2 "); + failureAction.CompleteOrchestration.FailureDetails.ErrorMessage.Should().NotContain("with id 3 "); + } + + [Fact] + public async Task LargePayloadsCapability_SingleOversizedAction_IsSentInsteadOfFailing() + { + // Arrange - same oversized action, but with LargePayloads capability announced. The action + // must be allowed through (as its own chunk) rather than failing the orchestration. + P.OrchestratorAction small = BuildScheduleTaskAction(0, 16); + P.OrchestratorAction large = BuildScheduleTaskAction(1, 2048); + int largeSize = large.CalculateSize(); + int maxChunkBytes = largeSize - 1; + + P.OrchestratorResponse response = BuildResponse("instance-6", small, large); + using Fixture fixture = Fixture.Create(largePayloads: true); + + // Act + await fixture.InvokeAsync(response, maxChunkBytes); + + // Assert - both actions are sent (each too big to share a chunk with the other), and + // neither call carries a failed CompleteOrchestration action. + fixture.Sent.Should().HaveCountGreaterOrEqualTo(1); + fixture.Sent.SelectMany(r => r.Actions).Should().Contain(a => a.Id == 0); + fixture.Sent.SelectMany(r => r.Actions).Should().Contain(a => a.Id == 1); + fixture.Sent.SelectMany(r => r.Actions).Should().NotContain(a => a.OrchestratorActionTypeCase == P.OrchestratorAction.OrchestratorActionTypeOneofCase.CompleteOrchestration); + } + + [Fact] + public async Task MultiChunkResponse_PreservesWireCompatibilityFieldsAcrossChunks() + { + // Arrange - three actions that must span three separate chunks, verifying InstanceId, + // CompletionToken, and CustomStatus repeat every chunk; OrchestrationTraceContext is only + // set on the first chunk; NumEventsProcessed is null on the first chunk and 0 afterward; + // RequiresHistory is preserved; and chunk indices/IsPartial sequence correctly. + P.OrchestratorAction action0 = BuildScheduleTaskAction(0, 16); + P.OrchestratorAction action1 = BuildScheduleTaskAction(1, 16); + P.OrchestratorAction action2 = BuildScheduleTaskAction(2, 16); + int maxChunkBytes = Math.Max(action0.CalculateSize(), Math.Max(action1.CalculateSize(), action2.CalculateSize())); + + P.OrchestratorResponse response = BuildResponse("instance-7", action0, action1, action2); + response.CustomStatus = "custom-status"; + response.RequiresHistory = true; + response.OrchestrationTraceContext = new P.OrchestrationTraceContext { SpanID = "span-7" }; + + using Fixture fixture = Fixture.Create(); + + // Act + await fixture.InvokeAsync(response, maxChunkBytes); + + // Assert + fixture.Sent.Should().HaveCount(3); + for (int i = 0; i < fixture.Sent.Count; i++) + { + P.OrchestratorResponse chunk = fixture.Sent[i]; + chunk.InstanceId.Should().Be("instance-7"); + chunk.CompletionToken.Should().Be(response.CompletionToken); + chunk.CustomStatus.Should().Be("custom-status"); + chunk.RequiresHistory.Should().BeTrue(); +#pragma warning disable CS0612 + chunk.ChunkIndex.Should().Be(i); + chunk.IsPartial.Should().Be(i < fixture.Sent.Count - 1); +#pragma warning restore CS0612 + + if (i == 0) + { + chunk.NumEventsProcessed.Should().BeNull(); + chunk.OrchestrationTraceContext.Should().NotBeNull(); + chunk.OrchestrationTraceContext.SpanID.Should().Be("span-7"); + } + else + { + chunk.NumEventsProcessed.Should().Be(0); + chunk.OrchestrationTraceContext.Should().BeNull(); + } + } + + // All three actions were sent, in order, across the chunks, with none duplicated or dropped. + fixture.Sent.SelectMany(r => r.Actions).Select(a => a.Id).Should().Equal(0, 1, 2); + } + + [Fact] + public async Task TransientRpcError_DuringSend_RetriesAndEventuallySucceeds() + { + // Arrange - a response that fits in a single chunk (fast path), whose first send attempt + // fails with a transient gRPC error. The method must still rely on ExecuteWithRetryAsync to + // retry the same request and eventually succeed. + P.OrchestratorResponse response = BuildResponse("instance-8", BuildScheduleTaskAction(0, 16)); + int maxChunkBytes = response.CalculateSize(); + + using Fixture fixture = Fixture.Create(transientRetryBackoffBase: TimeSpan.FromMilliseconds(1)); + fixture.FailNextAttempts(1); + + // Act + await fixture.InvokeAsync(response, maxChunkBytes); + + // Assert + fixture.AttemptCount.Should().Be(2); + fixture.Sent.Should().HaveCount(1); + fixture.Sent[0].Should().BeSameAs(response); + } + + static P.OrchestratorResponse BuildResponse(string instanceId, params P.OrchestratorAction[] actions) + { + P.OrchestratorResponse response = new() + { + InstanceId = instanceId, + CompletionToken = Guid.NewGuid().ToString("N"), + }; + response.Actions.AddRange(actions); + return response; + } + + static P.OrchestratorAction BuildScheduleTaskAction(int id, int payloadBytes) + { + return new P.OrchestratorAction + { + Id = id, + ScheduleTask = new P.ScheduleTaskAction + { + Name = "Echo", + Input = new string('x', payloadBytes), + }, + }; + } + + /// + /// Test fixture that constructs a real GrpcDurableTaskWorker.Processor via reflection (it + /// is a private nested type) wired to a strictly-mocked gRPC client, and exposes a helper to + /// invoke the private CompleteOrchestratorTaskWithChunkingAsync method directly. + /// + sealed class Fixture : IDisposable + { + static readonly MethodInfo Method = FindMethod(); + + readonly object processor; + readonly object gate = new(); + int attemptsToFail; + + Fixture(object processor, Mock clientMock) + { + this.processor = processor; + this.ClientMock = clientMock; + } + + public Mock ClientMock { get; } + + public List Sent { get; } = new(); + + public int AttemptCount { get; private set; } + + public static Fixture Create( + bool largePayloads = false, + TimeSpan? transientRetryBackoffBase = null) + { + GrpcDurableTaskWorkerOptions grpcOptionsValue = new(); + if (largePayloads) + { + grpcOptionsValue.Capabilities.Add(P.WorkerCapability.LargePayloads); + } + + if (transientRetryBackoffBase.HasValue) + { + grpcOptionsValue.Internal.TransientRetryBackoffBase = transientRetryBackoffBase.Value; + } + + OptionsMonitorStub grpcOptions = new(grpcOptionsValue); + OptionsMonitorStub workerOptions = new(new DurableTaskWorkerOptions()); + Mock factoryMock = new(MockBehavior.Strict); + + GrpcDurableTaskWorker worker = new( + name: "Test", + factory: factoryMock.Object, + grpcOptions: grpcOptions, + workerOptions: workerOptions, + services: Mock.Of(), + loggerFactory: NullLoggerFactory.Instance, + orchestrationFilter: null, + exceptionPropertiesProvider: null); + + CallInvoker callInvoker = Mock.Of(); + Mock clientMock = new( + MockBehavior.Strict, new object[] { callInvoker }); + + Type processorType = typeof(GrpcDurableTaskWorker).GetNestedType("Processor", BindingFlags.NonPublic)!; + object processorInstance = Activator.CreateInstance( + processorType, + BindingFlags.Public | BindingFlags.Instance, + binder: null, + args: new object?[] { worker, clientMock.Object, null, null }, + culture: null)!; + + Fixture fixture = new(processorInstance, clientMock); + + clientMock + .Setup(c => c.CompleteOrchestratorTaskAsync( + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns((P.OrchestratorResponse r, Metadata h, DateTime? d, CancellationToken ct) => + fixture.HandleSend(r)); + + return fixture; + } + + /// + /// Configures the next send attempts (across all chunks) to fail + /// with a transient (Unavailable) gRPC error before subsequent attempts succeed. + /// + public void FailNextAttempts(int count) + { + lock (this.gate) + { + this.attemptsToFail = count; + } + } + + public Task InvokeAsync(P.OrchestratorResponse response, int maxChunkBytes, CancellationToken cancellationToken = default) + { + return (Task)Method.Invoke(this.processor, new object?[] { response, maxChunkBytes, cancellationToken })!; + } + + public void Dispose() + { + } + + AsyncUnaryCall HandleSend(P.OrchestratorResponse response) + { + bool shouldFail; + lock (this.gate) + { + this.AttemptCount++; + shouldFail = this.attemptsToFail > 0; + if (shouldFail) + { + this.attemptsToFail--; + } + } + + if (shouldFail) + { + return RpcExceptionAsyncUnaryCall(StatusCode.Unavailable); + } + + this.Sent.Add(response); + return CompletedAsyncUnaryCall(new P.CompleteTaskResponse()); + } + + static MethodInfo FindMethod() + { + Type processorType = typeof(GrpcDurableTaskWorker).GetNestedType("Processor", BindingFlags.NonPublic)!; + return processorType.GetMethod("CompleteOrchestratorTaskWithChunkingAsync", BindingFlags.Instance | BindingFlags.NonPublic)!; + } + + static AsyncUnaryCall CompletedAsyncUnaryCall(T response) + { + Task respTask = Task.FromResult(response); + return new AsyncUnaryCall( + respTask, + Task.FromResult(new Metadata()), + () => new Status(StatusCode.OK, string.Empty), + () => new Metadata(), + () => { }); + } + + static AsyncUnaryCall RpcExceptionAsyncUnaryCall(StatusCode statusCode, string detail = "transient error") + { + RpcException ex = new(new Status(statusCode, detail)); + Task respTask = Task.FromException(ex); + return new AsyncUnaryCall( + respTask, + Task.FromResult(new Metadata()), + () => new Status(statusCode, detail), + () => new Metadata(), + () => { }); + } + + sealed class OptionsMonitorStub : IOptionsMonitor + where T : class, new() + { + readonly T value; + + public OptionsMonitorStub(T value) => this.value = value; + + public T CurrentValue => this.value; + + public T Get(string? name) => this.value; + + public IDisposable OnChange(Action listener) => NullDisposable.Instance; + + sealed class NullDisposable : IDisposable + { + public static readonly NullDisposable Instance = new(); + + public void Dispose() + { + } + } + } + } +}