Fix lookup for current Thread in async signal handler

src/coreclr/nativeaot/Runtime/CMakeLists.txt

@@ -55,6 +55,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/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);
+
+#ifdef TARGET_UNIX
+    if (!InsertThreadIntoAsyncSafeMap(pAttachingThread->m_threadId, pAttachingThread))
+    {
+        ASSERT_UNCONDITIONALLY("Failed to insert thread into async-safe map due to OOM.");
+        RhFailFast();
+    }
+#endif
 }
 
 // static
@@ -188,6 +198,9 @@ void ThreadStore::DetachCurrentThread()
         pTS->m_ThreadList.RemoveFirst(pDetachingThread);
         // tidy up GC related stuff (release allocation context, etc..)
         pDetachingThread->Detach();
+#ifdef TARGET_UNIX
+        RemoveThreadFromAsyncSafeMap(pDetachingThread->m_threadId, pDetachingThread);
+#endif
     }
 
     // post-mortem clean up.
@@ -352,6 +365,13 @@ EXTERN_C RuntimeThreadLocals* RhpGetThread()
     return &tls_CurrentThread;
 }
 
+#ifdef TARGET_UNIX
+Thread * ThreadStore::GetCurrentThreadIfAvailableAsyncSafe()
+{
+    return (Thread*)FindThreadInAsyncSafeMap(minipal_get_current_thread_id_no_cache());
+}
+#endif // TARGET_UNIX
+
 #endif // !DACCESS_COMPILE
 
 #ifdef _WIN32

src/coreclr/nativeaot/Runtime/threadstore.h

@@ -47,6 +47,7 @@ class ThreadStore
     static Thread *         RawGetCurrentThread();
     static Thread *         GetCurrentThread();
     static Thread *         GetCurrentThreadIfAvailable();
+    static Thread *         GetCurrentThreadIfAvailableAsyncSafe();
     static PTR_Thread       GetSuspendingThread();
     static void             AttachCurrentThread();
     static void             AttachCurrentThread(bool fAcquireThreadStoreLock);

src/coreclr/nativeaot/Runtime/threadstore.inl

@@ -14,6 +14,7 @@ inline Thread * ThreadStore::RawGetCurrentThread()
     return (Thread *) &tls_CurrentThread;
 }
 
+#if defined(TARGET_UNIX) && !defined(DACCESS_COMPILE)
 // static
 inline Thread * ThreadStore::GetCurrentThread()
 {
@@ -24,6 +25,7 @@ inline Thread * ThreadStore::GetCurrentThread()
     ASSERT(pCurThread->IsInitialized());
     return pCurThread;
 }
+#endif // TARGET_UNIX && !DACCESS_COMPILE
 
 // static
 inline Thread * ThreadStore::GetCurrentThreadIfAvailable()

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

@@ -1032,7 +1032,7 @@ static void ActivationHandler(int code, siginfo_t* siginfo, void* context)
         errno = savedErrNo;
     }
 
-    Thread* pThread = ThreadStore::GetCurrentThreadIfAvailable();
+    Thread* pThread = ThreadStore::GetCurrentThreadIfAvailableAsyncSafe();
     if (pThread)
     {
         pThread->SetActivationPending(false);

src/coreclr/runtime/asyncsafethreadmap.cpp

@@ -0,0 +1,125 @@
+// 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 = *ppSegment;
+            }
+
+            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 insert_thread_to_async_safe_map
+            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

