WIP [NTOS:KE/x64] Handle NewThread == OldThread in KiIdleLoop

ntoskrnl/ke/amd64/stubs.c

@@ -158,6 +158,9 @@ KiIdleLoop(VOID)
             /* The thread is now running */
             NewThread->State = Running;
 
+            /* Check if we're actually running a different thread */
+            if (NewThread != OldThread)
+            {
                 /* Clear idle summary bit */
                 InterlockedBitTestAndResetAffinity(&KiIdleSummary, Prcb->Number);
 
@@ -166,6 +169,11 @@ KiIdleLoop(VOID)
 
                 /* Set idle summary bit */
                 InterlockedBitTestAndSetAffinity(&KiIdleSummary, Prcb->Number);
+            }
+            else
+            {
+                NewThread->SwapBusy = FALSE;
+            }
 
 #ifdef CONFIG_SMP
             /* Go back to DISPATCH_LEVEL */

ntoskrnl/ke/amd64/traphandler.c

@@ -66,6 +66,17 @@ KiDpcInterruptHandler(VOID)
         /* Acquire the PRCB lock */
         KiAcquirePrcbLock(Prcb);
 
+        if (Prcb->NextThread == Prcb->CurrentThread)
+        {
+            __debugbreak();
+            /* This can happen, when the idle thread is running and a different
+               processor reschedules the thread */
+            ASSERT(Prcb->NextThread == Prcb->IdleThread);
+            Prcb->NextThread = NULL;
+            KiReleasePrcbLock(Prcb);
+            goto Exit;
+        }
+
         /* Capture current thread data */
         OldThread = Prcb->CurrentThread;
         NewThread = Prcb->NextThread;
@@ -78,7 +89,7 @@ KiDpcInterruptHandler(VOID)
         NewThread->State = Running;
         OldThread->WaitReason = WrDispatchInt;
 
-        /* Make the old thread ready */
+        /* Make the old thread ready (this releases the PRCB lock) */
         KxQueueReadyThread(OldThread, Prcb);
 
         /* Swap to the new thread */

ntoskrnl/ke/thrdschd.c

@@ -485,8 +485,27 @@ KiSwapThread(IN PKTHREAD CurrentThread,
     /* Save the wait IRQL */
     WaitIrql = CurrentThread->WaitIrql;
 
+#ifdef CONFIG_SMP
+    /* On SMP builds it is possible that the new thread is the old thread. */
+    if (NextThread == CurrentThread)
+    {
+        /* Unset SwapBusy */
+        CurrentThread->SwapBusy = FALSE;
+
+        /* Check for pending APCs */
+        if ((NextThread->ApcState.KernelApcPending) &&
+            (NextThread->SpecialApcDisable == FALSE) &&
+            (WaitIrql == PASSIVE_LEVEL))
+        {
+            ApcState = TRUE;
+        }
+    }
+    else
+#endif
+    {
         /* Swap contexts */
         ApcState = KiSwapContext(WaitIrql, CurrentThread);
+    }
 
     /* Get the wait status */
     WaitStatus = CurrentThread->WaitStatus;