uthread.cpp

#include "gocpparallel.hpp"
#include <sys/ucontext.h>
#define UCONTEXT_STACK_SIZE 16384
#define CLONE_STACK_SIZE 16384
#define MAX_NUM_UTHREADS 1000
#define gettid() (syscall(SYS_gettid))
#define tgkill(tgid, tid, sig) (syscall(SYS_tgkill, tgid, tid, sig))

bool waiting_uthread_has_priority_over(uthread_t *ut);
void uthread_system_shutdown();
void uthread_init(uthread_t *ut, void(run_func)());

Uthread::Uthread() : ucontext() {}
Uthread::Uthread(const Uthread &other) : ucontext(other.ucontext) {}

void system_init(int max_num_kthreads) {
  uthread_system_init(max_num_kthreads);
}

void uthread_system_init(int max_num_kthreads) {
  // Basic assertions
  assert(_shutdown == false);
  assert(1 <= max_num_kthreads && max_num_kthreads <= MAX_NUM_UTHREADS);

  // Initialize global variables
  _num_kthreads = 0;
  _max_num_kthreads = max_num_kthreads;
  getcontext(&_system_initializer_context);

  _waiting_uthreads = new std::queue<uthread_t>;
  _kthreads = new kthread_t[max_num_kthreads];
  assert(_kthreads != NULL);

  for (int i = 0; i < max_num_kthreads; ++i) {
    kthread_init(&_kthreads[i]);
  }
}

std::thread::id uthread_create(void(run_func)()) {
  std::thread::id tid;

  _mutex.lock();
  assert(_shutdown == false);
  // Lock the system from shutting down while there is a uthread running.
  if (_num_kthreads == 0) {
    _shutdown_mutex.lock();
  }
  uthread_t *uthread = new uthread_t;
  assert(uthread != NULL);
  uthread_init(uthread, run_func);

  if (_num_kthreads == _max_num_kthreads) {
    // Add the new `uthread` to the heap.
    _waiting_uthreads->push(*uthread);
  } else {
    // Make a new `kthread` to run this new `uthread` immediately.

    assert(_waiting_uthreads->size() == 0); // There must not be waiting
                                            // uthreads if `_num_kthreads` is
                                            // less than `_max_num_kthreads`.

    kthread_t *kthread = find_inactive_kthread();
    assert(kthread !=
           NULL); // There must be an inactive `kthread` if
                  // `_num_kthreads` is less than `_max_num_kthreads`.

    tid = kthread_create(kthread, uthread);
    _num_kthreads += 1;
  }
  _mutex.unlock();
  return tid;
}

void uthread_yield() {
  _mutex.lock();
  assert(_shutdown == false);
  kthread_t *self_kthread = kthread_self();
  assert(self_kthread != NULL);
  uthread_t *cur = self_kthread->running;
  // transfer_elapsed_time(self_kthread, cur);
  if (waiting_uthread_has_priority_over(cur)) {
    // Get another `uthread` from the heap to be run.
    uthread_t next;
    // HEAPextract(_waiting_uthreads, (void **) &next);
    next = Uthread(_waiting_uthreads->front());
    _waiting_uthreads->pop();
    self_kthread->running = &next;

    // Save the current `uthread` to the heap.
    _waiting_uthreads->push(*cur);

    // TODO: consider possibility of race conditions!
    _mutex.unlock();
    kthread_handoff(cur, &next);
  } else {
    _mutex.unlock();
  }
}

void uthread_exit() {
  _mutex.lock();
  kthread_t *self = kthread_self();
  int self_tid;
  int self_tgid;
  _mutex.unlock();

  // If the calling thread is not a `kthread` created by the system, block on a
  // mutex until there are no running `kthreads`.
  if (self == NULL) {
    _shutdown_mutex.lock();
    uthread_system_shutdown();
    _shutdown_mutex.unlock();
    return;
  }

  _mutex.lock();
  assert(_shutdown == false);

  self = kthread_self();
  uthread_t *prev = self->running;
  uthread_t *next = nullptr;

  // TODO: print prev->running_time for debug.

  // Stop running the `prev` `uthread`, and clean up any references to it.
  self->running = nullptr;
  uthread_destroy(prev);
  free(prev);

  // Check if a `uthread` can use this kthread.
  if (_waiting_uthreads->size() > 0) {
    // Use this `kthread` to run a different `uthread`.
    uthread_t next;
    next = Uthread(_waiting_uthreads->front());
    self->running = &next;

    // kthread_update_timestamps(self);

    // TODO: consider possibility of race conditions.
    _mutex.unlock();
    setcontext(&((&next)->ucontext));
  } else {
    // There are no `uthread`s that might use `self`. Kill `self`.
    // Find out more about `self` first.
    // self_tid = self->tid;
    // Clean up `kthread`-associated system data structures.
    _num_kthreads--;

    // If this was the last `kthread`, then the system-shutdown mutex is
    // unlocked. Free lock and kill self.
    if (_num_kthreads == 0) {
      _shutdown_mutex.unlock();
    }

    _mutex.unlock();
    // tgkill(self_tid, self_tgid, SIGKILL);
    self->kthread.~thread();
    assert(false); // Control should never reach here.
  }
}

void uthread_init(uthread_t *ut, void(run_func)()) {
  assert(ut != NULL);
  assert(run_func != NULL);

  // Initialize the `ucontext`.
  ucontext_t *uc = &(ut->ucontext);
  *uc = _system_initializer_context;
  uc->uc_stack.ss_sp = malloc(UCONTEXT_STACK_SIZE);
  uc->uc_stack.ss_size = UCONTEXT_STACK_SIZE;
  makecontext(uc, run_func, 0);
}

void uthread_destroy(uthread_t *ut) {
  assert(ut != NULL);
  free(ut->ucontext.uc_stack.ss_sp);
}

void uthread_system_shutdown() {
  if (_shutdown != true) {
    _shutdown = true;

    delete[] _waiting_uthreads;
    _waiting_uthreads = nullptr;

    // for (int i = 0; i < _max_num_kthreads; ++i) {
    //   kthread_destroy(&_kthreads[i]);
    // }
    delete[] _kthreads;
    _kthreads = nullptr;

    // Note that there is nothing to free from _system_initializer_context,
    // because its stack was never allocated.
  }
  // Otherwise, this function was already called, so there's nothing to do.
}

bool waiting_uthread_has_priority_over(uthread_t *ut) {
  if (_waiting_uthreads->size() > 0) {
    return true;
    // const uthread_t* highest_priority_waiting = &(_waiting_uthreads->top());
    // return uthread_priority(ut, highest_priority_waiting) < 0;
  } else {
    return false;
  }
}