switch_yield.c

# include <stdio.h>
# include <stdlib.h>
# include <stdbool.h>
# include <setjmp.h>

# undef _FORTIFY_SOURCE
# define _FORTIFY_SOURCE 0


typedef int64_t PID_t;
typedef void (*action_t)();


// Process

typedef struct process_t {
  PID_t    pid;
  jmp_buf* context;
  char*    stack;
} process_t;

process_t* process_new(PID_t pid, jmp_buf* context, char* stack) {
  process_t* process = malloc(sizeof(process_t));
  process->pid     = pid;
  process->context = context;
  process->stack   = stack;

  return process;
}

void process_free(process_t* process) {
  free(process->context);
  free(process->stack);
  free(process);
}



// Tasks queue

typedef struct tasks_queue_node_t {
  process_t*                 head;
  struct tasks_queue_node_t* tail;
} tasks_queue_node_t;

typedef struct tasks_queue_t {
  tasks_queue_node_t* beg;
  tasks_queue_node_t* end;
} tasks_queue_t;

tasks_queue_t* tasks_queue_new() {
  tasks_queue_t* queue = malloc(sizeof(tasks_queue_t));
  queue->beg = NULL;
  queue->end = NULL;
  return queue;
}

bool tasks_queue_is_empty(tasks_queue_t* queue) {
  return queue->beg == NULL;
}

void tasks_queue_enqueue(tasks_queue_t* queue, process_t* task) {
  tasks_queue_node_t* node = malloc(sizeof(tasks_queue_node_t));
  node->head = task;
  node->tail = NULL;

  if (tasks_queue_is_empty(queue)) {
    queue->beg = node;
    queue->end = node;
  } else {
    queue->end->tail = node;
    queue->end = node;
  }
}

process_t* tasks_queue_dequeue(tasks_queue_t* queue) {
  if (tasks_queue_is_empty(queue)) return NULL;

  tasks_queue_node_t* node = queue->beg;
  process_t* process = node->head;

  queue->beg = node->tail;
  if (node->tail == NULL)
    queue->end = NULL;

  free(node);
  return process;
}



// Scheduler

struct scheduler_t {
  tasks_queue_t* queue;
  tasks_queue_t* to_free;
  process_t*     current;
  jmp_buf*       exit;
  int64_t        fuel;
  int64_t        gen;
  int64_t        available_fuel;
}* scheduler;

void scheduler_init(int64_t available_fuel) {
  scheduler = malloc(sizeof(struct scheduler_t));
  scheduler->queue          = tasks_queue_new();
  scheduler->to_free        = tasks_queue_new();
  scheduler->current        = NULL;
  scheduler->exit           = malloc(sizeof(jmp_buf));
  scheduler->fuel           = 0;
  scheduler->gen            = 0;
  scheduler->available_fuel = available_fuel;
}

static inline PID_t self() {
  return scheduler->current->pid;
}

void collect_free_processes() {
  process_t* process;
  while ((process = tasks_queue_dequeue(scheduler->to_free))) {
    process_free(process);
  }
}

_Noreturn void switch_context() {
  scheduler->current = tasks_queue_dequeue(scheduler->queue);
  scheduler->fuel    = scheduler->available_fuel;

  if (scheduler->current) {
    longjmp(*scheduler->current->context, 1);
  } else {
    longjmp(*scheduler->exit, 1);
  }
}

_Noreturn void action_wrapper(action_t action, PID_t pid, jmp_buf* spawner, char* stack) {
  process_t* process = process_new(pid, malloc(sizeof(jmp_buf)), stack);
  tasks_queue_enqueue(scheduler->queue, process);

  if (setjmp(*process->context) == 0)
    longjmp(*spawner, 1);

  action();
  tasks_queue_enqueue(scheduler->to_free, process);
  switch_context();
}

# define STACK_SIZE (8 * 1024 * 1024)
PID_t spawn(action_t action) {
  jmp_buf* spawner    = malloc(sizeof(jmp_buf));
  PID_t    pid        = scheduler->gen++;
  char*    stack      = aligned_alloc(16, STACK_SIZE);
  char*    stack_head = stack + STACK_SIZE - 8;

  if (setjmp(*spawner) == 0) asm(
    // use new stack
    "movq  %[stack_head], %%rsp \n"

    // jump to wrapper to create context
    "movq  %[action]    , %%rdi \n"
    "movq  %[pid]       , %%rsi \n"
    "movq  %[spawner]   , %%rdx \n"
    "movq  %[stack]     , %%rcx \n"
    "jmpq *%[action_wrapper]    \n"
    :
    : [stack_head]     "r" (stack_head)
    , [action]         "r" (action)
    , [pid]            "r" (pid)
    , [spawner]        "r" (spawner)
    , [stack]          "r" (stack)
    , [action_wrapper] "r" (action_wrapper)
    : "rdi", "rsi", "rdx", "rcx", "memory"
  );

  return pid;
}

void yield() {
  scheduler->fuel -= 1;
  if (scheduler->fuel <= 0) {
    tasks_queue_enqueue(scheduler->queue, scheduler->current);
    if (setjmp(*scheduler->current->context) == 0) {
      switch_context();
    }
  }

  collect_free_processes();
}

void run(action_t action) {
  if (setjmp(*scheduler->exit) == 0) {
    spawn(action);
    switch_context();
  }

  collect_free_processes();
}



// Example

void printer() {
  printf("<%lld> start\n", self());
  for (int i = 0; i < 5; i++) {
    printf("<%lld> %d\n", self(), i);
    yield();
  }
  printf("<%lld> end\n", self());
}

void example() {
  printf("<%lld> main start\n", self());
  for (int i = 0; i < 5; i++) {
    spawn(printer);
    yield();
  }
  printf("<%lld> main end\n", self());
}

int main() {
  scheduler_init(2);
  run(example);
  puts("done");
}