Kernel: Restructure process and thread termination

When we want to kill a process, we mark its threads as Terminating
or Terminated. If the thread is in critical section that has to be
finished, it will be in Terminating state until done. Once Scheduler
is trying to execute Terminated thread it will instead delete it.

Once processes last thread is marked Terminated, the processes will
turn it into a cleanup thread, that will allow blocks and memory
cleanup to be done.

kernel/include/kernel/Process.h

@@ -41,8 +41,9 @@ namespace Kernel
 		static Process* create_kernel(entry_t, void*);
 		static BAN::ErrorOr<Process*> create_userspace(const Credentials&, BAN::StringView);
 		~Process();
+		void cleanup_function();
 
-		[[noreturn]] void exit(int status);
+		void exit(int status);
 
 		void add_thread(Thread*);
 		void on_thread_exit(Thread&);

kernel/include/kernel/Scheduler.h

@@ -20,9 +20,6 @@ namespace Kernel
 		void reschedule_if_idling();
 
 		void set_current_thread_sleeping(uint64_t);
-		[[noreturn]] void set_current_thread_done();
-
-		[[noreturn]] void set_current_process_done();
 
 		void block_current_thread(Semaphore*);
 		void unblock_threads(Semaphore*);
@@ -33,6 +30,8 @@ namespace Kernel
 
 		static bool is_valid_tid(pid_t tid);
 
+		[[noreturn]] void delete_current_process_and_thread();
+
 	private:
 		Scheduler() = default;
 
@@ -76,6 +75,7 @@ namespace Kernel
 
 		uint64_t m_last_reschedule = 0;
 
+		friend class Thread;
 		friend class Process;
 	};
 

kernel/include/kernel/Thread.h

@@ -27,6 +27,16 @@ namespace Kernel
 			NotStarted,
 			Executing,
 			Terminating,
+			Terminated
+		};
+
+		class TerminateBlocker
+		{
+		public:
+			TerminateBlocker(Thread&);
+			~TerminateBlocker();
+		private:
+			Thread& m_thread;
 		};
 
 	public:
@@ -36,6 +46,7 @@ namespace Kernel
 
 		BAN::ErrorOr<Thread*> clone(Process*, uintptr_t rsp, uintptr_t rip);
 		void setup_exec();
+		void setup_process_cleanup();
 
 		bool has_signal_to_execute() const;
 		void set_signal_done(int signal);
@@ -55,8 +66,8 @@ namespace Kernel
 		uintptr_t rip() const { return m_rip; }
 
 		void set_started() { ASSERT(m_state == State::NotStarted); m_state = State::Executing; }
+		void set_terminating();
 		State state() const { return m_state; }
-		void terminate() { m_state = State::Terminating; }
 
 		vaddr_t stack_base() const { return m_stack->vaddr(); }
 		size_t stack_size() const { return m_stack->size(); }
@@ -103,6 +114,9 @@ namespace Kernel
 		int							m_handling_signal	{ 0 };
 		static_assert(_SIGMAX < 64);
 
+		uint64_t m_terminate_blockers { 0 };
+
+		friend class TerminateBlocker;
 		friend class Scheduler;
 	};
 	

kernel/kernel/Process.cpp

@@ -2,6 +2,7 @@
 #include <BAN/StringView.h>
 #include <kernel/CriticalScope.h>
 #include <kernel/FS/VirtualFileSystem.h>
+#include <kernel/InterruptController.h>
 #include <kernel/LockGuard.h>
 #include <kernel/Memory/Heap.h>
 #include <kernel/Memory/PageTableScope.h>
@@ -108,6 +109,7 @@ namespace Kernel
 		ASSERT(m_fixed_width_allocators.empty());
 		ASSERT(!m_general_allocator);
 		ASSERT(m_mapped_ranges.empty());
+		ASSERT(m_exit_status.waiting == 0);
 		ASSERT(&PageTable::current() != m_page_table.ptr());
 	}
 
@@ -117,20 +119,15 @@ namespace Kernel
 		MUST(m_threads.push_back(thread));
 	}
 
-	void Process::on_thread_exit(Thread& thread)
+	void Process::cleanup_function()
 	{
-		LockGuard _(m_lock);
-		for (size_t i = 0; i < m_threads.size(); i++)
-			if (m_threads[i] == &thread)
-				m_threads.remove(i);
-		if (m_threads.empty())
-			exit(0);
-	}
+		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();
 
-	void Process::exit(int status)
-	{
 		m_lock.lock();
-		m_exit_status.exit_code = status;
 		m_exit_status.exited = true;
 		while (m_exit_status.waiting > 0)
 		{
@@ -141,7 +138,6 @@ namespace Kernel
 			m_lock.lock();
 		}
 
-		m_threads.clear();
 		m_open_file_descriptors.close_all();
 
 		// NOTE: We must unmap ranges while the page table is still alive
@@ -153,21 +149,47 @@ namespace Kernel
 
 		if (m_tty && m_tty->foreground_process() == pid())
 			m_tty->set_foreground_process(0);
+	}
 
