From f924ac9265257735c7a0baed03fd7f61e9201472 Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Sun, 9 Apr 2023 22:18:12 +0300
Subject: [PATCH] Kernel: Threads can now be terminated mid execution

---
 kernel/arch/i386/Thread.S       |  1 -
 kernel/arch/x86_64/Thread.S     |  1 -
 kernel/include/kernel/Process.h |  6 ++++--
 kernel/include/kernel/Thread.h  | 16 +++++++++++++---
 kernel/kernel/Process.cpp       | 22 +++++++++++++++++++---
 kernel/kernel/Scheduler.cpp     | 21 ++++++++++++---------
 6 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/kernel/arch/i386/Thread.S b/kernel/arch/i386/Thread.S
index e31f378729..af2cb91c22 100644
--- a/kernel/arch/i386/Thread.S
+++ b/kernel/arch/i386/Thread.S
@@ -15,7 +15,6 @@ start_thread:
 	movl 8(%esp), %ecx
 	movl 4(%esp), %esp
 	movl $0, %ebp
-	movl $0, (%esp)
 	pushl $exit_thread_trampoline
 	sti
 	jmp *%ecx
diff --git a/kernel/arch/x86_64/Thread.S b/kernel/arch/x86_64/Thread.S
index a75600c930..3c72cb179d 100644
--- a/kernel/arch/x86_64/Thread.S
+++ b/kernel/arch/x86_64/Thread.S
@@ -14,7 +14,6 @@ start_thread:
 	movq %rdi, %rsp
 	popq %rdi
 	movq $0, %rbp
-	movq $0, (%rsp)
 	pushq $exit_thread_trampoline
 	sti
 	jmp *%rsi
diff --git a/kernel/include/kernel/Process.h b/kernel/include/kernel/Process.h
index 7713afaf4e..11a45d01d3 100644
--- a/kernel/include/kernel/Process.h
+++ b/kernel/include/kernel/Process.h
@@ -23,9 +23,11 @@ namespace Kernel
 
 	public:
 		static BAN::ErrorOr<BAN::RefPtr<Process>> create_kernel(entry_t, void*);
-		~Process() {}
+		~Process();
 
-		BAN::ErrorOr<void> add_thread(entry_t, void*);
+		void exit();
+
+		BAN::ErrorOr<Thread*> add_thread(entry_t, void*);
 		void on_thread_exit(Thread&);
 
 		BAN::ErrorOr<void> init_stdio();
diff --git a/kernel/include/kernel/Thread.h b/kernel/include/kernel/Thread.h
index 5e4139e3f6..c401dcc3b9 100644
--- a/kernel/include/kernel/Thread.h
+++ b/kernel/include/kernel/Thread.h
@@ -18,6 +18,13 @@ namespace Kernel
 	public:
 		using entry_t = void(*)(void*);
 
+		enum class State
+		{
+			NotStarted,
+			Executing,
+			Terminating,
+		};
+
 	public:
 		static BAN::ErrorOr<Thread*> create(entry_t, void* = nullptr, BAN::RefPtr<Process> = nullptr);
 		~Thread();
@@ -29,8 +36,9 @@ namespace Kernel
 		uintptr_t rsp() const { return m_rsp; }
 		uintptr_t rip() const { return m_rip; }
 
-		void set_started() { m_started = true; }
-		bool started() const { return m_started; }
+		void set_started() { ASSERT(m_state == State::NotStarted); m_state = State::Executing; }
+		State state() const { return m_state; }
+		void terminate() { m_state = State::Terminating; }
 
 		static Thread& current() ;
 		BAN::RefPtr<Process> process();
@@ -46,8 +54,10 @@ namespace Kernel
 		uintptr_t	m_rip			= 0;
 		uintptr_t	m_rsp			= 0;
 		const pid_t	m_tid			= 0;
-		bool		m_started		= false;
+		State		m_state { State::NotStarted };
 		BAN::RefPtr<Process> m_process;
+
+		friend class Scheduler;
 	};
 	
 }
\ No newline at end of file
diff --git a/kernel/kernel/Process.cpp b/kernel/kernel/Process.cpp
index 62f2ceeeee..e6f95156fa 100644
--- a/kernel/kernel/Process.cpp
+++ b/kernel/kernel/Process.cpp
@@ -23,7 +23,12 @@ namespace Kernel
 		, m_tty(TTY::current())
 	{ }
 
-	BAN::ErrorOr<void> Process::add_thread(entry_t entry, void* data)
+	Process::~Process()
+	{
+		exit();
+	}
+
+	BAN::ErrorOr<Thread*> Process::add_thread(entry_t entry, void* data)
 	{
 		Thread* thread = TRY(Thread::create(entry, data, this));
 
@@ -32,10 +37,10 @@ namespace Kernel
 		if (auto res = Scheduler::get().add_thread(thread); res.is_error())
 		{
 			m_threads.pop_back();
-			return res;
+			return res.release_error();
 		}
 
-		return {};
+		return thread;
 	}
 
 	void Process::on_thread_exit(Thread& thread)
@@ -46,6 +51,17 @@ namespace Kernel
 				m_threads.remove(i);
 	}
 
+	void Process::exit()
+	{
+		LockGuard _(m_lock);
+		for (auto* thread : m_threads)
+			thread->terminate();
+		while (!m_threads.empty())
+			PIT::sleep(1);
+		for (auto& open_fd : m_open_files)
+			open_fd.inode = nullptr;
+	}
+
 	BAN::ErrorOr<void> Process::init_stdio()
 	{
 		if (!m_open_files.empty())
diff --git a/kernel/kernel/Scheduler.cpp b/kernel/kernel/Scheduler.cpp
index df089066db..728503cfac 100644
--- a/kernel/kernel/Scheduler.cpp
+++ b/kernel/kernel/Scheduler.cpp
@@ -16,8 +16,8 @@
 namespace Kernel
 {
 
-	extern "C" void start_thread(uintptr_t rsp, uintptr_t rip);
-	extern "C" void continue_thread(uintptr_t rsp, uintptr_t rip);
+	extern "C" [[noreturn]] void start_thread(uintptr_t rsp, uintptr_t rip);
+	extern "C" [[noreturn]] void continue_thread(uintptr_t rsp, uintptr_t rip);
 	extern "C" uintptr_t read_rip();
 
 	static Scheduler* s_instance = nullptr;
@@ -166,14 +166,17 @@ namespace Kernel
 		
 		Thread& current = current_thread();
 
-		if (current.started())
+		switch (current.state())
 		{
-			continue_thread(current.rsp(), current.rip());
-		}
-		else
-		{
-			current.set_started();
-			start_thread(current.rsp(), current.rip());
+			case Thread::State::NotStarted:
+				current.set_started();
+				start_thread(current.rsp(), current.rip());
+			case Thread::State::Executing:
+				continue_thread(current.rsp(), current.rip());
+			case Thread::State::Terminating:
+				ENABLE_INTERRUPTS();
+				current.on_exit();
+				ASSERT_NOT_REACHED();
 		}
 
 		ASSERT_NOT_REACHED();