Do not reset size of TailCallArgBuffer (#118365)

* TailCallArgBuffer:State is 32-bit int. The code incorrectly used native int operation to set. It lead to TailCallArgBuffer::Size to be reset as well and needlessly re-allocating the tail call buffer.

* Switch to preemptive mode around malloc/free in TailCallArgBuffer allocation

Move the fast path of the TailCallArgBuffer allocation helper to C# to avoid perf overload of switching to preemptive mode. Also, aggressively inline it to allow zero initialization to be unrolled by the JIT.

* Skip early GC reporting de-action if there is nothing to report

Fixes #118166

Author: jkotas

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeHelpers.CoreCLR.cs

@@ -488,15 +488,41 @@ public static IntPtr AllocateTypeAssociatedMemory(Type type, int size)
         private static partial IntPtr AllocateTypeAssociatedMemory(QCallTypeHandle type, uint size);
 
         [MethodImpl(MethodImplOptions.InternalCall)]
-        private static extern IntPtr AllocTailCallArgBufferWorker(int size, IntPtr gcDesc);
+        private static extern unsafe TailCallArgBuffer* GetTailCallArgBuffer();
 
-        private static IntPtr AllocTailCallArgBuffer(int size, IntPtr gcDesc)
+        [LibraryImport(QCall, EntryPoint = "TailCallHelp_AllocTailCallArgBufferInternal")]
+        private static unsafe partial TailCallArgBuffer* AllocTailCallArgBufferInternal(int size);
+
+        private const int TAILCALLARGBUFFER_ACTIVE = 0;
+        // private const int TAILCALLARGBUFFER_INSTARG_ONLY = 1;
+        private const int TAILCALLARGBUFFER_INACTIVE = 2;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)] // To allow unrolling of Span.Clear
+        private static unsafe TailCallArgBuffer* AllocTailCallArgBuffer(int size, IntPtr gcDesc)
         {
-            IntPtr buffer = AllocTailCallArgBufferWorker(size, gcDesc);
-            if (buffer == IntPtr.Zero)
+            TailCallArgBuffer* buffer = GetTailCallArgBuffer();
+            if (buffer != null && buffer->Size >= size)
+            {
+                buffer->State = TAILCALLARGBUFFER_INACTIVE;
+            }
+            else
             {
-                throw new OutOfMemoryException();
+                buffer = AllocTailCallArgBufferWorker(size);
+
+                [MethodImpl(MethodImplOptions.NoInlining)]
+                static TailCallArgBuffer* AllocTailCallArgBufferWorker(int size) => AllocTailCallArgBufferInternal(size);
             }
+            Debug.Assert(buffer != null);
+            Debug.Assert(buffer->Size >= size);
+            Debug.Assert(buffer->State == TAILCALLARGBUFFER_INACTIVE);
+
+            buffer->GCDesc = gcDesc;
+
+            new Span<byte>(buffer + 1, size - sizeof(TailCallArgBuffer)).Clear();
+
+            // The buffer is now ready to be used.
+            buffer->State = TAILCALLARGBUFFER_ACTIVE;
+
             return buffer;
         }
 