-		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();
+	void Process::on_thread_exit(Thread& thread)
+	{
+		ASSERT(!interrupts_enabled());
 
-		// FIXME: we can't assume this is the current process
-		ASSERT(&Process::current() == this);
-		Scheduler::get().set_current_process_done();
+		ASSERT(m_threads.size() > 0);
+
+		if (m_threads.size() == 1)
+		{
+			ASSERT(m_threads.front() == &thread);
+			m_threads.clear();
+
+			thread.setup_process_cleanup();
+			Scheduler::get().execute_current_thread();
+		}
+
+		for (size_t i = 0; i < m_threads.size(); i++)
+		{
+			if (m_threads[i] == &thread)
+			{
+				m_threads.remove(i);
+				return;
+			}
+		}
+
+		ASSERT_NOT_REACHED();
+	}
+
+	void Process::exit(int status)
+	{
+		LockGuard _(m_lock);
+		m_exit_status.exit_code = status;
+		for (auto* thread : m_threads)
+			thread->set_terminating();
 	}
 
 	BAN::ErrorOr<long> Process::sys_exit(int status)
 	{
 		exit(status);
+		Thread::TerminateBlocker _(Thread::current());
 		ASSERT_NOT_REACHED();
 	}
 
@@ -558,6 +580,7 @@ namespace Kernel
 
 	BAN::ErrorOr<long> Process::sys_read(int fd, void* buffer, size_t count)
 	{
+		Thread::TerminateBlocker blocker(Thread::current());
 		LockGuard _(m_lock);
 		return TRY(m_open_file_descriptors.read(fd, buffer, count));
 	}

kernel/kernel/Scheduler.cpp

@@ -192,32 +192,67 @@ namespace Kernel
 		return false;
 	}
 
+	void Scheduler::delete_current_process_and_thread()
+	{
+		DISABLE_INTERRUPTS();
+
+		load_temp_stack();
+		PageTable::kernel().load();
+
+		Thread* thread = m_current_thread->thread;
+
+		ASSERT(thread->has_process());
+		delete &thread->process();
+		
+		remove_and_advance_current_thread();
+
+		delete thread;
+
+		execute_current_thread();
+	}
+
 	void Scheduler::execute_current_thread()
 	{
 		VERIFY_CLI();
 
-		Thread& current = current_thread();
+		load_temp_stack();
+		PageTable::kernel().load();
 
-		if (current.has_process())
+		Thread* current = &current_thread();
+
+		while (current->state() == Thread::State::Terminated)
 		{
-			current.process().page_table().load();
-			GDT::set_tss_stack(current.interrupt_stack_base() + current.interrupt_stack_size());
+			Thread* thread = m_current_thread->thread;
+			if (thread->has_process())
+				thread->process().on_thread_exit(*thread);
+			
+			remove_and_advance_current_thread();
+
+			delete thread;
+
+			current = &current_thread();
+		}
+
+		if (current->has_process())
+		{
+			current->process().page_table().load();
+			GDT::set_tss_stack(current->interrupt_stack_base() + current->interrupt_stack_size());
 		}
 		else
 			PageTable::kernel().load();
 
-		switch (current.state())
+		switch (current->state())
 		{
 			case Thread::State::NotStarted:
-				current.set_started();
-				start_thread(current.rsp(), current.rip());
+				current->set_started();
+				start_thread(current->rsp(), current->rip());
 			case Thread::State::Executing:
-				while (current.has_signal_to_execute())
-					current.handle_next_signal();
-				continue_thread(current.rsp(), current.rip());
+				while (current->has_signal_to_execute() && current->state() == Thread::State::Executing)
+					current->handle_next_signal();
+				// fall through
 			case Thread::State::Terminating:
-				ENABLE_INTERRUPTS();
-				current.on_exit();
+				continue_thread(current->rsp(), current->rip());
+			case Thread::State::Terminated:
 				ASSERT_NOT_REACHED();
 		}
 
@@ -253,60 +288,6 @@ namespace Kernel
 		ASSERT_NOT_REACHED();
 	}
 
