[NTOS:KE] Implement KiFindIdealProcessor

ntoskrnl/include/internal/ke.h

@@ -1079,6 +1079,14 @@ KeBugCheckUnicodeToAnsi(
     IN ULONG Length
 );
 
+#ifdef CONFIG_SMP
+ULONG
+NTAPI
+KiFindIdealProcessor(
+    _In_ KAFFINITY ProcessorSet,
+    _In_ UCHAR OriginalIdealProcessor);
+#endif // CONFIG_SMP
+
 #ifdef __cplusplus
 } // extern "C"
 

ntoskrnl/ke/thrdobj.c

@@ -1119,29 +1119,10 @@ KeSetSystemAffinityThread(IN KAFFINITY Affinity)
     CurrentThread->Affinity = Affinity;
     CurrentThread->SystemAffinityActive = TRUE;
 
-    /* Check if the ideal processor is part of the affinity */
 #ifdef CONFIG_SMP
-    if (!(Affinity & AFFINITY_MASK(CurrentThread->IdealProcessor)))
-    {
-        KAFFINITY AffinitySet, NodeMask;
-        ULONG IdealProcessor;
-
-        /* It's not! Get the PRCB */
-        Prcb = KiProcessorBlock[CurrentThread->IdealProcessor];
-
-        /* Calculate the affinity set */
-        AffinitySet = KeActiveProcessors & Affinity;
-        NodeMask = Prcb->ParentNode->ProcessorMask & AffinitySet;
-        if (NodeMask)
-        {
-            /* Use the Node set instead */
-            AffinitySet = NodeMask;
-        }
-
-        /* Calculate the ideal CPU from the affinity set */
-        BitScanReverseAffinity(&IdealProcessor, AffinitySet);
-        CurrentThread->IdealProcessor = (UCHAR)IdealProcessor;
-    }
+    /* Calculate the ideal processor from the affinity set */
+    CurrentThread->IdealProcessor =
+        KiFindIdealProcessor(Affinity, CurrentThread->IdealProcessor);
 #endif
 
     /* Get the current PRCB and check if it doesn't match this affinity */

ntoskrnl/ke/thrdschd.c

@@ -75,6 +75,42 @@ KiQueueReadyThread(IN PKTHREAD Thread,
 }
 
 #ifdef CONFIG_SMP
+ULONG
+NTAPI
+KiFindIdealProcessor(
+    _In_ KAFFINITY ProcessorSet,
+    _In_ UCHAR OriginalIdealProcessor)
+{
+    PKPRCB OriginalIdealPrcb;
+    KAFFINITY NodeMask;
+    ULONG Processor;
+
+    /* Check if we can use the original ideal processor */
+    if (ProcessorSet & AFFINITY_MASK(OriginalIdealProcessor))
+    {
+        /* We can, so use it */
+        return OriginalIdealProcessor;
+    }
+
+    /* Only use active processors */
+    ProcessorSet &= KeActiveProcessors;
+
+    /* Get the original ideal PRCB */
+    OriginalIdealPrcb = KiProcessorBlock[OriginalIdealProcessor];
+
+    /* Check if we can use the original node */
+    NodeMask = OriginalIdealPrcb->ParentNode->ProcessorMask & ProcessorSet;
+    if (NodeMask)
+    {
+        /* Use the node set instead */
+        ProcessorSet = NodeMask;
+    }
+
+    /* Calculate the ideal CPU from the affinity set */
+    BitScanReverseAffinity(&Processor, ProcessorSet);
+    return Processor;
+}
+
 static
 ULONG
 KiSelectNextProcessor(
@@ -740,6 +776,10 @@ KiSetAffinityThread(IN PKTHREAD Thread,
     if (!Thread->SystemAffinityActive)
     {
 #ifdef CONFIG_SMP
+        /* Calculate the new ideal processor from the affinity set */
+        Thread->UserIdealProcessor =
+            KiFindIdealProcessor(Affinity, Thread->UserIdealProcessor);
+
         /* FIXME: TODO */
         DPRINT1("Affinity support disabled!\n");
 #endif