From c53108487331a579a3be9b43ec60d6e73cee26ed Mon Sep 17 00:00:00 2001
From: Tom <tomut@yahoo.com>
Date: Wed, 28 Oct 2020 16:06:16 -0600
Subject: [PATCH] Kernel: Track processor idle state and wake processors when
 waking threads

Attempt to wake idle processors to get threads to be scheduled more quickly.
We don't want to wait until the next timer tick if we have processors that
aren't doing anything.
---
 Kernel/Arch/i386/CPU.cpp | 51 ++++++++++++++++++++++++++++++++++++++++
 Kernel/Arch/i386/CPU.h   | 12 ++++++++++
 Kernel/Scheduler.cpp     |  8 +++++--
 Kernel/Thread.cpp        |  4 +++-
 4 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/Kernel/Arch/i386/CPU.cpp b/Kernel/Arch/i386/CPU.cpp
index 89f2c48a26..7dc7f082c6 100644
--- a/Kernel/Arch/i386/CPU.cpp
+++ b/Kernel/Arch/i386/CPU.cpp
@@ -1734,6 +1734,57 @@ ProcessorMessage& Processor::smp_get_from_pool()
     return *msg;
 }
 
+Atomic<u32> Processor::s_idle_cpu_mask{ 0 };
+
+u32 Processor::smp_wake_n_idle_processors(u32 wake_count)
+{
+    ASSERT(Processor::current().in_critical());
+    ASSERT(wake_count > 0);
+    if (!s_smp_enabled)
+        return 0;
+
+    // Wake at most N - 1 processors
+    if (wake_count >= Processor::count()) {
+        wake_count = Processor::count() - 1;
+        ASSERT(wake_count > 0);
+    }
+
+    u32 current_id = Processor::current().id();
+
+    u32 did_wake_count = 0;
+    auto& apic = APIC::the();
+    while (did_wake_count < wake_count) {
+        // Try to get a set of idle CPUs and flip them to busy
+        u32 idle_mask = s_idle_cpu_mask.load(AK::MemoryOrder::memory_order_relaxed) & ~(1u << current_id);
+        u32 idle_count = __builtin_popcountl(idle_mask);
+        if (idle_count == 0)
+            break; // No (more) idle processor available
+
+        u32 found_mask = 0;
+        for (u32 i = 0; i < idle_count; i++) {
+            u32 cpu = __builtin_ffsl(idle_mask) - 1;
+            idle_mask &= ~(1u << cpu);
+            found_mask |= 1u << cpu;
+        }
+
+        idle_mask = s_idle_cpu_mask.fetch_and(~found_mask, AK::MemoryOrder::memory_order_acq_rel) & found_mask;
+        if (idle_mask == 0)
+            continue; // All of them were flipped to busy, try again
+        idle_count = __builtin_popcountl(idle_mask);
+        for (u32 i = 0; i < idle_count; i++) {
+            u32 cpu = __builtin_ffsl(idle_mask) - 1;
+            idle_mask &= ~(1u << cpu);
+
+            // Send an IPI to that CPU to wake it up. There is a possibility
+            // someone else woke it up as well, or that it woke up due to
+            // a timer interrupt. But we tried hard to avoid this...
+            apic.send_ipi(cpu);
+            did_wake_count++;
+        }
+    }
+    return did_wake_count;
+}
+
 void Processor::smp_enable()
 {
     size_t msg_pool_size = Processor::count() * 100u;
diff --git a/Kernel/Arch/i386/CPU.h b/Kernel/Arch/i386/CPU.h
index 77ed93552b..325f637bed 100644
--- a/Kernel/Arch/i386/CPU.h
+++ b/Kernel/Arch/i386/CPU.h
@@ -709,6 +709,7 @@ class Processor {
     u32 m_cpu;
     u32 m_in_irq;
     Atomic<u32, AK::MemoryOrder::memory_order_relaxed> m_in_critical;
+    static Atomic<u32> s_idle_cpu_mask;
 
     TSS32 m_tss;
     static FPUState s_clean_fpu_state;
@@ -763,6 +764,16 @@ public:
     void early_initialize(u32 cpu);
     void initialize(u32 cpu);
 
+    void idle_begin()
+    {
+        s_idle_cpu_mask.fetch_or(1u << m_cpu, AK::MemoryOrder::memory_order_relaxed);
+    }
+
+    void idle_end()
+    {
+        s_idle_cpu_mask.fetch_and(~(1u << m_cpu), AK::MemoryOrder::memory_order_relaxed);
+    }
+
     static u32 count()
     {
         // NOTE: because this value never changes once all APs are booted,
@@ -1010,6 +1021,7 @@ public:
     static void smp_unicast(u32 cpu, void (*callback)(), bool async);
     static void smp_unicast(u32 cpu, void (*callback)(void*), void* data, void (*free_data)(void*), bool async);
     static void smp_broadcast_flush_tlb(const PageDirectory*, VirtualAddress, size_t);
+    static u32 smp_wake_n_idle_processors(u32 wake_count);
 
     template<typename Callback>
     static void deferred_call_queue(Callback callback)
diff --git a/Kernel/Scheduler.cpp b/Kernel/Scheduler.cpp
index 1fd17e1ae7..5c0eb58506 100644
--- a/Kernel/Scheduler.cpp
+++ b/Kernel/Scheduler.cpp
@@ -544,13 +544,17 @@ void Scheduler::notify_finalizer()
 
 void Scheduler::idle_loop(void*)
 {
-    dbgln("Scheduler[{}]: idle loop running", Processor::id());
+    auto& proc = Processor::current();
+    dbgln("Scheduler[{}]: idle loop running", proc.get_id());
     ASSERT(are_interrupts_enabled());
 
     for (;;) {
+        proc.idle_begin();
         asm("hlt");
 
-        if (Processor::id() == 0)
+        proc.idle_end();
+        ASSERT_INTERRUPTS_ENABLED();
+        if (Processor::current().id() == 0)
             yield();
     }
 }
diff --git a/Kernel/Thread.cpp b/Kernel/Thread.cpp
index b2b86979fe..0931fc11c3 100644
--- a/Kernel/Thread.cpp
+++ b/Kernel/Thread.cpp
@@ -919,7 +919,9 @@ void Thread::set_state(State new_state, u8 stop_signal)
         }
     }
 
-    if (m_state == Stopped) {
+    if (m_state == Runnable) {
+        Processor::smp_wake_n_idle_processors(1);
+    } else if (m_state == Stopped) {
         // We don't want to restore to Running state, only Runnable!
         m_stop_state = previous_state != Running ? previous_state : Runnable;
         auto& process = this->process();