From 7d003da9609a25672afd0a44bf661034d2d90101 Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Tue, 28 Mar 2023 21:39:59 +0100
Subject: [PATCH] sched: add current_task(), current_thread(), and
 preempt_disable/enable()

---
 include/socks/cpu.h       |  6 +++++-
 include/socks/sched.h     |  7 +++++++
 kernel/cpu.c              | 30 ++++++++++++++++++++++++++++++
 sched/{sched.c => core.c} |  4 +++-
 sched/task.c              |  5 +++++
 sched/thread.c            | 16 +++++++++++++++-
 6 files changed, 65 insertions(+), 3 deletions(-)
 rename sched/{sched.c => core.c} (82%)

diff --git a/include/socks/cpu.h b/include/socks/cpu.h
index 375840d..a0e8045 100644
--- a/include/socks/cpu.h
+++ b/include/socks/cpu.h
@@ -15,8 +15,9 @@ typedef enum cpu_flags {
 typedef struct cpu_data {
 	cpu_flags_t c_flags;
 	unsigned int c_id;
+	unsigned int c_preempt_count;
 
-	task_t *c_current_task;
+	thread_t *c_current_thread;
 	runqueue_t c_rq;
 } cpu_data_t;
 
@@ -35,6 +36,9 @@ extern void cpu_set_online(unsigned int cpu_id);
 #define irq_enable() ml_int_enable()
 #define irq_disable() ml_int_disable()
 
+extern void preempt_disable(void);
+extern void preempt_enable(void);
+
 extern unsigned int cpu_get_highest_available(void);
 
 #ifdef __cplusplus
diff --git a/include/socks/sched.h b/include/socks/sched.h
index 8197a59..a2294c6 100644
--- a/include/socks/sched.h
+++ b/include/socks/sched.h
@@ -73,6 +73,9 @@ typedef struct runqueue {
 } runqueue_t;
 
 extern kern_status_t sched_init(void);
+extern void schedule(void);
+extern void preempt_disable(void);
+extern void preempt_enable(void);
 
 extern void runqueue_init(runqueue_t *rq);
 
@@ -82,6 +85,10 @@ static inline void task_deref(task_t *task) { object_deref(object_header(task));
 extern task_t *task_from_pid(unsigned int pid);
 extern task_t *kernel_task(void);
 
+extern bool need_resched(void);
+extern task_t *current_task(void);
+extern thread_t *current_thread(void);
+
 static inline void task_lock_irqsave(task_t *task, unsigned long *flags)
 {
 	object_lock(object_header(task), flags);
diff --git a/kernel/cpu.c b/kernel/cpu.c
index c90a0ea..ceed9da 100644
--- a/kernel/cpu.c
+++ b/kernel/cpu.c
@@ -35,6 +35,36 @@ void cpu_set_online(unsigned int cpu_id)
 	bitmap_set(cpu_online, cpu_id);
 }
 
+void preempt_disable(void)
+{
+	ml_cpu_block *ml_cpu = ml_this_cpu();
+	cpu_data_t *cpu_data = ml_cpu_block_get_data(ml_cpu);
+	if (!cpu_data) {
+		return;
+	}
+
+	unsigned int preempt = READ_ONCE(cpu_data->c_preempt_count);
+	preempt++;
+	WRITE_ONCE(cpu_data->c_preempt_count, preempt);
+}
+
+void preempt_enable(void)
+{
+	ml_cpu_block *ml_cpu = ml_this_cpu();
+	cpu_data_t *cpu_data = ml_cpu_block_get_data(ml_cpu);
+	if (!cpu_data) {
+		return;
+	}
+
+	unsigned int preempt = READ_ONCE(cpu_data->c_preempt_count);
+
+	if (preempt > 0) {
+		preempt--;
+	}
+
+	WRITE_ONCE(cpu_data->c_preempt_count, preempt);
+}
+
 unsigned int cpu_get_highest_available(void)
 {
 	return bitmap_highest_set(cpu_available, CPU_MAX);
diff --git a/sched/sched.c b/sched/core.c
similarity index 82%
rename from sched/sched.c
rename to sched/core.c
index e6dcb8b..5ae953a 100644
--- a/sched/sched.c
+++ b/sched/core.c
@@ -26,9 +26,11 @@ kern_status_t sched_init(void)
 		return status;
 	}
 
+	thread_t *this_thread = QUEUE_CONTAINER(thread_t, tr_threads, queue_first(&kernel_task()->t_threads));
+
 	cpu_data_t *this_cpu = get_this_cpu();
 	runqueue_init(&this_cpu->c_rq);
-	this_cpu->c_current_task = kernel_task();
+	this_cpu->c_current_thread = this_thread;
 	put_cpu(this_cpu);
 
 	printk("sched: initialised");
diff --git a/sched/task.c b/sched/task.c
index 1bcd960..104226e 100644
--- a/sched/task.c
+++ b/sched/task.c
@@ -77,3 +77,8 @@ task_t *task_from_pid(unsigned int pid)
 	spin_unlock_irqrestore(&task_list_lock, flags);
 	return t;
 }
+
+task_t *current_task(void)
+{
+	return current_thread()->tr_parent;
+}
diff --git a/sched/thread.c b/sched/thread.c
index fcf0ed4..79f6ba0 100644
--- a/sched/thread.c
+++ b/sched/thread.c
@@ -1,5 +1,6 @@
 #include <socks/sched.h>
 #include <socks/object.h>
+#include <socks/cpu.h>
 
 static object_type_t thread_type = {
 	.ob_name = "thread",
@@ -17,7 +18,7 @@ thread_t *thread_alloc(void)
 	if (!thread_obj) {
 		return NULL;
 	}
-	
+
 	thread_t *t = object_data(thread_obj);
 	memset(t, 0x00, sizeof *t);
 	return t;
@@ -27,3 +28,16 @@ void thread_free(thread_t *thr)
 {
 
 }
+
+thread_t *current_thread(void)
+{
+	cpu_data_t *cpu = get_this_cpu();
+	thread_t *out = cpu->c_current_thread;
+	put_cpu(cpu);
+	return out;
+}
+
+bool need_resched(void)
+{
+	return (current_thread()->tr_flags & THREAD_F_NEED_RESCHED) != 0;
+}