From 5001fa58e02432cd029edc0aa06e48486b1238a2 Mon Sep 17 00:00:00 2001
From: Bananymous <oskari.alaranta@bananymous.com>
Date: Wed, 24 Jan 2024 14:32:52 +0200
Subject: [PATCH] Kernel: Fix wait syscall with atomics

---
 kernel/include/kernel/Process.h | 11 +++--
 kernel/kernel/Process.cpp       | 82 +++++++++++++++------------------
 2 files changed, 42 insertions(+), 51 deletions(-)

diff --git a/kernel/include/kernel/Process.h b/kernel/include/kernel/Process.h
index f9595689..ae794da5 100644
--- a/kernel/include/kernel/Process.h
+++ b/kernel/include/kernel/Process.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <BAN/Atomic.h>
 #include <BAN/Iteration.h>
 #include <BAN/String.h>
 #include <BAN/StringView.h>
@@ -173,7 +174,7 @@ namespace Kernel
 		// Load elf from a file
 		static BAN::ErrorOr<BAN::UniqPtr<LibELF::LoadableELF>> load_elf_for_exec(const Credentials&, BAN::StringView file_path, const BAN::String& cwd, Kernel::PageTable&);
 
-		int block_until_exit();
+		BAN::ErrorOr<int> block_until_exit(pid_t pid);
 
 		BAN::ErrorOr<BAN::String> absolute_path_of(BAN::StringView) const;
 
@@ -183,10 +184,10 @@ namespace Kernel
 	private:
 		struct ExitStatus
 		{
-			Semaphore semaphore;
-			int exit_code { 0 };
-			bool exited { false };
-			int waiting { 0 };
+			Semaphore			semaphore;
+			int					exit_code	{ 0 };
+			BAN::Atomic<bool>	exited		{ false };
+			BAN::Atomic<int>	waiting		{ 0 };
 		};
 
 		Credentials m_credentials;
diff --git a/kernel/kernel/Process.cpp b/kernel/kernel/Process.cpp
index da34db89..63e02f95 100644
--- a/kernel/kernel/Process.cpp
+++ b/kernel/kernel/Process.cpp
@@ -202,24 +202,22 @@ namespace Kernel
 
 	void Process::cleanup_function()
 	{
-		s_process_lock.lock();
-		for (size_t i = 0; i < s_processes.size(); i++)
-			if (s_processes[i] == this)
-				s_processes.remove(i);
-		s_process_lock.unlock();
+		{
+			LockGuard _(s_process_lock);
+			for (size_t i = 0; i < s_processes.size(); i++)
+				if (s_processes[i] == this)
+					s_processes.remove(i);
+		}
 
 		ProcFileSystem::get().on_process_delete(*this);
 
-		m_lock.lock();
 		m_exit_status.exited = true;
+		m_exit_status.semaphore.unblock();
+
 		while (m_exit_status.waiting > 0)
-		{
-			m_exit_status.semaphore.unblock();
-			while (m_lock.is_locked())
-				m_lock.unlock();
 			Scheduler::get().reschedule();
-			m_lock.lock();
-		}
+
+		m_lock.lock();
 
 		m_open_file_descriptors.close_all();
 
@@ -241,6 +239,7 @@ namespace Kernel
 
 			thread.setup_process_cleanup();
 			Scheduler::get().execute_current_thread();
+			ASSERT_NOT_REACHED();
 		}
 
 		for (size_t i = 0; i < m_threads.size(); i++)
@@ -547,44 +546,17 @@ namespace Kernel
 		ASSERT_NOT_REACHED();
 	}
 
-	int Process::block_until_exit()
+	BAN::ErrorOr<int> Process::block_until_exit(pid_t pid)
 	{
-		ASSERT(this != &Process::current());
+		ASSERT(this->pid() != pid);
 
-		m_lock.lock();
-		m_exit_status.waiting++;
-		while (!m_exit_status.exited)
-		{
-			m_lock.unlock();
-			m_exit_status.semaphore.block();
-			m_lock.lock();
-		}
-
-		int ret = m_exit_status.exit_code;
-		m_exit_status.waiting--;
-		m_lock.unlock();
-
-		return ret;
-	}
-
-	BAN::ErrorOr<long> Process::sys_wait(pid_t pid, int* stat_loc, int options)
-	{
 		Process* target = nullptr;
-
-		{
-			LockGuard _(m_lock);
-			TRY(validate_pointer_access(stat_loc, sizeof(int)));
-		}
-
-		// FIXME: support options
-		if (options)
-			return BAN::Error::from_errno(EINVAL);
-
 		for_each_process(
-			[&](Process& process)
+			[pid, &target](Process& process)
 			{
 				if (process.pid() == pid)
 				{
+					process.m_exit_status.waiting++;
 					target = &process;
 					return BAN::Iteration::Break;
 				}
@@ -595,13 +567,31 @@ namespace Kernel
 		if (target == nullptr)
 			return BAN::Error::from_errno(ECHILD);
 
-		pid_t ret = target->pid();
+		while (!target->m_exit_status.exited)
+			TRY(Thread::current().block_or_eintr(target->m_exit_status.semaphore));
 
-		int stat = target->block_until_exit();
+		int exit_status = target->m_exit_status.exit_code;
+		target->m_exit_status.waiting--;
+
+		return exit_status;
+	}
+
+	BAN::ErrorOr<long> Process::sys_wait(pid_t pid, int* stat_loc, int options)
+	{
+		{
+			LockGuard _(m_lock);
+			TRY(validate_pointer_access(stat_loc, sizeof(int)));
+		}
+
+		// FIXME: support options
+		if (options)
+			return BAN::Error::from_errno(EINVAL);
+
+		int stat = TRY(block_until_exit(pid));
 		if (stat_loc)
 			*stat_loc = stat;
 
-		return ret;
+		return pid;
 	}
 
 	BAN::ErrorOr<long> Process::sys_sleep(int seconds)