From 2574384c487a60be72ee4fe9150fc458c0bef12f Mon Sep 17 00:00:00 2001 From: Jakob Botsch Nielsen Date: Fri, 26 Jun 2026 14:45:27 +0200 Subject: [PATCH] Save async contexts for ValueTask methods We were not properly saving/restoring async contexts for ValueTask methods because of an incorrect flag check. --- src/coreclr/jit/compiler.h | 1 + src/coreclr/jit/importer.cpp | 33 ++++++++++++++++++- src/coreclr/jit/jitconfigvalues.h | 2 ++ src/coreclr/vm/method.hpp | 9 ++++- .../tests/FunctionalTests/DiagnosticsTests.cs | 1 - .../execution-context/execution-context.cs | 18 ++++++++++ 6 files changed, 61 insertions(+), 3 deletions(-) diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h index c2d5676a4fe219..f5387077ecf46f 100644 --- a/src/coreclr/jit/compiler.h +++ b/src/coreclr/jit/compiler.h @@ -5468,6 +5468,7 @@ class Compiler bool impMatchIsInstBooleanConversion(const BYTE* codeAddr, const BYTE* codeEndp, int* consumed); const BYTE* impMatchTaskAwaitPattern(const BYTE* codeAddr, const BYTE* codeEndp, int* configVal, IL_OFFSET* awaitOffset); + bool impCheckOptimizeAwait(IL_OFFSET awaitOffset); GenTree* impCastClassOrIsInstToTree( GenTree* op1, GenTree* op2, CORINFO_RESOLVED_TOKEN* pResolvedToken, bool isCastClass, bool* booleanCheck, IL_OFFSET ilOffset); diff --git a/src/coreclr/jit/importer.cpp b/src/coreclr/jit/importer.cpp index 1b811f91084775..791d2d162eea92 100644 --- a/src/coreclr/jit/importer.cpp +++ b/src/coreclr/jit/importer.cpp @@ -9091,7 +9091,7 @@ void Compiler::impImportBlockCode(BasicBlock* block) if (isAwait) { _impResolveToken(CORINFO_TOKENKIND_Await); - if (resolvedToken.hMethod == NO_METHOD_HANDLE) + if ((resolvedToken.hMethod == NO_METHOD_HANDLE) || !impCheckOptimizeAwait(awaitOffset)) { // This can happen in cases when the Task-returning method is not a runtime Async // function. For example "T M1(T arg) => arg" when called with a Task argument. @@ -11155,6 +11155,37 @@ void Compiler::impImportBlockCode(BasicBlock* block) #undef _impResolveToken } +//------------------------------------------------------------------------ +// impCheckOptimizeAwait: +// Check if an await at a specified offset should be optimized. +// +// Arguments: +// awaitOffset - The offset of the await +// +// Returns: +// True if we should use the async variant. +// +bool Compiler::impCheckOptimizeAwait(IL_OFFSET awaitOffset) +{ +#ifdef DEBUG + static ConfigMethodRange s_jitOptimizeAwaitRange; + s_jitOptimizeAwaitRange.EnsureInit(JitConfig.JitOptimizeAwaitRange()); + + unsigned hash = impInlineRoot()->info.compMethodHash(); + hash = ((hash << 5) + hash) ^ (compIsForInlining() ? compInlineContext->GetOrdinal() : 0); + hash = ((hash << 5) + hash) ^ awaitOffset; + + if ((JitConfig.JitAwaitHashBreak() != -1) && (hash == (unsigned)JitConfig.JitAwaitHashBreak())) + { + assert(!"JitAwaitHashBreak reached"); + } + + return s_jitOptimizeAwaitRange.Contains(hash); +#else + return true; +#endif +} + //------------------------------------------------------------------------ // impCreateLocal: create a GT_LCL_VAR node to access a local that might need to be normalized on load // diff --git a/src/coreclr/jit/jitconfigvalues.h b/src/coreclr/jit/jitconfigvalues.h index 1cc7e8aec86497..f67938cf9d1dee 100644 --- a/src/coreclr/jit/jitconfigvalues.h +++ b/src/coreclr/jit/jitconfigvalues.h @@ -134,6 +134,8 @@ RELEASE_CONFIG_INTEGER(JitInlineBudget, "JitInlineBudget", DEFAULT_INLINE_BUDGET CONFIG_INTEGER(JitForceInlineDepth, "JitForceInlineDepth", DEFAULT_MAX_FORCE_INLINE_DEPTH) RELEASE_CONFIG_INTEGER(JitInlineMethodsWithEH, "JitInlineMethodsWithEH", 1) CONFIG_STRING(JitInlineMethodsWithEHRange, "JitInlineMethodsWithEHRange") +CONFIG_STRING(JitOptimizeAwaitRange, "JitOptimizeAwaitRange") // Optimize awaits in this hash range +CONFIG_INTEGER(JitAwaitHashBreak, "JitAwaitHashBreak", -1) // Break on jitting await with specified hash CONFIG_INTEGER(JitLongAddress, "JitLongAddress", 0) // Force using the large pseudo instruction form for long address CONFIG_INTEGER(JitMaxUncheckedOffset, "JitMaxUncheckedOffset", 8) diff --git a/src/coreclr/vm/method.hpp b/src/coreclr/vm/method.hpp index 539b464b339147..653549db2e37fb 100644 --- a/src/coreclr/vm/method.hpp +++ b/src/coreclr/vm/method.hpp @@ -75,6 +75,8 @@ enum class AsyncMethodFlags Thunk = 16, // A special thunk to drop return value in covariant return scenario ReturnDroppingThunk = 32, + // Note: If adding more flags make sure to modify RequiresAsyncContextSaveAndRestore + // The rest of the methods that are not in any of the above groups. // Such methods are not interesting to the Runtime Async feature. // Note: Generic T-returning methods are classified as "None", even if T could be a Task. @@ -117,6 +119,11 @@ inline AsyncMethodFlags operator&(AsyncMethodFlags lhs, AsyncMethodFlags rhs) return (AsyncMethodFlags)((int)lhs & (int)rhs); } +inline AsyncMethodFlags operator~(AsyncMethodFlags flags) +{ + return (AsyncMethodFlags)~(int)flags; +} + inline AsyncMethodFlags& operator|=(AsyncMethodFlags& lhs, AsyncMethodFlags rhs) { lhs = lhs | rhs; @@ -2161,7 +2168,7 @@ class MethodDesc AsyncMethodFlags asyncFlags = GetAddrOfAsyncMethodData()->flags; // asynccall that is also async variant, but not a thunk - return asyncFlags == (AsyncMethodFlags::AsyncCall | AsyncMethodFlags::IsAsyncVariant); + return (asyncFlags & ~AsyncMethodFlags::IsAsyncVariantForValueTask) == (AsyncMethodFlags::AsyncCall | AsyncMethodFlags::IsAsyncVariant); } #ifdef FEATURE_METADATA_UPDATER diff --git a/src/libraries/System.Net.Http/tests/FunctionalTests/DiagnosticsTests.cs b/src/libraries/System.Net.Http/tests/FunctionalTests/DiagnosticsTests.cs index a61aac0f808451..f6a843cf21f530 100644 --- a/src/libraries/System.Net.Http/tests/FunctionalTests/DiagnosticsTests.cs +++ b/src/libraries/System.Net.Http/tests/FunctionalTests/DiagnosticsTests.cs @@ -577,7 +577,6 @@ protected static void VerifyTag(KeyValuePair[] tags, string } } - [ActiveIssue("https://github.com/dotnet/runtime/issues/129771")] [ConditionalTheory(typeof(RemoteExecutor), nameof(RemoteExecutor.IsSupported))] [InlineData(false)] [InlineData(true)] diff --git a/src/tests/async/execution-context/execution-context.cs b/src/tests/async/execution-context/execution-context.cs index c925759c3db41f..389f33eaf43a80 100644 --- a/src/tests/async/execution-context/execution-context.cs +++ b/src/tests/async/execution-context/execution-context.cs @@ -152,4 +152,22 @@ private static async Task LoopWithOsrTransition() s_osrLocal.Value = val; } + + [Fact] + public static void TestValueTask() + { + TestValueTaskAsync().GetAwaiter().GetResult(); + } + + private static async ValueTask TestValueTaskAsync() + { + s_local.Value = 42; + await ChangeThenReturnValueTask(); + Assert.Equal(42, s_local.Value); + } + + private static async ValueTask ChangeThenReturnValueTask() + { + s_local.Value = 123; + } }