diff --git a/src/WebJobs.Extensions.DurableTask/OutOfProcMiddleware.cs b/src/WebJobs.Extensions.DurableTask/OutOfProcMiddleware.cs
index b8aafa31b..e45102218 100644
--- a/src/WebJobs.Extensions.DurableTask/OutOfProcMiddleware.cs
+++ b/src/WebJobs.Extensions.DurableTask/OutOfProcMiddleware.cs
@@ -24,6 +24,7 @@ internal class OutOfProcMiddleware
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 @@ await this.LifeCycleNotificationHelper.OrchestratorStartingAsync(
workItemMetadata.IncludeState);
bool workerRequiresHistory = false;
+ bool isWorkerDraining = false;
var input = new TriggeredFunctionData
{
@@ -153,6 +155,7 @@ await this.LifeCycleNotificationHelper.OrchestratorStartingAsync(
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 @@ await this.LifeCycleNotificationHelper.OrchestratorStartingAsync(
{
// 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 @@ await this.LifeCycleNotificationHelper.OrchestratorStartingAsync(
{
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 @@ void SetErrorResult(FailureDetails failureDetails)
workItemMetadata.IncludeState);
bool workerRequiresEntityState = false;
+ bool isWorkerDraining = false;
var input = new TriggeredFunctionData
{
@@ -393,6 +405,7 @@ void SetErrorResult(FailureDetails failureDetails)
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 @@ void SetErrorResult(FailureDetails failureDetails)
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 @@ void SetErrorResult(FailureDetails failureDetails)
{
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 @@ public async Task CallActivityAsync(DispatchMiddlewareContext dispatchContext, F
// 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 @@ private static bool IsWorkerNotFullyInitializedException(Exception? exception)
{
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)
diff --git a/src/Worker.Extensions.DurableTask/Exceptions/WorkerDrainingException.cs b/src/Worker.Extensions.DurableTask/Exceptions/WorkerDrainingException.cs
new file mode 100644
index 000000000..bdd9e8b66
--- /dev/null
+++ b/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;
+
+///
+/// Thrown by the worker when a Function completes while the worker is draining (shutting down).
+///
+internal sealed class WorkerDrainingException : Exception
+{
+ ///
+ /// 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.
+ ///
+ 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}")
+ {
+ }
+}
diff --git a/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Activity.cs b/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Activity.cs
index 163941855..b4fe8cb84 100644
--- a/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Activity.cs
+++ b/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)
{
diff --git a/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Entity.cs b/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Entity.cs
index 32e198421..b8de5443d 100644
--- a/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Entity.cs
+++ b/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 @@ private async ValueTask RunEntityAsync(FunctionContext context, BindingMetadata
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 @@ private async Task RunDirectEntityAsync(
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());
+ }
}
diff --git a/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Orchestration.cs b/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Orchestration.cs
index dc64c2872..c328d66a2 100644
--- a/src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Orchestration.cs
+++ b/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 @@ private async ValueTask RunOrchestrationAsync(FunctionContext context, BindingMe
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