Kernel: Don't use peridic interrupts for scheduling

Scheduler will now set a timer interrupt deadline for when it should get
interrupted next. This fixes a big issue I've had for a long time where
sleeping on an idle core would have to wait for the next periodic
interrupt to wake up. This could cause 1 ms sleeps to be close to 10 ms.

This currently only supports lapic timers and will still fallback to
periodic interrupts when running with PIC. I should add deadline support
to PIT/HPET but why would you run with APIC disabled ;)

kernel/arch/i686/Yield.S

@@ -12,7 +12,7 @@ asm_yield_trampoline:
 	pushl %ebp
 
 	pushl %esp
-	call scheduler_on_yield
+	call scheduler_on_yield_trampoline
 	addl $4, %esp
 
 	popl %ebp

kernel/arch/x86_64/Yield.S

@@ -15,7 +15,7 @@ asm_yield_trampoline:
 	pushq %r15
 
 	movq %rsp, %rdi
-	call scheduler_on_yield
+	call scheduler_on_yield_trampoline
 
 	popq %r15
 	popq %r14

kernel/include/kernel/APIC.h

@@ -26,6 +26,7 @@ namespace Kernel
 		BAN::ErrorOr<uint8_t> reserve_gsi(uint32_t gsi);
 
 		void initialize_timer();
+		void set_timer_dealine(uint64_t ns);
 
 	private:
 		uint32_t read_from_local_apic(ptrdiff_t);

kernel/include/kernel/Scheduler.h

@@ -23,7 +23,7 @@ namespace Kernel
 
 	public:
 		void add_thread_to_back(Node*);
-		void add_thread_with_wake_time(Node*);
+		bool add_thread_with_wake_time(Node*); // return true if node was inserted as the first element
 		template<typename F>
 		Node* remove_with_condition(F callback);
 		void remove_node(Node*);
@@ -60,7 +60,8 @@ namespace Kernel
 		void reschedule(YieldRegisters*);
 		void reschedule_if_idle();
 
-		void timer_interrupt();
+		void on_timer_interrupt();
+		void on_yield(YieldRegisters*);
 
 		static BAN::ErrorOr<void> bind_thread_to_processor(Thread*, ProcessorID);
 		// if thread is already bound, this will never fail
@@ -82,6 +83,7 @@ namespace Kernel
 		void update_most_loaded_node_queue(SchedulerQueue::Node*, SchedulerQueue* target_queue);
 		void remove_node_from_most_loaded(SchedulerQueue::Node*);
 
+		void update_wake_up_deadline();
 		void wake_up_sleeping_threads();
 
 		void do_load_balancing();

kernel/kernel/APIC.cpp

@@ -437,8 +437,27 @@ namespace Kernel
 
 		dprintln("CPU {}: lapic timer frequency: {} Hz", Kernel::Processor::current_id(), m_lapic_timer_frequency_hz);
 
-		write_to_local_apic(LAPIC_TIMER_LVT,         TimerModePeriodic | IRQ_TIMER);
-		write_to_local_apic(LAPIC_TIMER_INITIAL_REG, m_lapic_timer_frequency_hz / 2 / 100);
+		write_to_local_apic(LAPIC_TIMER_LVT,         TimerModeOneShot | IRQ_TIMER);
+		write_to_local_apic(LAPIC_TIMER_INITIAL_REG, 0);
+	}
+
+	void APIC::set_timer_dealine(uint64_t target_ns)
+	{
+		// TODO: don't hardcode divide by 2
+		const uint64_t effective_freq = m_lapic_timer_frequency_hz / 2;
+
+		const uint64_t delta_ns = [&]() -> uint64_t {
+			const uint64_t current_ns = SystemTimer::get().ns_since_boot();
+			if (current_ns >= target_ns)
+				return 1;
+			const uint64_t delta_ns = target_ns - current_ns;
+			if (BAN::Math::will_multiplication_overflow(delta_ns, effective_freq))
+				return BAN::numeric_limits<uint64_t>::max() / effective_freq;
+			return delta_ns;
+		}();
+
+		const uint64_t ticks = BAN::Math::div_round_up<uint64_t>(delta_ns * effective_freq, 1'000'000'000);
+		write_to_local_apic(LAPIC_TIMER_INITIAL_REG, BAN::Math::min<uint64_t>(ticks, BAN::numeric_limits<uint32_t>::max()));
 	}
 
 	uint32_t APIC::read_from_local_apic(ptrdiff_t offset)

kernel/kernel/IDT.cpp

@@ -435,7 +435,7 @@ namespace Kernel
 		if (Processor::current_is_bsp())
 			Process::update_alarm_queue();
 
-		Processor::scheduler().timer_interrupt();
+		Processor::scheduler().on_timer_interrupt();
 	}
 
 	extern "C" void cpp_irq_handler(uint32_t irq)

