Fix to not reuse preallocated jump stubs for dynamic methods (#16941)

Fix to not reuse preallocated jump stubs for dynamic methods

Fixes https://github.com/dotnet/coreclr/issues/16940
- Allocate an extra jump stub per temporary entry points chunk that is shared by all precodes in the chunk. This jump stub always points to PrecodeFixupThunk.
- Use that for PrecodeFixupThunk, and use the precode-associated jump stub for pointing to the jitted code
- Considered allocating the extra jump stub only if coreclr is far away, but it involves reallocation which may be common in some environments/scenarios. Figured 12 extra bytes per dynamic type is not too significant.

7 files changed, 104 insertions, 17 deletions

diff --git a/src/vm/amd64/cgenamd64.cpp b/src/vm/amd64/cgenamd64.cpp
index 56e3bfa738..6d11c7f0fa 100644
--- a/src/vm/amd64/cgenamd64.cpp
+++ b/src/vm/amd64/cgenamd64.cpp
@@ -689,7 +689,7 @@ INT32 rel32UsingJumpStub(INT32 UNALIGNED * pRel32, PCODE target, MethodDesc *pMe
     return static_cast<INT32>(offset);
 }
 
-INT32 rel32UsingPreallocatedJumpStub(INT32 UNALIGNED * pRel32, PCODE target, PCODE jumpStubAddr)
+INT32 rel32UsingPreallocatedJumpStub(INT32 UNALIGNED * pRel32, PCODE target, PCODE jumpStubAddr, bool emitJump)
 {
     CONTRACTL
     {
@@ -711,7 +711,14 @@ INT32 rel32UsingPreallocatedJumpStub(INT32 UNALIGNED * pRel32, PCODE target, PCO
             EEPOLICY_HANDLE_FATAL_ERROR(COR_E_EXECUTIONENGINE);
         }
 
-        emitBackToBackJump((LPBYTE)jumpStubAddr, (LPVOID)target);
+        if (emitJump)
+        {
+            emitBackToBackJump((LPBYTE)jumpStubAddr, (LPVOID)target);
+        }
+        else
+        {
+            _ASSERTE(decodeBackToBackJump(jumpStubAddr) == target);
+        }
     }
 
     _ASSERTE(FitsInI4(offset));
diff --git a/src/vm/amd64/cgencpu.h b/src/vm/amd64/cgencpu.h
index 9136b168aa..ab049a49e2 100644
--- a/src/vm/amd64/cgencpu.h
+++ b/src/vm/amd64/cgencpu.h
@@ -383,7 +383,7 @@ INT32 rel32UsingJumpStub(INT32 UNALIGNED * pRel32, PCODE target, MethodDesc *pMe
     LoaderAllocator *pLoaderAllocator = NULL, bool throwOnOutOfMemoryWithinRange = true);
 
 // Get Rel32 destination, emit jumpStub if necessary into a preallocated location
-INT32 rel32UsingPreallocatedJumpStub(INT32 UNALIGNED * pRel32, PCODE target, PCODE jumpStubAddr);
+INT32 rel32UsingPreallocatedJumpStub(INT32 UNALIGNED * pRel32, PCODE target, PCODE jumpStubAddr, bool emitJump);
 
 void emitCOMStubCall (ComCallMethodDesc *pCOMMethod, PCODE target);
 
diff --git a/src/vm/crossgencompile.cpp b/src/vm/crossgencompile.cpp
index 790b0d817c..cea93bff57 100644
--- a/src/vm/crossgencompile.cpp
+++ b/src/vm/crossgencompile.cpp
@@ -270,7 +270,7 @@ INT32 rel32UsingJumpStub(INT32 UNALIGNED * pRel32, PCODE target, MethodDesc *pMe
     return 0;
 }
 
-INT32 rel32UsingPreallocatedJumpStub(INT32 UNALIGNED * pRel32, PCODE target, PCODE jumpStubAddr)
+INT32 rel32UsingPreallocatedJumpStub(INT32 UNALIGNED * pRel32, PCODE target, PCODE jumpStubAddr, bool emitJump)
 {
     // crossgen does not have jump stubs
     return 0;
diff --git a/src/vm/i386/stublinkerx86.cpp b/src/vm/i386/stublinkerx86.cpp
index ff3ec48d60..0b63e68942 100644
--- a/src/vm/i386/stublinkerx86.cpp
+++ b/src/vm/i386/stublinkerx86.cpp
@@ -6454,16 +6454,43 @@ TADDR FixupPrecode::GetMethodDesc()
 #endif
 
 #ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+PCODE FixupPrecode::GetDynamicMethodPrecodeFixupJumpStub()
+{
+    WRAPPER_NO_CONTRACT;
+    _ASSERTE(((PTR_MethodDesc)GetMethodDesc())->IsLCGMethod());
+
+    // The precode fixup jump stub is shared by all fixup precodes in a chunk, and immediately follows the MethodDesc. Jump
+    // stubs cannot be reused currently for the same method:
+    //   - The jump stub's target would change separately from the precode being updated from "call Func" to "jmp Func", both
+    //     changes would have to be done atomically with runtime suspension, which is not done currently
+    //   - When changing the entry point from one version of jitted code to another, the jump stub's target pointer is not
+    //     aligned to 8 bytes in order to be able to do an interlocked update of the target address
+    // So, when initially the precode intends to be of the form "call PrecodeFixupThunk", if the target address happens to be
+    // too far for a relative 32-bit jump, it will use the shared precode fixup jump stub. When changing the entry point to
+    // jitted code, the jump stub associated with the precode is patched, and the precode is updated to use that jump stub.
+    //
+    // Notes:
+    // - Dynamic method descs, and hence their precodes and preallocated jump stubs, may be reused for a different method
+    //   (along with reinitializing the precode), but only with a transition where the original method is no longer accessible
+    //   to user code
+    // - Concurrent calls to a dynamic method that has not yet been jitted may trigger multiple writes to the jump stub
+    //   associated with the precode, but only to the same target address (and while the precode is still pointing to
+    //   PrecodeFixupThunk)
+    return GetBase() + sizeof(PTR_MethodDesc);
+}
+
 PCODE FixupPrecode::GetDynamicMethodEntryJumpStub()
 {
+    WRAPPER_NO_CONTRACT;
     _ASSERTE(((PTR_MethodDesc)GetMethodDesc())->IsLCGMethod());
 
     // m_PrecodeChunkIndex has a value inverted to the order of precodes in memory (the precode at the lowest address has the
     // highest index, and the precode at the highest address has the lowest index). To map a precode to its jump stub by memory
-    // order, invert the precode index to get the jump stub index.
+    // order, invert the precode index to get the jump stub index. Also skip the precode fixup jump stub (see
+    // GetDynamicMethodPrecodeFixupJumpStub()).
     UINT32 count = ((PTR_MethodDesc)GetMethodDesc())->GetMethodDescChunk()->GetCount();
     _ASSERTE(m_PrecodeChunkIndex < count);
-    SIZE_T jumpStubIndex = count - 1 - m_PrecodeChunkIndex;
+    SIZE_T jumpStubIndex = count - m_PrecodeChunkIndex;
 
     return GetBase() + sizeof(PTR_MethodDesc) + jumpStubIndex * BACK_TO_BACK_JUMP_ALLOCATE_SIZE;
 }
@@ -6607,7 +6634,7 @@ void FixupPrecode::Init(MethodDesc* pMD, LoaderAllocator *pLoaderAllocator, int
 #ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
     if (pMD->IsLCGMethod())
     {
-        m_rel32 = rel32UsingPreallocatedJumpStub(&m_rel32, target, GetDynamicMethodEntryJumpStub());
+        m_rel32 = rel32UsingPreallocatedJumpStub(&m_rel32, target, GetDynamicMethodPrecodeFixupJumpStub(), false /* emitJump */);
         return;
     }
 #endif // FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
@@ -6632,12 +6659,15 @@ void FixupPrecode::ResetTargetInterlocked()
 
     PCODE target = (PCODE)GetEEFuncEntryPoint(PrecodeFixupThunk);
     MethodDesc* pMD = (MethodDesc*)GetMethodDesc();
-    newValue.m_rel32 =
 #ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
-        pMD->IsLCGMethod() ?
-            rel32UsingPreallocatedJumpStub(&m_rel32, target, GetDynamicMethodEntryJumpStub()) :
+    // The entry point of LCG methods cannot revert back to the original entry point, as their jump stubs would have to be
+    // reused, which is currently not supported. This method is intended for resetting the entry point while the method is
+    // callable, which implies that the entry point may later be changed again to something else. Currently, this is not done
+    // for LCG methods. See GetDynamicMethodPrecodeFixupJumpStub() for more.
+    _ASSERTE(!pMD->IsLCGMethod());
 #endif // FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
-            rel32UsingJumpStub(&m_rel32, target, pMD);
+
+    newValue.m_rel32 = rel32UsingJumpStub(&m_rel32, target, pMD);
 
     _ASSERTE(IS_ALIGNED(this, sizeof(INT64)));
     EnsureWritableExecutablePages(this, sizeof(INT64));
@@ -6659,6 +6689,11 @@ BOOL FixupPrecode::SetTargetInterlocked(TADDR target, TADDR expected)
     MethodDesc * pMD = (MethodDesc*)GetMethodDesc();
     g_IBCLogger.LogMethodPrecodeWriteAccess(pMD);
     
+#ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+    // A different jump stub is used for this case, see Init(). This call is unexpected for resetting the entry point.
+    _ASSERTE(!pMD->IsLCGMethod() || target != (TADDR)GetEEFuncEntryPoint(PrecodeFixupThunk));
+#endif // FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+
     INT64 newValue = oldValue;
     BYTE* pNewValue = (BYTE*)&newValue;
 
@@ -6687,7 +6722,7 @@ BOOL FixupPrecode::SetTargetInterlocked(TADDR target, TADDR expected)
     *(INT32*)(&pNewValue[offsetof(FixupPrecode, m_rel32)]) =
 #ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
         pMD->IsLCGMethod() ?
-            rel32UsingPreallocatedJumpStub(&m_rel32, target, GetDynamicMethodEntryJumpStub()) :
+            rel32UsingPreallocatedJumpStub(&m_rel32, target, GetDynamicMethodEntryJumpStub(), true /* emitJump */) :
 #endif // FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
             rel32UsingJumpStub(&m_rel32, target, pMD);
 
diff --git a/src/vm/i386/stublinkerx86.h b/src/vm/i386/stublinkerx86.h
index 229ab0767b..c5378422a1 100644
--- a/src/vm/i386/stublinkerx86.h
+++ b/src/vm/i386/stublinkerx86.h
@@ -720,6 +720,7 @@ struct FixupPrecode {
 #endif // HAS_FIXUP_PRECODE_CHUNKS
 
 #ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+    PCODE GetDynamicMethodPrecodeFixupJumpStub();
     PCODE GetDynamicMethodEntryJumpStub();
 #endif // FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
 
diff --git a/src/vm/precode.cpp b/src/vm/precode.cpp
index 103fc03959..444b4dd1fc 100644
--- a/src/vm/precode.cpp
+++ b/src/vm/precode.cpp
@@ -318,8 +318,10 @@ SIZE_T Precode::SizeOfTemporaryEntryPoints(PrecodeType t, bool preallocateJumpSt
         if (preallocateJumpStubs)
         {
             // For dynamic methods, space for jump stubs is allocated along with the precodes as part of the temporary entry
-            // points block. The first jump stub begins immediately after the PTR_MethodDesc.
-            size += count * BACK_TO_BACK_JUMP_ALLOCATE_SIZE;
+            // points block. The first jump stub begins immediately after the PTR_MethodDesc. Aside from a jump stub per
+            // precode, an additional shared precode fixup jump stub is also allocated (see
+            // GetDynamicMethodPrecodeFixupJumpStub()).
+            size += ((SIZE_T)count + 1) * BACK_TO_BACK_JUMP_ALLOCATE_SIZE;
         }
 #else // !FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
         _ASSERTE(!preallocateJumpStubs);
@@ -571,12 +573,32 @@ TADDR Precode::AllocateTemporaryEntryPoints(MethodDescChunk *  pChunk,
 #ifdef HAS_FIXUP_PRECODE_CHUNKS
     if (t == PRECODE_FIXUP)
     {
+#ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+        PCODE precodeFixupJumpStub = NULL;
+        if (preallocateJumpStubs)
+        {
+            // Emit the jump for the precode fixup jump stub now. This jump stub immediately follows the MethodDesc (see
+            // GetDynamicMethodPrecodeFixupJumpStub()).
+            precodeFixupJumpStub = temporaryEntryPoints + count * sizeof(FixupPrecode) + sizeof(PTR_MethodDesc);
+#ifndef CROSSGEN_COMPILE
+            emitBackToBackJump((LPBYTE)precodeFixupJumpStub, (LPVOID)GetEEFuncEntryPoint(PrecodeFixupThunk));
+#endif // !CROSSGEN_COMPILE
+        }
+#endif // FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+
         TADDR entryPoint = temporaryEntryPoints;
         MethodDesc * pMD = pChunk->GetFirstMethodDesc();
         for (int i = 0; i < count; i++)
         {
             ((FixupPrecode *)entryPoint)->Init(pMD, pLoaderAllocator, pMD->GetMethodDescIndex(), (count - 1) - i);
 
+#ifdef FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+            _ASSERTE(
+                !preallocateJumpStubs ||
+                !pMD->IsLCGMethod() ||
+                ((FixupPrecode *)entryPoint)->GetDynamicMethodPrecodeFixupJumpStub() == precodeFixupJumpStub);
+#endif // FIXUP_PRECODE_PREALLOCATE_DYNAMIC_METHOD_JUMP_STUBS
+
             _ASSERTE((Precode *)entryPoint == GetPrecodeForTemporaryEntryPoint(temporaryEntryPoints, i));
             entryPoint += sizeof(FixupPrecode);
     
diff --git a/tests/src/CoreMangLib/cti/system/reflection/emit/DynMethodJumpStubTests/DynMethodJumpStubTests.cs b/tests/src/CoreMangLib/cti/system/reflection/emit/DynMethodJumpStubTests/DynMethodJumpStubTests.cs
index 1dc717d884..d50d3cedbf 100644
--- a/tests/src/CoreMangLib/cti/system/reflection/emit/DynMethodJumpStubTests/DynMethodJumpStubTests.cs
+++ b/tests/src/CoreMangLib/cti/system/reflection/emit/DynMethodJumpStubTests/DynMethodJumpStubTests.cs
@@ -28,10 +28,11 @@ public static class DynamicMethodJumpStubTests
         // framework libraries.
         ReserveMemoryAround(new Action(ExecutionContext.RestoreFlow).Method.MethodHandle);
 
-        for (int i = 0; i < 64; ++i)
+        var dynamicMethodDelegates = new Action[64];
+        for (int i = 0; i < dynamicMethodDelegates.Length; ++i)
         {
             DynamicMethod dynamicMethod = CreateDynamicMethod("DynMethod" + i);
-            Action dynamicMethodDelegate = (Action)dynamicMethod.CreateDelegate(typeof(Action));
+            dynamicMethodDelegates[i] = (Action)dynamicMethod.CreateDelegate(typeof(Action));
 
             // Before compiling the dynamic method, reserve memory around its current entry point, which should be its
             // precode. Then, when compiling the method, there would be a good chance that the code will be located far from
@@ -44,9 +45,30 @@ public static class DynamicMethodJumpStubTests
                     null,
                     dynamicMethod,
                     null));
+        }
 
-            dynamicMethodDelegate();
+        // Call each dynamic method concurrently from several threads to validate jump stub usage
+        int threadCount = 64;
+        var barrier = new Barrier(threadCount);
+        ThreadStart threadStart = () =>
+        {
+            var dynamicMethodDelegatesLocal = dynamicMethodDelegates;
+            for (int i = 0; i < dynamicMethodDelegatesLocal.Length; ++i)
+            {
+                var dynamicMethodDelegate = dynamicMethodDelegatesLocal[i];
+                barrier.SignalAndWait();
+                dynamicMethodDelegate();
+            }
+        };
+        var threads = new Thread[threadCount];
+        for (int i = 0; i < threads.Length; ++i)
+        {
+            threads[i] = new Thread(threadStart);
+            threads[i].IsBackground = true;
+            threads[i].Start();
         }
+        foreach (var t in threads)
+            t.Join();
 
         // This test does not release reserved pages because they may have been committed by other components on the system
     }