Kernel: Initial work on new scheduler with queues

Sleeping is definately broken

kernel/include/kernel/Attributes.h

@@ -0,0 +1,3 @@
+#pragma once
+
+#define ALWAYS_INLINE inline __attribute__((always_inline))

kernel/include/kernel/Scheduler.h

@@ -1,7 +1,7 @@
 #pragma once
 
-#include <BAN/Function.h>
 #include <BAN/LinkedList.h>
+#include <BAN/Memory.h>
 #include <kernel/Thread.h>
 
 namespace Kernel
@@ -9,33 +9,48 @@ namespace Kernel
 
 	class Scheduler
 	{
-		BAN_NON_COPYABLE(Scheduler);
-		BAN_NON_MOVABLE(Scheduler);
-		
 	public:
-		static void initialize();
+		static BAN::ErrorOr<void> initialize();
 		static Scheduler& get();
 
-		const Thread& current_thread() const;
-		
-		BAN::ErrorOr<void> add_thread(const BAN::Function<void()>& function);
-
-		void reschedule();
-		void set_current_thread_sleeping();
 		void start();
+		void reschedule();
 
-		static constexpr size_t ms_between_switch = 4;
+		BAN::ErrorOr<void> add_thread(BAN::RefCounted<Thread>);
+
+		void set_current_thread_sleeping(uint64_t);
+		[[noreturn]] void set_current_thread_done();
 
 	private:
-		Scheduler() {}
-		void switch_thread();
+		Scheduler() = default;
+
+		BAN::RefCounted<Thread> current_thread();
+
+		void wake_threads();
+		[[nodiscard]] bool save_current_thread();
+		void get_next_thread();
+		[[noreturn]] void execute_current_thread();
 
 	private:
-		BAN::LinkedList<Thread>				m_threads;
-		BAN::LinkedList<Thread>::iterator	m_current_iterator;
+		struct ActiveThread
+		{
+			BAN::RefCounted<Thread> thread;
+			uint64_t padding;
+		};
+
+		struct SleepingThread
+		{
+			BAN::RefCounted<Thread> thread;
+			uint64_t wake_delta;
+		};
+
+		BAN::RefCounted<Thread> m_idle_thread;
+		BAN::LinkedList<ActiveThread> m_active_threads;
+		BAN::LinkedList<SleepingThread> m_sleeping_threads;
+
+		BAN::LinkedList<ActiveThread>::iterator m_current_thread;
+
 		uint64_t m_last_reschedule = 0;
-
-		friend class Thread;
 	};
 
 }

kernel/include/kernel/Thread.h

@@ -12,41 +12,35 @@ namespace Kernel
 		BAN_NON_MOVABLE(Thread);
 
 	public:
-		enum class State
-		{
-			NotStarted,
-			Running,
-			Paused,
-			Sleeping,
-			Done,
-		};
-
-	public:
-		Thread(const BAN::Function<void()>&);
+		static BAN::ErrorOr<BAN::RefCounted<Thread>> create(const BAN::Function<void()>&);
 		~Thread();
 
 		uint32_t tid() const { return m_tid; }
 
 		void set_rsp(uintptr_t rsp) { m_rsp = rsp; }
 		void set_rip(uintptr_t rip) { m_rip = rip; }
-		void set_state(State state) { m_state = state; }
 		uintptr_t rsp() const { return m_rsp; }
 		uintptr_t rip() const { return m_rip; }
-		State state() const { return m_state; }
+
+		void set_started() { m_started = true; }
+		bool started() const { return m_started; }
 
 		const BAN::Function<void()>* function() const { return &m_function; }
 
 	private:
+		Thread(const BAN::Function<void()>&);
 		void on_exit();
 		
 	private:
 		void*			m_stack_base	= nullptr;
-		State			m_state			= State::NotStarted;
 		uintptr_t		m_rip			= 0;
 		uintptr_t		m_rsp			= 0;
 		const uint32_t	m_tid			= 0;
+		bool			m_started		= false;
 		
 		BAN::Function<void()> m_function;
+
+		friend class BAN::RefCounted<Thread>;
 	};
 
 }

kernel/kernel/PIT.cpp

