Fix class construction cycle in Lock on NativeAOT (dotnet#100374)

When a thread reenters class construction through accessing NativeRuntimeEventSource.Log, Log would return null. Checks on IsFullyInitialized were added to ensure that the normal path would not be taken in that case, to avoid null checks in several places and in different files. That doesn't work when a different thread sees the initialization stage as Complete, as it would try to initialize NativeRuntimeEventSource and run into a class construction cycle.

Fixed by removing the IsFullyInitialized checks, introducing a new initialization stage PartiallyCompelte, and not setting the stage to Complete until it has been verified that Log does not return null. When the stage is PartiallyCompelte, a thread would retry the relevant initialization. This again guarantees that there would be at most one attempt at initialization through Lock at any given time, and prevents the class construction cycle.

Fixes dotnet#99663

Author: kouvel

src/coreclr/nativeaot/System.Private.CoreLib/src/System/Threading/Lock.NativeAot.cs

@@ -92,18 +92,6 @@ internal void Reenter(uint previousRecursionCount)
             _recursionCount = previousRecursionCount;
         }
 
-        private static bool IsFullyInitialized
-        {
-            get
-            {
-                // If NativeRuntimeEventSource is already being class-constructed by this thread earlier in the stack, Log can
-                // be null. This property is used to avoid going down the wait path in that case to avoid null checks in several
-                // other places.
-                Debug.Assert((StaticsInitializationStage)s_staticsInitializationStage == StaticsInitializationStage.Complete);
-                return NativeRuntimeEventSource.Log != null;
-            }
-        }
-
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private TryLockResult LazyInitializeOrEnter()
         {
@@ -113,10 +101,6 @@ private TryLockResult LazyInitializeOrEnter()
                 case StaticsInitializationStage.Complete:
                     if (_spinCount == SpinCountNotInitialized)
                     {
-                        if (!IsFullyInitialized)
-                        {
-                            goto case StaticsInitializationStage.Started;
-                        }
                         _spinCount = s_maxSpinCount;
                     }
                     return TryLockResult.Spin;
@@ -137,7 +121,7 @@ private TryLockResult LazyInitializeOrEnter()
                         }
 
                         stage = (StaticsInitializationStage)Volatile.Read(ref s_staticsInitializationStage);
-                        if (stage == StaticsInitializationStage.Complete && IsFullyInitialized)
+                        if (stage == StaticsInitializationStage.Complete)
                         {
                             goto case StaticsInitializationStage.Complete;
                         }
@@ -155,7 +139,9 @@ private TryLockResult LazyInitializeOrEnter()
                     }
 
                 default:
-                    Debug.Assert(stage == StaticsInitializationStage.NotStarted);
+                    Debug.Assert(
+                        stage == StaticsInitializationStage.NotStarted ||
+                        stage == StaticsInitializationStage.PartiallyComplete);
                     if (TryInitializeStatics())
                     {
                         goto case StaticsInitializationStage.Complete;
@@ -169,29 +155,49 @@ private static bool TryInitializeStatics()
         {
             // Since Lock is used to synchronize class construction, and some of the statics initialization may involve class
             // construction, update the stage first to avoid infinite recursion
-            switch (
-                (StaticsInitializationStage)
-                Interlocked.CompareExchange(
-                    ref s_staticsInitializationStage,
-                    (int)StaticsInitializationStage.Started,
-                    (int)StaticsInitializationStage.NotStarted))
+            var oldStage = (StaticsInitializationStage)s_staticsInitializationStage;
+            while (true)
             {
-                case StaticsInitializationStage.Started:
-                    return false;
-                case StaticsInitializationStage.Complete:
+                if (oldStage == StaticsInitializationStage.Complete)
+                {
                     return true;
+                }
+
+                var stageBeforeUpdate =
+                    (StaticsInitializationStage)Interlocked.CompareExchange(
+                        ref s_staticsInitializationStage,
+                        (int)StaticsInitializationStage.Started,
+                        (int)oldStage);
+                if (stageBeforeUpdate == StaticsInitializationStage.Started)
+                {
+                    return false;
+                }
+                if (stageBeforeUpdate == oldStage)
+                {
+                    Debug.Assert(
+                        oldStage == StaticsInitializationStage.NotStarted ||
+                        oldStage == StaticsInitializationStage.PartiallyComplete);
+                    break;
+                }
+
+                oldStage = stageBeforeUpdate;
             }
 
             bool isFullyInitialized;
             try
             {
-                s_isSingleProcessor = Environment.IsSingleProcessor;
-                s_maxSpinCount = DetermineMaxSpinCount();
-                s_minSpinCount = DetermineMinSpinCount();
+                if (oldStage == StaticsInitializationStage.NotStarted)
+                {
+                    // If the stage is PartiallyComplete, these will have already been initialized
+                    s_isSingleProcessor = Environment.IsSingleProcessor;
+                    s_maxSpinCount = DetermineMaxSpinCount();
+                    s_minSpinCount = DetermineMinSpinCount();
+                }
 
                 // Also initialize some types that are used later to prevent potential class construction cycles. If
                 // NativeRuntimeEventSource is already being class-constructed by this thread earlier in the stack, Log can be
-                // null. Avoid going down the wait path in that case to avoid null checks in several other places.
+                // null. Avoid going down the wait path in that case to avoid null checks in several other places. If not fully
+                // initialized, the stage will also be set to PartiallyComplete to try again.
                 isFullyInitialized = NativeRuntimeEventSource.Log != null;
             }
             catch
@@ -200,7 +206,11 @@ private static bool TryInitializeStatics()
                 throw;
             }
 
-            Volatile.Write(ref s_staticsInitializationStage, (int)StaticsInitializationStage.Complete);
+            Volatile.Write(
+                ref s_staticsInitializationStage,
+                isFullyInitialized
+                    ? (int)StaticsInitializationStage.Complete
+                    : (int)StaticsInitializationStage.PartiallyComplete);
             return isFullyInitialized;
         }
 
@@ -242,6 +252,7 @@ private enum StaticsInitializationStage
         {
             NotStarted,
             Started,
+            PartiallyComplete,
             Complete
         }
     }