[NativeAOT] Making thread local storage completely managed. (#84060)

* move m_pThreadLocalModuleStatics to managed

* typo - should use newSize once computed

src/coreclr/nativeaot/Runtime/AsmOffsets.h

@@ -52,7 +52,6 @@ ASM_OFFSET(   48,    80, Thread, m_pExInfoStackHead)
 ASM_OFFSET(   4c,    88, Thread, m_threadAbortException)
 
 ASM_OFFSET(   50,    90, Thread, m_pThreadLocalModuleStatics)
-ASM_OFFSET(   54,    98, Thread, m_numThreadLocalModuleStatics)
 
 ASM_SIZEOF(   14,    20, EHEnum)
 

src/coreclr/nativeaot/Runtime/amd64/MiscStubs.S

@@ -62,19 +62,25 @@ NESTED_ENTRY RhpGetThreadStaticBaseForType, _TEXT, NoHandler
         INLINE_GETTHREAD
 
         mov     r8d, [rbx + 8]         // Get ModuleIndex out of the TypeManagerSlot
-        cmp     r8d, [rax + OFFSETOF__Thread__m_numThreadLocalModuleStatics]
-        jae     LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
 
-        mov     r9, [rax + OFFSETOF__Thread__m_pThreadLocalModuleStatics]
-        mov     rax, [r9 + r8 * 8]     // Index into the array of modules
+        // get per-thread storage
+        mov     rax, [rax + OFFSETOF__Thread__m_pThreadLocalModuleStatics]
+
+        // get per-module storage
         test    rax, rax
         jz      LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
+        cmp     r8d, [rax + OFFSETOF__Array__m_Length]
+        jae     LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
+        mov     rax, [rax + r8 * 8 + 0x10]
 
-        mov     r8, [rax]              // Get the managed array from the handle
-        cmp     r12d, [r8 + OFFSETOF__Array__m_Length]
+        // get the actual per-type storage
+        test    rax, rax
+        jz      LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
+        cmp     r12d, [rax + OFFSETOF__Array__m_Length]
         jae     LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
-        mov     rax, [r8 + r12 * 8 + 0x10]
+        mov     rax, [rax + r12 * 8 + 0x10]
 
+        // if have storage, return it
         test    rax, rax
         jz      LOCAL_LABEL(RhpGetThreadStaticBaseForType_RarePath)
 

src/coreclr/nativeaot/Runtime/amd64/MiscStubs.asm

@@ -192,29 +192,30 @@ LEAF_ENTRY RhpGetThreadStaticBaseForType, _TEXT
         INLINE_GETTHREAD rax, r8
 
         mov     r8d, [rcx + 8]         ; Get ModuleIndex out of the TypeManagerSlot
-        cmp     r8d, [rax + OFFSETOF__Thread__m_numThreadLocalModuleStatics]
-        jae     RhpGetThreadStaticBaseForType_RarePath
 
-        mov     r9, [rax + OFFSETOF__Thread__m_pThreadLocalModuleStatics]
-        mov     rax, [r9 + r8 * 8]     ; Index into the array of modules
+        ;; get per-thread storage
+        mov     rax, [rax + OFFSETOF__Thread__m_pThreadLocalModuleStatics]
+
+        ;; get per-module storage
         test    rax, rax
-        jz      RhpGetThreadStaticBaseForType_RarePath
+        jz      RhpGetThreadStaticBaseForTypeSlow
+        cmp     r8d, [rax + OFFSETOF__Array__m_Length]
+        jae     RhpGetThreadStaticBaseForTypeSlow
+        mov     rax, [rax + r8 * 8 + 10h]
 
-        mov     r8, [rax]              ; Get the managed array from the handle
-        cmp     edx, [r8 + OFFSETOF__Array__m_Length]
-        jae     RhpGetThreadStaticBaseForType_RarePath
-        mov     rax, [r8 + rdx * 8 + 10h]
+        ;; get the actual per-type storage
+        test    rax, rax
+        jz      RhpGetThreadStaticBaseForTypeSlow
+        cmp     edx, [rax + OFFSETOF__Array__m_Length]
+        jae     RhpGetThreadStaticBaseForTypeSlow
+        mov     rax, [rax + rdx * 8 + 10h]
 
+        ;; if have storage, return it
         test    rax, rax
-        jz      RhpGetThreadStaticBaseForType_RarePath
+        jz      RhpGetThreadStaticBaseForTypeSlow
 
         ret
 
-RhpGetThreadStaticBaseForType_RarePath:
-        ;; We kept the arguments in their appropriate registers
-        ;; and we can tailcall right away.
-        jmp     RhpGetThreadStaticBaseForTypeSlow
-
 LEAF_END RhpGetThreadStaticBaseForType, _TEXT
 
 end

src/coreclr/nativeaot/Runtime/gcrhscan.cpp

@@ -28,6 +28,7 @@
 #ifndef DACCESS_COMPILE
 
 void GcEnumObjectsConservatively(PTR_PTR_Object ppLowerBound, PTR_PTR_Object ppUpperBound, EnumGcRefCallbackFunc * fnGcEnumRef, EnumGcRefScanContext * pSc);
+void GcEnumObject(PTR_PTR_Object ppObj, uint32_t flags, EnumGcRefCallbackFunc* fnGcEnumRef, EnumGcRefScanContext* pSc);
 
 /*
  * Scan all stack and statics roots
@@ -53,6 +54,9 @@ void GCToEEInterface::GcScanRoots(EnumGcRefCallbackFunc * fn,  int condemned, in
         else
 #endif
         {
+            STRESS_LOG1(LF_GC | LF_GCROOTS, LL_INFO100, "{ Scanning Thread's %p thread statics root. \n", pThread);
+            GcEnumObject(pThread->GetThreadStaticStorage(), 0 /*flags*/, fn, sc);
+
             STRESS_LOG1(LF_GC|LF_GCROOTS, LL_INFO100, "{ Starting scan of Thread %p\n", pThread);
             sc->thread_under_crawl = pThread;
 #if defined(FEATURE_EVENT_TRACE) && !defined(DACCESS_COMPILE)

