Add worker drain signal

src/WebJobs.Extensions.DurableTask/OutOfProcMiddleware.cs

@@ -24,6 +24,7 @@
         private const string NoProcessAssociatedMessage = "No process is associated";
         private const string NoWorkerInitializedMessage = "Did not find any initialized language workers";
         private const string AssemblyNotLoadedMessage = "Could not load file or assembly";
+        private const string WorkerDrainingMessageMarker = "[DurableTask:WorkerDraining]";
 
         private readonly DurableTaskExtension extension;
 
@@ -124,6 +125,7 @@
                 workItemMetadata.IncludeState);
 
             bool workerRequiresHistory = false;
+            bool isWorkerDraining = false;
 
             var input = new TriggeredFunctionData
             {
@@ -153,6 +155,7 @@
                     P.OrchestratorResponse response = P.OrchestratorResponse.Parser.ParseFrom(triggerReturnValueBytes);
 
                     workerRequiresHistory = response.RequiresHistory;
+                    isWorkerDraining = response.IsWorkerDraining;
 
                     // TrySetResult may throw if a platform-level error is encountered (like an out of memory exception).
                     context.SetResult(
@@ -177,9 +180,15 @@
                 {
                     // Shutdown can surface as a completed invocation in a failed state.
                     // Re-throw so we can abort this invocation.
-                    this.HostLifetimeService.OnStopping.ThrowIfCancellationRequested();
-                }
-
+                    this.HostLifetimeService.OnStopping.ThrowIfCancellationRequested();
+
+                    // The Function may have failed because the worker is draining, so we want to retry its execution in this case
+                    if (isWorkerDraining)
+                    {
+                        throw new Exception("Worker is shutting down. The orchestration execution will be aborted and retried.");
+                    }
+                }
+
                 // we abort the invocation on "platform level errors" such as:
                 // - a timeout
                 // - an out of memory exception
@@ -193,7 +202,9 @@
             {
                 string reason = this.HostLifetimeService.OnStopping.IsCancellationRequested ?
                     "The Functions/WebJobs runtime is shutting down!" :
-                    $"Unhandled exception in the Functions/WebJobs runtime: {hostRuntimeException}";
+                        isWorkerDraining ?
+                             "The worker is shutting down. The orchestration execution will be aborted and retried." :
+                                $"Unhandled exception in the Functions/WebJobs runtime: {hostRuntimeException}";
 
                 this.TraceHelper.FunctionAborted(
                     this.Options.HubName,
@@ -365,6 +376,7 @@
                 workItemMetadata.IncludeState);
 
             bool workerRequiresEntityState = false;
+            bool isWorkerDraining = false;
 
             var input = new TriggeredFunctionData
             {
@@ -393,6 +405,7 @@
                     byte[] triggerReturnValueBytes = Convert.FromBase64String(triggerReturnValue);
                     P.EntityBatchResult response = P.EntityBatchResult.Parser.ParseFrom(triggerReturnValueBytes);
                     workerRequiresEntityState = response.RequiresState;
+                    isWorkerDraining = response.IsWorkerDraining;
                     context.Result = response.ToEntityBatchResult(this.Options.DefaultVersion);
 
                     context.ThrowIfFailed();
@@ -418,6 +431,12 @@
                         throw functionResult.Exception;
                     }
 
+                    // The Function may have failed because the worker is draining, so we want to retry its execution in this case
+                    if (isWorkerDraining)
+                    {
+                        throw new Exception("Worker is shutting down. The entity execution will be aborted and retried.");
+                    }
+
                     // Shutdown can surface as a completed invocation in a failed state.
                     // Re-throw so we can abort this invocation.
                     this.HostLifetimeService.OnStopping.ThrowIfCancellationRequested();
@@ -427,7 +446,9 @@
             {
                 string reason = this.HostLifetimeService.OnStopping.IsCancellationRequested ?
                     "The Functions/WebJobs runtime is shutting down!" :
-                    $"Unhandled exception in the Functions/WebJobs runtime: {hostRuntimeException}";
+                        isWorkerDraining ?
+                             "The worker is shutting down. The entity execution will be aborted and retried." :
+                                $"Unhandled exception in the Functions/WebJobs runtime: {hostRuntimeException}";
 
                 this.TraceHelper.FunctionAborted(
                     this.Options.HubName,
@@ -596,6 +617,19 @@
 
                 // This will abort the current execution and force an durable retry
                 throw new SessionAbortedException(reason);
+            }
+
+            // The Function may have failed because the worker is draining, so we want to retry its execution in this case
+            if (!result.Succeeded && IsWorkerDrainingException(result.Exception))
+            {
+                this.TraceHelper.FunctionAborted(
+                    this.Options.HubName,
+                    functionName.Name,
+                    instance.InstanceId,
+                    "The worker is shutting down. The activity execution will be aborted and retried.",
+                    functionType: FunctionType.Activity);
+
+                throw new SessionAbortedException("The worker is shutting down. The activity execution will be aborted and retried.");
             }
 
             ActivityExecutionResult activityResult;
@@ -681,6 +715,19 @@
         {
             return (exception?.InnerException is InvalidOperationException ioe && ioe.Message.Contains(NoWorkerInitializedMessage, StringComparison.Ordinal))
                 || (exception?.InnerException is FileNotFoundException fnfe && fnfe.Message.Contains(AssemblyNotLoadedMessage, StringComparison.Ordinal));
+        }
+
+        private static bool IsWorkerDrainingException(Exception? exception)
+        {
+            for (Exception? current = exception; current != null; current = current.InnerException)
+            {
+                if (current.Message != null && current.Message.Contains(WorkerDrainingMessageMarker, StringComparison.Ordinal))
+                {
+                    return true;
+                }
+            }
+
+            return false;
         }
 
         private static FailureDetails GetFailureDetails(Exception e, out bool fromSerializedException)

src/Worker.Extensions.DurableTask/Exceptions/WorkerDrainingException.cs

@@ -0,0 +1,24 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+
+#nullable enable
+namespace Microsoft.Azure.Functions.Worker.Extensions.DurableTask.Exceptions;
+
+/// <summary>
+/// Thrown by the worker when a Function completes while the worker is draining (shutting down).
+/// </summary>
+internal sealed class WorkerDrainingException : Exception
+{
+    /// <summary>
+    /// Stable marker the host matches on to detect this exception after it crosses the gRPC boundary
+    /// as a serialized string. The host keeps an identical copy of this literal; keep them in sync.
+    /// </summary>
+    internal const string Sentinel = "[DurableTask:WorkerDraining]";
+
+    public WorkerDrainingException()
+        : base($"The worker is shutting down and will not commit this result; the work item should be retried on another worker. {Sentinel}")
+    {
+    }
+}

src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Activity.cs

@@ -28,11 +28,16 @@ private async ValueTask RunActivityAsync(FunctionContext context, BindingMetadat
             if (context.FunctionDefinition.EntryPoint == ActivityEntryPoint)
             {
                 await this.RunDirectActivityAsync(context, triggerBinding);
-                return;
+            }
+            else
+            {
+                await inner.ExecuteAsync(context);
             }
 
-            await inner.ExecuteAsync(context);
-            return;
+            if (this.IsWorkerDraining)
+            {
+                throw new WorkerDrainingException();
+            }
         }
         catch (Exception ex)
         {

src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Entity.cs

@@ -3,9 +3,11 @@
 
 using System;
 using System.Threading.Tasks;
+using Google.Protobuf;
 using Microsoft.DurableTask.Entities;
 using Microsoft.DurableTask.Worker;
 using Microsoft.DurableTask.Worker.Grpc;
+using P = Microsoft.DurableTask.Protobuf;
 
 namespace Microsoft.Azure.Functions.Worker.Extensions.DurableTask.Execution;
 
@@ -40,7 +42,15 @@
         TaskEntityDispatcher dispatcher = new(encodedEntityBatch, context.InstanceServices, extendedSessionsCache);
         triggerInputData.Value = dispatcher;
         await inner.ExecuteAsync(context);
-        context.GetInvocationResult().Value = dispatcher.Result;
+
+        string entityResult = dispatcher.Result;
+
+        if (this.IsWorkerDraining)
+        {
+            entityResult = FlagEntityDraining(entityResult);
+        }
+
+        context.GetInvocationResult().Value = entityResult;
     }
 
     private async Task RunDirectEntityAsync(
@@ -61,6 +71,20 @@
 
         string result = await GrpcEntityRunner.LoadAndRunAsync(
             encodedEntityBatch, entity, context.InstanceServices);
+
+        if (this.IsWorkerDraining)
+        {
+            result = FlagEntityDraining(result);
+        }
+
         context.GetInvocationResult().Value = result;
     }
+
+    private static string FlagEntityDraining(string encodedEntityBatchResult)
+    {
+        byte[] resultBytes = Convert.FromBase64String(encodedEntityBatchResult);
+        P.EntityBatchResult result = P.EntityBatchResult.Parser.ParseFrom(resultBytes);
+        result.IsWorkerDraining = true;
+        return Convert.ToBase64String(result.ToByteArray());
+    }
 }

src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Orchestration.cs

@@ -3,8 +3,10 @@
 
 using System;
 using System.Threading.Tasks;
+using Google.Protobuf;
 using Microsoft.DurableTask;
 using Microsoft.DurableTask.Worker.Grpc;
+using P = Microsoft.DurableTask.Protobuf;
 
 namespace Microsoft.Azure.Functions.Worker.Extensions.DurableTask.Execution;
 
@@ -34,10 +36,23 @@
         string orchestratorOutput = GrpcOrchestrationRunner.LoadAndRun(
             encodedOrchestratorState, orchestrator, extendedSessionsCache, context.InstanceServices);
 
+        if (this.IsWorkerDraining)
+        {
+            orchestratorOutput = FlagOrchestratorDraining(orchestratorOutput);
+        }
+
         // Send the encoded orchestrator output as the return value seen by the functions host extension
         context.GetInvocationResult().Value = orchestratorOutput;
     }
 
+    private static string FlagOrchestratorDraining(string encodedOrchestratorOutput)
+    {
+        byte[] responseBytes = Convert.FromBase64String(encodedOrchestratorOutput);
+        P.OrchestratorResponse response = P.OrchestratorResponse.Parser.ParseFrom(responseBytes);
+        response.IsWorkerDraining = true;
+        return Convert.ToBase64String(response.ToByteArray());
+    }
+
     private IDisposableOrchestrator CreateOrchestrator(
         FunctionContext context, InputBindingData<object> triggerInputData)
     {

src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.cs

@@ -6,6 +6,7 @@
 using Microsoft.Azure.Functions.Worker.Invocation;
 using Microsoft.DurableTask;
 using Microsoft.DurableTask.Worker;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Options;
 
 namespace Microsoft.Azure.Functions.Worker.Extensions.DurableTask.Execution;
@@ -15,11 +16,15 @@ internal partial class DurableFunctionExecutor(
     ExtendedSessionsCache extendedSessionsCache,
     IDurableTaskFactory factory,
     IOptions<DurableTaskWorkerOptions> options,
+    IHostApplicationLifetime? hostApplicationLifetime = null,
     IExceptionPropertiesProvider? exceptionPropertiesProvider = null)
     : IFunctionExecutor
 {
     private DataConverter Converter => options.Value.DataConverter;
 
+    private bool IsWorkerDraining =>
+        hostApplicationLifetime?.ApplicationStopping.IsCancellationRequested ?? false;
+
     public virtual ValueTask ExecuteAsync(FunctionContext context)
     {
         if (context is null)