@@ -23,6 +23,7 @@
 
 namespace PIT
 {
+
 	static uint64_t s_system_time = 0;
 
 	void irq_handler()
@@ -52,9 +53,9 @@ namespace PIT
 
 	void sleep(uint64_t ms)
 	{
-		uint64_t end = s_system_time + ms;
-		while (s_system_time < end)
-			Kernel::Scheduler::get().set_current_thread_sleeping();
+		if (ms == 0)
+			return;
+		Kernel::Scheduler::get().set_current_thread_sleeping(ms);
 	}
 
 }

kernel/kernel/Scheduler.cpp

@@ -1,22 +1,42 @@
 #include <kernel/Arch.h>
+#include <kernel/Attributes.h>
 #include <kernel/InterruptController.h>
 #include <kernel/Scheduler.h>
 
+#include <kernel/PCI.h>
+
+#define DISABLE_INTERRUPTS() asm volatile("cli")
+#define ENABLE_INTERRUPTS() asm volatile("sti")
+
+#if 1
+	#define VERIFY_CLI() ASSERT(interrupts_disabled())
+#else
+	#define VERIFY_CLI()
+#endif
+
 namespace Kernel
 {
 
-	static Scheduler* s_instance = nullptr;
-
 	extern "C" void start_thread(const BAN::Function<void()>* function, uintptr_t rsp, uintptr_t rip);
 	extern "C" void continue_thread(uintptr_t rsp, uintptr_t rip);
 	extern "C" uintptr_t read_rip();
 
-	void Scheduler::initialize()
+	static Scheduler* s_instance = nullptr;
+
+	static bool interrupts_disabled()
 	{
-		ASSERT(!s_instance);
+		uintptr_t flags;
+		asm volatile("pushf; pop %0" : "=r"(flags));
+		return !(flags & (1 << 9));
+	}
+
+	BAN::ErrorOr<void> Scheduler::initialize()
+	{
+		ASSERT(s_instance == nullptr);
 		s_instance = new Scheduler();
 		ASSERT(s_instance);
-		MUST(s_instance->add_thread(BAN::Function<void()>([] { for(;;) asm volatile("hlt"); })));
+		s_instance->m_idle_thread = TRY(Thread::create([] { for (;;) asm volatile("hlt"); }));
+		return {};
 	}
 
 	Scheduler& Scheduler::get()
@@ -25,21 +45,18 @@ namespace Kernel
 		return *s_instance;
 	}
 
-	const Thread& Scheduler::current_thread() const
+	void Scheduler::start()
 	{
-		return *m_current_iterator;
+		VERIFY_CLI();
+		ASSERT(!m_active_threads.empty());
+		m_current_thread = m_active_threads.begin();
+		execute_current_thread();
+		ASSERT_NOT_REACHED();
 	}
 
-
-	BAN::ErrorOr<void> Scheduler::add_thread(const BAN::Function<void()>& function)
+	BAN::RefCounted<Thread> Scheduler::current_thread()
 	{
-		uintptr_t flags;
-		asm volatile("pushf; pop %0" : "=r"(flags));
-		asm volatile("cli");
-		TRY(m_threads.emplace_back(function));
-		if (flags & (1 << 9))
-			asm volatile("sti");
-		return {};
+		return m_current_thread ? m_current_thread->thread : m_idle_thread;
 	}
 
 	void Scheduler::reschedule()
@@ -47,94 +64,151 @@ namespace Kernel
 		ASSERT(InterruptController::get().is_in_service(PIT_IRQ));
 		InterruptController::get().eoi(PIT_IRQ);
 
-		uint64_t current_time = PIT::ms_since_boot();
-		if (m_last_reschedule + ms_between_switch > current_time)
+		if (PIT::ms_since_boot() <= m_last_reschedule)
 			return;
-		m_last_reschedule = current_time;
+		m_last_reschedule = PIT::ms_since_boot();
 		
-		for (Thread& thread : m_threads)
-			if (thread.state() == Thread::State::Sleeping)
-				thread.set_state(Thread::State::Paused);
+		if (!m_sleeping_threads.empty())
+			m_sleeping_threads.front().wake_delta--;
+		wake_threads();
 
-		switch_thread();
+		if (save_current_thread())
+			return;
+		get_next_thread();
+		execute_current_thread();
+		ASSERT_NOT_REACHED();
 	}
 