src/coreclr/nativeaot/Runtime/thread.cpp

@@ -285,7 +285,6 @@ void Thread::Construct()
     // Everything else should be initialized to 0 via the static initialization of tls_CurrentThread.
 
     ASSERT(m_pThreadLocalModuleStatics == NULL);
-    ASSERT(m_numThreadLocalModuleStatics == 0);
 
     ASSERT(m_pGCFrameRegistrations == NULL);
 
@@ -349,18 +348,6 @@ void Thread::Destroy()
     if (m_hPalThread != INVALID_HANDLE_VALUE)
         PalCloseHandle(m_hPalThread);
 
-    if (m_pThreadLocalModuleStatics != NULL)
-    {
-        for (uint32_t i = 0; i < m_numThreadLocalModuleStatics; i++)
-        {
-            if (m_pThreadLocalModuleStatics[i] != NULL)
-            {
-                RhHandleFree(m_pThreadLocalModuleStatics[i]);
-            }
-        }
-        delete[] m_pThreadLocalModuleStatics;
-    }
-
 #ifdef STRESS_LOG
     ThreadStressLog* ptsl = reinterpret_cast<ThreadStressLog*>(GetThreadStressLog());
     StressLog::ThreadDetach(ptsl);
@@ -1269,78 +1256,15 @@ COOP_PINVOKE_HELPER(Object *, RhpGetThreadAbortException, ())
     return pCurThread->GetThreadAbortException();
 }
 
-Object* Thread::GetThreadStaticStorageForModule(uint32_t moduleIndex)
-{
-    // Return a pointer to the TLS storage if it has already been
-    // allocated for the specified module.
-    if (moduleIndex < m_numThreadLocalModuleStatics)
-    {
-        Object** threadStaticsStorageHandle = (Object**)m_pThreadLocalModuleStatics[moduleIndex];
-        if (threadStaticsStorageHandle != NULL)
-        {
-            return *threadStaticsStorageHandle;
-        }
-    }
-
-    return NULL;
-}
-
-bool Thread::SetThreadStaticStorageForModule(Object * pStorage, uint32_t moduleIndex)
-{
-    // Grow thread local storage if needed.
-    if (m_numThreadLocalModuleStatics <= moduleIndex)
-    {
-        uint32_t newSize = moduleIndex + 1;
-        if (newSize < moduleIndex)
-        {
-            return false;
-        }
-
-        PTR_PTR_VOID pThreadLocalModuleStatics = new (nothrow) PTR_VOID[newSize];
-        if (pThreadLocalModuleStatics == NULL)
-        {
-            return false;
-        }
-
-        memset(&pThreadLocalModuleStatics[m_numThreadLocalModuleStatics], 0, sizeof(PTR_VOID) * (newSize - m_numThreadLocalModuleStatics));
-
-        if (m_pThreadLocalModuleStatics != NULL)
-        {
-            memcpy(pThreadLocalModuleStatics, m_pThreadLocalModuleStatics, sizeof(PTR_VOID) * m_numThreadLocalModuleStatics);
-            delete[] m_pThreadLocalModuleStatics;
-        }
-
-        m_pThreadLocalModuleStatics = pThreadLocalModuleStatics;
-        m_numThreadLocalModuleStatics = newSize;
-    }
-
-    if (m_pThreadLocalModuleStatics[moduleIndex] != NULL)
-    {
-        RhHandleSet(m_pThreadLocalModuleStatics[moduleIndex], pStorage);
-    }
-    else
-    {
-        void* threadStaticsStorageHandle = RhpHandleAlloc(pStorage, 2 /* Normal */);
-        if (threadStaticsStorageHandle == NULL)
-        {
-            return false;
-        }
-        m_pThreadLocalModuleStatics[moduleIndex] = threadStaticsStorageHandle;
-    }
-
-    return true;
-}
-
-COOP_PINVOKE_HELPER(Object*, RhGetThreadStaticStorageForModule, (uint32_t moduleIndex))
+Object** Thread::GetThreadStaticStorage()
 {
-    Thread * pCurrentThread = ThreadStore::RawGetCurrentThread();
-    return pCurrentThread->GetThreadStaticStorageForModule(moduleIndex);
+    return &m_pThreadLocalModuleStatics;
 }
 
