Luna/kernel/src/thread/Scheduler.cpp

#include "thread/Scheduler.h"
#include "ELF.h"
#include "Log.h"
#include "arch/CPU.h"
#include "arch/MMU.h"
#include "memory/MemoryManager.h"
#include <luna/Alignment.h>
#include <luna/ScopeGuard.h>
#include <luna/Stack.h>

static Thread g_idle;
static Thread* g_current = nullptr;

static const usize TICKS_PER_TIMESLICE = 20;

namespace Scheduler
{
    void init()
    {
        g_idle.id = 0;
        g_idle.init_regs_kernel();
        g_idle.set_ip((u64)CPU::idle_loop);
        g_idle.state = ThreadState::Idle;
        g_idle.is_kernel = true;

        g_idle.ticks_left = 1;

        // Map some stack for the idle task
        u64 idle_stack_vm = MemoryManager::alloc_for_kernel(1, MMU::NoExecute | MMU::ReadWrite)
                                .expect_value("Error while setting up the idle task, cannot continue");

        Stack idle_stack { idle_stack_vm, ARCH_PAGE_SIZE };
        g_idle.set_sp(idle_stack.top());

        g_idle.stack = idle_stack;

        kinfoln("CREATED IDLE THREAD: id %lu with ip %#lx and sp %#lx", g_idle.id, g_idle.ip(), g_idle.sp());

        g_current = &g_idle;
    }

    Thread* current()
    {
        return g_current;
    }

    Thread* idle()
    {
        return &g_idle;
    }

    Result<void> new_kernel_thread_impl(Thread* thread)
    {
        // If anything fails, make sure to clean up.
        auto guard = make_scope_guard([&] { delete thread; });

        const u64 thread_stack_vm = TRY(MemoryManager::alloc_for_kernel(4, MMU::NoExecute | MMU::ReadWrite));

        guard.deactivate();

        const Stack thread_stack { thread_stack_vm, ARCH_PAGE_SIZE * 4 };
        thread->set_sp(thread_stack.top());

        thread->stack = thread_stack;

        thread->is_kernel = true;

        g_threads.append(thread);

        kinfoln("CREATED THREAD: id %lu with ip %#lx and sp %#lx", thread->id, thread->ip(), thread->sp());

        return {};
    }

    Result<void> new_kernel_thread(u64 address)
    {
        Thread* const thread = TRY(new_thread());
        thread->init_regs_kernel();
        thread->set_ip(address);

        return new_kernel_thread_impl(thread);
    }

    Result<void> new_kernel_thread(void (*func)(void))
    {
        Thread* const thread = TRY(new_thread());
        thread->init_regs_kernel();
        thread->set_ip((u64)func);

        return new_kernel_thread_impl(thread);
    }

    Result<void> new_kernel_thread(void (*func)(void*), void* arg)
    {
        Thread* const thread = TRY(new_thread());
        thread->init_regs_kernel();
        thread->set_ip((u64)func);
        thread->set_arguments((u64)arg, 0, 0, 0);

        return new_kernel_thread_impl(thread);
    }

    static Result<void> create_stacks(Stack& user_stack, Stack& kernel_stack)
    {
        const u64 THREAD_STACK_BASE = 0x10000;

        TRY(MemoryManager::alloc_at(THREAD_STACK_BASE, 4, MMU::ReadWrite | MMU::NoExecute | MMU::User));

        auto guard = make_scope_guard([&] { MemoryManager::unmap_owned(THREAD_STACK_BASE, 4); });

        const u64 kernel_stack_base = TRY(MemoryManager::alloc_for_kernel(4, MMU::ReadWrite | MMU::NoExecute));

        guard.deactivate();

        user_stack = { THREAD_STACK_BASE, 4 * ARCH_PAGE_SIZE };
        kernel_stack = { kernel_stack_base, 4 * ARCH_PAGE_SIZE };

        return {};
    }

