[release/10.0] Fix lookup for current Thread in async signal handler

src/coreclr/nativeaot/Runtime/CMakeLists.txt

@@ -56,6 +56,12 @@ set(COMMON_RUNTIME_SOURCES
     ${CLR_SRC_NATIVE_DIR}/minipal/xoshiro128pp.c
 )
 
+if (CLR_CMAKE_TARGET_UNIX AND NOT CLR_CMAKE_TARGET_ARCH_WASM)
+    list(APPEND COMMON_RUNTIME_SOURCES
+        ${RUNTIME_DIR}/asyncsafethreadmap.cpp
+    )
+endif()
+
 set(SERVER_GC_SOURCES
     ${GC_DIR}/gceesvr.cpp
     ${GC_DIR}/gcsvr.cpp

src/coreclr/nativeaot/Runtime/thread.cpp

@@ -625,7 +625,7 @@ void Thread::Hijack()
     PalHijack(this);
 }
 
-void Thread::HijackCallback(NATIVE_CONTEXT* pThreadContext, Thread* pThreadToHijack)
+void Thread::HijackCallback(NATIVE_CONTEXT* pThreadContext, Thread* pThreadToHijack, bool doInlineSuspend)
 {
     // If we are no longer trying to suspend, no need to do anything.
     // This is just an optimization. It is ok to race with the setting the trap flag here.
@@ -694,9 +694,8 @@ void Thread::HijackCallback(NATIVE_CONTEXT* pThreadContext, Thread* pThreadToHij
         ASSERT(codeManager->IsUnwindable(pvAddress) || runtime->IsConservativeStackReportingEnabled());
 #endif
 
-        // if we are not given a thread to hijack
         // perform in-line wait on the current thread
-        if (pThreadToHijack == NULL)
+        if (doInlineSuspend)
         {
             ASSERT(pThread->m_interruptedContext == NULL);
             pThread->InlineSuspend(pThreadContext);

src/coreclr/nativeaot/Runtime/thread.h

@@ -280,7 +280,7 @@ class Thread : private RuntimeThreadLocals
     void*               GetHijackedReturnAddress();
     static bool         IsHijackTarget(void * address);
 
-    static void HijackCallback(NATIVE_CONTEXT* pThreadContext, Thread* pThreadToHijack);
+    static void HijackCallback(NATIVE_CONTEXT* pThreadContext, Thread* pThreadToHijack, bool doInlineSuspend);
 #else // FEATURE_HIJACK
     void                Unhijack() { }
     bool                IsHijacked() { return false; }

src/coreclr/nativeaot/Runtime/threadstore.cpp

@@ -22,6 +22,8 @@
 #include "TargetPtrs.h"
 #include "yieldprocessornormalized.h"
 #include <minipal/time.h>
+#include <minipal/thread.h>
+#include "asyncsafethreadmap.h"
 
 #include "slist.inl"
 
@@ -143,6 +145,14 @@ void ThreadStore::AttachCurrentThread(bool fAcquireThreadStoreLock)
     pAttachingThread->m_ThreadStateFlags = Thread::TSF_Attached;
 
     pTS->m_ThreadList.PushHead(pAttachingThread);
+
+#if defined(TARGET_UNIX) && !defined(TARGET_WASM)
+    if (!InsertThreadIntoAsyncSafeMap(pAttachingThread->m_threadId, pAttachingThread))
+    {
+        PalPrintFatalError("\nFailed to insert thread into async-safe map due to out of memory.\n");
+        RhFailFast();
+    }
+#endif // TARGET_UNIX && !TARGET_WASM
 }
 
 // static
@@ -188,6 +198,9 @@ void ThreadStore::DetachCurrentThread()
         pTS->m_ThreadList.RemoveFirst(pDetachingThread);
         // tidy up GC related stuff (release allocation context, etc..)
         pDetachingThread->Detach();
+#if defined(TARGET_UNIX) && !defined(TARGET_WASM)
+        RemoveThreadFromAsyncSafeMap(pDetachingThread->m_threadId, pDetachingThread);
+#endif
     }
 
     // post-mortem clean up.
@@ -426,6 +439,13 @@ EXTERN_C RuntimeThreadLocals* RhpGetThread()
     return &tls_CurrentThread;
 }
 