-	void Scheduler::set_current_thread_done()
-	{
-		DISABLE_INTERRUPTS();
-
-		load_temp_stack();
-
-		ASSERT(m_current_thread);
-
-		Thread* thread = m_current_thread->thread;
-		remove_and_advance_current_thread();
-		delete thread;
-
-		execute_current_thread();
-		ASSERT_NOT_REACHED();
-	}
-
-#define remove_threads(list, condition)				\
-	for (auto it = list.begin(); it != list.end();) \
-	{												\
-		if (condition)								\
-		{											\
-			delete it->thread;						\
-			it = list.remove(it);					\
-		}											\
-		else										\
-		{											\
-			it++;									\
-		}											\
-	}
-
-	void Scheduler::set_current_process_done()
-	{
-		DISABLE_INTERRUPTS();
-
-		load_temp_stack();
-		PageTable::kernel().load();
-
-		ASSERT(m_current_thread);
-		Thread* thread = m_current_thread->thread;
-		Process* process = &thread->process();
-		remove_threads(m_blocking_threads, it->thread->process().pid() == process->pid());
-		remove_threads(m_sleeping_threads, it->thread->process().pid() == process->pid());
-		remove_threads(m_active_threads, it != m_current_thread && it->thread->process().pid() == process->pid());
-
-		remove_and_advance_current_thread();
-
-		delete thread;
-		delete process;
-
-		execute_current_thread();
-
-		ASSERT_NOT_REACHED();
-	}
-
 	void Scheduler::block_current_thread(Semaphore* semaphore)
 	{
 		VERIFY_STI();

kernel/kernel/Thread.cpp

@@ -25,6 +25,52 @@ namespace Kernel
 			memcpy((void*)rsp, (void*)&value, sizeof(uintptr_t));
 	}
 
+
+	Thread::TerminateBlocker::TerminateBlocker(Thread& thread)
+		: m_thread(thread)
+	{
+		{
+			CriticalScope _;
+
+			if (m_thread.state() == State::Executing || m_thread.m_terminate_blockers > 0)
+			{
+				m_thread.m_terminate_blockers++;
+				return;
+			}
+
+			if (m_thread.state() == State::Terminating && m_thread.m_terminate_blockers == 0)
+				m_thread.m_state = State::Terminated;
+		}
+
+		while (true)
+			Scheduler::get().reschedule();
+	}
+
+	Thread::TerminateBlocker::~TerminateBlocker()
+	{
+		{
+			CriticalScope _;
+	
+			m_thread.m_terminate_blockers--;
+
+			if (m_thread.state() == State::Executing || m_thread.m_terminate_blockers > 0)
+				return;
+
+			if (m_thread.state() == State::Terminating && m_thread.m_terminate_blockers == 0)
+				m_thread.m_state = State::Terminated;
+		}
+
+		while (true)
+			Scheduler::get().reschedule();
+	}
+
+	void Thread::set_terminating()
+	{
+		CriticalScope _;
+		m_state = m_terminate_blockers == 0 ? State::Terminated : State::Terminating;
+		Scheduler::get().unblock_thread(tid());
+	}
+
 	static pid_t s_next_tid = 1;
 
 	BAN::ErrorOr<Thread*> Thread::create_kernel(entry_t entry, void* data, Process* process)
@@ -144,6 +190,32 @@ namespace Kernel
 		}
 	}
 
+	void Thread::setup_process_cleanup()
+	{
+		m_state = State::NotStarted;
+		static entry_t entry(
+			[](void* process)
+			{
+				((Process*)process)->cleanup_function();
+				Scheduler::get().delete_current_process_and_thread();
+				ASSERT_NOT_REACHED();				
+			}
+		);
+		m_rsp = stack_base() + stack_size();
+		m_rip = (uintptr_t)entry;
+
+		m_signal_mask = ~0ull;
+
+		// Setup stack for returning
+		{
+			// FIXME: don't use PageTableScope
+			PageTableScope _(m_process->page_table());
+			write_to_stack(m_rsp, this);
+			write_to_stack(m_rsp, &Thread::on_exit);
+			write_to_stack(m_rsp, m_process);
+		}
+	}
+
 	bool Thread::has_signal_to_execute() const
 	{
 		return !m_signal_queue.empty() && !m_handling_signal;
@@ -227,12 +299,8 @@ namespace Kernel
 				case SIGPROF:
 				case SIGVTALRM:
 				{
-					// NOTE: we cannot have schedulers temporary stack when
-					//       enabling interrupts
-					asm volatile("movq %0, %%rsp" :: "r"(m_return_rsp));
-					ENABLE_INTERRUPTS();
 					process().exit(128 + signal);
-					ASSERT_NOT_REACHED();
+					break;
 				}
 
 				// Ignore the signal
@@ -280,9 +348,8 @@ namespace Kernel
 
 	void Thread::on_exit()
 	{
-		if (m_process)
-			m_process->on_thread_exit(*this);
-		Scheduler::get().set_current_thread_done();
+		set_terminating();
+		TerminateBlocker(*this);
 		ASSERT_NOT_REACHED();
 	}