kernel/kernel/Scheduler.cpp

@@ -1,5 +1,6 @@
 #include <BAN/Optional.h>
 #include <BAN/Sort.h>
+#include <kernel/APIC.h>
 #include <kernel/InterruptController.h>
 #include <kernel/Lock/Mutex.h>
 #include <kernel/Process.h>
@@ -45,12 +46,15 @@ namespace Kernel
 		m_tail = node;
 	}
 
-	void SchedulerQueue::add_thread_with_wake_time(Node* node)
+	bool SchedulerQueue::add_thread_with_wake_time(Node* node)
 	{
 		ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
 
 		if (m_tail == nullptr || node->wake_time_ns >= m_tail->wake_time_ns)
-			return add_thread_to_back(node);
+		{
+			add_thread_to_back(node);
+			return node == m_head;
+		}
 
 		Node* next = m_head;
 		Node* prev = nullptr;
@@ -64,6 +68,8 @@ namespace Kernel
 		node->prev = prev;
 		(next ? next->prev : m_tail) = node;
 		(prev ? prev->next : m_head) = node;
+
+		return node == m_head;
 	}
 
 	template<typename F>
@@ -122,15 +128,9 @@ namespace Kernel
 		m_idle_thread = TRY(Thread::create_kernel([](void*) { asm volatile("1: hlt; jmp 1b"); }, nullptr));
 		ASSERT(m_idle_thread);
 
-		size_t processor_index = 0;
-		for (; processor_index < Processor::count(); processor_index++)
-			if (Processor::id_from_index(processor_index) == Processor::current_id())
-				break;
-		ASSERT(processor_index < Processor::count());
-
 		// each CPU does load balance at different times. This calulates the offset to other CPUs
-		m_last_load_balance_ns = s_load_balance_interval_ns * processor_index / Processor::count();
-		m_idle_ns              = -m_last_load_balance_ns;
+		m_last_load_balance_ns = s_load_balance_interval_ns * Processor::current_index() / Processor::count();
+		m_idle_ns              = m_last_load_balance_ns;
 
 		s_schedulers_initialized++;
 		while (s_schedulers_initialized < Processor::count())
@@ -139,8 +139,6 @@ namespace Kernel
 		if (Processor::count() > 1)
 			Processor::set_smp_enabled();
 
-		m_next_reschedule_ns = SystemTimer::get().ns_since_boot();
-
 		return {};
 	}
 
@@ -232,17 +230,14 @@ namespace Kernel
 					m_thread_count--;
 					break;
 				case Thread::State::Executing:
-				{
-					const uint64_t current_ns = SystemTimer::get().ns_since_boot();
 					m_current->thread->yield_registers() = *yield_registers;
-					m_current->time_used_ns += current_ns - m_current->last_start_ns;
+					m_current->time_used_ns += SystemTimer::get().ns_since_boot() - m_current->last_start_ns;
 					add_current_to_most_loaded(m_current->blocked ? &m_block_queue : &m_run_queue);
 					if (!m_current->blocked)
 						m_run_queue.add_thread_to_back(m_current);
-					else
-						m_block_queue.add_thread_with_wake_time(m_current);
+					else if (m_block_queue.add_thread_with_wake_time(m_current))
+						update_wake_up_deadline();
 					break;
-				}
 				case Thread::State::NotStarted:
 					ASSERT(!m_current->blocked);
 					m_current->time_used_ns = 0;
