Add SetThreadIdealProcessorEx PAL API (#13606)

* add SetThreadIdealProcessorEx PAL API

* fix lpPreviousIdealProcessor

2 files changed, 121 insertions, 0 deletions

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