From 20bc9f4a64081fab249de66b8a3cc7382eadc6c5 Mon Sep 17 00:00:00 2001
From: Edgaru089 <sjs333@outlook.com>
Date: Thu, 11 Nov 2021 15:44:59 +0800
Subject: [PATCH] smp: proper cleanup on (non-main) thread exit

---
 kernel/kmain.c    |  5 +++--
 smp/kthread.c     | 21 ++++++++++++++++++++-
 smp/kthread_asm.S | 19 +++++++++++++++++++
 3 files changed, 42 insertions(+), 3 deletions(-)
 create mode 100644 smp/kthread_asm.S

diff --git a/kernel/kmain.c b/kernel/kmain.c
index 0f90b84..9f1d257 100644
--- a/kernel/kmain.c
+++ b/kernel/kmain.c
@@ -39,10 +39,11 @@ static inline int minmax(int val, int min, int max) {
 
 static SYSV_ABI void kThreader(int a, int b, int c, int d, int e, int f) {
 	io_Printf("kThreader: args[%d,%d,%d,%d,%d,%d]\n", a, b, c, d, e, f);
-	for (;;) {
-		smp_thread_Sleep(768);
+	for (int i = 0; i < 10; i++) {
 		io_Printf("kThreader: Not sleeping!\n");
+		smp_thread_Sleep(768);
 	}
+	io_Printf("kThreader: byebye!\n");
 }
 
 
diff --git a/smp/kthread.c b/smp/kthread.c
index 24c3eff..f965e26 100644
--- a/smp/kthread.c
+++ b/smp/kthread.c
@@ -3,6 +3,7 @@
 #include "kthread_switch.h"
 #include "internal.h"
 
+#include "../runtime/panic_assert.h"
 #include "../util/tree.h"
 #include "../interrupt/interrupt.h"
 #include "../driver/irq/pic/rtc/rtc.h"
@@ -10,6 +11,23 @@
 #include "../memory/paging_internal.h"
 
 
+// defined in assembly
+SYSV_ABI void __smp_thread_EntryPoint();
+
+// for __smp_thread_EntryPoint cleanups
+// assumes interrupt is disabled
+SYSV_ABI NORETURN void __smp_thread_Cleanup() {
+	__smp_Thread *t = __smp_Current[0];
+	tree_Delete(__smp_Threads, tree_FindNode(__smp_Threads, t->id));
+	__smp_Current[0] = 0;
+
+	// Now that we're pretending to be a idle thread, this should not return
+	// thread_Switch() sets the interrupt flag on idle for us
+	smp_thread_Yield();
+	Panic("__smp_thread_Cleanup: Yield not switching to another thread");
+}
+
+
 smp_thread_ID smp_thread_Init() {
 	INTERRUPT_DISABLE;
 	__smp_Threads        = tree_Create(sizeof(__smp_Thread));
@@ -51,8 +69,9 @@ smp_thread_ID smp_thread_Start(void *entry, const smp_thread_Arguments *args, un
 
 	t->state.cs  = GDT_EXEC_SELECTOR;
 	t->state.ss  = 0;
-	t->state.rip = (uint64_t)entry;
+	t->state.rip = (uint64_t)__smp_thread_EntryPoint;
 
+	t->state.rax = (uint64_t)entry;
 	t->state.rdi = args->a;
 	t->state.rsi = args->b;
 	t->state.rdx = args->c;
diff --git a/smp/kthread_asm.S b/smp/kthread_asm.S
new file mode 100644
index 0000000..7ad7360
--- /dev/null
+++ b/smp/kthread_asm.S
@@ -0,0 +1,19 @@
+format elf64
+
+extrn  __smp_thread_Cleanup
+public __smp_thread_EntryPoint
+
+
+; sysvx64abi void __smp_thread_EntryPoint(uintptr_t a,b,c,d,e,f)
+;
+; Calls thread entry point in RAX, and, after it returns,
+; removes the current thread from the structure.
+section '.text' executable
+__smp_thread_EntryPoint:
+	call rax
+
+	; Thread returned; delete the current thread
+	; The current thread: In __smp_Threads, not in __smp_ThreadsWaiting
+
+	cli  ; Disable interrupt on thread structure operations
+	jmp __smp_thread_Cleanup  ; This will not return