sched: add current_task(), current_thread(), and preempt_disable/enable()

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

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);

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);

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");

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;
+}

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",
@@ -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;
+}