1 files changed, 41 insertions, 3 deletions

diff --git a/kc/core/task.c b/kc/core/task.c
index 79fae5f..13263d8 100644
--- a/kc/core/task.c
+++ b/kc/core/task.c
@@ -6,9 +6,16 @@
 #include "kprint.h"
 #include "cpu/irq.h"
 
+#include "pit8253.h"
+
+#include <stdatomic.h>
+#include <stdbool.h>
+
 struct kc_thread *current_task = NULL;
 struct kc_thread *first_ready_task = NULL;
 struct kc_thread *last_ready_task = NULL;
+struct kc_thread *sleeping_tasks = NULL;
+struct kc_thread *blocked_tasks = NULL;
 
 static void lock_scheduler(void);
 static void unlock_scheduler(void);
@@ -18,9 +25,14 @@ static void idle_task_thread(void);
 extern uint64_t *get_tss_rsp0(void);
 static uint64_t last_count = 0;
 
+const struct timer_source * timesource;
+
+static atomic_int preempt_switch_count = 0;
+static volatile atomic_bool preempt_switch_flag = false;
+
 void update_time_used(void)
 {
-    uint64_t current_count = timer_ticks();
+    uint64_t current_count = timesource->nanoseconds_elapsed();
     uint64_t elapsed = current_count - last_count;
     last_count = current_count;
     current_task->time_elapsed += elapsed;
@@ -29,13 +41,23 @@ void update_time_used(void)
 void task_set_thread(struct kc_thread *task)
 {
     // wrapper around cpu_set_thread
+    if (preempt_switch_count)
+    {
+        preempt_switch_flag = true;
+        return;
+    }
     uint64_t *rsp0 = get_tss_rsp0();
     cpu_set_thread(&current_task->state, &task->state, rsp0);
 }
 
 void task_schedule(void)
 {
-    kprintf("task_schedule() called @%lu ticks\n", last_count);
+    if (preempt_switch_count)
+    {
+        preempt_switch_flag = true;
+        return;
+    }
+
     update_time_used();
     if (first_ready_task)
     {
@@ -49,10 +71,22 @@ void task_schedule(void)
 static void lock_scheduler(void)
 {
     irq_lock();
+    preempt_switch_count++;
 }
 
 static void unlock_scheduler(void)
 {
+    if (preempt_switch_count >= 1)
+    {
+        preempt_switch_count--;
+    }
+
+    if (!preempt_switch_count && preempt_switch_flag)
+    {
+        preempt_switch_flag = false;
+        task_schedule();
+    }
+
     irq_unlock();
 }
 
@@ -94,8 +128,11 @@ static void idle_task_thread(void)
 noreturn void task_init(void)
 {
     lock_scheduler();
+    timesource = &pit8253_timer_source;
+    timesource->start();
     char *kernel_rsp0 = vm_alloc(4096);
     // TODO: make it so this isn't demand-allocated.
+    // TODO: decouple the task management code from cpu-dependent task code
     kernel_rsp0[1] = 0;
     *get_tss_rsp0() = (uintptr_t)kernel_rsp0;
     char *first_task = vm_alloc(16384);
@@ -104,6 +141,7 @@ noreturn void task_init(void)
     idle_thread->state.stack = (uintptr_t)idle_thread;
     idle_thread->state.stack_top = (uintptr_t)kernel_rsp0;
 
+    // XXX: this is kindof a train wreck
     __asm__ volatile
         (
          "mov %%cr3, %%rax\n\t"
@@ -118,6 +156,6 @@ noreturn void task_init(void)
         );
 
     // shouldn't ever get here
-    panic(GENERAL_PANIC);
+    PANIC(DEAD_END);
 }