+#if defined(TARGET_UNIX) && !defined(TARGET_WASM)
+Thread * ThreadStore::GetCurrentThreadIfAvailableAsyncSafe()
+{
+    return (Thread*)FindThreadInAsyncSafeMap(minipal_get_current_thread_id_no_cache());
+}
+#endif // TARGET_UNIX && !TARGET_WASM
+
 #endif // !DACCESS_COMPILE
 
 #ifdef _WIN32

src/coreclr/nativeaot/Runtime/threadstore.h

@@ -48,6 +48,9 @@ class ThreadStore
     static Thread *         RawGetCurrentThread();
     static Thread *         GetCurrentThread();
     static Thread *         GetCurrentThreadIfAvailable();
+#if defined(TARGET_UNIX) && !defined(TARGET_WASM)
+    static Thread *         GetCurrentThreadIfAvailableAsyncSafe();
+#endif
     static PTR_Thread       GetSuspendingThread();
     static void             AttachCurrentThread();
     static void             AttachCurrentThread(bool fAcquireThreadStoreLock);

src/coreclr/nativeaot/Runtime/unix/PalUnix.cpp

@@ -1016,24 +1016,24 @@ static struct sigaction g_previousActivationHandler;
 
 static void ActivationHandler(int code, siginfo_t* siginfo, void* context)
 {
-    // Only accept activations from the current process
-    if (siginfo->si_pid == getpid()
+    Thread* pThread = ThreadStore::GetCurrentThreadIfAvailableAsyncSafe();
+    if (pThread)
+    {
+        // Only accept activations from the current process
+        if (siginfo->si_pid == getpid()
 #ifdef HOST_APPLE
-        // On Apple platforms si_pid is sometimes 0. It was confirmed by Apple to be expected, as the si_pid is tracked at the process level. So when multiple
-        // signals are in flight in the same process at the same time, it may be overwritten / zeroed.
-        || siginfo->si_pid == 0
+            // On Apple platforms si_pid is sometimes 0. It was confirmed by Apple to be expected, as the si_pid is tracked at the process level. So when multiple
+            // signals are in flight in the same process at the same time, it may be overwritten / zeroed.
+            || siginfo->si_pid == 0
 #endif
-        )
-    {
-        // Make sure that errno is not modified
-        int savedErrNo = errno;
-        Thread::HijackCallback((NATIVE_CONTEXT*)context, NULL);
-        errno = savedErrNo;
-    }
+            )
+        {
+            // Make sure that errno is not modified
+            int savedErrNo = errno;
+            Thread::HijackCallback((NATIVE_CONTEXT*)context, pThread, true /* doInlineSuspend */);
+            errno = savedErrNo;
+        }
 
-    Thread* pThread = ThreadStore::GetCurrentThreadIfAvailable();
-    if (pThread)
-    {
         pThread->SetActivationPending(false);
     }
 

src/coreclr/nativeaot/Runtime/windows/PalMinWin.cpp

@@ -673,9 +673,9 @@ static void* g_returnAddressHijackTarget = NULL;
 static void NTAPI ActivationHandler(ULONG_PTR parameter)
 {
     CLONE_APC_CALLBACK_DATA* data = (CLONE_APC_CALLBACK_DATA*)parameter;
-    Thread::HijackCallback((NATIVE_CONTEXT*)data->ContextRecord, NULL);
-
     Thread* pThread = (Thread*)data->Parameter;
+    Thread::HijackCallback((NATIVE_CONTEXT*)data->ContextRecord, pThread, true /* doInlineSuspend */);
+
     pThread->SetActivationPending(false);
 }
 
@@ -833,7 +833,7 @@ void PalHijack(Thread* pThreadToHijack)
 
         if (isSafeToRedirect)
         {
-            Thread::HijackCallback((NATIVE_CONTEXT*)&win32ctx, pThreadToHijack);
+            Thread::HijackCallback((NATIVE_CONTEXT*)&win32ctx, pThreadToHijack, false /* doInlineSuspend */);
         }
     }
 

src/coreclr/pal/src/exception/signal.cpp

@@ -936,22 +936,20 @@ static void inject_activation_handler(int code, siginfo_t *siginfo, void *contex
             CONTEXTToNativeContext(&winContext, ucontext);
         }
     }
