thread.cc

#include "core/thread.h"

#include <cxx/linked_list.h>
#include <cxx/sprintf.h>
#include <cxx/string.h>
#include <cxx/units.h>
#include <elf/elf.h>
#include <libkernel/assert.h>
#include <libkernel/status.h>

#include "core/constants.h"
#include "core/event.h"
#include "core/loader.h"
#include "core/slab.h"
#include "core/vm.h"
#include "platform/arch.h"
#include "platform/exception.h"
#include "platform/mmu.h"

using namespace core::thread;
using libkernel::PerfMonitor;
using loader::Range;

namespace core::thread {
   Thread* __current_thread = nullptr;
} // namespace core::thread

static Thread* thread_pool = nullptr;
static cxx::IntrusiveList<Thread> thread_freelist;
static cxx::IntrusiveList<Thread> thread_ready_queue[core::MaxPriorityNum];
constexpr char const* ThreadStateStr[] = {
      "INVALID", "CREATING", "READY", "RUNNING", "IPC_WAIT", "IPC_DEFERRED", "EVENT_WAIT",
};
constexpr char const* IpcStateStr[] = {
      "INITIAL", "AWAIT_SEND", "AWAIT_RECEIVE", "AWAIT_REPLY", "REPLIED", "SENT",
};

DECLARE_KCOUNTERS(all_schedule_counts, core::MaxThreadsNum)
DECLARE_KCOUNTERS(all_runtime_perfs, core::MaxThreadsNum)
DECLARE_KCOUNTERS(all_ipc_block_perfs, core::MaxThreadsNum)
DECLARE_KCOUNTERS(all_event_block_perfs, core::MaxThreadsNum)
DECLARE_KCOUNTERS(all_reschedule_perfs, core::MaxThreadsNum)
DECLARE_KCOUNTERS(pqueue_active_perfs, core::MaxPriorityNum)
DECLARE_KCOUNTERS(idle_thread_time, 1)
DECLARE_KCOUNTERS(total_thread_time, 1)
tid_t global_idle_thread;

//===----------------------------------------------------------------------===//
// Static helper functions
//===----------------------------------------------------------------------===//

static inline Result<int64_t> destroy_and_return(E err, Thread* thread) {
   // FIXME(frank): Do we need anything other than putting the thread back to freelist?
   thread_freelist.push_back(thread);
   return err;
}

static void finalize_thread_create(Thread* thread) {
   // 1. Activate the address space mappings, note that we're already in this
   // thread's address space when this is called!
   thread->aspace.activate();
   // 3. Populate the ELF contents of the address space
   auto img = reinterpret_cast<uint8_t const*>(thread->image.img);
   for (auto& ph : thread->image.get_pheaders()) {
      if (ph.p_type != PT_LOAD) continue;
      const auto p_start = img + ph.p_offset;
      const auto v_start = reinterpret_cast<void*>(ph.p_vaddr);
      cxx::memset(v_start, 0, ph.p_memsz);
      cxx::memcpy(v_start, p_start, ph.p_filesz);
   }
   // 4. Find out where the _Exit() symbol is
   thread->exit_address = thread->image.get_func_addr("_Exit");
}

//===----------------------------------------------------------------------===//
// Class Thread:: methods
//===----------------------------------------------------------------------===//

Thread::Thread(mem::Page* l0, uint64_t priority, elf::Parser& parser)
      : cxx::IntrusiveListNode<Thread>{},
        user_tid{-1},
        user_pid{-1},
        priority{priority},
        state{State::CREATING},
        entry_address{0},
        saved_context{},
        aspace{l0},
        image{parser},
        ipc{},
        mailbox{} {}

