From 84372edd58aee075fa386e043e07d222e724fe6c Mon Sep 17 00:00:00 2001 From: Ben Adams Date: Wed, 12 Aug 2020 16:05:18 +0100 Subject: [PATCH] ExecutionContext.Restore (#40322) Add ExecutionContext.Restore --- .../src/System/Threading/ExecutionContext.cs | 35 +++++++++++++++++++++- .../Sources/ManualResetValueTaskSourceCore.cs | 8 ++--- .../System.Threading/ref/System.Threading.cs | 1 + .../tests/ExecutionContextTests.cs | 23 ++++++++++++++ 4 files changed, 62 insertions(+), 5 deletions(-) diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs index 40b6ca1..9f5b9cd 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs @@ -67,6 +67,21 @@ namespace System.Threading return executionContext; } + // Allows capturing asynclocals for a FlowSuppressed ExecutionContext rather than returning null. + internal static ExecutionContext? CaptureForRestore() + { + // This is a short cut for: + // + // ExecutionContext.RestoreFlow() + // var ec = ExecutionContext.Capture() + // ExecutionContext.SuppressFlow(); + // ... + // ExecutionContext.Restore(ec) + // ExecutionContext.SuppressFlow(); + + return Thread.CurrentThread._executionContext; + } + private ExecutionContext? ShallowClone(bool isFlowSuppressed) { Debug.Assert(isFlowSuppressed != m_isFlowSuppressed); @@ -199,7 +214,25 @@ namespace System.Threading edi?.Throw(); } - internal static void Restore(ExecutionContext? executionContext) + /// + /// Restores a captured execution context to on the current thread. + /// + /// + /// To revert to the current execution context; capture it before Restore, and Restore it again. + /// It will not automatically be reverted unlike . + /// + /// The ExecutionContext to set. + public static void Restore(ExecutionContext executionContext) + { + if (executionContext == null) + { + ThrowNullContext(); + } + + RestoreInternal(executionContext); + } + + internal static void RestoreInternal(ExecutionContext? executionContext) { Thread currentThread = Thread.CurrentThread; diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs index b9caa53..647e4d7 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs @@ -244,7 +244,7 @@ namespace System.Threading.Tasks.Sources Debug.Assert(_continuation != null); Debug.Assert(_executionContext != null); - ExecutionContext? currentContext = ExecutionContext.Capture(); + ExecutionContext? currentContext = ExecutionContext.CaptureForRestore(); // Restore the captured ExecutionContext before executing anything. ExecutionContext.Restore(_executionContext); @@ -259,7 +259,7 @@ namespace System.Threading.Tasks.Sources finally { // Restore the current ExecutionContext. - ExecutionContext.Restore(currentContext); + ExecutionContext.RestoreInternal(currentContext); } } else @@ -284,7 +284,7 @@ namespace System.Threading.Tasks.Sources // Set sync context back to what it was prior to coming in SynchronizationContext.SetSynchronizationContext(syncContext); // Restore the current ExecutionContext. - ExecutionContext.Restore(currentContext); + ExecutionContext.RestoreInternal(currentContext); } // Now rethrow the exception; if there is one. @@ -301,7 +301,7 @@ namespace System.Threading.Tasks.Sources finally { // Restore the current ExecutionContext. - ExecutionContext.Restore(currentContext); + ExecutionContext.RestoreInternal(currentContext); } } diff --git a/src/libraries/System.Threading/ref/System.Threading.cs b/src/libraries/System.Threading/ref/System.Threading.cs index 888abda..d2f60ac 100644 --- a/src/libraries/System.Threading/ref/System.Threading.cs +++ b/src/libraries/System.Threading/ref/System.Threading.cs @@ -128,6 +128,7 @@ namespace System.Threading public void Dispose() { } public void GetObjectData(System.Runtime.Serialization.SerializationInfo info, System.Runtime.Serialization.StreamingContext context) { } public static bool IsFlowSuppressed() { throw null; } + public static void Restore(System.Threading.ExecutionContext executionContext) { } public static void RestoreFlow() { } public static void Run(System.Threading.ExecutionContext executionContext, System.Threading.ContextCallback callback, object? state) { } public static System.Threading.AsyncFlowControl SuppressFlow() { throw null; } diff --git a/src/libraries/System.Threading/tests/ExecutionContextTests.cs b/src/libraries/System.Threading/tests/ExecutionContextTests.cs index 6db03a6..96261fb 100644 --- a/src/libraries/System.Threading/tests/ExecutionContextTests.cs +++ b/src/libraries/System.Threading/tests/ExecutionContextTests.cs @@ -31,6 +31,29 @@ namespace System.Threading.Tests } [Fact] + public static void RestoreTest() + { + ExecutionContext defaultEC = ExecutionContext.Capture(); + var asyncLocal = new AsyncLocal(); + Assert.Equal(0, asyncLocal.Value); + + asyncLocal.Value = 1; + ExecutionContext oneEC = ExecutionContext.Capture(); + Assert.Equal(1, asyncLocal.Value); + + ExecutionContext.Restore(defaultEC); + Assert.Equal(0, asyncLocal.Value); + + ExecutionContext.Restore(oneEC); + Assert.Equal(1, asyncLocal.Value); + + ExecutionContext.Restore(defaultEC); + Assert.Equal(0, asyncLocal.Value); + + Assert.Throws(() => ExecutionContext.Restore(null!)); + } + + [Fact] public static void DisposeTest() { ExecutionContext executionContext = ExecutionContext.Capture(); -- 2.7.4