From 53315683e9be6c87814149bdea05b06a412cbfd2 Mon Sep 17 00:00:00 2001
From: Vladimir Sadov <vsadov@microsoft.com>
Date: Fri, 29 Jul 2022 02:57:56 -0700
Subject: [PATCH] [NativeAOT] Fix assert in Thread.SpinWait(0) (#73033)

* do not inline long wait
---
 src/coreclr/nativeaot/Runtime/MiscHelpers.cpp      |  6 ++++++
 .../src/System/Runtime/RuntimeImports.cs           |  5 +++++
 .../src/System/Threading/Thread.NativeAot.cs       | 25 +++++++++++++++++++++-
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/coreclr/nativeaot/Runtime/MiscHelpers.cpp b/src/coreclr/nativeaot/Runtime/MiscHelpers.cpp
index 064da7e..381079d 100644
--- a/src/coreclr/nativeaot/Runtime/MiscHelpers.cpp
+++ b/src/coreclr/nativeaot/Runtime/MiscHelpers.cpp
@@ -48,6 +48,12 @@ COOP_PINVOKE_HELPER(void, RhDebugBreak, ())
 // Busy spin for the given number of iterations.
 EXTERN_C NATIVEAOT_API void __cdecl RhSpinWait(int32_t iterations)
 {
+    ASSERT(iterations > 0);
+
+    // limit the spin count in coop mode.
+    ASSERT_MSG(iterations <= 10000 || !ThreadStore::GetCurrentThread()->IsCurrentThreadInCooperativeMode(),
+        "This is too long wait for coop mode. You must p/invoke with GC transition.");
+
     YieldProcessorNormalizationInfo normalizationInfo;
     YieldProcessorNormalizedForPreSkylakeCount(normalizationInfo, iterations);
 }
diff --git a/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/RuntimeImports.cs b/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/RuntimeImports.cs
index 287e80b..5489553 100644
--- a/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/RuntimeImports.cs
+++ b/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/RuntimeImports.cs
@@ -388,6 +388,11 @@ namespace System.Runtime
         [UnmanagedCallConv(CallConvs = new Type[] { typeof(CallConvCdecl) })]
         internal static partial void RhSpinWait(int iterations);
 
+        // Call RhSpinWait with a GC transition
+        [LibraryImport(RuntimeLibrary, EntryPoint = "RhSpinWait")]
+        [UnmanagedCallConv(CallConvs = new Type[] { typeof(CallConvCdecl) })]
+        internal static partial void RhLongSpinWait(int iterations);
+
         // Yield the cpu to another thread ready to process, if one is available.
         [LibraryImport(RuntimeLibrary, EntryPoint = "RhYield")]
         [UnmanagedCallConv(CallConvs = new Type[] { typeof(CallConvCdecl) })]
diff --git a/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Threading/Thread.NativeAot.cs b/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Threading/Thread.NativeAot.cs
index 2564eec..543ba77 100644
--- a/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Threading/Thread.NativeAot.cs
+++ b/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Threading/Thread.NativeAot.cs
@@ -322,7 +322,30 @@ namespace System.Threading
         /// </summary>
         internal const int OptimalMaxSpinWaitsPerSpinIteration = 64;
 
-        public static void SpinWait(int iterations) => RuntimeImports.RhSpinWait(iterations);
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private static void LongSpinWait(int iterations)
+        {
+            RuntimeImports.RhLongSpinWait(iterations);
+        }
+
+        public static void SpinWait(int iterations)
+        {
+            if (iterations <= 0)
+                return;
+
+            // Max iterations to be done in RhSpinWait.
+            // RhSpinWait does not switch GC modes and we want to avoid native spinning in coop mode for too long.
+            const int spinWaitCoopThreshold = 10000;
+
+            if (iterations > spinWaitCoopThreshold)
+            {
+                LongSpinWait(iterations);
+            }
+            else
+            {
+                RuntimeImports.RhSpinWait(iterations);
+            }
+        }
 
         [MethodImpl(MethodImplOptions.NoInlining)] // Slow path method. Make sure that the caller frame does not pay for PInvoke overhead.
         public static bool Yield() => RuntimeImports.RhYield();
-- 
2.7.4