Use red-black trees to store process run queue and process list

kernel/amd64/mm.c

@@ -77,6 +77,14 @@ static uint64_t* amd64_mm_next_table (uint64_t* table, uint64_t entry_idx, bool
   return (uint64_t*)((uintptr_t)hhdm->offset + (uintptr_t)paddr);
 }
 
+static bool amd64_mm_is_table_empty (uint64_t* table) {
+  for (size_t i = 0; i < 512; i++) {
+    if (table[i] & AMD64_PG_PRESENT)
+      return false;
+  }
+  return true;
+}
+
 /// Convert generic memory management subsystem flags into AMD64-specific flags
 static uint64_t amd64_mm_resolve_flags (uint32_t generic) {
   uint64_t flags = 0;
@@ -171,8 +179,28 @@ void mm_unmap_page (struct pd* pd, uintptr_t vaddr, uint32_t flags) {
 
   uint64_t* pte = &pml1[pg_index.pml1];
 
-  *pte &= ~AMD64_PG_PRESENT;
+  if ((*pte) & AMD64_PG_PRESENT) {
+    *pte = 0;
     do_reload = true;
+  }
+
+  if (amd64_mm_is_table_empty (pml1)) {
+    uintptr_t pml1_phys = pml2[pg_index.pml2] & ~0xFFFULL;
+    pmm_free (pml1_phys, 1);
+    pml2[pg_index.pml2] = 0;
+
+    if (amd64_mm_is_table_empty (pml2)) {
+      uintptr_t pml2_phys = pml3[pg_index.pml3] & ~0xFFFULL;
+      pmm_free (pml2_phys, 1);
+      pml3[pg_index.pml3] = 0;
+
+      if (amd64_mm_is_table_empty (pml3)) {
+        uintptr_t pml3_phys = pml4[pg_index.pml4] & ~0xFFFULL;
+        pmm_free (pml3_phys, 1);
+        pml4[pg_index.pml4] = 0;
+      }
+    }
+  }
 
 done:
   if (do_reload && (flags & MM_PD_RELOAD))

kernel/amd64/smp.h

@@ -4,6 +4,7 @@
 #include <amd64/gdt.h>
 #include <amd64/tss.h>
 #include <aux/compiler.h>
+#include <libk/rbtree.h>
 #include <libk/std.h>
 #include <proc/proc.h>
 
@@ -30,9 +31,9 @@ struct cpu {
 
   spin_lock_t lock;
 
-  struct proc* proc_run_q;
+  struct rb_node_link* proc_run_q;
   struct proc* proc_current;
-} PACKED;
+};
 
 struct cpu* cpu_make (void);
 struct cpu* cpu_get (void);

kernel/libk/rbtree.h

@@ -0,0 +1,281 @@
+#ifndef _KERNEL_LIBK_RBTREE_H
+#define _KERNEL_LIBK_RBTREE_H
+
+struct rb_node_link {
+  struct rb_node_link* left;
+  struct rb_node_link* right;
+  struct rb_node_link* parent;
+  int color;
+};
+
+#define RBTREE_RED   0
+#define RBTREE_BLACK 1
+
+#define rbtree_parent(x) ((x)->parent)
+#define rbtree_left(x)   ((x)->left)
+#define rbtree_right(x)  ((x)->right)
+#define rbtree_color(x)  ((x)->color)
+
+#define rbtree_container_of(ptr, type, member) ((type*)((char*)(ptr) - offsetof (type, member)))
+
+#define rbtree_entry(node, type, member) rbtree_container_of (node, type, member)
+
+#define rbtree_rotate_left(root_ptr, x)                                                            \
+  do {                                                                                             \
+    struct rb_node_link* __y = (x)->right;                                                         \
+    (x)->right = __y->left;                                                                        \
+    if (__y->left)                                                                                 \
+      __y->left->parent = (x);                                                                     \
+    __y->parent = (x)->parent;                                                                     \
+    if (!(x)->parent)                                                                              \
+      *(root_ptr) = __y;                                                                           \
+    else if ((x) == (x)->parent->left)                                                             \
+      (x)->parent->left = __y;                                                                     \
+    else                                                                                           \
+      (x)->parent->right = __y;                                                                    \
+    __y->left = (x);                                                                               \
+    (x)->parent = __y;                                                                             \
+  } while (0)
+
+#define rbtree_rotate_right(root_ptr, y)                                                           \
+  do {                                                                                             \
+    struct rb_node_link* __x = (y)->left;                                                          \
+    (y)->left = __x->right;                                                                        \
+    if (__x->right)                                                                                \
+      __x->right->parent = (y);                                                                    \
+    __x->parent = (y)->parent;                                                                     \
+    if (!(y)->parent)                                                                              \
+      *(root_ptr) = __x;                                                                           \
+    else if ((y) == (y)->parent->right)                                                            \
+      (y)->parent->right = __x;                                                                    \
+    else                                                                                           \
+      (y)->parent->left = __x;                                                                     \
+    __x->right = (y);                                                                              \
+    (y)->parent = __x;                                                                             \
+  } while (0)
+
+#define rbtree_insert_fixup(root_ptr, z_node)                                                      \
+  do {                                                                                             \
+    struct rb_node_link* __z = (z_node);                                                           \
+    while (__z->parent && __z->parent->color == RBTREE_RED) {                                      \
+      if (__z->parent == __z->parent->parent->left) {                                              \
+        struct rb_node_link* __y = __z->parent->parent->right;                                     \
+        if (__y && __y->color == RBTREE_RED) {                                                     \
+          __z->parent->color = RBTREE_BLACK;                                                       \
+          __y->color = RBTREE_BLACK;                                                               \
+          __z->parent->parent->color = RBTREE_RED;                                                 \
+          __z = __z->parent->parent;                                                               \
+        } else {                                                                                   \
+          if (__z == __z->parent->right) {                                                         \
+            __z = __z->parent;                                                                     \
+            rbtree_rotate_left ((root_ptr), __z);                                                  \
+          }                                                                                        \
+          __z->parent->color = RBTREE_BLACK;                                                       \
+          __z->parent->parent->color = RBTREE_RED;                                                 \
+          rbtree_rotate_right ((root_ptr), __z->parent->parent);                                   \
+        }                                                                                          \
+      } else {                                                                                     \
+        struct rb_node_link* __y = __z->parent->parent->left;                                      \
+        if (__y && __y->color == RBTREE_RED) {                                                     \
+          __z->parent->color = RBTREE_BLACK;                                                       \
+          __y->color = RBTREE_BLACK;                                                               \
+          __z->parent->parent->color = RBTREE_RED;                                                 \
+          __z = __z->parent->parent;                                                               \
+        } else {                                                                                   \
+          if (__z == __z->parent->left) {                                                          \
+            __z = __z->parent;                                                                     \
+            rbtree_rotate_right ((root_ptr), __z);                                                 \
+          }                                                                                        \
+          __z->parent->color = RBTREE_BLACK;                                                       \
+          __z->parent->parent->color = RBTREE_RED;                                                 \
+          rbtree_rotate_left ((root_ptr), __z->parent->parent);                                    \
+        }                                                                                          \
+      }                                                                                            \
+    }                                                                                              \
+    (*(root_ptr))->color = RBTREE_BLACK;                                                           \
+  } while (0)
+
+#define rbtree_insert(type, root_ptr, node, member, keyfield)                                      \
+  do {                                                                                             \
+    struct rb_node_link** __link = (root_ptr);                                                     \
+    struct rb_node_link* __parent = NULL;                                                          \
+    struct rb_node_link* __new = (node);                                                           \
+    type* __nobj = rbtree_entry (__new, type, member);                                             \
+    while (*__link) {                                                                              \
+      __parent = *__link;                                                                          \
+      type* __xobj = rbtree_entry (*__link, type, member);                                         \
+      if (__nobj->keyfield < __xobj->keyfield)                                                     \
+        __link = &((*__link)->left);                                                               \
+      else                                                                                         \
+        __link = &((*__link)->right);                                                              \
+    }                                                                                              \
+    __new->parent = __parent;                                                                      \
+    __new->left = NULL;                                                                            \
+    __new->right = NULL;                                                                           \
+    __new->color = RBTREE_RED;                                                                     \
+    *__link = __new;                                                                               \
+    rbtree_insert_fixup (root_ptr, __new);                                                         \
+  } while (0)
+
+#define rbtree_find(type, root_ptr, keyval, out, member, keyfield)                                 \
+  do {                                                                                             \
+    (out) = NULL;                                                                                  \
+    struct rb_node_link* __cur = *(root_ptr);                                                      \
+    while (__cur) {                                                                                \
+      type* __obj = rbtree_entry (__cur, type, member);                                            \
+      if ((keyval) == __obj->keyfield) {                                                           \
+        (out) = __cur;                                                                             \
+        break;                                                                                     \
+      } else if ((keyval) < __obj->keyfield)                                                       \
+        __cur = __cur->left;                                                                       \
+      else                                                                                         \
+        __cur = __cur->right;                                                                      \
+    }                                                                                              \
+  } while (0)
+
+#define rbtree_min(node, out)                                                                      \
+  do {                                                                                             \
+    (out) = NULL;                                                                                  \
+    if ((node)) {                                                                                  \
+      struct rb_node_link* __n = (node);                                                           \
+      while (__n->left)                                                                            \
+        __n = __n->left;                                                                           \
+      (out) = __n;                                                                                 \
+    }                                                                                              \
+  } while (0)
+
+#define rbtree_transplant(root_ptr, u, v)                                                          \
+  do {                                                                                             \
+    if (!(u)->parent)                                                                              \
+      *(root_ptr) = (v);                                                                           \
+    else if ((u) == (u)->parent->left)                                                             \
+      (u)->parent->left = (v);                                                                     \
+    else                                                                                           \
+      (u)->parent->right = (v);                                                                    \
+    if (v)                                                                                         \
+      (v)->parent = (u)->parent;                                                                   \
+  } while (0)
+
+#define rbtree_delete_fixup(root_ptr, x_node, xparent_node)                                        \
+  do {                                                                                             \
+    struct rb_node_link* __x = (x_node);                                                           \
+    struct rb_node_link* __xparent = (xparent_node);                                               \
+    while (__x != *(root_ptr) && (__x == NULL || __x->color == RBTREE_BLACK)) {                    \
+      if (__x == __xparent->left) {                                                                \
+        struct rb_node_link* __w = __xparent->right;                                               \
+        if (__w && __w->color == RBTREE_RED) {                                                     \
+          __w->color = RBTREE_BLACK;                                                               \
+          __xparent->color = RBTREE_RED;                                                           \
+          rbtree_rotate_left (root_ptr, __xparent);                                                \
+          __w = __xparent->right;                                                                  \
+        }                                                                                          \
+        if ((!__w->left || __w->left->color == RBTREE_BLACK) &&                                    \
+            (!__w->right || __w->right->color == RBTREE_BLACK)) {                                  \
+          __w->color = RBTREE_RED;                                                                 \
+          __x = __xparent;                                                                         \
+          __xparent = __x->parent;                                                                 \
+        } else {                                                                                   \
+          if (!__w->right || __w->right->color == RBTREE_BLACK) {                                  \
+            if (__w->left)                                                                         \
+              __w->left->color = RBTREE_BLACK;                                                     \
+            __w->color = RBTREE_RED;                                                               \
+            rbtree_rotate_right (root_ptr, __w);                                                   \
+            __w = __xparent->right;                                                                \
+          }                                                                                        \
+          __w->color = __xparent->color;                                                           \
+          __xparent->color = RBTREE_BLACK;                                                         \
+          if (__w->right)                                                                          \
+            __w->right->color = RBTREE_BLACK;                                                      \
+          rbtree_rotate_left (root_ptr, __xparent);                                                \
+          __x = *(root_ptr);                                                                       \
+        }                                                                                          \
+      } else {                                                                                     \
+        struct rb_node_link* __w = __xparent->left;                                                \
+        if (__w && __w->color == RBTREE_RED) {                                                     \
+          __w->color = RBTREE_BLACK;                                                               \
+          __xparent->color = RBTREE_RED;                                                           \
+          rbtree_rotate_right (root_ptr, __xparent);                                               \
+          __w = __xparent->left;                                                                   \
+        }                                                                                          \
+        if ((!__w->right || __w->right->color == RBTREE_BLACK) &&                                  \
+            (!__w->left || __w->left->color == RBTREE_BLACK)) {                                    \
+          __w->color = RBTREE_RED;                                                                 \
+          __x = __xparent;                                                                         \
+          __xparent = __x->parent;                                                                 \
+        } else {                                                                                   \
+          if (!__w->left || __w->left->color == RBTREE_BLACK) {                                    \
+            if (__w->right)                                                                        \
+              __w->right->color = RBTREE_BLACK;                                                    \
+            __w->color = RBTREE_RED;                                                               \
+            rbtree_rotate_left (root_ptr, __w);                                                    \
+            __w = __xparent->left;                                                                 \
+          }                                                                                        \
+          __w->color = __xparent->color;                                                           \
+          __xparent->color = RBTREE_BLACK;                                                         \
+          if (__w->left)                                                                           \
+            __w->left->color = RBTREE_BLACK;                                                       \
+          rbtree_rotate_right (root_ptr, __xparent);                                               \
+          __x = *(root_ptr);                                                                       \
+        }                                                                                          \
+      }                                                                                            \
+    }                                                                                              \
+    if (__x)                                                                                       \
+      __x->color = RBTREE_BLACK;                                                                   \
+  } while (0)
+
+#define rbtree_next(node, out)                                                                     \
+  do {                                                                                             \
+    (out) = NULL;                                                                                  \
+    if (node) {                                                                                    \
+      if ((node)->right) {                                                                         \
+        struct rb_node_link* __n = (node)->right;                                                  \
+        while (__n->left)                                                                          \
+          __n = __n->left;                                                                         \
+        (out) = __n;                                                                               \
+      } else {                                                                                     \
+        struct rb_node_link* __n = (node);                                                         \
+        struct rb_node_link* __p = (node)->parent;                                                 \
+        while (__p && __n == __p->right) {                                                         \
+          __n = __p;                                                                               \
+          __p = __p->parent;                                                                       \
+        }                                                                                          \
+        (out) = __p;                                                                               \
+      }                                                                                            \
+    }                                                                                              \
+  } while (0)
+
+#define rbtree_prev(node, out)                                                                     \
+  do {                                                                                             \
+    (out) = NULL;                                                                                  \
+    if (node) {                                                                                    \
+      if ((node)->left) {                                                                          \
+        struct rb_node_link* __n = (node)->left;                                                   \
+        while (__n->right)                                                                         \
+          __n = __n->right;                                                                        \
+        (out) = __n;                                                                               \
+      } else {                                                                                     \
+        struct rb_node_link* __n = (node);                                                         \
+        struct rb_node_link* __p = (node)->parent;                                                 \
+        while (__p && __n == __p->left) {                                                          \
+          __n = __p;                                                                               \
+          __p = __p->parent;                                                                       \
+        }                                                                                          \
+        (out) = __p;                                                                               \
+      }                                                                                            \
+    }                                                                                              \
+  } while (0)
+
+#define rbtree_first(root_ptr, out) rbtree_min (*(root_ptr), out)
+
+#define rbtree_last(root_ptr, out)                                                                 \
+  do {                                                                                             \
+    (out) = NULL;                                                                                  \
+    struct rb_node_link* __n = *(root_ptr);                                                        \
+    if (__n) {                                                                                     \
+      while (__n->right)                                                                           \
+        __n = __n->right;                                                                          \
+      (out) = __n;                                                                                 \
+    }                                                                                              \
+  } while (0)
+
+#endif // _KERNEL_LIBK_RBTREE_H

kernel/proc/proc.c

@@ -3,6 +3,7 @@
 #include <irq/irq.h>
 #include <libk/align.h>
 #include <libk/list.h>
+#include <libk/rbtree.h>
 #include <libk/std.h>
 #include <libk/string.h>
 #include <limine/requests.h>
@@ -22,8 +23,8 @@
 #include <amd64/intr_defs.h>
 #endif
 
-static struct procw* procs;
-static spin_lock_t procs_lock = SPIN_LOCK_INIT;
+static struct rb_node_link* proc_tree = NULL;
+static spin_lock_t proc_tree_lock = SPIN_LOCK_INIT;
 
 static bool proc_check_elf (uint8_t* elf) {
   if (!((elf[0] == 0x7F) && (elf[1] == 'E') && (elf[2] == 'L') && (elf[3] == 'F')))
@@ -108,52 +109,51 @@ static struct proc* proc_spawn_rd (char* name) {
   return proc_from_elf (rd_file->content);
 }
 
-static void proc_register_for_cpu (struct proc* proc, struct cpu* cpu) {
-  /* make available globally. */
-  struct procw* procw = malloc (sizeof (*procw));
-  if (procw == NULL)
-    return;
-  procw->proc = proc;
-  proc->procw = procw;
+static void proc_register (struct proc* proc, struct cpu* cpu) {
   proc->cpu = cpu;
 
-  spin_lock (&procs_lock);
+  spin_lock (&proc_tree_lock);
 
   spin_lock (&cpu->lock);
 
-  linklist_append (struct procw*, procs, procw);
-  linklist_append (struct proc*, cpu->proc_run_q, proc);
+  rbtree_insert (struct proc, &cpu->proc_run_q, &proc->cpu_run_q_link, cpu_run_q_link, pid);
+  rbtree_insert (struct proc, &proc_tree, &proc->proc_tree_link, proc_tree_link, pid);
 
   if (cpu->proc_current == NULL)
     cpu->proc_current = proc;
 
   spin_unlock (&cpu->lock);
 
-  spin_unlock (&procs_lock);
+  spin_unlock (&proc_tree_lock);
 }
 
 static struct proc* proc_find_sched (void) {
+  struct rb_node_link* node = NULL;
   struct proc* start = thiscpu->proc_current;
-  struct proc* proc = start->next;
+  struct proc* proc = NULL;
 
-  for (;;) {
-    if (proc == NULL) {
-      proc = thiscpu->proc_run_q;
-    }
+  if (start)
+    node = &start->cpu_run_q_link;
 
-    if (atomic_load (&proc->state) == PROC_READY) {
+  if (!node)
+    rbtree_first (&thiscpu->proc_run_q, node);
+
+  struct rb_node_link* first = node;
+  while (node) {
+    proc = rbtree_entry (node, struct proc, cpu_run_q_link);
+
+    if (atomic_load (&proc->state) == PROC_READY)
       return proc;
+
+    rbtree_next (node, node);
+
+    if (node == first)
+      break;
   }
 
-    /* No runnable processes found. */
-    if (proc == start) {
   return NULL;
 }
 
-    proc = proc->next;
-  }
-}
-
 void proc_sched (void) {
   struct proc* next = NULL;
 
@@ -190,7 +190,7 @@ static void proc_irq_sched (void* arg, void* regs) {
 
 void proc_init (void) {
   struct proc* init = proc_spawn_rd ("init.exe");
-  proc_register_for_cpu (init, thiscpu);
+  proc_register (init, thiscpu);
 
 #if defined(__x86_64__)
   irq_attach (&proc_irq_sched, NULL, SCHED_PREEMPT_TIMER, IRQ_INTERRUPT_SAFE);

kernel/proc/proc.h

@@ -3,6 +3,7 @@
 
 #include <aux/compiler.h>
 #include <aux/elf.h>
+#include <libk/rbtree.h>
 #include <libk/std.h>
 #include <sync/spin_lock.h>
 #include <sys/mm.h>
@@ -29,13 +30,16 @@ struct proc_mapping {
 struct procw;
 
 struct proc {
-  struct proc* next;
+  int pid;
+  struct rb_node_link proc_tree_link;
+  struct rb_node_link cpu_run_q_link;
+
   struct proc_mapping* mappings; /* pd.lock implicitly protects this field */
   struct proc_platformdata pdata;
   struct pd pd;
   spin_lock_t lock;
   struct cpu* cpu;
-  struct procw* procw; /* link to it's global struct */
+  // struct procw* procw; /* link to it's global struct */
   atomic_int state;
 };
 
@@ -44,10 +48,10 @@ struct proc {
  * struct procw is a process wrapper that is a member of
  * a global process list.
  */
-struct procw {
-  struct procw* next;
-  struct proc* proc;
-};
+/* struct procw { */
+/*   struct procw* next; */
+/*   struct proc* proc; */
+/* }; */
 
 void proc_sched (void);
 void proc_kill (struct proc* proc);