From 664379b35ca1c9a75b860974c43e67d2b0f6467b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thomas=20Lov=C3=A9n?= <thomasloven@gmail.com>
Date: Fri, 18 Nov 2016 10:36:21 +0100
Subject: [PATCH] [MULTITASKING] Preemptive multitasking

---
 kernel/boot/kmain.c        |  2 --
 kernel/include/registers.h |  3 +++
 kernel/include/thread.h    |  3 +++
 kernel/interrupts/apic.c   |  6 ++----
 kernel/proc/scheduler.c    |  8 +++++++-
 kernel/proc/swtch.S        |  2 ++
 kernel/proc/thread.c       | 10 +++++++++-
 7 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/kernel/boot/kmain.c b/kernel/boot/kmain.c
index 0392381..7898299 100644
--- a/kernel/boot/kmain.c
+++ b/kernel/boot/kmain.c
@@ -18,8 +18,6 @@ void thread_function()
   while(1)
   {
     debug((char *)0x10000);
-    schedule();
-    pit_delay(100);
   }
 }
 
diff --git a/kernel/include/registers.h b/kernel/include/registers.h
index e250a9f..1b7dbce 100644
--- a/kernel/include/registers.h
+++ b/kernel/include/registers.h
@@ -14,3 +14,6 @@ uint64_t read_cr4();
 
 #define CR4_OSFXSR (1<<9)
 #define CR4_OSXMMEXCPT (1<<10)
+
+#define RFLAGS_IF (1<<9)
+
diff --git a/kernel/include/thread.h b/kernel/include/thread.h
index 75d321a..8405d68 100644
--- a/kernel/include/thread.h
+++ b/kernel/include/thread.h
@@ -3,11 +3,13 @@ typedef struct thread_st thread_t;
 #include <stdint.h>
 #include <list.h>
 #include <process.h>
+#include <int.h>
 
 #define THREAD_STACK_SIZE 0x1000-sizeof(thread_t)
 
 typedef struct thread_st
 {
+  registers_t r;
   uint64_t stack_pointer; // Top of the kernel stack for thread
   uint64_t tid;
   uint64_t state;
@@ -32,6 +34,7 @@ typedef struct thread_stack_st
   uint64_t R13;
   uint64_t R14;
   uint64_t R15;
+  uint64_t thread;
   uint64_t zero_frame;
   uint64_t function_address;
   thread_t tcb;
diff --git a/kernel/interrupts/apic.c b/kernel/interrupts/apic.c
index d682c5d..6fab6b7 100644
--- a/kernel/interrupts/apic.c
+++ b/kernel/interrupts/apic.c
@@ -3,6 +3,7 @@
 #include <msr.h>
 #include <debug.h>
 #include <mem.h>
+#include <scheduler.h>
 
 #define APIC_MSR_ENABLE 0x800
 
@@ -62,8 +63,7 @@ registers_t *apic_timer_handler(registers_t *r)
 {
   // APIC timer timeout occurred
   APIC(R_EOI) = 0;
-  debug_putch('-');
-  apic_timer(1000000);
+  schedule();
   return r;
 }
 void apic_timer(uint64_t us)
@@ -112,6 +112,4 @@ void apic_init()
 
   // Register temporary timer handler to go off every 10 ms
   register_int_handler(INT_APIC_TIMER, apic_timer_handler);
-
-  apic_timer(1000000);
 }
diff --git a/kernel/proc/scheduler.c b/kernel/proc/scheduler.c
index 1c78bf8..7a6312c 100644
--- a/kernel/proc/scheduler.c
+++ b/kernel/proc/scheduler.c
@@ -2,6 +2,7 @@
 #include <list.h>
 #include <debug.h>
 #include <sse.h>
+#include <apic.h>
 
 LIST(thread_t, ready_queue);
 
@@ -39,6 +40,7 @@ void scheduler()
 {
   while(1)
   {
+    apic_timer_stop();
     thread_t *old = 0, *new = 0;
     if((old = get_last_thread()))
     {
@@ -51,7 +53,8 @@ void scheduler()
       }
     }
 
-    while(!(new = scheduler_next()));
+    while(!(new = scheduler_next()))asm("sti");
+    asm("cli");
 
     if(!process_alive(new->process))
     {
@@ -62,6 +65,7 @@ void scheduler()
     set_last_thread(new);
     switch_process(new->process);
     sse_restore(new->sse_registers);
+    apic_timer(1000000);
     switch_thread(scheduler_th, new);
   }
 }
@@ -71,6 +75,8 @@ void schedule()
   // This function handles swithing to the next thread in the ready queue
 
   thread_t *old = get_current_thread();
+  if(old == scheduler_th)
+    return;
 
   if(old)
   {
diff --git a/kernel/proc/swtch.S b/kernel/proc/swtch.S
index 84d5226..09dee32 100644
--- a/kernel/proc/swtch.S
+++ b/kernel/proc/swtch.S
@@ -7,6 +7,7 @@ swtch:
 # Switches stacks preserving callee preserved registers according to System V ABI
   push rbp
   mov rbp, rsp
+  push rdi
   push r15
   push r14
   push r13
@@ -24,5 +25,6 @@ swtch:
   pop r13
   pop r14
   pop r15
+  pop rdi
   leaveq
   ret
diff --git a/kernel/proc/thread.c b/kernel/proc/thread.c
index 221cdb4..16a6665 100644
--- a/kernel/proc/thread.c
+++ b/kernel/proc/thread.c
@@ -4,6 +4,7 @@
 #include <mem.h>
 #include <list.h>
 #include <debug.h>
+#include <registers.h>
 
 thread_t *current_thread = 0;
 
@@ -16,9 +17,16 @@ thread_t *new_thread(void (*func)(void))
   thread_stack_t *stack = kcalloc(1, sizeof(thread_stack_t));
   thread_t *th = &stack->tcb;
 
-  stack->function_address = (uint64_t)func;
+  stack->thread = (uint64_t)th;
+  stack->function_address = (uint64_t)isr_return;
   stack->RBP = (uint64_t)&stack->zero_frame;
 
+  th->r.rip = (uint64_t)func;
+  th->r.rflags = RFLAGS_IF;
+  th->r.cs = 0x8;
+  th->r.ss = 0x10;
+  th->r.rsp = (uint64_t)&th->stack_pointer;
+
   th->tid = tid++;
   th->state = THREAD_STATE_READY;
   th->stack_pointer = (uint64_t)&stack->RBP;