-COOP_PINVOKE_HELPER(FC_BOOL_RET, RhSetThreadStaticStorageForModule, (Array * pStorage, uint32_t moduleIndex))
+COOP_PINVOKE_HELPER(Object**, RhGetThreadStaticStorage, ())
 {
     Thread * pCurrentThread = ThreadStore::RawGetCurrentThread();
-    FC_RETURN_BOOL(pCurrentThread->SetThreadStaticStorageForModule((Object*)pStorage, moduleIndex));
+    return pCurrentThread->GetThreadStaticStorage();
 }
 
 // This is function is used to quickly query a value that can uniquely identify a thread

src/coreclr/nativeaot/Runtime/thread.h

@@ -88,8 +88,7 @@ struct ThreadBuffer
     uintptr_t               m_uHijackedReturnValueFlags;            
     PTR_ExInfo              m_pExInfoStackHead;
     Object*                 m_threadAbortException;                 // ThreadAbortException instance -set only during thread abort
-    PTR_PTR_VOID            m_pThreadLocalModuleStatics;
-    uint32_t                m_numThreadLocalModuleStatics;
+    Object*                 m_pThreadLocalModuleStatics;
     GCFrameRegistration*    m_pGCFrameRegistrations;
     PTR_VOID                m_pStackLow;
     PTR_VOID                m_pStackHigh;
@@ -283,11 +282,10 @@ class Thread : private ThreadBuffer
     void InlinePInvoke(PInvokeTransitionFrame * pFrame);
     void InlinePInvokeReturn(PInvokeTransitionFrame * pFrame);
 
-    Object * GetThreadAbortException();
+    Object* GetThreadAbortException();
     void SetThreadAbortException(Object *exception);
 
-    Object* GetThreadStaticStorageForModule(uint32_t moduleIndex);
-    bool SetThreadStaticStorageForModule(Object* pStorage, uint32_t moduleIndex);
+    Object** GetThreadStaticStorage();
 
     NATIVE_CONTEXT* GetInterruptedContext();
 

src/coreclr/nativeaot/System.Private.CoreLib/src/Internal/Runtime/ThreadStatics.cs

