From 5faf6ed3390573fa51dec32e4fc54e8bac7478a4 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Tue, 10 Dec 2024 19:26:02 +0200
Subject: [PATCH] WIP [NTOS:KE/x64] Handle NewThread == OldThread in KiIdleLoop

---
 ntoskrnl/ke/amd64/stubs.c       | 20 ++++++++++++++------
 ntoskrnl/ke/amd64/traphandler.c | 13 ++++++++++++-
 ntoskrnl/ke/thrdschd.c          | 23 +++++++++++++++++++++--
 3 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/ntoskrnl/ke/amd64/stubs.c b/ntoskrnl/ke/amd64/stubs.c
index a32d3022fe7..3d343d5f667 100644
--- a/ntoskrnl/ke/amd64/stubs.c
+++ b/ntoskrnl/ke/amd64/stubs.c
@@ -158,14 +158,22 @@ KiIdleLoop(VOID)
             /* The thread is now running */
             NewThread->State = Running;
 
-            /* Clear idle summary bit */
-            InterlockedBitTestAndResetAffinity(&KiIdleSummary, Prcb->Number);
+            /* Check if we're actually running a different thread */
+            if (NewThread != OldThread)
+            {
+                /* Clear idle summary bit */
+                InterlockedBitTestAndResetAffinity(&KiIdleSummary, Prcb->Number);
 
-            /* Switch away from the idle thread */
-            KiSwapContext(APC_LEVEL, OldThread);
+                /* Switch away from the idle thread */
+                KiSwapContext(APC_LEVEL, OldThread);
 
-            /* Set idle summary bit */
-            InterlockedBitTestAndSetAffinity(&KiIdleSummary, Prcb->Number);
+                /* Set idle summary bit */
+                InterlockedBitTestAndSetAffinity(&KiIdleSummary, Prcb->Number);
+            }
+            else
+            {
+                NewThread->SwapBusy = FALSE;
+            }
 
 #ifdef CONFIG_SMP
             /* Go back to DISPATCH_LEVEL */
diff --git a/ntoskrnl/ke/amd64/traphandler.c b/ntoskrnl/ke/amd64/traphandler.c
index b96fb5323b9..cdd16d82b5b 100644
--- a/ntoskrnl/ke/amd64/traphandler.c
+++ b/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 */
diff --git a/ntoskrnl/ke/thrdschd.c b/ntoskrnl/ke/thrdschd.c
index 10dc41488c3..cd280cf518d 100644
--- a/ntoskrnl/ke/thrdschd.c
+++ b/ntoskrnl/ke/thrdschd.c
@@ -485,8 +485,27 @@ KiSwapThread(IN PKTHREAD CurrentThread,
     /* Save the wait IRQL */
     WaitIrql = CurrentThread->WaitIrql;
 
-    /* Swap contexts */
-    ApcState = KiSwapContext(WaitIrql, CurrentThread);
+#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;