From 7a533242ad07933832837a837f575630d85f4bbe Mon Sep 17 00:00:00 2001
From: Matheus <matheusuky@gmail.com>
Date: Wed, 16 Sep 2026 00:00:56 -0300
Subject: [PATCH] scheduler: make idle cpus steal work

---
 kernel-src/include/kernel/scheduler.h |  1 +
 kernel-src/sys/scheduler/queue.c      | 50 +++++++++++++++++++++++++++
 kernel-src/sys/scheduler/scheduler.c  |  8 +++--
 recipes/astral                        |  2 +-
 4 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/kernel-src/include/kernel/scheduler.h b/kernel-src/include/kernel/scheduler.h
index c4ae3f92..c78936c5 100644
--- a/kernel-src/include/kernel/scheduler.h
+++ b/kernel-src/include/kernel/scheduler.h
@@ -58,6 +58,7 @@ thread_t *sched_select_next_thread(void);
 void sched_insert_in_cpu_queue(struct cpu_t *cpu, thread_t *thread);
 void sched_preempt_cpu(struct cpu_t *cpu);
 thread_t *sched_steal_work_from_cpu(struct cpu_t *cpu);
+bool sched_idle_steal_work(void);
 
 static inline int sched_thread_run_queue_index(int interactivity) {
 	return min(interactivity * SCHED_RUN_QUEUE_SIZE / SCHED_MAX_INTERACTIVITY, SCHED_RUN_QUEUE_SIZE - 1);
diff --git a/kernel-src/sys/scheduler/queue.c b/kernel-src/sys/scheduler/queue.c
index 93dc7e3a..5fd352bd 100644
--- a/kernel-src/sys/scheduler/queue.c
+++ b/kernel-src/sys/scheduler/queue.c
@@ -195,11 +195,61 @@ thread_t *sched_steal_work_from_cpu(cpu_t *cpu) {
 	thread = steal_from_run_queue(&cpu->idle_queue);
 	found:
 
+	if (thread == NULL)
+		return NULL;
+
 	--cpu->thread_count;
 	--cpu->stealable_thread_count;
 	return thread;
 }
 
+bool sched_idle_steal_work(void) {
+	bool intstatus = interrupt_set(false);
+	cpu_t *destination = current_cpu();
+	bool work = false;
+
+	spinlock_acquire(&destination->sched_lock);
+	work = destination->thread_count != 0;
+	spinlock_release(&destination->sched_lock);
+	if (work)
+		goto leave;
+
+	size_t cpu_count = arch_smp_get_cpu_count();
+	if (cpu_count == 1 || unlikely(arch_smp_cpusawake != cpu_count))
+		goto leave;
+
+	for (size_t offset = 1; offset < cpu_count; ++offset) {
+		cpu_t *source = smp_cpus[(destination->internal_id + offset) % cpu_count];
+		cpu_t *first = source->internal_id < destination->internal_id ? source : destination;
+		cpu_t *second = first == source ? destination : source;
+
+		if (!spinlock_try(&first->sched_lock))
+			continue;
+
+		if (!spinlock_try(&second->sched_lock)) {
+			spinlock_release(&first->sched_lock);
+			continue;
+		}
+
+		work = destination->thread_count != 0;
+		if (!work && source->stealable_thread_count) {
+			thread_t *thread = sched_steal_work_from_cpu(source);
+			__assert(thread);
+			sched_insert_in_cpu_queue(destination, thread);
+			work = true;
+		}
+
+		spinlock_release(&second->sched_lock);
+		spinlock_release(&first->sched_lock);
+		if (work)
+			break;
+	}
+
+leave:
+	interrupt_set(intstatus);
+	return work;
+}
+
 // requires cpu queue to be locked
 thread_t *sched_select_next_thread(void) {
 	thread_t *thread = pop_from_run_queue(&current_cpu()->rt_queue);
diff --git a/kernel-src/sys/scheduler/scheduler.c b/kernel-src/sys/scheduler/scheduler.c
index 7d7a1ee5..5b0b6a0b 100644
--- a/kernel-src/sys/scheduler/scheduler.c
+++ b/kernel-src/sys/scheduler/scheduler.c
@@ -274,8 +274,12 @@ void sched_preempt_cpu(cpu_t *cpu) {
 
 static void idle_thread(void) {
 	interrupt_set(true);
-	while (1)
-		CPU_HALT();
+	while (1) {
+		if (sched_idle_steal_work())
+			sched_yield();
+		else
+			CPU_HALT();
+	}
 }
 
 void sched_target_cpu(cpu_t *cpu) {
diff --git a/recipes/astral b/recipes/astral
index 09623da2..dd28b71c 100644
--- a/recipes/astral
+++ b/recipes/astral
@@ -1,5 +1,5 @@
 name=astral
-revision=32
+revision=33
 from_source=astral
 imagedeps="build-essential nasm"
 hostdeps="xgcc xbinutils limine"