Fix a series of possible race conditions in thread static variable initialization

src/coreclr/vm/threadstatics.cpp

@@ -164,6 +164,7 @@ PTR_MethodTable LookupMethodTableAndFlagForThreadStatic(TLSIndex index, bool *pI
     }
     CONTRACTL_END;
 
+    // This method is called without the g_TLSCrst lock being held
     PTR_MethodTable retVal;
     if (index.GetTLSIndexType() == TLSIndexType::NonCollectible)
     {
@@ -173,7 +174,7 @@ PTR_MethodTable LookupMethodTableAndFlagForThreadStatic(TLSIndex index, bool *pI
     {
         *pIsGCStatic = false;
         *pIsCollectible = false;
-        retVal = g_pMethodTablesForDirectThreadLocalData[IndexOffsetToDirectThreadLocalIndex(index.GetIndexOffset())];
+        retVal = VolatileLoadWithoutBarrier(&g_pMethodTablesForDirectThreadLocalData[IndexOffsetToDirectThreadLocalIndex(index.GetIndexOffset())]);
     }
     else
     {
@@ -230,8 +231,8 @@ void TLSIndexToMethodTableMap::Set(TLSIndex index, PTR_MethodTable pMT, bool isG
             memcpy(newMap, pMap, m_maxIndex * sizeof(TADDR));
             // Don't delete the old map in case some other thread is reading from it, this won't waste significant amounts of memory, since this map cannot grow indefinitely
         }
-        pMap = newMap;
-        m_maxIndex = newSize;
+        VolatileStore(&pMap, newMap); // Use a volatile store, so that the memcpy is guaranteed to be visible before the new pMap is visible on other threads.  This is paired with VolatileLoad in the Lookup functions although it probably could be a VolatileLoadWithoutBarrier and take advantage of data dependency.
+        VolatileStore(&m_maxIndex, newSize); // Use a volatile store so that any thread that reads the size will also have a consistent view of a map at least as big as newSize. This is paired with VolatileLoad in the Lookup functions.
     }
 
     TADDR rawValue = dac_cast<TADDR>(pMT);
@@ -245,7 +246,12 @@ void TLSIndexToMethodTableMap::Set(TLSIndex index, PTR_MethodTable pMT, bool isG
         m_collectibleEntries++;
     }
     _ASSERTE(pMap[index.GetIndexOffset()] == 0 || IsClearedValue(pMap[index.GetIndexOffset()]));
-    pMap[index.GetIndexOffset()] = rawValue;
+    // This VolatileStore does not have a clear consumer-producer relationship. Notably, the MethodTable is always considered to have been
+    // fully published by the time its pointer is stored in the map, so there is no need for a memory barrier to ensure visibility of the MethodTable's contents.
+    // However, we have had issues with race safety in this codebase, and we fixed them by adding several VolatileStore which we believe are actually critical
+    // but this is a scenario which AI indicated should also have a VolatileStore. We've chosen to add the VolatileStore as a defensive measure, as this is very
+    // rarely run code.
+    VolatileStore(&pMap[index.GetIndexOffset()], rawValue); 
 }
 
 void TLSIndexToMethodTableMap::Clear(TLSIndex index, uint8_t whenCleared)
@@ -265,7 +271,7 @@ void TLSIndexToMethodTableMap::Clear(TLSIndex index, uint8_t whenCleared)
     {
         m_collectibleEntries--;
     }
-    pMap[index.GetIndexOffset()] = (whenCleared << 2) | 0x3;
+    VolatileStore(&pMap[index.GetIndexOffset()], (TADDR)((whenCleared << 2) | 0x3));
     _ASSERTE(GetClearedMarker(pMap[index.GetIndexOffset()]) == whenCleared);
     _ASSERTE(IsClearedValue(pMap[index.GetIndexOffset()]));
 }
@@ -749,7 +755,7 @@ void GetTLSIndexForThreadStatic(MethodTable* pMT, bool gcStatic, TLSIndex* pInde
                 usedDirectOnThreadLocalDataPath = true;
             }
             if (usedDirectOnThreadLocalDataPath)