@@ -290,6 +290,7 @@ set(VM_SOURCES_WKS
     assemblyspec.cpp
     baseassemblyspec.cpp
     bundle.cpp
+    ${RUNTIME_DIR}/asyncsafethreadmap.cpp
     ${RUNTIME_DIR}/CachedInterfaceDispatch.cpp
     CachedInterfaceDispatch_Coreclr.cpp
     cachelinealloc.cpp

src/coreclr/vm/codeman.cpp

@@ -858,7 +858,7 @@ IJitManager::IJitManager()
 // been stopped when we suspend the EE so they won't be touching an element that is about to be deleted.
 // However for pre-emptive mode threads, they could be stalled right on top of the element we want
 // to delete, so we need to apply the reader lock to them and wait for them to drain.
-ExecutionManager::ScanFlag ExecutionManager::GetScanFlags()
+ExecutionManager::ScanFlag ExecutionManager::GetScanFlags(Thread *pThread)
 {
     CONTRACTL {
         NOTHROW;
@@ -869,7 +869,10 @@ ExecutionManager::ScanFlag ExecutionManager::GetScanFlags()
 #if !defined(DACCESS_COMPILE)
 
 
-    Thread *pThread = GetThreadNULLOk();
+    if (!pThread)
+    {
+        pThread = GetThreadNULLOk();
+    }
 
     if (!pThread)
         return ScanNoReaderLock;
@@ -5229,6 +5232,24 @@ BOOL ExecutionManager::IsManagedCode(PCODE currentPC)
     return IsManagedCodeWorker(currentPC, &lockState);
 }
 
+//**************************************************************************
+BOOL ExecutionManager::IsManagedCodeNoLock(PCODE currentPC)
+{
+    CONTRACTL {
+        NOTHROW;
+        GC_NOTRIGGER;
+    } CONTRACTL_END;
+
+    if (currentPC == (PCODE)NULL)
+        return FALSE;
+
+    _ASSERTE(GetScanFlags() != ScanReaderLock);
+
+    // Since ScanReaderLock is not set, then we must assume that the ReaderLock is effectively taken.
+    RangeSectionLockState lockState = RangeSectionLockState::ReaderLocked;
+    return IsManagedCodeWorker(currentPC, &lockState);
+}
+
 //**************************************************************************
 NOINLINE // Make sure that the slow path with lock won't affect the fast path
 BOOL ExecutionManager::IsManagedCodeWithLock(PCODE currentPC)

src/coreclr/vm/codeman.h

@@ -2293,11 +2293,15 @@ class ExecutionManager
     };
 
     // Returns default scan flag for current thread
-    static ScanFlag GetScanFlags();
+    static ScanFlag GetScanFlags(Thread *pThread = NULL);
 
     // Returns whether currentPC is in managed code. Returns false for jump stubs on WIN64.
     static BOOL IsManagedCode(PCODE currentPC);
 
+    // Returns whether currentPC is in managed code. Returns false for jump stubs on WIN64.
+    // Does not acquire the reader lock. Caller must ensure it is safe.
+    static BOOL IsManagedCodeNoLock(PCODE currentPC);
+
     // Returns true if currentPC is ready to run codegen
     static BOOL IsReadyToRunCode(PCODE currentPC);
 

src/coreclr/vm/threads.cpp

@@ -33,6 +33,8 @@
 #include "vmholder.h"
 #include "exceptmacros.h"
 #include "minipal/time.h"
+#include "minipal/thread.h"
+#include "asyncsafethreadmap.h"
 
 #ifdef FEATURE_COMINTEROP
 #include "runtimecallablewrapper.h"
@@ -62,6 +64,13 @@
 #include "interpexec.h"
 #endif // FEATURE_INTERPRETER
 
+#if defined(TARGET_UNIX) && !defined(TARGET_WASM)
+Thread* GetThreadAsyncSafe()
+{
+    return (Thread*)FindThreadInAsyncSafeMap(minipal_get_current_thread_id_no_cache());
+}
+#endif // TARGET_UNIX && !TARGET_WASM
+
 static const PortableTailCallFrame g_sentinelTailCallFrame = { NULL, NULL };
 
 TailCallTls::TailCallTls()
@@ -371,6 +380,21 @@ void SetThread(Thread* t)
 
     // Clear or set the app domain to the one domain based on if the thread is being nulled out or set
     t_CurrentThreadInfo.m_pAppDomain = t == NULL ? NULL : AppDomain::GetCurrentDomain();
+
+#if defined(TARGET_UNIX) && !defined(TARGET_WASM)
+    if (t != NULL)
+    {
+        if (!InsertThreadIntoAsyncSafeMap(t->GetOSThreadId64(), t))
+        {
+            // TODO: can we handle this OOM more gracefully?
+            EEPOLICY_HANDLE_FATAL_ERROR_WITH_MESSAGE(COR_E_EXECUTIONENGINE, W("Failed to insert thread into async-safe map due to OOM."));
+        }
+    }
+    else if (origThread != NULL)
+    {
+        RemoveThreadFromAsyncSafeMap(origThread->GetOSThreadId64(), origThread);
+    }
+#endif // TARGET_UNIX && !TARGET_WASM
 }
 
 BOOL Thread::Alert ()

src/coreclr/vm/threads.h

@@ -5512,6 +5512,10 @@ class StackWalkerWalkingThreadHolder
     Thread* m_PreviousValue;
 };
 
+#ifdef TARGET_UNIX
+EXTERN_C Thread* GetThreadAsyncSafe();
+#endif
+
 #ifndef DACCESS_COMPILE
 #if defined(TARGET_WINDOWS) && defined(TARGET_AMD64)
 EXTERN_C void STDCALL ClrRestoreNonvolatileContextWorker(PCONTEXT ContextRecord, DWORD64 ssp);

src/coreclr/vm/threadsuspend.cpp

@@ -5731,16 +5731,18 @@ void ThreadSuspend::SuspendEE(SUSPEND_REASON reason)
 //          It is unsafe to use blocking APIs or allocate in this method.
 BOOL CheckActivationSafePoint(SIZE_T ip)
 {
-    Thread *pThread = GetThreadNULLOk();
+    Thread *pThread = GetThreadAsyncSafe();
+    _ASSERTE(pThread != NULL);
 
     // The criteria for safe activation is to be running managed code.
     // Also we are not interested in handling interruption if we are already in preemptive mode nor if we are single stepping
     BOOL isActivationSafePoint = pThread != NULL &&
         (pThread->m_StateNC & Thread::TSNC_DebuggerIsStepping) == 0 &&
         pThread->PreemptiveGCDisabled() &&
-        ExecutionManager::IsManagedCode(ip);
+        (ExecutionManager::GetScanFlags(pThread) != ExecutionManager::ScanReaderLock) &&
+        ExecutionManager::IsManagedCodeNoLock(ip);
 
-    if (!isActivationSafePoint)
+    if (!isActivationSafePoint && pThread != NULL)
     {
         pThread->m_hasPendingActivation = false;
     }