@@ -23,13 +23,22 @@ internal static class ThreadStatics
         /// </summary>
         internal static unsafe object GetThreadStaticBaseForType(TypeManagerSlot* pModuleData, int typeTlsIndex)
         {
-            // Get the array that holds thread static memory blocks for each type in the given module
-            object[] storage = RuntimeImports.RhGetThreadStaticStorageForModule(pModuleData->ModuleIndex);
+            Debug.Assert(typeTlsIndex >= 0);
+            int moduleIndex = pModuleData->ModuleIndex;
+            Debug.Assert(moduleIndex >= 0);
 
-            // Check whether thread static storage has already been allocated for this module and type.
-            if ((storage != null) && ((uint)typeTlsIndex < (uint)storage.Length) && (storage[typeTlsIndex] != null))
+            object[][] threadStorage = RuntimeImports.RhGetThreadStaticStorage();
+            if (threadStorage != null && threadStorage.Length > moduleIndex)
             {
-                return storage[typeTlsIndex];
+                object[] moduleStorage = threadStorage[moduleIndex];
+                if (moduleStorage != null && moduleStorage.Length > typeTlsIndex)
+                {
+                    object threadStaticBase = moduleStorage[typeTlsIndex];
+                    if (threadStaticBase != null)
+                    {
+                        return threadStaticBase;
+                    }
+                }
             }
 
             return GetThreadStaticBaseForTypeSlow(pModuleData, typeTlsIndex);
@@ -39,49 +48,42 @@ internal static unsafe object GetThreadStaticBaseForType(TypeManagerSlot* pModul
         [MethodImpl(MethodImplOptions.NoInlining)]
         internal static unsafe object GetThreadStaticBaseForTypeSlow(TypeManagerSlot* pModuleData, int typeTlsIndex)
         {
-            // Get the array that holds thread static memory blocks for each type in the given module
-            object[] storage = RuntimeImports.RhGetThreadStaticStorageForModule(pModuleData->ModuleIndex);
-
-            // This the first access to the thread statics of the type corresponding to typeTlsIndex.
-            // Make sure there is enough storage allocated to hold it.
-            storage = EnsureThreadStaticStorage(pModuleData->ModuleIndex, storage, requiredSize: typeTlsIndex + 1);
-
-            // Allocate an object that will represent a memory block for all thread static fields of the type
-            object threadStaticBase = AllocateThreadStaticStorageForType(pModuleData->TypeManager, typeTlsIndex);
-
-            Debug.Assert(storage[typeTlsIndex] == null);
-
-            storage[typeTlsIndex] = threadStaticBase;
-
-            return threadStaticBase;
-        }
-
-        /// <summary>
-        /// if it is required, this method extends thread static storage of the given module
-        /// to the specified size and then registers the memory with the runtime.
-        /// </summary>
-        private static object[] EnsureThreadStaticStorage(int moduleIndex, object[] existingStorage, int requiredSize)
-        {
-            if ((existingStorage != null) && (requiredSize < existingStorage.Length))
+            Debug.Assert(typeTlsIndex >= 0);
+            int moduleIndex = pModuleData->ModuleIndex;
+            Debug.Assert(typeTlsIndex >= 0);
+
+            // Get the array that holds thread statics for the current thread, if none present
+            // allocate a new one big enough to hold the current module data
+            ref object[][] threadStorage = ref RuntimeImports.RhGetThreadStaticStorage();
+            if (threadStorage == null)
             {
-                return existingStorage;
+                threadStorage = new object[moduleIndex + 1][];
             }
-
-            object[] newStorage = new object[requiredSize];
-            if (existingStorage != null)
+            else if (moduleIndex >= threadStorage.Length)
             {
-                Array.Copy(existingStorage, newStorage, existingStorage.Length);
+                Array.Resize(ref threadStorage, moduleIndex + 1);
             }
 
-            // Install the newly created array as thread static storage for the given module
-            // on the current thread. This call can fail due to a failure to allocate/extend required
-            // internal thread specific resources.
-            if (!RuntimeImports.RhSetThreadStaticStorageForModule(newStorage, moduleIndex))
+            // Get the array that holds thread static memory blocks for each type in the given module
+            ref object[] moduleStorage = ref threadStorage[moduleIndex];
+            if (moduleStorage == null)
             {
-                throw new OutOfMemoryException();
+                moduleStorage = new object[typeTlsIndex + 1];
+            }
+            else if (typeTlsIndex >= moduleStorage.Length)
+            {
+                // typeTlsIndex could have a big range, we do not want to reallocate every time we see +1 index
+                // so we double up from previous size to guarantee a worst case linear complexity
+                int newSize = Math.Max(typeTlsIndex + 1, moduleStorage.Length * 2);
+                Array.Resize(ref moduleStorage, newSize);
             }
 
-            return newStorage;
+            // Allocate an object that will represent a memory block for all thread static fields of the type
+            object threadStaticBase = AllocateThreadStaticStorageForType(pModuleData->TypeManager, typeTlsIndex);
+
+            Debug.Assert(moduleStorage[typeTlsIndex] == null);
+            moduleStorage[typeTlsIndex] = threadStaticBase;
+            return threadStaticBase;
         }
 
         /// <summary>

src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/RuntimeImports.cs

@@ -561,12 +561,8 @@ internal static IntPtr RhGetModuleSection(TypeManagerHandle module, ReadyToRunSe
         internal static extern IntPtr RhGetOSModuleFromEEType(IntPtr pEEType);
 
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
-        [RuntimeImport(RuntimeLibrary, "RhGetThreadStaticStorageForModule")]
-        internal static extern unsafe object[] RhGetThreadStaticStorageForModule(int moduleIndex);
-
-        [MethodImplAttribute(MethodImplOptions.InternalCall)]
-        [RuntimeImport(RuntimeLibrary, "RhSetThreadStaticStorageForModule")]
-        internal static extern unsafe bool RhSetThreadStaticStorageForModule(object[] storage, int moduleIndex);
+        [RuntimeImport(RuntimeLibrary, "RhGetThreadStaticStorage")]
+        internal static extern ref object[][] RhGetThreadStaticStorage();
 
         [MethodImplAttribute(MethodImplOptions.InternalCall)]
         [RuntimeImport(RuntimeLibrary, "RhCurrentNativeThreadId")]