-                VolatileStoreWithoutBarrier(&g_pMethodTablesForDirectThreadLocalData[IndexOffsetToDirectThreadLocalIndex(newTLSIndex.GetIndexOffset())], pMT);
+                VolatileStore(&g_pMethodTablesForDirectThreadLocalData[IndexOffsetToDirectThreadLocalIndex(newTLSIndex.GetIndexOffset())], pMT);
         }
 
         if (!usedDirectOnThreadLocalDataPath)
@@ -772,7 +778,7 @@ void GetTLSIndexForThreadStatic(MethodTable* pMT, bool gcStatic, TLSIndex* pInde
         pMT->GetLoaderAllocator()->GetTLSIndexList().Append(newTLSIndex);
     }
 
-    *pIndex = newTLSIndex;
+    pIndex->VolatileStore(newTLSIndex); // Use a volatile store so that any other thread that sees the allocated index will also see the writes throughout this path.
 }
 
 void FreeTLSIndicesForLoaderAllocator(LoaderAllocator *pLoaderAllocator)

src/coreclr/vm/threadstatics.h

@@ -63,6 +63,7 @@ struct TLSIndex
     static TLSIndex Unallocated() { LIMITED_METHOD_DAC_CONTRACT; return TLSIndex(0xFFFFFFFF); }
     bool operator == (TLSIndex index) const { LIMITED_METHOD_DAC_CONTRACT; return TLSIndexRawIndex == index.TLSIndexRawIndex; }
     bool operator != (TLSIndex index) const { LIMITED_METHOD_DAC_CONTRACT; return TLSIndexRawIndex != index.TLSIndexRawIndex; }
+    void VolatileStore(const TLSIndex &other) { LIMITED_METHOD_CONTRACT; ::VolatileStore(&TLSIndexRawIndex, other.TLSIndexRawIndex); }
 };
 
 // Used to store access to TLS data for a single index when the TLS is accessed while the class constructor is running
@@ -123,6 +124,7 @@ class TLSIndexToMethodTableMap
     PTR_MethodTable Lookup(TLSIndex index, bool *isGCStatic, bool *isCollectible) const
     {
         LIMITED_METHOD_CONTRACT;
+        // This method is called without the g_TLSCrst lock being held
         *isGCStatic = false;
         *isCollectible = false;
         if (index.GetIndexOffset() < VolatileLoad(&m_maxIndex))
@@ -142,6 +144,7 @@ class TLSIndexToMethodTableMap
     PTR_MethodTable LookupTlsIndexKnownToBeAllocated(TLSIndex index) const
     {
         LIMITED_METHOD_CONTRACT;
+        // This method is called without the g_TLSCrst lock being held
         if (index.GetIndexOffset() < VolatileLoad(&m_maxIndex))
         {
             TADDR rawValue = VolatileLoadWithoutBarrier(&VolatileLoad(&pMap)[index.GetIndexOffset()]);
@@ -165,9 +168,14 @@ class TLSIndexToMethodTableMap
 
     entry Lookup(TLSIndex index) const
     {
+        // This method is called with the g_TLSCrst lock held, so we don't actually
+        // need all of these volatile loads, but using VolatileLoad is more similar to the other
+        // paths, and the performance penalty is negligible, so we keep them.
+
         LIMITED_METHOD_CONTRACT;
         entry e(index);
-        if (index.GetIndexOffset() < VolatileLoad(&m_maxIndex))
+        int32_t maxIndex = VolatileLoad(&m_maxIndex); // This VolatileLoad pairs with a VolatileStore in TLSIndexToMethodTableMap::Set to ensure that if we read a maxIndex that is large enough to contain our index
+        if (index.GetIndexOffset() < maxIndex)
         {
             TADDR rawValue = VolatileLoadWithoutBarrier(&VolatileLoad(&pMap)[index.GetIndexOffset()]);
             if (!IsClearedValue(rawValue))
@@ -184,7 +192,7 @@ class TLSIndexToMethodTableMap
         }
         else
         {
-            e.TlsIndex = TLSIndex(m_indexType, m_maxIndex);
+            e.TlsIndex = TLSIndex(m_indexType, maxIndex);
         }
         return e;
     }