Hide some task-switching internals inside task_yield.

Signed-off-by: Amy Ringo <me@remexre.com>

Author: remexre

src/kernel/arch/riscv64/task.S

@@ -48,6 +48,9 @@ task_switch:
   li t0, 1
   bne a2, t0, 1f
 
+  // Save the new task to the hart locals.
+  sd a0, 8(a1)
+
   // Get the f_off and f_len. Jump to the appropriate place based on the f_len,
   // to swap the callee-save floating-point registers.
   ld a5, 16(a4) // a5 = f_off
@@ -134,8 +137,10 @@ task_switch:
   la t0, task_resume
   sd t0, 40(a5)
 
-  // Save SSTATUS to the old task.
+  // Save SSTATUS to the old task, setting its SPP bit so that we return to
+  // supervisor mode.
   csrr t0, sstatus
+  ori t0, t0, 1<<8
   sd t0, 296(a5)
 
   // Load SEPC and SSTATUS from the new task.
@@ -207,4 +212,4 @@ task_switch:
 .p2align 2
 .type task_resume, @function
 task_resume:
-	ebreak
+	ret

src/kernel/drivers/riscv_cpu.c

@@ -85,7 +85,7 @@ constexpr u64 TASK_MAGIC = 0x636967614d6f6b75;
 static void task_main(void *main_arg, void (*main)(void *), u64 magic) {
   assert(TASK_MAGIC == magic);
   main(main_arg);
-  TODO("task exit");
+  task_exit();
 }
 
 static void riscv_cpu_set_task_regs(struct hart *hart, struct task *task,

src/kernel/include/scheduler.h

@@ -14,13 +14,12 @@
  *
  * Ensures:
  *
- * - The task is in the list of running tasks.
+ * - If the task is non-`nullptr`, the task is in the list of running tasks.
  */
-struct task *scheduler_get(void);
+struct task *scheduler_get();
 
 /**
- * Puts a task back in the scheduler, allowing other harts to run it later. The
- * task must not be accessed again after calling this function.
+ * Puts a task back in the scheduler, allowing other harts to run it later.
  *
  * Requires:
  *
@@ -29,8 +28,7 @@ struct task *scheduler_get(void);
 void scheduler_put(struct task *task);
 
 /**
- * Puts a new task into the scheduler's list of runnable tasks. The task must
- * not be accessed again afer calling this function.
+ * Puts a new task into the scheduler's list of runnable tasks.
  *
  * Requires:
  *

src/kernel/include/task.h

@@ -106,10 +106,17 @@ static_assert(offsetof(struct task, register_save) == 40);
 struct task *task_new(struct hart *hart, void (*main)(void *), void *main_arg);
 
 /**
- * Switches to the given task.
+ * Considers whether to switch to another task.
  *
- * TODO: This should require that a hartlock be held.
+ * There is no guarantee that a switch will occur; in particular, this does
+ * _not_ necessarily give up the current task's quantum.
  */
-void task_switch(struct task *task);
+void task_yield();
+
+/**
+ * Exits from the current task.
+ */
+[[noreturn]]
+void task_exit();
 
 #endif // UKO_OS_KERNEL__TASK_H

src/kernel/main.c

@@ -71,16 +71,19 @@ void main(u64 hart_id, paddr devicetree_start, paddr kernel_start,
   run_selftests();
 
   // Create another task and switch to it.
-  struct task *task = task_new(get_hart_locals()->hart, gm, (void *)0x12345);
-  scheduler_put_new(task);
-  get_hart_locals()->task->flags = TASK_STATE_RUNNABLE;
-  task->flags = TASK_STATE_RUNNING;
-  task_switch(task);
+  scheduler_put_new(task_new(get_hart_locals()->hart, gm, (void *)0x12345));
+  task_yield();
+  print("back in main");
+  task_yield();
+  print("back in main again");
+  task_yield();
+  print("back in main a last time -- nothing should've switched this time");
 
   TODO();
 }
 
 static void gm(void *ptr) {
   print("GM {uptr}", ptr);
+  task_yield();
   print("GN {uptr}", ptr);
 }

src/kernel/scheduler.c

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: GPL-3.0-or-later
  */
 
+#include <print.h>
 #include <scheduler.h>
 
 /**
@@ -24,6 +25,24 @@ static struct list_head scheduler_runnable = LIST_INIT(scheduler_runnable);
  */
 static struct list_head scheduler_blocked = LIST_INIT(scheduler_blocked);
 
+struct task *scheduler_get() {
+  struct task *task =
+      container_of(list_shift(&scheduler_runnable), struct task, list);
+  assert(task->flags == TASK_STATE_RUNNABLE);
+  task->flags = TASK_STATE_RUNNING;
+  list_push(&scheduler_running, &task->list);
+  return task;
+}
+
+void scheduler_put(struct task *task) {
+  assert(!list_is_empty(&task->list));
+  assert(task->flags == TASK_STATE_RUNNING);
+
+  list_remove(&task->list);
+  task->flags = TASK_STATE_RUNNABLE;
+  list_push(&scheduler_runnable, &task->list);
+}
+
 void scheduler_put_new(struct task *task) {
   assert(list_is_empty(&task->list));
   assert(task->flags == TASK_STATE_CONSTRUCTING);

src/kernel/task.c

@@ -7,9 +7,17 @@
 #include <mm/alloc.h>
 #include <mm/paging.h>
 #include <mm/physical_alloc.h>
+#include <scheduler.h>
 #include <stdatomic.h>
 #include <task.h>
 
+/**
+ * Switches to the given task.
+ *
+ * TODO: This should require that a hartlock be held.
+ */
+void task_switch(struct task *task);
+
 struct task *task_new(struct hart *hart, void (*main)(void *), void *main_arg) {
   assert(hart);
   assert(hart->hart_group);
@@ -54,3 +62,35 @@ struct task *task_new(struct hart *hart, void (*main)(void *), void *main_arg) {
 
   return task;
 }
+
+void task_yield() {
+  // Ask the scheduler for another task to run. If there was nothing to run,
+  // don't try to switch.
+  struct task *new_task = scheduler_get();
+  if (!new_task)
+    return;
+  assert(new_task->flags == TASK_STATE_RUNNING);
+
+  // Get the current task.
+  struct task *old_task = get_hart_locals()->task;
+
+  // Hand the current task back to the scheduler.
+  //
+  // TODO: This _really_ ought to be under some kind of lock that gets released
+  // during the context-switch.
+  assert(old_task->flags == TASK_STATE_RUNNING);
+  scheduler_put(old_task);
+  assert(old_task->flags == TASK_STATE_RUNNABLE);
+
+  // Switch to the new task.
+  task_switch(new_task);
+
+  // Make sure we came back all right.
+  assert(get_hart_locals()->task == old_task);
+  assert(old_task->flags == TASK_STATE_RUNNING);
+}
+
+[[noreturn]]
+void task_exit() {
+  TODO();
+}