-	void Scheduler::switch_thread()
+	void Scheduler::wake_threads()
 	{
+		VERIFY_CLI();
+
+		while (!m_sleeping_threads.empty() && m_sleeping_threads.front().wake_delta == 0)
+		{
+			auto thread = m_sleeping_threads.front().thread;
+			m_sleeping_threads.remove(m_sleeping_threads.begin());
+
+			// This should work as we released enough memory from sleeping thread
+			static_assert(sizeof(ActiveThread) == sizeof(SleepingThread));
+			MUST(m_active_threads.push_back({ thread, 0 }));
+		}
+	}
+
+	BAN::ErrorOr<void> Scheduler::add_thread(BAN::RefCounted<Thread> thread)
+	{
+		if (interrupts_disabled())
+		{
+			TRY(m_active_threads.push_back({ thread, 0 }));
+		}
+		else
+		{
+			DISABLE_INTERRUPTS();
+			TRY(m_active_threads.push_back({ thread, 0 }));
+			ENABLE_INTERRUPTS();
+		}
+		return {};
+	}
+
+	void Scheduler::get_next_thread()
+	{
+		VERIFY_CLI();
+
+		if (m_active_threads.empty())
+		{
+			m_current_thread = {};
+			return;
+		}
+
+		if (!m_current_thread || ++m_current_thread == m_active_threads.end())
+			m_current_thread = m_active_threads.begin();
+	}
+
+	// NOTE: this is declared always inline, so we don't corrupt the stack
+	//       after getting the rsp
+	ALWAYS_INLINE bool Scheduler::save_current_thread()
+	{
+		VERIFY_CLI();
+
 		uintptr_t rsp, rip;
 		push_callee_saved();
 		if (!(rip = read_rip()))
 		{
 			pop_callee_saved();
-			return;
+			return true;
 		}
 		read_rsp(rsp);
 
-		ASSERT(m_threads.size() > 1);
+		auto current = current_thread();
+		current->set_rip(rip);
+		current->set_rsp(rsp);
+		return false;
+	}
 
-		Thread& current = *m_current_iterator;
-
-		if (current.state() == Thread::State::Done)
+	void Scheduler::execute_current_thread()
 	{
-			m_threads.remove(m_current_iterator);
-			m_current_iterator = m_threads.begin();
+		VERIFY_CLI();
+
+		auto current = current_thread();
+
+		if (current->started())
+		{
+			continue_thread(current->rsp(), current->rip());
 		}
 		else
 		{
-			current.set_rsp(rsp);
-			current.set_rip(rip);
-			if (current.state() != Thread::State::Sleeping)
-				current.set_state(Thread::State::Paused);
+			current->set_started();
+			start_thread(current->function(), current->rsp(), current->rip());
 		}
 
-		auto next_iterator = m_current_iterator;
-		if (++next_iterator == m_threads.end())
-			next_iterator = ++m_threads.begin();
-		if (next_iterator->state() == Thread::State::Sleeping)
-			next_iterator = m_threads.begin();
-		Thread& next = *next_iterator;
+		ASSERT_NOT_REACHED();
+	}
 
-		m_current_iterator = next_iterator;
-
-		switch (next.state())
+	void Scheduler::set_current_thread_sleeping(uint64_t wake_delta)
 	{
-			case Thread::State::NotStarted:
-				next.set_state(Thread::State::Running);
-				start_thread(next.function(), next.rsp(), next.rip());
+		DISABLE_INTERRUPTS();
+
+		ASSERT(m_current_thread);
+
+		auto current = m_current_thread->thread;
+
+		auto temp = m_current_thread;
+		if (save_current_thread())
+			return;
+		get_next_thread();
+		m_active_threads.remove(temp);
+
+		auto it = m_sleeping_threads.begin();
+
+		for (; it != m_sleeping_threads.end(); it++)
+		{
+			if (wake_delta <= it->wake_delta)
 				break;
-			case Thread::State::Paused:
-				next.set_state(Thread::State::Running);
-				continue_thread(next.rsp(), next.rip());
-				break;
-			case Thread::State::Sleeping:	ASSERT(false);
-			case Thread::State::Running:	ASSERT(false);
-			case Thread::State::Done:		ASSERT(false);
+			wake_delta -= it->wake_delta;
 		}
 
-		ASSERT(false);
+		if (it != m_sleeping_threads.end())
+			it->wake_delta -= wake_delta;
+
+		// This should work as we released enough memory from active thread
+		static_assert(sizeof(ActiveThread) == sizeof(SleepingThread));
+		MUST(m_sleeping_threads.insert(it, { current, wake_delta }));
+
+		execute_current_thread();
+		ASSERT_NOT_REACHED();
 	}
 