void core::thread::Thread::destory() {
   remove_from_scheduler();
   state = Thread::State::INVALID;
   thread_freelist.push_back(this);
   for (auto thread : mailbox) {
      // FIXME(kevin): Do something about threads still waiting here.
      // Otherwise, they will become zombies x_x. Also, if a thread dies in
      // the Receive() syscall, the caller becomes a zombie.
      (void)thread;
   }
   if (cur_thread() == this) [[likely]] {
      __current_thread = nullptr;
   }
   all_schedule_counts[this->user_tid].reset();
   all_runtime_perfs[this->user_tid].reset();
   all_ipc_block_perfs[this->user_tid].reset();
   all_event_block_perfs[this->user_tid].reset();
   all_reschedule_perfs[this->user_tid].reset();
}

PerfMonitor& core::thread::Thread::get_event_block_perf() const {
   return all_event_block_perfs[this->user_tid];
}

PerfMonitor& core::thread::Thread::get_ipc_block_perf() const {
   return all_ipc_block_perfs[this->user_tid];
}

PerfMonitor& core::thread::Thread::get_runtime_perf() const {
   return all_runtime_perfs[this->user_tid];
}

PerfMonitor& core::thread::Thread::get_reschedule_perf() const {
   return all_reschedule_perfs[this->user_tid];
}

size_t core::thread::Thread::receive_one_mailbox() {
   auto tid = ipc.receive.user_tid_buf;
   auto msglen = ipc.receive.expected_send_len;
   auto msg = ipc.receive.user_send_buf;
   assert(!mailbox.empty());
   auto sender = mailbox.pop_front();
   libkernel::log(LL::VERBOSE, "%s\r\n", sender->dump());
   const auto actual_size = cxx::min(msglen, sender->ipc.send.msg_len);
   assert(sender->state == Thread::State::IPC_WAIT);
   assert(sender->ipc.state == Ipc::State::AWAIT_RECEIVE);
   cxx::memcpy(msg, sender->ipc.buf, actual_size);
   *tid = sender->user_tid;
   sender->ipc.state = Ipc::State::AWAIT_REPLY;
   return actual_size;
}

void core::thread::Thread::add_to_scheduler() {
   libkernel::log(LL::VERBOSE, "%s\r\n", dump(true));
   if (state == Thread::State::RUNNING) {
      state = Thread::State::READY;
   }
   thread_ready_queue[priority].push_back(this);

   // This queue just activate
   if (thread_ready_queue[priority].size() == 1) {
      pqueue_active_perfs[priority].start_measure();
   }
}

void core::thread::Thread::remove_from_scheduler() {
   libkernel::log(LL::VERBOSE, "%s\r\n", dump(true));
   if (!thread_ready_queue[priority].empty()) {
      thread_ready_queue[priority].remove(this);
      if (thread_ready_queue[priority].empty()) {
         pqueue_active_perfs[priority].end_measure();
      }
   }
   if (this == cur_thread()) {
      __current_thread = nullptr;
   }
}

const char* core::thread::Thread::dump(bool print_state) const {
   static char buf[128];
   int eos = 0;
   if (print_state) {
      eos = cxx::sprintf(buf, "[\"%s\", %lld, %s/%s]", name, user_tid, ThreadStateStr[(int)state],
                         IpcStateStr[(int)ipc.state]);
   } else {
      eos = cxx::sprintf(buf, "[\"%s\", %lld]", name, user_tid);
   }
   buf[eos] = 0;
   return buf;
}

//===----------------------------------------------------------------------===//
// Public thread:: API
//===----------------------------------------------------------------------===//

void core::thread::init() {
   auto thread_pool_r = slab::alloc<Thread[MaxThreadsNum]>();
   assert(thread_pool_r, "Failed to allocate thread pool");
   thread_pool = reinterpret_cast<Thread*>(thread_pool_r.unwrap());
   for (int i = 0; i < MaxPriorityNum; ++i) {
      thread_freelist.push_back(thread_pool + i);
      thread_pool[i].state = Thread::State::INVALID;
   }
   libkernel::log("Thread pool initialized with %d threads\r\n", MaxThreadsNum);
   global_idle_thread = -1;
}