    Result<void> new_userspace_thread(const TarStream::Entry& entry, const TarStream& stream)
    {
        Thread* const thread = TRY(new_thread());

        thread->is_kernel = false;

        auto guard = make_scope_guard([&] { delete thread; });

        thread->vm_allocator = TRY(UserVM::try_create());

        PageDirectory* const directory = TRY(MMU::create_page_directory_for_userspace());

        auto directory_guard = make_scope_guard([&] {
            MMU::switch_page_directory(MMU::kernel_page_directory());
            MemoryManager::free_frame((u64)directory);
        });

        MMU::switch_page_directory(directory);

        thread->init_regs_user();

        const ELFData data = TRY(ELFLoader::load(entry, stream));

        thread->set_ip(data.entry);

        TRY(create_stacks(thread->stack, thread->kernel_stack));
        thread->set_sp(thread->stack.top());

        thread->directory = directory;

        guard.deactivate();
        directory_guard.deactivate();

        kinfoln("CREATED USERSPACE THREAD: id %lu with ip %#.16lx and sp %#.16lx (ksp %#lx)", thread->id, thread->ip(),
                thread->sp(), thread->kernel_stack.top());

        g_threads.append(thread);

        return {};
    }

    void reap_thread(Thread* thread)
    {
        kinfoln("reap: reaping thread with id %zu", thread->id);

        if (thread->is_kernel)
        {
            auto stack = thread->stack;
            // FIXME: Propagate errors I guess?
            kinfoln("deleting stack @ %#lx", stack.bottom());
            MemoryManager::unmap_owned_and_free_vm(stack.bottom(), stack.bytes() / ARCH_PAGE_SIZE).release_value();
        }
        else
        {
            auto stack = thread->kernel_stack;
            kinfoln("deleting kstack @ %#lx", stack.bottom());
            // FIXME: Propagate errors I guess?
            MemoryManager::unmap_owned_and_free_vm(stack.bottom(), stack.bytes() / ARCH_PAGE_SIZE).release_value();
        }

        if (!thread->is_kernel) MMU::delete_userspace_page_directory(thread->directory);

        delete thread;
    }

    Thread* pick_task()
    {
        Thread* old = g_current;
        if (old->is_idle())
        {
            auto maybe_last = g_threads.last();
            if (!maybe_last.has_value()) // No threads!!
                return &g_idle;
            g_current = old = maybe_last.value();
        }

        bool has_found_thread = false;

        do {
            auto maybe_next = g_threads.next(g_current);
            if (!maybe_next.has_value()) g_current = g_threads.expect_first();
            else
                g_current = maybe_next.value();

            if (g_current->state == ThreadState::Runnable)
            {
                has_found_thread = true;
                break;
            }
        } while (g_current != old);

        if (!has_found_thread) g_current = &g_idle;

        return g_current;
    }

    void generic_switch_context(Thread* old_thread, Thread* new_thread, Registers* regs)
    {
        if (old_thread != new_thread)
        {
            switch_context(old_thread, new_thread, regs);
            if (!new_thread->is_kernel)
            {
                MMU::switch_page_directory(new_thread->directory);
                CPU::switch_kernel_stack(new_thread->kernel_stack.top());
            }
        }

        if (new_thread->is_idle())
        {
            new_thread->ticks_left = 1; // The idle task only runs for 1 tick so we can check for new runnable tasks
                                        // as fast as possible.
        }
        else
            new_thread->ticks_left = TICKS_PER_TIMESLICE;
    }

    void switch_task(Registers* regs)
    {
        Thread* old_thread = g_current;
        Thread* new_thread = pick_task();
        generic_switch_context(old_thread, new_thread, regs);
    }

    void invoke(Registers* regs)
    {
        CPU::disable_interrupts();

        g_current->ticks++;

        if (is_in_kernel(regs)) g_current->ticks_in_kernel++;
        else
            g_current->ticks_in_user++;

        g_current->ticks_left--;

        g_threads.for_each([](Thread* thread) {
            if (thread->state == ThreadState::Sleeping)
            {
                if (--thread->sleep_ticks_left == 0) thread->state = ThreadState::Runnable;
            }
        });

        if (!g_current->ticks_left) switch_task(regs);
    }

    LinkedList<Thread> check_for_dying_threads()
    {
        LinkedList<Thread> result;

        g_threads.delayed_for_each([&](Thread* thread) {
            if (thread->state == ThreadState::Dying)
            {
                g_threads.remove(thread);
                result.append(thread);
            }
        });

        return result;
    }
}

void kernel_sleep(u64 ms)
{
    g_current->sleep_ticks_left = ms;
    g_current->state = ThreadState::Sleeping;
    kernel_yield();
}

[[noreturn]] void kernel_exit()
{
    g_current->state = ThreadState::Dying;
    kernel_yield();
    unreachable();
}