Mutexes and process suspension in kernels

#include <stdatomic.h>

#define SPIN_LOCK_INIT ATOMIC_FLAG_INIT

typedef atomic_flag spin_lock_t;

void spin_lock (spin_lock_t* sl) {
  while (atomic_flag_test_and_set_explicit (sl, memory_order_acquire))
    ;
}

void spin_unlock (spin_lock_t* sl) {
  atomic_flag_clear_explicit (sl, memory_order_release);
}

// List same as before...

// the lock
spin_lock_t the_lock = SPIN_LOCK_INIT;

void thread_A (void) {
  struct my_list_node* new_node = malloc (sizeof (*new_node));
  new_node->value = 67;

  spin_lock (&the_lock);
  LIST_ADD (my_list, new_node);
  spin_unlock (&the_lock);
}

void thread_B (void) {
  spin_lock (&the_lock);
  LIST_REMOVE (my_list, my_list); /* remove head */
  spin_unlock (&the_lock);
}

/* process states */
#define PROC_READY     0
#define PROC_DEAD      1
#define PROC_SUSPENDED 2 /* NEW */

// ...

static struct proc* proc_find_sched (struct cpu* cpu) {
  if (!cpu->proc_run_q)
    return NULL;

  struct list_node_link *current, *start;

  if (cpu->proc_current)
    current = cpu->proc_current->cpu_run_q_link.next;
  else
    current = cpu->proc_run_q;

  if (!current)
    current = cpu->proc_run_q;

  start = current;

  do {
    struct proc* proc = list_entry (current, struct proc, cpu_run_q_link);

    if (atomic_load (&proc->state) == PROC_READY)
      return proc;

    current = current->next ? current->next : cpu->proc_run_q;
  } while (current != start);

  return NULL;
}

// ...

void proc_sched (void) {
  spin_lock_ctx_t ctxcpu;

  int s_cycles = atomic_fetch_add (&sched_cycles, 1);

  if (s_cycles % SCHED_REAP_FREQ == 0)
    proc_reap ();

  struct proc* next = NULL;
  struct cpu* cpu = thiscpu;

  spin_lock (&cpu->lock, &ctxcpu);

  next = proc_find_sched (cpu);

  if (next) {
    cpu->proc_current = next;

    do_sched (next, &cpu->lock, &ctxcpu);
  } else {
    cpu->proc_current = NULL;
    spin_unlock (&cpu->lock, &ctxcpu);

    spin ();
  }
}

struct proc {
  int pid;
  struct rb_node_link proc_tree_link;
  struct rb_node_link procgroup_memb_tree_link;
  struct list_node_link cpu_run_q_link;
  struct list_node_link reap_link;
  struct list_node_link* sq_entries; /* list of intermediate structs */
  struct procgroup* procgroup;
  struct proc_platformdata pdata;
  uint32_t flags;
  spin_lock_t lock;
  struct cpu* cpu;
  atomic_int state;
  uintptr_t uvaddr_argument;
};

struct proc_suspension_q {
  struct list_node_link* proc_list; /* list of processes via proc_sq_entry */
  spin_lock_t lock;
};

/* the intermediate struct */
struct proc_sq_entry {
  struct list_node_link sq_link; /* list link */
  struct list_node_link proc_link; /* list link */
  struct proc* proc; /* pointer back to the process */
  struct proc_suspension_q* sq; /* pointer back to the suspension queue */
};

bool proc_sq_suspend (struct proc* proc, struct proc_suspension_q* sq, spin_lock_t* resource_lock,
                      spin_lock_ctx_t* ctxrl) {
  spin_lock_ctx_t ctxpr, ctxcpu, ctxsq;
  struct cpu* cpu = proc->cpu;

  /* allocate intermediate struct */
  struct proc_sq_entry* sq_entry = malloc (sizeof (*sq_entry));
  if (!sq_entry) {
    spin_unlock (resource_lock, ctxrl);
    return PROC_NO_RESCHEDULE;
  }

  /* set links on both ends */
  sq_entry->proc = proc;
  sq_entry->sq = sq;

  /* lock with compliance to the lock hierachy */
  spin_lock (&cpu->lock, &ctxcpu);
  spin_lock (&proc->lock, &ctxpr);
  spin_lock (&sq->lock, &ctxsq);

  spin_unlock (resource_lock, ctxrl);

  /* transition the state to PROC_SUSPENDED */
  atomic_store (&proc->state, PROC_SUSPENDED);

  /* append to sq's list */
  list_append (sq->proc_list, &sq_entry->sq_link);

  /* append to proc's list */
  list_append (proc->sq_entries, &sq_entry->proc_link);

  /* remove from CPU's run list and decrement CPU load counter */
  list_remove (cpu->proc_run_q, &proc->cpu_run_q_link);
  atomic_fetch_sub (&cpu->proc_run_q_count, 1);

  if (cpu->proc_current == proc)
    cpu->proc_current = NULL;

  /* a process is now unschedulable, so it doesn't need a CPU */
  proc->cpu = NULL;

  /* release the locks */
  spin_unlock (&sq->lock, &ctxsq);
  spin_unlock (&proc->lock, &ctxpr);
  spin_unlock (&cpu->lock, &ctxcpu);

  return PROC_NEED_RESCHEDULE;
}