Result<int64_t> core::thread::create(const char* name, void* image_ptr, int priority) {
   if (priority < 0 || priority >= MaxPriorityNum) return E::INVALID_ARGUMENT;

   // 1. Parse the ELF image
   auto image_r = elf::Parser::parse(image_ptr);
   if (!image_r) return image_r.status();
   elf::Parser image = image_r.unwrap();

   // 2. Allocate the thread struct, and grab the address space struct out of it
   Thread* thread;
   {
      // 2a. Allocate the L0 page table
      mem::PageList l0list{};
      auto res = slab::alloc(4_KiB, 4_KiB);
      if (!res) return res.status();

      // 2b. Grab the physical address from virtual
      auto l0 = reinterpret_cast<platform::mmu::TableDescriptor*>(res.unwrap());
      auto pfn = core::mem::traverse(reinterpret_cast<uint64_t>(l0)).phys_addr;
      auto l0table_page = core::mem::phys_to_pfn(pfn * 4_KiB);

      // 2c. Clear the table and set up the recursive mapping
      cxx::memset(l0, 0, 4_KiB);
      l0[511] = platform::mmu::UnifyDescriptors(
            platform::mmu::TableDescriptor{.valid = 1, .table_phys_addr = pfn});

      // 2d. Grab a free thread from the thread slab array
      thread = thread_freelist.pop_front();
      if (!thread) return E::OUT_OF_MEMORY;
      new ((void*)thread) Thread(l0table_page, priority, image);
   }
   auto& aspace = thread->aspace;

   // 3. Populate the new thread structure
   int64_t new_tid = thread - thread_pool;
   cxx::strncpy(thread->name, name, 64);
   thread->name[63] = '\0';
   thread->user_tid = new_tid;
   thread->user_pid = cur_thread() != nullptr ? cur_thread()->user_tid : -1;

   // 4. Populate the new address space regions with the ELF sections
   uint64_t elf_end_addr = 0;
   for (auto& ph : image.get_pheaders()) {
      // Set the thread's TLS pointer
      if (ph.p_type == PT_TLS) thread->tls_address = ph.p_vaddr - 16;
      // All the code below is for LOAD sections
      if (ph.p_type != PT_LOAD) continue;
      core::mem::Region::Flags flags{};
      flags.readable = 1;
      if (ph.p_flags & 1) flags.executable = 1;
      if (ph.p_flags & 0b10) flags.writable = 1;
      const Range v(ph.p_vaddr, cxx::align_up(ph.p_vaddr + ph.p_memsz, 4_KiB));
      assert(v.start == cxx::align_down(v.start, 4_KiB),
             "Program section not aligned - all sections must be page-aligned");
      const auto name = image.get_section_name(image.find_section_by_va(v.start));
      // FIXME(kevin): Allow for cached ELF regions
      auto result = aspace.add_region(name, v, flags);
      if (!result) return destroy_and_return(result.status(), thread);
      elf_end_addr = cxx::max(elf_end_addr, v.end);
   }
   libkernel::log(LL::DEBUG, "TLS pointer %llX\r\n", thread->tls_address);
   assert(elf_end_addr < UserHeapRange.start);

   // 5. Add the following sections to the task address space, allocating as needed
   //  - Stack, grows down, r/w
   //  - Heap, grows up, r/w
   //  - Device, shared memory, r/w
   //  - Initrd, shared memory, r
   {
      Result<core::mem::Region*> result{E::INVALID};
      result = aspace.add_region("stack", UserStackRange,
                                 core::mem::Region::Flags{
                                       .readable = 1,
                                       .writable = 1,
                                       .executable = 0,
                                       .shared = 0,
                                       .lazy_commit = 1,
                                       .grows_down = 1,
                                 });
      if (!result) return destroy_and_return(result.status(), thread);
      result = aspace.add_region("heap", UserHeapRange,
                                 core::mem::Region::Flags{
                                       .readable = 1,
                                       .writable = 1,
                                       .executable = 0,
                                       .shared = 0,
                                       .lazy_commit = 1,
                                       .grows_down = 0,
                                 });
      if (!result) return destroy_and_return(result.status(), thread);
      result = aspace.add_region("initrd", UserInitrdAddress, core::mem::get_initrd_vmo(),
                                 core::mem::Region::Flags{
                                       .readable = 1,
                                       .writable = 0,
                                       .executable = 0,
                                       .shared = 1,
                                       .lazy_commit = 0,
                                       .grows_down = 0,
                                 });
      if (!result) return destroy_and_return(result.status(), thread);
      result = aspace.add_region("device", UserDeviceRange,
                                 core::mem::Region::Flags{
                                       .readable = 1,
                                       .writable = 1,
                                       .shared = 1,
                                       .lazy_commit = 1,
                                       .is_device = 1,
                                 });
      if (!result) return destroy_and_return(result.status(), thread);
   }

   // 6. Stuff the task's register states
   thread->entry_address = image.img->e_entry;
   thread->saved_context.elr_elx = image.img->e_entry;
   thread->saved_context.sp_elx = UserStackRange.end;
   thread->image = image;

   // 7. Add the thread to the ready queue and wait for being sheduled
   thread->add_to_scheduler();
   thread->state = Thread::State::CREATING;

   // 8. The remaining creation steps will be in finalize_thread_create()
   return new_tid;
}

