diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bdaf23b..4309a5a2 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/src/Extensions/AzureBlobPayloads/DependencyInjection/DurableTaskWorkerBuilderExtensions.AzureBlobPayloads.cs b/src/Extensions/AzureBlobPayloads/DependencyInjection/DurableTaskWorkerBuilderExtensions.AzureBlobPayloads.cs index a5564431..62ec0f19 100644 --- a/src/Extensions/AzureBlobPayloads/DependencyInjection/DurableTaskWorkerBuilderExtensions.AzureBlobPayloads.cs +++ b/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); diff --git a/src/Extensions/AzureBlobPayloads/Interceptors/AzureBlobPayloadsSideCarInterceptor.cs b/src/Extensions/AzureBlobPayloads/Interceptors/AzureBlobPayloadsSideCarInterceptor.cs index 29fd76f5..21b65099 100644 --- a/src/Extensions/AzureBlobPayloads/Interceptors/AzureBlobPayloadsSideCarInterceptor.cs +++ b/src/Extensions/AzureBlobPayloads/Interceptors/AzureBlobPayloadsSideCarInterceptor.cs @@ -27,6 +27,50 @@ public sealed class AzureBlobPayloadsSideCarInterceptor(PayloadStore payloadStor Action? SharedMessageLockAcquiredForTest { get; set; } + /// + /// Externalizes a response before the worker finalizes terminal entity cleanup. + /// + /// The orchestration response. + /// Cancellation token. + /// A task representing response preparation. + 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, + }, + }, + }); + } + } + /// protected override async Task ExternalizeRequestPayloadsAsync(TRequest request, CancellationToken cancellation) { @@ -76,37 +120,7 @@ protected override async Task ExternalizeRequestPayloadsAsync(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); diff --git a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs index 3274229c..76b9fb5f 100644 --- a/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs +++ b/src/Extensions/AzureBlobPayloads/Interceptors/PayloadInterceptor.cs @@ -174,7 +174,7 @@ public override AsyncServerStreamingCall AsyncServerStreamingCall(TResponse response, CancellationToken cancellation); /// - /// Externalizes a payload if it exceeds the threshold. + /// Externalizes a payload if it exceeds the threshold and is not already a known payload token. /// /// The value to potentially externalize. /// Cancellation token. @@ -232,11 +232,11 @@ private protected async Task ExternalizePayloadAsync( /// Calculates a value's UTF-8 size and determines whether it requires externalization. /// /// The value to inspect. - /// The calculated UTF-8 size, or zero when is empty. - /// true when the value meets the configured externalization threshold. + /// The calculated UTF-8 size, or zero when is empty or a known payload token. + /// true when a non-token value meets the configured externalization threshold. 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; diff --git a/src/Shared/Grpc/ProtoUtils.cs b/src/Shared/Grpc/ProtoUtils.cs index 45779e67..981610ee 100644 --- a/src/Shared/Grpc/ProtoUtils.cs +++ b/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); } + /// + /// Inserts any outstanding entity releases before a terminal completion. + /// + /// The response containing the completion. + /// The completion action already added to the response. + 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, + }, + }, + }); + } + } + /// /// Returns the remaining unlock obligations, and clears the list. /// diff --git a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs b/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs index abfba18a..21206156 100644 --- a/src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs +++ b/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( diff --git a/src/Worker/Grpc/GrpcDurableTaskWorkerOptions.cs b/src/Worker/Grpc/GrpcDurableTaskWorkerOptions.cs index 0134a959..46cbcc6f 100644 --- a/src/Worker/Grpc/GrpcDurableTaskWorkerOptions.cs +++ b/src/Worker/Grpc/GrpcDurableTaskWorkerOptions.cs @@ -135,6 +135,12 @@ internal class InternalOptions /// public bool InsertEntityUnlocksOnCompletion { get; set; } + /// + /// 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. + /// + public Func? PrepareOrchestratorResponseAsync { get; set; } + /// /// 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. diff --git a/src/Worker/Grpc/Internal/InternalOptionsExtensions.cs b/src/Worker/Grpc/Internal/InternalOptionsExtensions.cs index bf075440..bd561b9d 100644 --- a/src/Worker/Grpc/Internal/InternalOptionsExtensions.cs +++ b/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)); } + + /// + /// 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. + /// + /// The gRPC worker options. + /// The callback receiving the mutable response and the worker cancellation token. + /// + /// 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. + /// + /// or is null. + public static void SetOrchestratorResponsePreparation( + this GrpcDurableTaskWorkerOptions options, + Func preparation) + { + Check.NotNull(options); + options.Internal.PrepareOrchestratorResponseAsync = preparation ?? throw new ArgumentNullException(nameof(preparation)); + } } diff --git a/test/Client/Grpc.Tests/ProtoUtilsTests.cs b/test/Client/Grpc.Tests/ProtoUtilsTests.cs index 5eb0d278..bf487e71 100644 --- a/test/Client/Grpc.Tests/ProtoUtilsTests.cs +++ b/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(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() { diff --git a/test/Extensions/AzureBlobPayloads.Tests/DependencyInjection/UseExternalizedPayloadsTests.cs b/test/Extensions/AzureBlobPayloads.Tests/DependencyInjection/UseExternalizedPayloadsTests.cs index 242a5fad..389d355f 100644 --- a/test/Extensions/AzureBlobPayloads.Tests/DependencyInjection/UseExternalizedPayloadsTests.cs +++ b/test/Extensions/AzureBlobPayloads.Tests/DependencyInjection/UseExternalizedPayloadsTests.cs @@ -15,6 +15,7 @@ using Microsoft.Extensions.Options; using Moq; using Xunit; +using P = Microsoft.DurableTask.Protobuf; namespace Microsoft.DurableTask.Extensions.AzureBlobPayloads.Tests.DependencyInjection; @@ -122,12 +123,69 @@ public void UseExternalizedPayloads_Worker_LeavesTheGrpcWorkerOptionsAlone() GrpcDurableTaskWorkerOptions grpcOptions = provider.GetRequiredService>().Get(string.Empty); - // Assert - the configured transport survives PostConfigure untouched, and the large-payload capability - // is the only thing this extension adds to the worker's gRPC options. + // Assert - the configured transport survives PostConfigure untouched. grpcOptions.CallInvoker.Should().BeSameAs(configured); grpcOptions.Channel.Should().BeNull(); } + [Theory] + [InlineData("")] + [InlineData("worker-hub")] + public async Task UseExternalizedPayloads_Worker_RegistersResponsePreparation(string name) + { + // Arrange + const string Payload = "payload"; + const string Token = "blob:v1:test:prepared"; + using CancellationTokenSource cancellation = new(); + Mock store = new(MockBehavior.Strict); + store.Setup(value => value.IsKnownPayloadToken(It.IsAny())) + .Returns((string value) => value == Token); + store.Setup(value => value.UploadAsync(Payload, cancellation.Token)).ReturnsAsync(Token); + ServiceCollection services = new(); + services.AddSingleton(store.Object); + Mock builder = new(); + builder.Setup(value => value.Services).Returns(services); + builder.Setup(value => value.Name).Returns(name); + builder.Object.UseExternalizedPayloads(options => + { + options.ThresholdBytes = 1; + options.MaxPayloadBytes = 1024; + }); + using ServiceProvider provider = services.BuildServiceProvider(); + IOptionsMonitor monitor = + provider.GetRequiredService>(); + P.OrchestratorResponse response = new() + { + InstanceId = "instance", + CompletionToken = "completion-token", + Actions = + { + new P.OrchestratorAction + { + CompleteOrchestration = new P.CompleteOrchestrationAction + { + OrchestrationStatus = P.OrchestrationStatus.Completed, + Result = Payload, + }, + }, + }, + }; + + // Act + GrpcDurableTaskWorkerOptions grpcOptions = monitor.Get(name); + Func preparation = + Assert.IsType>( + grpcOptions.Internal.PrepareOrchestratorResponseAsync); + await preparation(response, cancellation.Token); + + // Assert + preparation.Target.Should().BeSameAs(Assert.Single(grpcOptions.Interceptors)); + Assert.Single(response.Actions).CompleteOrchestration.Result.Should().Be(Token); + response.CompletionToken.Should().Be("completion-token"); + monitor.Get("other-hub").Internal.PrepareOrchestratorResponseAsync.Should().BeNull(); + store.Verify(value => value.UploadAsync(Payload, cancellation.Token), Times.Once); + } + [Theory] [InlineData(false)] [InlineData(true)] diff --git a/test/Extensions/AzureBlobPayloads.Tests/OrchestratorPayloadFailureTests.cs b/test/Extensions/AzureBlobPayloads.Tests/OrchestratorPayloadFailureTests.cs new file mode 100644 index 00000000..020b2ab4 --- /dev/null +++ b/test/Extensions/AzureBlobPayloads.Tests/OrchestratorPayloadFailureTests.cs @@ -0,0 +1,352 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Text.Json; +using System.Threading.Channels; +using Azure; +using Google.Protobuf.WellKnownTypes; +using Grpc.Core; +using Microsoft.DurableTask.Entities; +using Microsoft.DurableTask.Worker; +using Microsoft.DurableTask.Worker.Grpc; +using Microsoft.DurableTask.Worker.Grpc.Internal; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Hosting; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.Logging.Abstractions; +using Moq; +using Xunit; +using P = Microsoft.DurableTask.Protobuf; + +namespace Microsoft.DurableTask.Extensions.AzureBlobPayloads.Tests; + +public sealed class OrchestratorPayloadFailureTests +{ + enum PayloadFailure + { + None, + SizeLimit, + ForbiddenUpload, + } + + [Theory] + [InlineData(false, false, false)] + [InlineData(true, false, false)] + [InlineData(false, true, false)] + [InlineData(true, true, false)] + [InlineData(false, false, true)] + [InlineData(true, false, true)] + [InlineData(false, true, true)] + [InlineData(true, true, true)] + public Task PermanentPayloadFailure_ReleasesEntitiesBeforeSendingFailedCompletion( + bool completeOrchestration, + bool streamHistory, + bool failUpload) + => RunLockedOrchestrationAsync( + completeOrchestration, + streamHistory, + failUpload ? PayloadFailure.ForbiddenUpload : PayloadFailure.SizeLimit); + + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public Task SuccessfulCompletion_ReleasesEntitiesExactlyOnceBeforeCompletion( + bool useExternalizedPayloads, + bool explicitlyReleaseLocks) + => RunLockedOrchestrationAsync( + completeOrchestration: true, + streamHistory: false, + failure: PayloadFailure.None, + useExternalizedPayloads: useExternalizedPayloads, + explicitlyReleaseLocks: explicitlyReleaseLocks); + + static async Task RunLockedOrchestrationAsync( + bool completeOrchestration, + bool streamHistory, + PayloadFailure failure, + bool useExternalizedPayloads = true, + bool explicitlyReleaseLocks = true) + { + // Arrange + const string InstanceId = "locked-payload-instance"; + const string ExecutionId = "execution-1"; + const string OrchestratorName = "LockedPayloadOrchestrator"; + EntityInstanceId[] entities = [new("counter", "one"), new("counter", "two")]; + string payload = failure == PayloadFailure.None + ? "completed" + : new('x', GrpcDurableTaskWorkerOptions.MinCompleteOrchestrationWorkItemChunkSizeInBytes + 1024); + string serializedPayload = JsonSerializer.Serialize(payload); + Channel workItems = Channel.CreateUnbounded(); + Channel history = Channel.CreateUnbounded(); + Channel completions = Channel.CreateUnbounded(); + Mock invoker = new(MockBehavior.Strict); + invoker.Setup(value => value.AsyncUnaryCall( + It.IsAny>(), + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns(() => CreateUnaryCall(new Empty())); + invoker.Setup(value => value.AsyncServerStreamingCall( + It.IsAny>(), + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns(() => new AsyncServerStreamingCall( + new ChannelStreamReader(workItems.Reader), + Task.FromResult(new Metadata()), + () => Status.DefaultSuccess, + () => new Metadata(), + () => { })); + invoker.Setup(value => value.AsyncServerStreamingCall( + It.IsAny>(), + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Returns(() => new AsyncServerStreamingCall( + new ChannelStreamReader(history.Reader), + Task.FromResult(new Metadata()), + () => Status.DefaultSuccess, + () => new Metadata(), + () => { })); + invoker.Setup(value => value.AsyncUnaryCall( + It.IsAny>(), + It.IsAny(), + It.IsAny(), + It.IsAny())) + .Callback, string, CallOptions, P.OrchestratorResponse>( + (_, _, _, response) => completions.Writer.TryWrite(response.Clone())) + .Returns(() => CreateUnaryCall(new P.CompleteTaskResponse())); + Mock store = new(MockBehavior.Strict); + store.Setup(value => value.IsKnownPayloadToken(It.IsAny())).Returns(false); + if (failure == PayloadFailure.ForbiddenUpload) + { + store.Setup(value => value.UploadAsync(serializedPayload, It.IsAny())) + .ThrowsAsync(new RequestFailedException(403, "Simulated permanent Blob upload failure.")); + } + + ServiceCollection services = new(); + services.AddSingleton(NullLoggerFactory.Instance); + services.AddSingleton(typeof(ILogger<>), typeof(NullLogger<>)); + services.AddSingleton(store.Object); + services.AddDurableTaskWorker(builder => + { + builder.Configure(options => options.EnableEntitySupport = true); + builder.UseGrpc(options => + { + options.CallInvoker = invoker.Object; + options.ConfigureForAzureManaged(); + options.CompleteOrchestrationWorkItemChunkSizeInBytes = + GrpcDurableTaskWorkerOptions.MinCompleteOrchestrationWorkItemChunkSizeInBytes; + }); + if (useExternalizedPayloads) + { + builder.UseExternalizedPayloads(options => + { + options.ThresholdBytes = 1024; + options.MaxPayloadBytes = failure == PayloadFailure.ForbiddenUpload + ? 2 * GrpcDurableTaskWorkerOptions.MinCompleteOrchestrationWorkItemChunkSizeInBytes + : 2048; + }); + } + + builder.AddTasks(tasks => tasks.AddOrchestratorFunc( + OrchestratorName, + async (context, _) => + { + IAsyncDisposable locks = await context.Entities.LockEntitiesAsync(entities); + await using IAsyncDisposable? releasedLocks = explicitlyReleaseLocks ? locks : null; + if (!completeOrchestration) + { + await context.Entities.CallEntityAsync(entities[0], "set", payload); + } + + return payload; + })); + }); + await using ServiceProvider provider = services.BuildServiceProvider(); + IHostedService worker = Assert.Single(provider.GetServices()); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(30)); + DateTime timestamp = new(2026, 1, 1, 0, 0, 0, DateTimeKind.Utc); + P.OrchestratorRequest initialRequest = new() + { + InstanceId = InstanceId, + ExecutionId = ExecutionId, + EntityParameters = new P.OrchestratorEntityParameters + { + EntityMessageReorderWindow = Duration.FromTimeSpan(TimeSpan.Zero), + }, + NewEvents = + { + new P.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(timestamp), + OrchestratorStarted = new P.OrchestratorStartedEvent(), + }, + new P.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(timestamp), + ExecutionStarted = new P.ExecutionStartedEvent + { + Name = OrchestratorName, + OrchestrationInstance = new P.OrchestrationInstance + { + InstanceId = InstanceId, + ExecutionId = ExecutionId, + }, + }, + }, + }, + }; + await worker.StartAsync(timeout.Token); + try + { + await workItems.Writer.WriteAsync(new P.WorkItem + { + OrchestratorRequest = initialRequest, + CompletionToken = "acquire-token", + }, timeout.Token); + P.OrchestratorResponse acquiring = await completions.Reader.ReadAsync(timeout.Token); + P.OrchestratorAction acquireAction = Assert.Single(acquiring.Actions); + Assert.True(acquireAction.SendEntityMessage is not null, acquiring.ToString()); + P.EntityLockRequestedEvent acquire = acquireAction.SendEntityMessage.EntityLockRequested; + Assert.NotNull(acquire); + P.OrchestratorRequest resumedRequest = new() + { + InstanceId = InstanceId, + ExecutionId = ExecutionId, + EntityParameters = initialRequest.EntityParameters.Clone(), + PastEvents = + { + initialRequest.NewEvents, + new P.HistoryEvent + { + EventId = acquireAction.Id, + Timestamp = Timestamp.FromDateTime(timestamp), + EntityLockRequested = acquire.Clone(), + }, + new P.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(timestamp), + OrchestratorCompleted = new P.OrchestratorCompletedEvent(), + }, + }, + NewEvents = + { + new P.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(timestamp.AddSeconds(1)), + OrchestratorStarted = new P.OrchestratorStartedEvent(), + }, + new P.HistoryEvent + { + EventId = -1, + Timestamp = Timestamp.FromDateTime(timestamp.AddSeconds(1)), + EntityLockGranted = new P.EntityLockGrantedEvent + { + CriticalSectionId = acquire.CriticalSectionId, + }, + }, + }, + }; + if (streamHistory) + { + await history.Writer.WriteAsync( + new P.HistoryChunk { Events = { resumedRequest.PastEvents } }, timeout.Token); + history.Writer.Complete(); + resumedRequest.PastEvents.Clear(); + resumedRequest.RequiresHistoryStreaming = true; + } + + // Act + await workItems.Writer.WriteAsync(new P.WorkItem + { + OrchestratorRequest = resumedRequest, + CompletionToken = "resumed-token", + }, timeout.Token); + P.OrchestratorResponse response = await completions.Reader.ReadAsync(timeout.Token); + + // Assert + P.OrchestratorAction completionAction = Assert.Single( + response.Actions, action => action.CompleteOrchestration is not null); + Assert.Same(completionAction, response.Actions.Last()); + P.CompleteOrchestrationAction completion = completionAction.CompleteOrchestration; + if (failure == PayloadFailure.None) + { + Assert.Equal(P.OrchestrationStatus.Completed, completion.OrchestrationStatus); + Assert.Null(completion.FailureDetails); + Assert.Equal(serializedPayload, completion.Result); + } + else + { + Assert.Equal(P.OrchestrationStatus.Failed, completion.OrchestrationStatus); + Assert.Equal( + failure == PayloadFailure.ForbiddenUpload + ? typeof(RequestFailedException).FullName + : typeof(PayloadStorageException).FullName, + completion.FailureDetails.ErrorType); + Assert.True(completion.FailureDetails.IsNonRetriable); + Assert.Null(completion.Result); + } + + Assert.Equal(entities.Length + 1, response.Actions.Count); + P.EntityUnlockSentEvent[] unlocks = response.Actions.Take(entities.Length) + .Select(action => Assert.IsType(action.SendEntityMessage?.EntityUnlockSent)) + .ToArray(); + Assert.Equal(entities.Select(entity => entity.ToString()), unlocks.Select(unlock => unlock.TargetInstanceId)); + Assert.All(unlocks, unlock => + { + Assert.Equal(acquire.CriticalSectionId, unlock.CriticalSectionId); + Assert.Equal(InstanceId, unlock.ParentInstanceId); + }); + Assert.Equal(InstanceId, response.InstanceId); + Assert.Equal("resumed-token", response.CompletionToken); + Assert.Null(response.CustomStatus); +#pragma warning disable CS0612 // isPartial/chunkIndex are deprecated but still required for chunked response wire compatibility. + Assert.False(response.IsPartial); + Assert.Null(response.ChunkIndex); +#pragma warning restore CS0612 + invoker.Verify(value => value.AsyncUnaryCall( + It.IsAny>(), + It.IsAny(), + It.IsAny(), + It.IsAny()), Times.Exactly(2)); + store.Verify( + value => value.UploadAsync(It.IsAny(), It.IsAny()), + failure == PayloadFailure.ForbiddenUpload ? Times.Once() : Times.Never()); + } + finally + { + await worker.StopAsync(CancellationToken.None).WaitAsync(TimeSpan.FromSeconds(10)); + } + } + + static AsyncUnaryCall CreateUnaryCall(T response) + => new( + Task.FromResult(response), + Task.FromResult(new Metadata()), + () => Status.DefaultSuccess, + () => new Metadata(), + () => { }); + + sealed class ChannelStreamReader(ChannelReader reader) : IAsyncStreamReader + { + public T Current { get; private set; } = default!; + + public async Task MoveNext(CancellationToken cancellationToken) + { + if (await reader.WaitToReadAsync(cancellationToken)) + { + this.Current = await reader.ReadAsync(cancellationToken); + return true; + } + + return false; + } + } +} diff --git a/test/Grpc.IntegrationTests/AzureBlobPayloadsSideCarInterceptorTests.cs b/test/Grpc.IntegrationTests/AzureBlobPayloadsSideCarInterceptorTests.cs index c8135ef6..9c6a765e 100644 --- a/test/Grpc.IntegrationTests/AzureBlobPayloadsSideCarInterceptorTests.cs +++ b/test/Grpc.IntegrationTests/AzureBlobPayloadsSideCarInterceptorTests.cs @@ -456,6 +456,24 @@ public async Task ExternalizeRequestPayloadsAsync_OrchestratorResponse_Externali store.MaxObservedConcurrency.Should().BeLessOrEqualTo(8, "concurrency must be bounded to avoid Azure Storage throttling"); } + [Fact] + public async Task ExternalizeRequestPayloadsAsync_CreateInstanceRequest_KnownTokenBypassesSizeLimits() + { + // Arrange + TrackingPayloadStore store = new(); + string token = store.Seed("original input"); + LargePayloadStorageOptions options = new() { ThresholdBytes = 1, MaxPayloadBytes = 1 }; + AzureBlobPayloadsSideCarInterceptor interceptor = new(store, options); + P.CreateInstanceRequest request = new() { Input = token }; + + // Act + await ExternalizeAsync(interceptor, request, CancellationToken.None); + + // Assert + request.Input.Should().Be(token); + store.UploadCount.Should().Be(0); + } + [Fact] public async Task ExternalizeRequestPayloadsAsync_RewindInstanceRequest_ExternalizesReason() { @@ -520,6 +538,7 @@ public async Task ExternalizeRequestPayloadsAsync_OrchestratorResponse_Externali // Act await ExternalizeAsync(interceptor, response, CancellationToken.None); + await ExternalizeAsync(interceptor, response, CancellationToken.None); // Assert IReadOnlyList tokens = GetHistoryEventPayloads(historyEvent); @@ -587,6 +606,7 @@ public async Task ExternalizeRequestPayloadsAsync_EntityBatchRequest_Externalize // Act await ExternalizeAsync(interceptor, request, CancellationToken.None); + await ExternalizeAsync(interceptor, request, CancellationToken.None); // Assert store.GetUploadedValue(request.EntityState).Should().Be("state-0"); @@ -596,6 +616,8 @@ public async Task ExternalizeRequestPayloadsAsync_EntityBatchRequest_Externalize token.Should().NotBe(expectedInputs[i]); store.GetUploadedValue(token).Should().Be(expectedInputs[i]); } + + store.UploadCount.Should().Be(operationCount + 1); } [Fact] @@ -628,6 +650,7 @@ public async Task ExternalizeRequestPayloadsAsync_EntityBatchResult_Externalizes // Act await ExternalizeAsync(interceptor, result, CancellationToken.None); + await ExternalizeAsync(interceptor, result, CancellationToken.None); // Assert store.GetUploadedValue(result.EntityState).Should().Be("state-final"); @@ -638,6 +661,7 @@ public async Task ExternalizeRequestPayloadsAsync_EntityBatchResult_Externalizes store.GetUploadedValue(result.Actions[0].SendSignal.Input).Should().Be("signal-input"); store.GetUploadedValue(result.Actions[1].StartNewOrchestration.Input).Should().Be("start-input"); + store.UploadCount.Should().Be(resultCount + 3); } [Fact] @@ -699,6 +723,80 @@ public async Task ExternalizeRequestPayloadsAsync_OrchestratorResponse_Permanent #pragma warning restore CS0612 } + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ExternalizeRequestPayloadsAsync_OrchestratorResponse_PermanentFailure_PreservesEntityUnlocks( + bool useLegacyEntityEvents) + { + // Arrange + TrackingPayloadStore store = new(); + LargePayloadStorageOptions options = new() { ThresholdBytes = 1024, MaxPayloadBytes = 2048 }; + AzureBlobPayloadsSideCarInterceptor interceptor = new(store, options); + P.OrchestratorAction unlock = new() { Id = 7 }; + if (useLegacyEntityEvents) + { + unlock.SendEvent = new P.SendEventAction + { + Instance = new P.OrchestrationInstance { InstanceId = "@counter@one" }, + Name = "release", + Data = """{"id":"11111111-1111-1111-1111-111111111111","parent":"instance-1"}""", + }; + } + else + { + unlock.SendEntityMessage = new P.SendEntityMessageAction + { + EntityUnlockSent = new P.EntityUnlockSentEvent + { + CriticalSectionId = "11111111-1111-1111-1111-111111111111", + ParentInstanceId = "instance-1", + TargetInstanceId = "@counter@one", + }, + }; + } + + P.OrchestratorResponse response = new() + { + InstanceId = "instance-1", + CompletionToken = "completion-token", + Actions = + { + unlock, + new P.OrchestratorAction + { + Id = 8, + SendEvent = new P.SendEventAction + { + Instance = new P.OrchestrationInstance { InstanceId = "another-orchestration" }, + Name = "release", + Data = "discarded", + }, + }, + new P.OrchestratorAction + { + Id = 9, + CompleteOrchestration = new P.CompleteOrchestrationAction + { + OrchestrationStatus = P.OrchestrationStatus.Completed, + Result = new string('x', 4096), + }, + }, + }, + }; + + // Act + await ExternalizeAsync(interceptor, response, CancellationToken.None); + + // Assert + response.Actions.Should().HaveCount(2); + response.Actions[0].Should().BeSameAs(unlock); + response.Actions[1].CompleteOrchestration.OrchestrationStatus.Should().Be(P.OrchestrationStatus.Failed); + response.Actions[1].CompleteOrchestration.FailureDetails.IsNonRetriable.Should().BeTrue(); + response.CompletionToken.Should().Be("completion-token"); + store.UploadCount.Should().Be(0); + } + [Fact] public async Task ExternalizeRequestPayloadsAsync_OrchestratorResponse_EarlierRetriableFailureWinsOverLaterPermanentFailure() { diff --git a/test/Grpc.IntegrationTests/LargePayloadTests.cs b/test/Grpc.IntegrationTests/LargePayloadTests.cs index 8502ab56..3b262c41 100644 --- a/test/Grpc.IntegrationTests/LargePayloadTests.cs +++ b/test/Grpc.IntegrationTests/LargePayloadTests.cs @@ -2,17 +2,76 @@ // Licensed under the MIT License. using System.Text.Json; +using FluentAssertions; +using Grpc.Core; +using Grpc.Core.Interceptors; using Microsoft.DurableTask.Client; using Microsoft.DurableTask; using Microsoft.DurableTask.Converters; using Microsoft.DurableTask.Worker; +using Microsoft.DurableTask.Worker.Grpc; using Microsoft.Extensions.DependencyInjection; using Xunit.Abstractions; +using P = Microsoft.DurableTask.Protobuf; namespace Microsoft.DurableTask.Grpc.Tests; public class LargePayloadTests(ITestOutputHelper output, GrpcSidecarFixture sidecarFixture) : IntegrationTestBase(output, sidecarFixture) { + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task CompletionRetry_DoesNotExternalizePayloadTokensAgain(bool retryActivity) + { + // Arrange + string payload = new('R', 1024); + TaskName orchestratorName = nameof(CompletionRetry_DoesNotExternalizePayloadTokensAgain); + TaskName activityName = "RetryPayloadActivity"; + InMemoryPayloadStore store = new(); + RetryCompletionInterceptor retry = new(retryActivity); + await using HostTestLifetime server = await this.StartWorkerAsync( + worker => + { + worker.AddTasks(tasks => tasks + .AddOrchestratorFunc(orchestratorName, async (context, _) => + { + string result = await context.CallActivityAsync(activityName); + context.SetCustomStatus(result); + return result; + }) + .AddActivityFunc(activityName, _ => Task.FromResult(payload))); + worker.UseExternalizedPayloads(); + worker.Services.AddSingleton(store); + worker.Services.PostConfigure( + options => options.Interceptors.Add(retry)); + }, + client => + { + client.UseExternalizedPayloads(); + client.Services.AddSingleton(store); + }, + services => services.AddExternalizedPayloadStore(options => + { + options.ThresholdBytes = 16; + options.ConnectionString = "UseDevelopmentStorage=true"; + })); + + // Act + string instanceId = await server.Client.ScheduleNewOrchestrationInstanceAsync(orchestratorName); + OrchestrationMetadata result = await server.Client.WaitForInstanceCompletionAsync( + instanceId, getInputsAndOutputs: true, this.TimeoutToken); + + // Assert + retry.Payloads.Should().HaveCount(2); + retry.Payloads[0].Length.Should().BeGreaterThan(16); + retry.Payloads[1].Should().Be(retry.Payloads[0]); + store.UploadCount.Should().Be(3, "only the activity result, orchestration result, and custom status need uploading"); + store.uploadedPayloads.Should().NotContain(value => store.IsKnownPayloadToken(value)); + result.RuntimeStatus.Should().Be(OrchestrationRuntimeStatus.Completed); + result.ReadOutputAs().Should().Be(payload); + result.ReadCustomStatusAs().Should().Be(payload); + } + // Validates client externalizes a large orchestration input and worker resolves it. [Fact] public async Task LargeOrchestrationInputAndOutputAndCustomStatus() @@ -792,6 +851,42 @@ public async Task LargeActivityOutput_ExceedsChunkSize_ExternalizesWithLP() Assert.Equal(largeOutput.Length, output?.Length); } + sealed class RetryCompletionInterceptor(bool retryActivity) : Interceptor + { + public List Payloads { get; } = []; + + public override AsyncUnaryCall AsyncUnaryCall( + TRequest request, + ClientInterceptorContext context, + AsyncUnaryCallContinuation continuation) + { + string? payload = request switch + { + P.ActivityResponse response when retryActivity => response.Result, + P.OrchestratorResponse response when !retryActivity => + response.Actions.FirstOrDefault(action => action.CompleteOrchestration is not null) + ?.CompleteOrchestration.Result, + _ => null, + }; + if (payload is not null) + { + this.Payloads.Add(payload); + if (this.Payloads.Count == 1) + { + Status status = new(StatusCode.Unavailable, "Simulated completion transport failure."); + return new AsyncUnaryCall( + Task.FromException(new RpcException(status)), + Task.FromResult(new Metadata()), + () => status, + () => new Metadata(), + () => { }); + } + } + + return continuation(request, context); + } + } + class InMemoryPayloadStore : PayloadStore { const string TokenPrefix = "blob:v1:"; diff --git a/test/Worker/Grpc.Tests/GrpcDurableTaskWorkerOptionsInternalTests.cs b/test/Worker/Grpc.Tests/GrpcDurableTaskWorkerOptionsInternalTests.cs index f7216c23..8606f269 100644 --- a/test/Worker/Grpc.Tests/GrpcDurableTaskWorkerOptionsInternalTests.cs +++ b/test/Worker/Grpc.Tests/GrpcDurableTaskWorkerOptionsInternalTests.cs @@ -2,6 +2,7 @@ // Licensed under the MIT License. using Microsoft.DurableTask.Worker.Grpc.Internal; +using P = Microsoft.DurableTask.Protobuf; namespace Microsoft.DurableTask.Worker.Grpc.Tests; @@ -26,6 +27,7 @@ public void InternalOptions_HasSafeDefaults() internalOptions.TransientRetryMaxAttempts.Should().Be(10); internalOptions.SilentDisconnectTimeout.Should().Be(TimeSpan.FromSeconds(120)); internalOptions.ChannelRecreator.Should().BeNull(); + internalOptions.PrepareOrchestratorResponseAsync.Should().BeNull(); internalOptions.VersioningExemptOrchestrations.Should().BeEmpty(); internalOptions.VersioningExemptActivities.Should().BeEmpty(); } @@ -99,4 +101,35 @@ public void SetChannelRecreator_StoresCallbackOnInternalOptions() options.Internal.ChannelRecreator!.Invoke(null!, CancellationToken.None); invoked.Should().BeTrue(); } + + [Theory] + [InlineData(true, "options")] + [InlineData(false, "preparation")] + public void SetOrchestratorResponsePreparation_NullArgument_Throws(bool nullOptions, string parameterName) + { + // Arrange + GrpcDurableTaskWorkerOptions options = new(); + Func preparation = (_, _) => Task.CompletedTask; + + // Act + Action act = () => (nullOptions ? null! : options) + .SetOrchestratorResponsePreparation(nullOptions ? preparation : null!); + + // Assert + act.Should().Throw().Which.ParamName.Should().Be(parameterName); + } + + [Fact] + public void SetOrchestratorResponsePreparation_StoresCallbackOnInternalOptions() + { + // Arrange + GrpcDurableTaskWorkerOptions options = new(); + Func preparation = (_, _) => Task.CompletedTask; + + // Act + options.SetOrchestratorResponsePreparation(preparation); + + // Assert + options.Internal.PrepareOrchestratorResponseAsync.Should().BeSameAs(preparation); + } }