+
+    // Call the original handler when it is not ignored or default (terminate).
+    if (g_previous_activation.sa_flags & SA_SIGINFO)
+    {
+        _ASSERTE(g_previous_activation.sa_sigaction != NULL);
+        g_previous_activation.sa_sigaction(code, siginfo, context);
+    }
     else
     {
-        // Call the original handler when it is not ignored or default (terminate).
-        if (g_previous_activation.sa_flags & SA_SIGINFO)
-        {
-            _ASSERTE(g_previous_activation.sa_sigaction != NULL);
-            g_previous_activation.sa_sigaction(code, siginfo, context);
-        }
-        else
+        if (g_previous_activation.sa_handler != SIG_IGN &&
+            g_previous_activation.sa_handler != SIG_DFL)
         {
-            if (g_previous_activation.sa_handler != SIG_IGN &&
-                g_previous_activation.sa_handler != SIG_DFL)
-            {
-                _ASSERTE(g_previous_activation.sa_handler != NULL);
-                g_previous_activation.sa_handler(code);
-            }
+            _ASSERTE(g_previous_activation.sa_handler != NULL);
+            g_previous_activation.sa_handler(code);
         }
     }
 }

src/coreclr/runtime/asyncsafethreadmap.cpp

@@ -0,0 +1,126 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#include "common.h"
+
+#include "asyncsafethreadmap.h"
+
+// Async safe lock free thread map for use in signal handlers
+
+struct ThreadEntry
+{
+    size_t osThread;
+    void* pThread;
+};
+
+#define MAX_THREADS_IN_SEGMENT 256
+
+struct ThreadSegment
+{
+    ThreadEntry entries[MAX_THREADS_IN_SEGMENT];
+    ThreadSegment* pNext;
+};
+
+static ThreadSegment *s_pAsyncSafeThreadMapHead = NULL;
+
+bool InsertThreadIntoAsyncSafeMap(size_t osThread, void* pThread)
+{
+    size_t startIndex = osThread % MAX_THREADS_IN_SEGMENT;
+
+    ThreadSegment* pSegment = s_pAsyncSafeThreadMapHead;
+    ThreadSegment** ppSegment = &s_pAsyncSafeThreadMapHead;
+    while (true)
+    {
+        if (pSegment == NULL)
+        {
+            // Need to add a new segment
+            ThreadSegment* pNewSegment = new (nothrow) ThreadSegment();
+            if (pNewSegment == NULL)
+            {
+                // Memory allocation failed
+                return false;
+            }
+
+            memset(pNewSegment, 0, sizeof(ThreadSegment));
+            ThreadSegment* pExpected = NULL;           
+            if (!__atomic_compare_exchange_n(
+                ppSegment,
+                &pExpected,
+                pNewSegment,
+                false /* weak */,
+                __ATOMIC_RELEASE  /* success_memorder */,
+                __ATOMIC_RELAXED /* failure_memorder */))
+            {
+                // Another thread added the segment first
+                delete pNewSegment;
+                pNewSegment = pExpected;
+            }
+
+            pSegment = pNewSegment;
+        }
+        for (size_t i = 0; i < MAX_THREADS_IN_SEGMENT; i++)
+        {
+            size_t index = (startIndex + i) % MAX_THREADS_IN_SEGMENT;
+
+            size_t expected = 0;
+            if (__atomic_compare_exchange_n(
+                    &pSegment->entries[index].osThread,
+                    &expected,
+                    osThread,
+                    false /* weak */,
+                    __ATOMIC_RELEASE  /* success_memorder */,
+                    __ATOMIC_RELAXED /* failure_memorder */))
+            {
+                // Successfully inserted
+                // Use atomic store with release to ensure proper ordering
+                __atomic_store_n(&pSegment->entries[index].pThread, pThread, __ATOMIC_RELEASE);
+                return true;
+            }
+        }
+
+        ppSegment = &pSegment->pNext;
+        pSegment = __atomic_load_n(&pSegment->pNext, __ATOMIC_ACQUIRE);
+    }
+}
+
+void RemoveThreadFromAsyncSafeMap(size_t osThread, void* pThread)
+{
+    size_t startIndex = osThread % MAX_THREADS_IN_SEGMENT;
+
+    ThreadSegment* pSegment = s_pAsyncSafeThreadMapHead;
+    while (pSegment)
+    {
+        for (size_t i = 0; i < MAX_THREADS_IN_SEGMENT; i++)
+        {
+            size_t index = (startIndex + i) % MAX_THREADS_IN_SEGMENT;
+            if (pSegment->entries[index].pThread == pThread)
+            {
+                // Found the entry, remove it
+                pSegment->entries[index].pThread = NULL;
+                __atomic_exchange_n(&pSegment->entries[index].osThread, (size_t)0, __ATOMIC_RELEASE);
+                return;
+            }
+        }
+        pSegment = __atomic_load_n(&pSegment->pNext, __ATOMIC_ACQUIRE);
+    }
+}
+
+void *FindThreadInAsyncSafeMap(size_t osThread)
+{
+    size_t startIndex = osThread % MAX_THREADS_IN_SEGMENT;
+    ThreadSegment* pSegment = s_pAsyncSafeThreadMapHead;
+    while (pSegment)
+    {
+        for (size_t i = 0; i < MAX_THREADS_IN_SEGMENT; i++)
+        {
+            size_t index = (startIndex + i) % MAX_THREADS_IN_SEGMENT;
+            // Use acquire to synchronize with release in InsertThreadIntoAsyncSafeMap
+            if (__atomic_load_n(&pSegment->entries[index].osThread, __ATOMIC_ACQUIRE) == osThread)
+            {
+                return pSegment->entries[index].pThread;
+            }
+        }
+        pSegment = __atomic_load_n(&pSegment->pNext, __ATOMIC_ACQUIRE);
+    }
+    return NULL;
+}

