From 5f3472f65155406a6b28f037c7016291280c5508 Mon Sep 17 00:00:00 2001
From: Microindole <microindole@gmail.com>
Date: Tue, 12 May 2026 23:30:32 +0800
Subject: [PATCH] sched: track wait queue ownership

---
 include/tianole/sched.h |  2 ++
 kernel/sched/thread.c   |  5 ++++
 kernel/sched/wait.c     | 57 +++++++++++++++++++++++++----------------
 kernel/selftest/sched.c |  4 +++
 4 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/include/tianole/sched.h b/include/tianole/sched.h
index 6941e68..b96134e 100644
--- a/include/tianole/sched.h
+++ b/include/tianole/sched.h
@@ -67,6 +67,7 @@ struct wait_queue {
  * @wake_tick: Timer tick deadline for sleeping threads.
  * @next: Run queue link owned by the scheduler.
  * @wait_next: Wait queue link owned by wait queue code.
+ * @wait_queue: Wait queue currently owning @wait_next, or NULL.
  * @name: Diagnostic thread name.
  *
  * This structure is public only as an early-stage compromise. Long term, most
@@ -84,6 +85,7 @@ struct thread {
 	uint64_t wake_tick;
 	struct thread *next;
 	struct thread *wait_next;
+	struct wait_queue *wait_queue;
 	char name[32];
 };
 
diff --git a/kernel/sched/thread.c b/kernel/sched/thread.c
index c217a71..4269e1e 100644
--- a/kernel/sched/thread.c
+++ b/kernel/sched/thread.c
@@ -86,6 +86,7 @@ struct thread *kernel_thread_create(
 	thread->wake_tick = 0;
 	thread->next = 0;
 	thread->wait_next = 0;
+	thread->wait_queue = 0;
 	copy_thread_name(thread->name, sizeof(thread->name), name);
 
 	spin_lock_irqsave(&scheduler_lock, &flags);
@@ -98,6 +99,10 @@ struct thread *kernel_thread_create(
 
 static void release_thread(struct thread *thread)
 {
+	if (thread->wait_queue != 0) {
+		panic("reaping thread still on wait queue");
+	}
+
 	early_log_puts("thread reaped ");
 	early_log_puts(thread->name);
 	early_log_puts("\n");
diff --git a/kernel/sched/wait.c b/kernel/sched/wait.c
index 0bb65db..8a680c9 100644
--- a/kernel/sched/wait.c
+++ b/kernel/sched/wait.c
@@ -18,7 +18,12 @@ void wait_queue_init(struct wait_queue *queue)
 static void wait_queue_enqueue_locked(
 	struct wait_queue *queue, struct thread *thread)
 {
+	if (thread->wait_queue != 0) {
+		panic("thread already queued on wait queue");
+	}
+
 	thread->wait_next = 0;
+	thread->wait_queue = queue;
 
 	if (queue->tail != 0) {
 		queue->tail->wait_next = thread;
@@ -29,14 +34,36 @@ static void wait_queue_enqueue_locked(
 	queue->tail = thread;
 }
 
-static void wait_queue_remove_locked(
+static struct thread *wait_queue_remove_head_locked(struct wait_queue *queue)
+{
+	struct thread *thread = queue->head;
+
+	if (thread == 0) {
+		return 0;
+	}
+
+	queue->head = thread->wait_next;
+	if (queue->head == 0) {
+		queue->tail = 0;
+	}
+
+	thread->wait_next = 0;
+	thread->wait_queue = 0;
+	return thread;
+}
+
+static int wait_queue_remove_locked(
 	struct wait_queue *queue, struct thread *target)
 {
 	struct thread *prev = 0;
 	struct thread *thread;
 
 	if (queue == 0 || target == 0) {
-		return;
+		return 0;
+	}
+
+	if (target->wait_queue != queue) {
+		return 0;
 	}
 
 	for (thread = queue->head; thread != 0; thread = thread->wait_next) {
@@ -52,11 +79,14 @@ static void wait_queue_remove_locked(
 			}
 
 			thread->wait_next = 0;
-			return;
+			thread->wait_queue = 0;
+			return 1;
 		}
 
 		prev = thread;
 	}
+
+	panic("wait queue membership is inconsistent");
 }
 
 static void wait_queue_mark_ready_locked(struct thread *thread)
@@ -178,18 +208,7 @@ void wait_queue_wake_one(struct wait_queue *queue)
 	}
 
 	spin_lock_irqsave(&queue->lock, &flags);
-	if (queue->head == 0) {
-		spin_unlock_irqrestore(&queue->lock, flags);
-		return;
-	}
-
-	thread = queue->head;
-	queue->head = thread->wait_next;
-	if (queue->head == 0) {
-		queue->tail = 0;
-	}
-
-	thread->wait_next = 0;
+	thread = wait_queue_remove_head_locked(queue);
 	wait_queue_mark_ready_locked(thread);
 	spin_unlock_irqrestore(&queue->lock, flags);
 }
@@ -205,13 +224,7 @@ void wait_queue_wake_all(struct wait_queue *queue)
 
 	spin_lock_irqsave(&queue->lock, &flags);
 	while (queue->head != 0) {
-		thread = queue->head;
-		queue->head = thread->wait_next;
-		if (queue->head == 0) {
-			queue->tail = 0;
-		}
-
-		thread->wait_next = 0;
+		thread = wait_queue_remove_head_locked(queue);
 		wait_queue_mark_ready_locked(thread);
 	}
 	spin_unlock_irqrestore(&queue->lock, flags);
diff --git a/kernel/selftest/sched.c b/kernel/selftest/sched.c
index 0fcfc40..20eb019 100644
--- a/kernel/selftest/sched.c
+++ b/kernel/selftest/sched.c
@@ -72,6 +72,10 @@ void sched_selftest(void)
 		panic("kernel thread selftest state failed");
 	}
 
+	if (first->wait_queue != 0 || second->wait_queue != 0) {
+		panic("kernel thread selftest wait queue ownership failed");
+	}
+
 	if ((first->stack_top & (STACK_ALIGNMENT - 1)) != 0 ||
 		(second->stack_top & (STACK_ALIGNMENT - 1)) != 0) {
 		panic("kernel thread selftest stack alignment failed");