Add SetThreadIdealProcessorEx PAL API (#13606)

* add SetThreadIdealProcessorEx PAL API

* fix lpPreviousIdealProcessor

5 files changed, 125 insertions, 3 deletions

diff --git a/src/dlls/mscoree/mscorwks_unixexports.src b/src/dlls/mscoree/mscorwks_unixexports.src
index 0da8942a2b..a3ca37c3f2 100644
--- a/src/dlls/mscoree/mscorwks_unixexports.src
+++ b/src/dlls/mscoree/mscorwks_unixexports.src
@@ -105,4 +105,5 @@ SetThreadGroupAffinity
 SetThreadAffinityMask
 GetThreadGroupAffinity
 GetCurrentProcessorNumberEx
+SetThreadIdealProcessorEx
 
diff --git a/src/gc/gc.cpp b/src/gc/gc.cpp
index 89454c6bb0..779b1b3259 100644
--- a/src/gc/gc.cpp
+++ b/src/gc/gc.cpp
@@ -13351,7 +13351,6 @@ try_again:
                     org_hp->alloc_context_count--;
                     max_hp->alloc_context_count++;
                     acontext->set_alloc_heap(GCHeap::GetHeap(max_hp->heap_number));
-#if !defined(FEATURE_PAL)
                     if (CPUGroupInfo::CanEnableGCCPUGroups())
                     {   //only set ideal processor when max_hp and org_hp are in the same cpu
                         //group. DO NOT MOVE THREADS ACROSS CPU GROUPS
@@ -13385,7 +13384,6 @@ try_again:
                                         org_hp->heap_number));
                         }
                     }
-#endif // !FEATURE_PAL
                     dprintf (3, ("Switching context %p (home heap %d) ", 
                                  acontext,
                         acontext->get_home_heap()->pGenGCHeap->heap_number));
diff --git a/src/pal/inc/pal.h b/src/pal/inc/pal.h
index 1f611d0922..1a1b6c8625 100644
--- a/src/pal/inc/pal.h
+++ b/src/pal/inc/pal.h
@@ -4847,6 +4847,15 @@ GetProcessAffinityMask(
   OUT PDWORD_PTR lpSystemAffinityMask
 );
 
+PALIMPORT
+BOOL
+PALAPI
+SetThreadIdealProcessorEx(
+  IN HANDLE hThread,
+  IN PPROCESSOR_NUMBER lpIdealProcessor,
+  OUT PPROCESSOR_NUMBER lpPreviousIdealProcessor
+);
+
 //
 // The types of events that can be logged.
 //
diff --git a/src/pal/src/numa/numa.cpp b/src/pal/src/numa/numa.cpp
index 91cf6871c1..383fb0e3b0 100644
--- a/src/pal/src/numa/numa.cpp
+++ b/src/pal/src/numa/numa.cpp
@@ -867,3 +867,115 @@ VirtualAllocExNuma(
 
     return result;
 }