src/coreclr/runtime/asyncsafethreadmap.h

@@ -0,0 +1,27 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#ifndef __ASYNCSAFETHREADMAP_H__
+#define __ASYNCSAFETHREADMAP_H__
+
+#if defined(TARGET_UNIX) && !defined(TARGET_WASM)
+
+//  Insert a thread into the async-safe map.
+//  * osThread - The OS thread ID to insert.
+//  * pThread - A pointer to the thread object to associate with the OS thread ID.
+//  * return true if the insertion was successful, false otherwise (OOM).
+bool InsertThreadIntoAsyncSafeMap(size_t osThread, void* pThread);
+
+// Remove a thread from the async-safe map.
+// * osThread - The OS thread ID to remove.
+// * pThread - A pointer to the thread object associated with the OS thread ID.
+void RemoveThreadFromAsyncSafeMap(size_t osThread, void* pThread);
+
+// Find a thread in the async-safe map.
+// * osThread - The OS thread ID to search for.
+// * return - A pointer to the thread object associated with the OS thread ID, or NULL if not found.
+void* FindThreadInAsyncSafeMap(size_t osThread);
+
+#endif // TARGET_UNIX && !TARGET_WASM
+
+#endif // __ASYNCSAFETHREADMAP_H__

src/coreclr/vm/CMakeLists.txt

@@ -384,6 +384,12 @@ set(VM_SOURCES_WKS
     ${VM_SOURCES_GDBJIT}
 )
 
+if (CLR_CMAKE_TARGET_UNIX AND NOT CLR_CMAKE_TARGET_ARCH_WASM)
+    list(APPEND VM_SOURCES_WKS
+        ${RUNTIME_DIR}/asyncsafethreadmap.cpp
+    )
+endif()
+
 # coreclr needs to compile codeman.cpp differently depending on flavor (i.e. dll vs. static lib))
 list(REMOVE_ITEM VM_SOURCES_WKS codeman.cpp)
 
@@ -482,6 +488,12 @@ set(VM_HEADERS_WKS
     ${VM_HEADERS_GDBJIT}
 )
 
+if (CLR_CMAKE_TARGET_UNIX AND NOT CLR_CMAKE_TARGET_ARCH_WASM)
+    list(APPEND VM_HEADERS_WKS
+        ${RUNTIME_DIR}/asyncsafethreadmap.h
+    )
+endif()
+
 set(GC_SOURCES_WKS
     ${GC_SOURCES_DAC_AND_WKS_COMMON}
     ../gc/gceventstatus.cpp