From b5d4fc07e2e4319efde004a01b45b29d6cee3c59 Mon Sep 17 00:00:00 2001
From: Matheus <matheusuky@gmail.com>
Date: Wed, 3 May 2023 19:11:40 -0300
Subject: [PATCH] scheduler: have a separate stack for yield

---
 kernel-src/arch/x86-64/context.asm       | 12 +++++++-----
 kernel-src/include/x86-64/arch/context.h |  2 +-
 kernel-src/include/x86-64/arch/cpu.h     |  1 +
 kernel-src/sys/scheduler.c               |  7 ++++++-
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/kernel-src/arch/x86-64/context.asm b/kernel-src/arch/x86-64/context.asm
index b0985bac..4c436d7a 100644
--- a/kernel-src/arch/x86-64/context.asm
+++ b/kernel-src/arch/x86-64/context.asm
@@ -34,17 +34,19 @@ arch_context_switch:
 	o64 iret
 
 ; saves context on the stack and passes as the argument to a function
-; since this uses the C system V abi, some registers can be clobbered (in this case, rsi and r11)
+; since this uses the C system V abi, some registers can be clobbered (in this case, r10 and r11)
 ; rdi has the pointer to the function to call
+; rsi has the pointer to the stack the function should have on call
 ; if the main function returns, just return normally
 global arch_context_saveandcall
 arch_context_saveandcall:
 	mov r11, [rsp] 	; save return address on scratch register
-	mov rsi, rsp
-	add rsi, 8 	; stack pointer before the call instruction
+	mov r10, rsp
+	add r10, 8 	; stack pointer before the call instruction
 
+	mov rsp, rsi	; use desired stack
 	push qword 0x10 ; SS
-	push rsi 	; RSP
+	push r10 	; RSP
 	pushf 		; movs and pushes don't affect flags
 	push qword 0x08 ; CS
 	push r11 	; return address
@@ -73,7 +75,7 @@ arch_context_saveandcall:
 
 	mov r11, rdi 	; save function to call
 	mov rdi, rsp 	; first argument is the context struct
-	push rsi     	; save old stack
+	push r10     	; save old stack
 	call r11 	; jump to the desired function
 
 	pop rsp		; restore old stack
diff --git a/kernel-src/include/x86-64/arch/context.h b/kernel-src/include/x86-64/arch/context.h
index 4220a114..62c11017 100644
--- a/kernel-src/include/x86-64/arch/context.h
+++ b/kernel-src/include/x86-64/arch/context.h
@@ -60,7 +60,7 @@ typedef struct {
 }
 
 void arch_context_switch(context_t *context);
-void arch_context_saveandcall(void (*fn)(context_t *context));
+void arch_context_saveandcall(void (*fn)(context_t *context), void *stack);
 
 #include <string.h>
 
diff --git a/kernel-src/include/x86-64/arch/cpu.h b/kernel-src/include/x86-64/arch/cpu.h
index 00bdf234..0c9cbeb5 100644
--- a/kernel-src/include/x86-64/arch/cpu.h
+++ b/kernel-src/include/x86-64/arch/cpu.h
@@ -24,6 +24,7 @@ typedef struct cpu_t {
 	thread_t *thread;
 	thread_t *idlethread;
 	timerentry_t schedtimerentry;
+	void *schedulerstack;
 } cpu_t;
 
 #define CPU_HALT() asm volatile("hlt")
diff --git a/kernel-src/sys/scheduler.c b/kernel-src/sys/scheduler.c
index 5bb06ef9..09aea85b 100644
--- a/kernel-src/sys/scheduler.c
+++ b/kernel-src/sys/scheduler.c
@@ -7,6 +7,7 @@
 #include <kernel/alloc.h>
 
 #define QUANTUM_US 100000
+#define SCHEDULER_STACK_SIZE 4096
 
 static scache_t *threadcache;
 static scache_t *processcache;
@@ -218,7 +219,7 @@ static void yield(context_t *context) {
 
 void sched_yield() {
 	bool old = interrupt_set(false);
-	arch_context_saveandcall(yield);
+	arch_context_saveandcall(yield, _cpu()->schedulerstack);
 	interrupt_set(old);
 }
 
@@ -258,6 +259,10 @@ void sched_init() {
 	processcache = slab_newcache(sizeof(proc_t), 0, NULL, NULL);
 	__assert(processcache);
 
+	_cpu()->schedulerstack = alloc(SCHEDULER_STACK_SIZE);
+	__assert(_cpu()->schedulerstack);
+	_cpu()->schedulerstack = (void *)((uintptr_t)_cpu()->schedulerstack + SCHEDULER_STACK_SIZE);
+
 	_cpu()->schedtimerentry.func = timerhook;
 	_cpu()->schedtimerentry.repeatus = QUANTUM_US;
 	SPINLOCK_INIT(runqueuelock);