+
+/*++
+Function:
+  SetThreadIdealProcessorEx
+
+See MSDN doc.
+--*/
+BOOL
+PALAPI
+SetThreadIdealProcessorEx(
+  IN HANDLE hThread,
+  IN PPROCESSOR_NUMBER lpIdealProcessor,
+  OUT PPROCESSOR_NUMBER lpPreviousIdealProcessor)
+{
+    PERF_ENTRY(SetThreadIdealProcessorEx);
+    ENTRY("SetThreadIdealProcessorEx(hThread=%p, lpIdealProcessor=%p)\n", hThread, lpIdealProcessor);
+
+    CPalThread *pCurrentThread = InternalGetCurrentThread();
+    CPalThread *pTargetThread = NULL;
+    IPalObject *pTargetThreadObject = NULL;
+
+    PAL_ERROR palErr =
+        InternalGetThreadDataFromHandle(pCurrentThread, hThread,
+                                        0, // THREAD_SET_CONTEXT
+                                        &pTargetThread, &pTargetThreadObject);
+
+    if (NO_ERROR != palErr)
+    {
+        ERROR("Unable to obtain thread data for handle %p (error %x)!\n", hThread,
+              palErr);
+        return 0;
+    }
+
+    pthread_t thread = pTargetThread->GetPThreadSelf();
+
+#if HAVE_PTHREAD_GETAFFINITY_NP
+    int cpu = -1;
+    if ((lpIdealProcessor->Group < g_groupCount) &&
+        (lpIdealProcessor->Number < MaxCpusPerGroup) &&
+        (lpIdealProcessor->Reserved == 0))
+    {
+        cpu = g_groupAndIndexToCpu[lpIdealProcessor->Group * MaxCpusPerGroup + lpIdealProcessor->Number];
+    }
+
+    if (cpu == -1)
+    {
+        SetLastError(ERROR_INVALID_PARAMETER);
+        return FALSE;
+    }
+
+    if (lpPreviousIdealProcessor != NULL)
+    {
+        cpu_set_t prevCpuSet;
+        CPU_ZERO(&prevCpuSet);
+        DWORD prevCpu = GetCurrentProcessorNumber();
+
+        int st = pthread_getaffinity_np(thread, sizeof(cpu_set_t), &prevCpuSet);
+
+        if (st == 0)
+        {
+            for (int i = 0; i < g_possibleCpuCount; i++)
+            {
+                if (CPU_ISSET(i, &prevCpuSet))
+                {
+                    prevCpu = i;
+                    break;
+                }
+            }
+        }
+
+        _ASSERTE(prevCpu < g_possibleCpuCount);
+        lpPreviousIdealProcessor->Group = g_cpuToAffinity[prevCpu].Group;
+        lpPreviousIdealProcessor->Number = g_cpuToAffinity[prevCpu].Number;
+        lpPreviousIdealProcessor->Reserved = 0;
+    }
+
+    cpu_set_t cpuSet;
+    CPU_ZERO(&cpuSet);
+    CPU_SET(cpu, &cpuSet);
+
+    int st = pthread_setaffinity_np(thread, sizeof(cpu_set_t), &cpuSet);
+
+    if (st != 0)
+    {
+        switch (st)
+        {
+        case EINVAL:
+            // There is no processor in the mask that is allowed to execute the
+            // process
+            SetLastError(ERROR_INVALID_PARAMETER);
+            break;
+        case ESRCH:
+            SetLastError(ERROR_INVALID_HANDLE);
+            break;
+        default:
+            SetLastError(ERROR_GEN_FAILURE);
+            break;
+        }
+    }
+
+    BOOL success = (st == 0);
+
+#else  // HAVE_PTHREAD_GETAFFINITY_NP
+    // There is no API to manage thread affinity, so let's ignore the request
+    BOOL success = FALSE;
+#endif // HAVE_PTHREAD_GETAFFINITY_NP
+
+    LOGEXIT("SetThreadIdealProcessorEx returns BOOL %d\n", success);
+    PERF_EXIT(SetThreadIdealProcessorEx);
+
+    return success;
+}
\ No newline at end of file
diff --git a/src/vm/gcenv.os.cpp b/src/vm/gcenv.os.cpp
index 22c8a97176..b1b9c32635 100644
--- a/src/vm/gcenv.os.cpp
+++ b/src/vm/gcenv.os.cpp
@@ -83,7 +83,7 @@ bool GCToOSInterface::SetCurrentThreadIdealAffinity(GCThreadAffinity* affinity)
 
 #if !defined(FEATURE_CORESYSTEM)
     SetThreadIdealProcessor(GetCurrentThread(), (DWORD)affinity->Processor);
-#elif !defined(FEATURE_PAL)
+#else
     PROCESSOR_NUMBER proc;
 
     if (affinity->Group != -1)
@@ -94,6 +94,7 @@ bool GCToOSInterface::SetCurrentThreadIdealAffinity(GCThreadAffinity* affinity)
         
         success = !!SetThreadIdealProcessorEx(GetCurrentThread(), &proc, NULL);
     }
+#if !defined(FEATURE_PAL)
     else
     {
         if (GetThreadIdealProcessorEx(GetCurrentThread(), &proc))
@@ -102,6 +103,7 @@ bool GCToOSInterface::SetCurrentThreadIdealAffinity(GCThreadAffinity* affinity)
             success = !!SetThreadIdealProcessorEx(GetCurrentThread(), &proc, &proc);
         }        
     }
+#endif // !defined(FEATURE_PAL)
 #endif
 
     return success;