@@ -506,7 +532,7 @@ private static IntPtr AllocTailCallArgBuffer(int size, IntPtr gcDesc)
         [StackTraceHidden]
         private static unsafe void DispatchTailCalls(
             IntPtr callersRetAddrSlot,
-            delegate*<IntPtr, ref byte, PortableTailCallFrame*, void> callTarget,
+            delegate*<TailCallArgBuffer*, ref byte, PortableTailCallFrame*, void> callTarget,
             ref byte retVal)
         {
             IntPtr callersRetAddr;
@@ -537,11 +563,8 @@ private static unsafe void DispatchTailCalls(
             {
                 tls->Frame = prevFrame;
 
-                // If the arg buffer is reporting inst argument, it is safe to abandon it now
-                if (tls->ArgBuffer != IntPtr.Zero && *(int*)tls->ArgBuffer == 1 /* TAILCALLARGBUFFER_INSTARG_ONLY */)
-                {
-                    *(int*)tls->ArgBuffer = 2 /* TAILCALLARGBUFFER_ABANDONED */;
-                }
+                // If the arg buffer is reporting inst argument (TAILCALLARGBUFFER_INSTARG_ONLY), it is safe to abandon it now.
+                tls->ArgBuffer->State = TAILCALLARGBUFFER_INACTIVE;
             }
         }
 
@@ -1217,13 +1240,22 @@ private static bool CanCastToWorker(TypeHandle srcTH, TypeHandle destTH, bool nu
     internal unsafe struct PortableTailCallFrame
     {
         public IntPtr TailCallAwareReturnAddress;
-        public delegate*<IntPtr, ref byte, PortableTailCallFrame*, void> NextCall;
+        public delegate*<TailCallArgBuffer*, ref byte, PortableTailCallFrame*, void> NextCall;
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    internal struct TailCallArgBuffer
+    {
+        public int State;
+        public int Size;
+        public IntPtr GCDesc;
+        // Args
     }
 
     [StructLayout(LayoutKind.Sequential)]
     internal unsafe struct TailCallTls
     {
         public PortableTailCallFrame* Frame;
-        public IntPtr ArgBuffer;
+        public TailCallArgBuffer* ArgBuffer;
     }
 }

src/coreclr/vm/corelib.h

@@ -709,7 +709,7 @@ DEFINE_METHOD(RUNTIME_HELPERS,      GET_RAW_DATA,            GetRawData,
 DEFINE_METHOD(RUNTIME_HELPERS,      GET_UNINITIALIZED_OBJECT, GetUninitializedObject, SM_Type_RetObj)
 DEFINE_METHOD(RUNTIME_HELPERS,      ENUM_EQUALS,            EnumEquals, NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS,      ENUM_COMPARE_TO,        EnumCompareTo, NoSig)
-DEFINE_METHOD(RUNTIME_HELPERS,      ALLOC_TAILCALL_ARG_BUFFER, AllocTailCallArgBuffer,  SM_Int_IntPtr_RetIntPtr)
+DEFINE_METHOD(RUNTIME_HELPERS,      ALLOC_TAILCALL_ARG_BUFFER, AllocTailCallArgBuffer,  NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS,      GET_TAILCALL_INFO,      GetTailCallInfo, NoSig)
 DEFINE_METHOD(RUNTIME_HELPERS,      DISPATCH_TAILCALLS,     DispatchTailCalls,          NoSig)
 #ifdef FEATURE_IJW

src/coreclr/vm/ecalllist.h

@@ -334,7 +334,7 @@ FCFuncStart(gRuntimeHelpers)
     FCFuncElement("TryGetHashCode", ObjectNative::TryGetHashCode)
     FCFuncElement("ContentEquals", ObjectNative::ContentEquals)
     FCFuncElement("TryEnsureSufficientExecutionStack", ReflectionInvocation::TryEnsureSufficientExecutionStack)
-    FCFuncElement("AllocTailCallArgBufferWorker", TailCallHelp::AllocTailCallArgBufferWorker)
+    FCFuncElement("GetTailCallArgBuffer", TailCallHelp::GetTailCallArgBuffer)
     FCFuncElement("GetTailCallInfo", TailCallHelp::GetTailCallInfo)
 FCFuncEnd()
 

src/coreclr/vm/gcenv.ee.cpp

@@ -235,7 +235,7 @@ static void ScanTailCallArgBufferRoots(Thread* pThread, promote_func* fn, ScanCo
     if (argBuffer == NULL || argBuffer->GCDesc == NULL)
         return;
 
-    if (argBuffer->State == TAILCALLARGBUFFER_ABANDONED)
+    if (argBuffer->State == TAILCALLARGBUFFER_INACTIVE)
         return;
 
     bool instArgOnly = argBuffer->State == TAILCALLARGBUFFER_INSTARG_ONLY;

src/coreclr/vm/metasig.h

@@ -330,7 +330,6 @@ DEFINE_METASIG_T(SM(RetMethodBase, _, C(METHOD_BASE)))
 DEFINE_METASIG(SM(RetVoid, _, v))
 DEFINE_METASIG(SM(Str_IntPtr_Int_RetVoid, s I i, v))
 DEFINE_METASIG(SM(Int_RetIntPtr, i, I))
-DEFINE_METASIG(SM(Int_IntPtr_RetIntPtr, i I, I))
 
 DEFINE_METASIG_T(SM(DateTime_RetDbl, g(DATE_TIME), d))
 DEFINE_METASIG(SM(Dbl_RetLong, d, l))

src/coreclr/vm/qcallentrypoints.cpp

@@ -548,6 +548,7 @@ static const Entry s_QCall[] =
     DllImportEntry(ThrowInvalidCastException)
     DllImportEntry(IsInstanceOf_NoCacheLookup)
     DllImportEntry(VersionResilientHashCode_TypeHashCode)
+    DllImportEntry(TailCallHelp_AllocTailCallArgBufferInternal)
 };
 
 const void* QCallResolveDllImport(const char* name)

src/coreclr/vm/tailcallhelp.cpp

@@ -11,14 +11,28 @@
 #include "threads.h"
 
 
-FCIMPL2(void*, TailCallHelp::AllocTailCallArgBufferWorker, INT32 size, void* gcDesc)
+FCIMPL0(void*, TailCallHelp::GetTailCallArgBuffer)
 {
     FCALL_CONTRACT;
-    _ASSERTE(size >= 0);
-    return GetThread()->GetTailCallTls()->AllocArgBuffer(size, gcDesc);
+    return GetThread()->GetTailCallTls()->GetArgBuffer();
 }
 FCIMPLEND
 
+extern "C" void* QCALLTYPE TailCallHelp_AllocTailCallArgBufferInternal(int size)
+{
+    QCALL_CONTRACT;
+
+    void* retVal = NULL;
+
+    BEGIN_QCALL;
+
+    retVal = GetThread()->GetTailCallTls()->AllocArgBuffer(size);
+
+    END_QCALL;
+
+    return retVal;
+}
+
 FCIMPL2(void*, TailCallHelp::GetTailCallInfo, void** retAddrSlot, void** retAddr)
 {
     FCALL_CONTRACT;
@@ -30,7 +44,6 @@ FCIMPL2(void*, TailCallHelp::GetTailCallInfo, void** retAddrSlot, void** retAddr
 }
 FCIMPLEND
 
-
 struct ArgBufferValue
 {
     TypeHandle TyHnd;
@@ -486,14 +499,17 @@ MethodDesc* TailCallHelp::CreateCallTargetStub(const TailCallInfo& info)
         EmitLoadTyHnd(pCode, arg.TyHnd);
     }
 
-    // All arguments are loaded on the stack, it is safe to disable the GC reporting of ArgBuffer now.
-    // This is optimization to avoid extending argument lifetime unnecessarily.
-    // We still need to report the inst argument of shared generic code to prevent it from being unloaded. The inst
-    // argument is just a regular IntPtr on the stack. It is safe to stop reporting it only after the target method
-    // takes over.
-    pCode->EmitLDARG(ARG_ARG_BUFFER);
-    pCode->EmitLDC(info.ArgBufLayout.HasInstArg ? TAILCALLARGBUFFER_INSTARG_ONLY : TAILCALLARGBUFFER_ABANDONED);
-    pCode->EmitSTIND_I();
+    if (info.HasGCDescriptor)
+    {
+        // All arguments are loaded on the stack, it is safe to disable the GC reporting of ArgBuffer now.
+        // This is optimization to avoid extending argument lifetime unnecessarily.
+        // We still need to report the inst argument of shared generic code to prevent it from being unloaded. The inst
+        // argument is just a regular IntPtr on the stack. It is safe to stop reporting it only after the target method
+        // takes over.
+        pCode->EmitLDARG(ARG_ARG_BUFFER);
+        pCode->EmitLDC(info.ArgBufLayout.HasInstArg ? TAILCALLARGBUFFER_INSTARG_ONLY : TAILCALLARGBUFFER_INACTIVE);
+        pCode->EmitSTIND_I4();
+    }
 
     int numRetVals = info.CallSiteSig->IsReturnTypeVoid() ? 0 : 1;
     // Normally there will not be any target and we just emit a normal

src/coreclr/vm/tailcallhelp.h

@@ -13,7 +13,7 @@ struct ArgBufferLayout;
 class TailCallHelp
 {
 public:
-    static FCDECL2(void*, AllocTailCallArgBufferWorker, INT32, void*);
+    static FCDECL0(void*, GetTailCallArgBuffer);
     static FCDECL2(void*, GetTailCallInfo, void**, void**);
 
     static void CreateTailCallHelperStubs(
@@ -47,4 +47,6 @@ class TailCallHelp
     static void* AllocateBlob(LoaderAllocator* alloc, const void* blob, size_t blobLen);
 };
 
+extern "C" void* QCALLTYPE TailCallHelp_AllocTailCallArgBufferInternal(int size);
+
 #endif

src/coreclr/vm/threads.cpp

@@ -81,39 +81,37 @@ Thread* STDCALL GetThreadHelper()
     return GetThreadNULLOk();
 }
 
-TailCallArgBuffer* TailCallTls::AllocArgBuffer(int size, void* gcDesc)
+TailCallArgBuffer* TailCallTls::AllocArgBuffer(int size)
 {
-    CONTRACTL
-    {
-        NOTHROW;
-        GC_NOTRIGGER;
-    }
-    CONTRACTL_END
+    STANDARD_VM_CONTRACT;
 
     _ASSERTE(size >= (int)offsetof(TailCallArgBuffer, Args));
 
-    if (m_argBuffer != NULL && m_argBuffer->Size < size)
-    {
-        FreeArgBuffer();
-    }
+    // Round the size up
+    size = max(size, (int)(8 * sizeof(void*)));
+
+    TailCallArgBuffer* pBuffer = (TailCallArgBuffer*)new BYTE[size];
+    TailCallArgBuffer* pOldBuffer;
+
+    pBuffer->Size = size;
+    pBuffer->State = TAILCALLARGBUFFER_INACTIVE;
 
-    if (m_argBuffer == NULL)
     {
-        m_argBuffer = (TailCallArgBuffer*)new (nothrow) BYTE[size];
-        if (m_argBuffer == NULL)
-            return NULL;
-        m_argBuffer->Size = size;
-    }
+        // We need to ensure that the GC does not run while we are switching the arg buffer.
+        GCX_COOP();
 
-    m_argBuffer->State = TAILCALLARGBUFFER_ACTIVE;
+        pOldBuffer = m_argBuffer;
+        m_argBuffer = pBuffer;
+    }
 
-    m_argBuffer->GCDesc = gcDesc;
-    if (gcDesc != NULL)
+    if (pOldBuffer != NULL)
     {
-        memset(m_argBuffer->Args, 0, size - offsetof(TailCallArgBuffer, Args));
+        _ASSERTE(pOldBuffer->Size < size);
+
+        delete[] (BYTE*)pOldBuffer;
     }
 
-    return m_argBuffer;
+    return pBuffer;
 }
 
 #if defined (_DEBUG_IMPL)

src/coreclr/vm/threads.h

@@ -162,7 +162,7 @@ class Module;
 // TailCallArgBuffer states
 #define TAILCALLARGBUFFER_ACTIVE       0
 #define TAILCALLARGBUFFER_INSTARG_ONLY 1
-#define TAILCALLARGBUFFER_ABANDONED    2
+#define TAILCALLARGBUFFER_INACTIVE     2
 
 struct TailCallArgBuffer
 {
@@ -414,7 +414,7 @@ class TailCallTls
 
 public:
     TailCallTls();
-    TailCallArgBuffer* AllocArgBuffer(int size, void* gcDesc);
+    TailCallArgBuffer* AllocArgBuffer(int size);
     void FreeArgBuffer() { delete[] (BYTE*)m_argBuffer; m_argBuffer = NULL; }
     TailCallArgBuffer* GetArgBuffer()
     {