@@ -321,26 +316,50 @@ namespace Kernel
 		}
 	}
 
+	void Scheduler::update_wake_up_deadline()
+	{
+		ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
+
+		auto& interrupt_controller = InterruptController::get();
+
+		// TODO: support timer deadlines on non-apic timers and abstract it :)
+		if (!interrupt_controller.is_using_apic())
+			return;
+
+		wake_up_sleeping_threads();
+
+		uint64_t deadline_ns = m_next_reschedule_ns;
+		if (!m_block_queue.empty())
+			deadline_ns = BAN::Math::min(deadline_ns, m_block_queue.front()->wake_time_ns);
+		if (Processor::is_smp_enabled())
+			deadline_ns = BAN::Math::min(deadline_ns, m_last_load_balance_ns + s_load_balance_interval_ns);
+
+		static_cast<APIC&>(interrupt_controller).set_timer_dealine(deadline_ns);
+	}
+
 	void Scheduler::reschedule_if_idle()
 	{
 		ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
 
-		if (m_current != nullptr)
-			return;
-
-		if (m_run_queue.empty())
-			wake_up_sleeping_threads();
-
-		if (!m_run_queue.empty())
+		if (m_current == nullptr && !m_run_queue.empty())
 			Processor::yield();
 	}
 
-	extern "C" void scheduler_on_yield(YieldRegisters* yield_registers)
+	extern "C" void scheduler_on_yield_trampoline(YieldRegisters* yield_registers)
 	{
-		Processor::scheduler().reschedule(yield_registers);
+		Processor::scheduler().on_yield(yield_registers);
 	}
 
-	void Scheduler::timer_interrupt()
+	void Scheduler::on_yield(YieldRegisters* yield_registers)
+	{
+		reschedule(yield_registers);
+		m_next_reschedule_ns = !is_idle()
+			? SystemTimer::get().ns_since_boot() + s_reschedule_interval_ns
+			: BAN::numeric_limits<uint64_t>::max();
+		update_wake_up_deadline();
+	}
+
+	void Scheduler::on_timer_interrupt()
 	{
 		ASSERT(Processor::get_interrupt_state() == InterruptState::Disabled);
 
@@ -349,11 +368,8 @@ namespace Kernel
 
 		wake_up_sleeping_threads();
 
-		if (SystemTimer::get().ns_since_boot() >= m_next_reschedule_ns)
-		{
-			m_next_reschedule_ns += s_reschedule_interval_ns;
+		if (is_idle() || SystemTimer::get().ns_since_boot() >= m_next_reschedule_ns)
 			Processor::yield();
-		}
 	}
 
 	void Scheduler::unblock_thread(SchedulerQueue::Node* node)
@@ -394,8 +410,8 @@ namespace Kernel
 
 		if (!node->blocked)
 			m_run_queue.add_thread_to_back(node);
-		else
-			m_block_queue.add_thread_with_wake_time(node);
+		else if (m_block_queue.add_thread_with_wake_time(node))
+			update_wake_up_deadline();
 
 		if (auto* thread = node->thread; thread->is_userspace() && thread->has_process())
 			thread->update_processor_index_address();

kernel/kernel/Timer/HPET.cpp

@@ -306,7 +306,7 @@ namespace Kernel
 		SystemTimer::get().update_tsc();
 
 		if (should_invoke_scheduler())
-			Processor::scheduler().timer_interrupt();
+			Processor::scheduler().on_timer_interrupt();
 	}
 
 	uint64_t HPET::ms_since_boot() const

kernel/kernel/Timer/PIT.cpp

@@ -61,7 +61,7 @@ namespace Kernel
 		SystemTimer::get().update_tsc();
 
 		if (should_invoke_scheduler())
-			Processor::scheduler().timer_interrupt();
+			Processor::scheduler().on_timer_interrupt();
 	}
 
 	uint64_t PIT::read_counter_ns() const