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/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();
 	}