bool proc_sq_resume (struct proc* proc, struct proc_sq_entry* sq_entry) {
  spin_lock_ctx_t ctxsq, ctxpr, ctxcpu;
  struct cpu* cpu = cpu_find_lightest (); /* find the least-loaded CPU */
  struct proc_suspension_q* sq = sq_entry->sq;

  /* acquire necessary locks */
  spin_lock (&cpu->lock, &ctxcpu);
  spin_lock (&proc->lock, &ctxpr);
  spin_lock (&sq->lock, &ctxsq);

  /* remove from sq's list */
  list_remove (sq->proc_list, &sq_entry->sq_link);

  /* remove from proc's list */
  list_remove (proc->sq_entries, &sq_entry->proc_link);

  /* Give the CPU to the process */
  proc->cpu = cpu;

  /*
   * update process state to PROC_READY, but only if it's not waiting inside
   * another suspension queue.
   */
  if (proc->sq_entries == NULL)
    atomic_store (&proc->state, PROC_READY);

  /* attach to CPU's run list and increment load counter */
  list_append (cpu->proc_run_q, &proc->cpu_run_q_link);
  atomic_fetch_add (&cpu->proc_run_q_count, 1);

  /* unlock */
  spin_unlock (&sq->lock, &ctxsq);
  spin_unlock (&proc->lock, &ctxpr);
  spin_unlock (&cpu->lock, &ctxcpu);

  /* intermediate struct is no longer needed, so free it */
  free (sq_entry);

  return PROC_NEED_RESCHEDULE;
}

struct proc_mutex {
  struct proc_resource* resource; /* pointer back to the resource */

  bool locked; /* taken or free */
  struct proc_suspension_q suspension_q; /* suspension queue of waiting processes */
  struct proc* owner; /* current owner of the mutex / process, which can enter the critical section */
};

bool proc_mutex_lock (struct proc* proc, struct proc_mutex* mutex) {
  spin_lock_ctx_t ctxmt;

  spin_lock (&mutex->resource->lock, &ctxmt);

  /* we can enter if: the mutex is NOT locked or we're already the owner of the mutex */
  if (!mutex->locked || mutex->owner == proc) {
    mutex->locked = true;
    mutex->owner = proc;
    spin_unlock (&mutex->resource->lock, &ctxmt);
    return PROC_NO_RESCHEDULE;
  }

  /* we failed to become an owner, thus we have to free up the CPU and suspend ourselves */
  return proc_sq_suspend (proc, &mutex->suspension_q, &mutex->resource->lock, &ctxmt);
}

bool proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex) {
  spin_lock_ctx_t ctxmt, ctxsq;

  spin_lock (&mutex->resource->lock, &ctxmt);

  /*
   * we can't unlock a mutex, which we don't already own,
   * otherwise we could just steal someone else's mutex.
   */
  if (mutex->owner != proc) {
    spin_unlock (&mutex->resource->lock, &ctxmt);
    return PROC_NO_RESCHEDULE;
  }

  spin_lock (&mutex->suspension_q.lock, &ctxsq);

  struct list_node_link* node = mutex->suspension_q.proc_list;

  /*
   * This is a little tricky. What we're doing here is "direct handoff" - we basically
   * transfer ownership, while keeping the mutex locked, instead of just unlocking and
   * resuming everyone.
   */

  if (node) {
    struct proc_sq_entry* sq_entry = list_entry (node, struct proc_sq_entry, sq_link);
    struct proc* resumed_proc = sq_entry->proc;
    /* transfer ownership */
    mutex->owner = resumed_proc;
    mutex->locked = true;

    spin_unlock (&mutex->suspension_q.lock, &ctxsq);
    spin_unlock (&mutex->resource->lock, &ctxmt);

    /* resume new owner - it can enter the critical section now */
    return proc_sq_resume (resumed_proc, sq_entry);
  }

  /* no possible future owner found, so we just unlock */
  mutex->locked = false;
  mutex->owner = NULL;

  spin_unlock (&mutex->suspension_q.lock, &ctxsq);
  spin_unlock (&mutex->resource->lock, &ctxmt);

  return PROC_NEED_RESCHEDULE;
}