void core::thread::schedule() {
   // Current thread stops running
   if (cur_thread()) {
      cur_thread()->get_runtime_perf().end_measure();
      cur_thread()->get_reschedule_perf().start_measure();
   }

   // 1. Run the scheduling algorithm to determine the next thread.
   //    Search for the next thread to run using the following steps.
   Thread* next_thread = nullptr;

   // 1a. Grab the current thread's priority. Use the highest priority
   //     if there are no current thread (cur thread must be running too).
   const auto cur_thread_running = cur_thread() && cur_thread()->state == Thread::State::RUNNING;
   const int prio_idx = cur_thread_running ? cur_thread()->priority : 0;

   // 1b. Search all priorities higher than this thread for a runnable thread
   for (int i = 0; i < prio_idx; ++i) {
      if (!thread_ready_queue[i].empty()) {
         next_thread = thread_ready_queue[i].pop_front();
         if (thread_ready_queue[i].empty()) {
            pqueue_active_perfs[i].end_measure();
         }
         break;
      }
   }

   // 1c. Have we found a thread? If not AND we have a non-blocked current
   //     thread, then schedule round-robin in this priority level.
   //     !! IMPORTANT NOTE !!
   //     !! If this priority is empty, we re-schedule the current thread.
   //     !! IMPORTANT NOTE !!
   if (!next_thread && cur_thread_running) {
      if (thread_ready_queue[prio_idx].empty()) {
         next_thread = cur_thread();
         next_thread->state = Thread::State::READY;
      } else {
         next_thread = thread_ready_queue[prio_idx].pop_front();
         if (thread_ready_queue[prio_idx].empty()) {
            pqueue_active_perfs[prio_idx].end_measure();
         }
      }
   }

   // 1d. If we still haven't found a thread, we move to a lower priority
   //     and find a thread. If after this, there are still no threads,
   //     then the ready queues are empty.
   if (!next_thread) {
      for (int i = prio_idx; i < MaxPriorityNum; ++i) {
         if (!thread_ready_queue[i].empty()) {
            next_thread = thread_ready_queue[i].pop_front();
            if (thread_ready_queue[i].empty()) {
               pqueue_active_perfs[i].end_measure();
            }
            break;
         }
      }
   }

   // 2. Make sure the thread to be scheduled exists and is valid
   if (!next_thread && !cur_thread()) {
      if (event::queues_empty()) {
         libkernel::log(LL::ERROR, "No more threads to run (or deadlock), halting...\r\n");
      } else {
         libkernel::log(LL::ERROR, "No more threads to run, all threads waiting on event...\r\n");
         event::dump(LL::NORMAL);
      }
      platform::arch::halt();

   } else if (!next_thread && cur_thread()) {
      assert(cur_thread()->state == Thread::State::RUNNING);
      next_thread = cur_thread();
      if (next_thread->state == Thread::State::RUNNING) {
         next_thread->state = Thread::State::READY;
      }
   }
   assert(next_thread->state != Thread::State::INVALID, "Attempt to schedule invalid thread");
   libkernel::log(LL::VERBOSE, "Next: %s\r\n", next_thread->dump(false));

   // 3. Change to this thread's address space
   Thread* old = nullptr;
   if (next_thread != cur_thread()) [[likely]] {
      old = switch_to(next_thread);
   }

   // 4. Take action depending on the thread's state
   bool curr_just_created = false;
   switch (cur_thread()->state) {
      case Thread::State::CREATING:
         finalize_thread_create(cur_thread());
         cur_thread()->state = Thread::State::READY;
         curr_just_created = true;
         break;
      case Thread::State::READY:
         break;
      case Thread::State::IPC_DEFERRED:
         // 1. Has the receiver replied to us?
         // 2. Has the sender sent us something?
         if (cur_thread()->ipc.state == Ipc::State::REPLIED) {
            auto sender = cur_thread();
            cxx::memcpy(sender->ipc.send.user_reply_buf, sender->ipc.buf, sender->ipc.send.msg_len);
            sender->state = Thread::State::IPC_WAIT;
            sender->ipc.state = Ipc::State::AWAIT_REPLY;
         } else if (cur_thread()->ipc.state == Ipc::State::SENT) {
            const auto rc = cur_thread()->receive_one_mailbox();
            cur_thread()->set_retcode(rc);
            // Remember to set the current thread to ready again!
            cur_thread()->state = Thread::State::READY;
            cur_thread()->ipc.state = Ipc::State::INITIAL;
         }
         break;
      default:
         libkernel::log(LL::ERROR, "Thread state is invalid: %s\r\n", cur_thread()->dump(true));
         panic("Invalid thread state");
         break;
   }
   // Now current thread is ready to run
   cur_thread()->state = Thread::State::RUNNING;

   // 5. Add the old thread back to the scheduler
   if (old) {
      old->add_to_scheduler();
   }

   // 6. Update the performance monitors
   all_schedule_counts[cur_thread()->user_tid].count();
   cur_thread()->get_runtime_perf().start_measure();
   if (!curr_just_created) {
      cur_thread()->get_reschedule_perf().end_measure();
   }
   core::thread::sample_idle_stats();
}

void core::thread::sample_idle_stats() {
   if (global_idle_thread != -1) {
      if (cur_thread()->user_tid == global_idle_thread) {
         idle_thread_time[0].start_measure();
         total_thread_time[0].end_measure();
      } else {
         total_thread_time[0].start_measure();
         idle_thread_time[0].end_measure();
      }
   }
}

Thread* core::thread::switch_to(Thread* other) {
   if (!other) [[unlikely]] {
      panic("Attempting to switch to a null thread");
   }
   Thread* old_thread = cur_thread();
   libkernel::log(LL::VERBOSE, "From: %s\r\n", old_thread ? old_thread->dump(true) : "NULL");
   libkernel::log(LL::VERBOSE, "To: %s\r\n", other->dump(true));
   __current_thread = other;
   platform::arch::write_ttbr0_el1(other->aspace.l0_address());
   platform::arch::write_tpidr_el0(other->tls_address);
   platform::arch::flush_tlb();
   return old_thread;
}

Result<Thread*> core::thread::get_thread_by_tid(int64_t tid) {
   if (tid < 0 || tid >= MaxThreadsNum) return E::INVALID;
   auto res = &thread_pool[tid];
   if (res->state == Thread::State::INVALID) return E::INVALID;
   return res;
}