-	void Scheduler::set_current_thread_sleeping()
+	void Scheduler::set_current_thread_done()
 	{
-		asm volatile("cli");
-		m_current_iterator->set_state(Thread::State::Sleeping);
-		switch_thread();
-		asm volatile("sti");
-	}
+		DISABLE_INTERRUPTS();
 
-	void Scheduler::start()
-	{
-		ASSERT(m_threads.size() > 1);
+		ASSERT(m_current_thread);
 
-		m_current_iterator = m_threads.begin();
+		auto temp = m_current_thread;
+		get_next_thread();
+		m_active_threads.remove(temp);
 	
-		Thread& current = *m_current_iterator;
-		ASSERT(current.state() == Thread::State::NotStarted);
-		current.set_state(Thread::State::Running);
-
-		start_thread(current.function(), current.rsp(), current.rip());
-
-		ASSERT(false);
+		execute_current_thread();
+		ASSERT_NOT_REACHED();
 	}
 
 }

kernel/kernel/Shell.cpp

@@ -59,7 +59,8 @@ namespace Kernel
 		TTY_PRINT("{}", m_prompt);
 		for (;;)
 		{
-			Scheduler::get().set_current_thread_sleeping();
+			//PIT::sleep(1); // sleep until next reschedule
+			asm volatile("hlt");
 			Input::update();
 		}
 	}
@@ -182,7 +183,7 @@ argument_done:
 
 			s_thread_spinlock.lock();
 
-			MUST(Scheduler::get().add_thread(Function<void()>(
+			MUST(Scheduler::get().add_thread(MUST(Thread::create(
 				[this, &arguments]
 				{
 					auto args = arguments;
@@ -191,7 +192,7 @@ argument_done:
 					PIT::sleep(5000);
 					process_command(args);
 				}
-			)));
+			))));
 
 			while (s_thread_spinlock.is_locked());
 		}

kernel/kernel/Thread.cpp

@@ -21,6 +21,12 @@ namespace Kernel
 		memcpy((void*)rsp, (void*)&value, size);
 	}
 
+
+	BAN::ErrorOr<BAN::RefCounted<Thread>> Thread::create(const BAN::Function<void()>& function)
+	{
+		return BAN::RefCounted<Thread>::create(function);
+	}
+
 	Thread::Thread(const BAN::Function<void()>& function)
 		: m_tid(s_next_tid++)
 		, m_function(function)
@@ -46,10 +52,8 @@ namespace Kernel
 
 	void Thread::on_exit()
 	{
-		asm volatile("cli");
-		m_state = State::Done;
-		Scheduler::get().switch_thread();
-		ASSERT(false);
+		Scheduler::get().set_current_thread_done();
+		ASSERT_NOT_REACHED();
 	}
 
 }

kernel/kernel/kernel.cpp

@@ -104,9 +104,10 @@ extern "C" void kernel_main()
 		dprintln("Could not initialize input drivers");
 	dprintln("Input initialized");
 
-	Scheduler::initialize();
+	MUST(Scheduler::initialize());
 	Scheduler& scheduler = Scheduler::get();
-	MUST(scheduler.add_thread(BAN::Function<void()>(
+	#if 1
+	MUST(scheduler.add_thread(MUST(Thread::create(
 		[terminal_driver]
 		{
 			if (auto error = VirtualFileSystem::initialize(); error.is_error())
@@ -121,15 +122,16 @@ extern "C" void kernel_main()
 			else
 				terminal_driver->set_font(font_or_error.release_value());
 		}
-	)));
-	MUST(scheduler.add_thread(BAN::Function<void()>(
+	))));
+	#endif
+	MUST(scheduler.add_thread(MUST(Thread::create(
 		[tty1]
 		{
 			Shell* shell = new Shell(tty1);
 			ASSERT(shell);
 			shell->run();
 		}
-	)));
+	))));
 	scheduler.start();
 	ASSERT(false);
 }