For transition profiler callbacks, always load the thread (#105104)

Author: jkoritzinsky

src/coreclr/vm/dllimport.cpp

@@ -2272,15 +2272,7 @@ DWORD NDirectStubLinker::EmitProfilerBeginTransitionCallback(ILCodeStream* pcsEm
         EmitLoadStubContext(pcsEmit, dwStubFlags);
     }
 
-    if (SF_IsForwardStub(dwStubFlags))
-    {
-        pcsEmit->EmitLDLOC(GetThreadLocalNum());
-    }
-    else
-    {
-        // we use a null pThread to indicate reverse interop
-        pcsEmit->EmitLoadNullPtr();
-    }
+    pcsEmit->EmitLDLOC(GetThreadLocalNum());
 
     // In the unmanaged delegate case, we need the "this" object to retrieve the MD
     // in StubHelpers::ProfilerEnterCallback().

src/coreclr/vm/stubhelpers.cpp

@@ -547,6 +547,7 @@ FCIMPL3(SIZE_T, StubHelpers::ProfilerBeginTransitionCallback, SIZE_T pSecretPara
     }
 
     {
+        _ASSERTE(pThread != nullptr);
         GCX_PREEMP_THREAD_EXISTS(pThread);
 
         ProfilerManagedToUnmanagedTransitionMD(pRealMD, COR_PRF_TRANSITION_CALL);
@@ -577,6 +578,7 @@ FCIMPL2(void, StubHelpers::ProfilerEndTransitionCallback, MethodDesc* pRealMD, T
     // and the transition requires us to set up a HMF.
     HELPER_METHOD_FRAME_BEGIN_0();
     {
+        _ASSERTE(pThread != nullptr);
         GCX_PREEMP_THREAD_EXISTS(pThread);
 
         ProfilerUnmanagedToManagedTransitionMD(pRealMD, COR_PRF_TRANSITION_RETURN);

src/tests/profiler/native/profiler.def

@@ -5,4 +5,5 @@ EXPORTS
             DllCanUnloadNow         PRIVATE
             PassCallbackToProfiler  PRIVATE
             DoPInvoke               PRIVATE
+            DoPInvokeWithCallbackOnOtherThread PRIVATE
 

src/tests/profiler/native/transitions/transitions.cpp

@@ -69,6 +69,14 @@ extern "C" EXPORT void STDMETHODCALLTYPE DoPInvoke(int(*callback)(int), int i)
     printf("PInvoke received i=%d\n", callback(i));
 }
 
+extern "C" EXPORT void STDMETHODCALLTYPE DoPInvokeWithCallbackOnOtherThread(int(*callback)(int), int i)
+{
+    int j = 0;
+    std::thread t([&j, callback, i] { j = callback(i); });
+    t.join();
+    printf("PInvoke with callback on other thread received i=%d\n", j);
+}
+
 
 HRESULT Transitions::UnmanagedToManagedTransition(FunctionID functionID, COR_PRF_TRANSITION_REASON reason)
 {
@@ -124,4 +132,4 @@ bool Transitions::FunctionIsTargetFunction(FunctionID functionID, TransitionInst
     }
 
     return true;
-}
+}

src/tests/profiler/transitions/transitions.cs

@@ -32,7 +32,7 @@ private static int DoDelegateReversePInvokeNonBlittable(bool b)
         public static int BlittablePInvokeToBlittableInteropDelegate()
         {
             InteropDelegate del = DoDelegateReversePInvoke;
-            
+
             DoPInvoke((delegate* unmanaged<int,int>)Marshal.GetFunctionPointerForDelegate(del), 13);
             GC.KeepAlive(del);
 
@@ -42,13 +42,23 @@ public static int BlittablePInvokeToBlittableInteropDelegate()
         public static int NonBlittablePInvokeToNonBlittableInteropDelegate()
         {
             InteropDelegateNonBlittable del = DoDelegateReversePInvokeNonBlittable;
-            
+
             DoPInvokeNonBlitable((delegate* unmanaged<int,int>)Marshal.GetFunctionPointerForDelegate(del), true);
             GC.KeepAlive(del);
 
             return 100;
         }
 
+        public static int BlittablePInvokeToBlittableInteropDelegateOnOtherThread()
+        {
+            InteropDelegate del = DoDelegateReversePInvoke;
+
+            DoPInvokeWithCallbackOnOtherThread((delegate* unmanaged<int,int>)Marshal.GetFunctionPointerForDelegate(del), 13);
+            GC.KeepAlive(del);
+
+            return 100;
+        }
+
         [UnmanagedCallersOnly]
         private static int DoReversePInvoke(int i)
         {
@@ -61,6 +71,9 @@ private static int DoReversePInvoke(int i)
         [DllImport("Profiler", EntryPoint = nameof(DoPInvoke))]
         public static extern void DoPInvokeNonBlitable(delegate* unmanaged<int,int> callback, bool i);
 
+        [DllImport("Profiler")]
+        public static extern void DoPInvokeWithCallbackOnOtherThread(delegate* unmanaged<int,int> callback, int i);
+
         public static int BlittablePInvokeToUnmanagedCallersOnly()
         {
             DoPInvoke(&DoReversePInvoke, 13);
@@ -75,6 +88,13 @@ public static int NonBlittablePInvokeToUnmanagedCallersOnly()
             return 100;
         }
 
+        public static int BlittablePInvokeToUnmanagedCallersOnlyOnOtherThread()
+        {
+            DoPInvokeWithCallbackOnOtherThread(&DoReversePInvoke, 13);
+
+            return 100;
+        }
+
         public static int Main(string[] args)
         {
             if (args.Length > 1 && args[0].Equals("RunTest", StringComparison.OrdinalIgnoreCase))
@@ -89,6 +109,10 @@ public static int Main(string[] args)
                         return NonBlittablePInvokeToUnmanagedCallersOnly();
                     case nameof(NonBlittablePInvokeToNonBlittableInteropDelegate):
                         return NonBlittablePInvokeToNonBlittableInteropDelegate();
+                    case nameof(BlittablePInvokeToUnmanagedCallersOnlyOnOtherThread):
+                        return BlittablePInvokeToUnmanagedCallersOnlyOnOtherThread();
+                    case nameof(BlittablePInvokeToBlittableInteropDelegateOnOtherThread):
+                        return BlittablePInvokeToBlittableInteropDelegateOnOtherThread();
                 }
             }
 
@@ -104,12 +128,22 @@ public static int Main(string[] args)
 
             if (!RunProfilerTest(nameof(NonBlittablePInvokeToUnmanagedCallersOnly), nameof(DoPInvokeNonBlitable), nameof(DoReversePInvoke)))
             {
-                return 101;
+                return 103;
             }
 
             if (!RunProfilerTest(nameof(NonBlittablePInvokeToNonBlittableInteropDelegate), nameof(DoPInvokeNonBlitable), "Invoke"))
             {
-                return 102;
+                return 104;
+            }
+
+            if (!RunProfilerTest(nameof(BlittablePInvokeToUnmanagedCallersOnlyOnOtherThread), nameof(DoPInvokeWithCallbackOnOtherThread), nameof(DoReversePInvoke)))
+            {
+                return 105;
+            }
+
+            if (!RunProfilerTest(nameof(BlittablePInvokeToBlittableInteropDelegateOnOtherThread), nameof(DoPInvokeWithCallbackOnOtherThread), "Invoke"))
+            {
+                return 106;
             }
 
             return 100;