diff --git a/src/WebJobs.Extensions.DurableTask/ContextImplementations/DurableOrchestrationContext.cs b/src/WebJobs.Extensions.DurableTask/ContextImplementations/DurableOrchestrationContext.cs index 0d5666c44..f8aa2cfc6 100644 --- a/src/WebJobs.Extensions.DurableTask/ContextImplementations/DurableOrchestrationContext.cs +++ b/src/WebJobs.Extensions.DurableTask/ContextImplementations/DurableOrchestrationContext.cs @@ -30,6 +30,20 @@ internal class DurableOrchestrationContext : DurableCommonContext, IDurableOrche DurableOrchestrationContextBase // for v1 legacy compatibility. #pragma warning restore 618 { + /// + /// Error message for multithreaded execution detection. + /// + private const string MultithreadedAccessErrorMessage = + "Multithreaded execution was detected. This can happen if the orchestrator function code awaits on a task that was not created by a DurableOrchestrationContext method. More details can be found in this article https://learn.microsoft.com/azure/azure-functions/durable/durable-functions-code-constraints."; + + /// + /// AsyncLocal to track the current orchestration context. This is used in addition to + /// OrchestrationContext.IsOrchestratorThread to detect illegal async usage patterns + /// like calling activities from within ContinueWith blocks. AsyncLocal properly flows + /// with async context, unlike [ThreadStatic] which doesn't. + /// + private static readonly AsyncLocal CurrentContext = new AsyncLocal(); + private readonly Dictionary pendingExternalEvents = new Dictionary(StringComparer.OrdinalIgnoreCase); @@ -1201,8 +1215,43 @@ internal void ThrowIfInvalidAccess() if (!OrchestrationContext.IsOrchestratorThread) { - throw new InvalidOperationException( - "Multithreaded execution was detected. This can happen if the orchestrator function code awaits on a task that was not created by a DurableOrchestrationContext method. More details can be found in this article https://learn.microsoft.com/azure/azure-functions/durable/durable-functions-code-constraints."); + throw new InvalidOperationException(MultithreadedAccessErrorMessage); + } + + // Additional check using AsyncLocal to detect illegal access from ContinueWith blocks. + // The AsyncLocal value properly flows with async context, unlike the [ThreadStatic] + // IsOrchestratorThread which can be incorrectly true if thread pool reuses the same thread. + if (CurrentContext.Value != this) + { + throw new InvalidOperationException(MultithreadedAccessErrorMessage); + } + } + + /// + /// Sets this context as the current context for the async flow. + /// This should be called at the start of orchestration execution. + /// + internal void SetAsCurrentContext() + { + CurrentContext.Value = this; + } + + /// + /// Clears the current context from the async flow. + /// This should be called at the end of orchestration execution. + /// + internal void ClearCurrentContext() + { + // Only clear the context if it matches this instance. + // If it doesn't match, it could indicate a programming error or nested orchestration scenario. + // We use Debug.Assert to help detect such issues during development. + System.Diagnostics.Debug.Assert( + CurrentContext.Value == null || CurrentContext.Value == this, + "ClearCurrentContext called when the current context does not match this instance. This may indicate a programming error."); + + if (CurrentContext.Value == this) + { + CurrentContext.Value = null; } } diff --git a/src/WebJobs.Extensions.DurableTask/Listener/TaskOrchestrationShim.cs b/src/WebJobs.Extensions.DurableTask/Listener/TaskOrchestrationShim.cs index 5c6198c4f..198bd312b 100644 --- a/src/WebJobs.Extensions.DurableTask/Listener/TaskOrchestrationShim.cs +++ b/src/WebJobs.Extensions.DurableTask/Listener/TaskOrchestrationShim.cs @@ -138,6 +138,10 @@ private async Task InvokeUserCodeAndHandleResults( RegisteredFunctionInfo orchestratorInfo, OrchestrationContext innerContext) { + // Set the current context for the async flow to enable detection of illegal + // ContinueWith usage. This uses AsyncLocal which properly flows with async context. + this.context.SetAsCurrentContext(); + try { Task invokeTask = this.FunctionInvocationCallback(); @@ -205,6 +209,9 @@ private async Task InvokeUserCodeAndHandleResults( finally { this.context.IsCompleted = true; + + // Clear the current context from the async flow + this.context.ClearCurrentContext(); } }