Fix CPU load balancer bugs, scheduling points support for remote CPUs

kernel/amd64/bootmain.c

@@ -32,7 +32,7 @@ ALIGNED (16) static uint8_t uacpi_memory_buffer[UACPI_MEMORY_BUFFER_MAX];
 void bootmain (void) {
   struct limine_mp_response* mp = limine_mp_request.response;
 
-  struct cpu* bsp_cpu = cpu_make (mp->bsp_lapic_id);
+  struct cpu* bsp_cpu = cpu_make (mp->bsp_lapic_id, 0);
 
   amd64_init (bsp_cpu, false);
   syscall_init ();

kernel/amd64/intr.c

@@ -157,7 +157,9 @@ static void amd64_intr_exception (struct saved_regs* regs) {
                regs->rbx);
 
   if (regs->cs == (GDT_UCODE | 0x03)) {
-    proc_kill (thiscpu->proc_current);
+    struct cpu* reschedule_cpu;
+    if (proc_kill (thiscpu->proc_current, &reschedule_cpu) == PROC_NEED_RESCHEDULE)
+      cpu_request_sched (reschedule_cpu);
   } else {
     spin ();
   }

kernel/amd64/smp.c

@@ -15,22 +15,19 @@
 #include <sys/smp.h>
 #include <sys/syscall.h>
 
-/// Cpu ID counter
-static atomic_uint cpu_counter = 0;
 /// The CPUs
 static struct cpu cpus[CPUS_MAX];
 
-static atomic_int cpu_init_count;
+static atomic_int last_cpu_index = 0;
+static atomic_int cpu_counter;
 
 /// Allocate a CPU structure
-struct cpu* cpu_make (uint64_t lapic_id) {
-  int id = atomic_fetch_add (&cpu_counter, 1);
-
-  struct cpu* cpu = &cpus[id];
+struct cpu* cpu_make (uint64_t lapic_id, uint64_t cpu_id) {
+  struct cpu* cpu = &cpus[cpu_id];
 
   memset (cpu, 0, sizeof (*cpu));
   cpu->lock = SPIN_LOCK_INIT;
-  cpu->id = id;
+  cpu->id = cpu_id;
   cpu->lapic_id = lapic_id;
 
   amd64_wrmsr (MSR_GS_BASE, (uint64_t)cpu);
@@ -53,27 +50,30 @@ void cpu_request_sched (struct cpu* cpu) {
 }
 
 struct cpu* cpu_find_lightest (void) {
-  struct cpu* cpu = &cpus[0];
+  struct limine_mp_response* mp = limine_mp_request.response;
 
-  int load = atomic_load (&cpu->proc_run_q_count);
+  int start = atomic_fetch_add (&last_cpu_index, 1) % mp->cpu_count;
+  struct cpu* best_cpu = &cpus[start];
+  int best_load = atomic_load (&best_cpu->proc_run_q_count);
 
-  for (unsigned int i = 1; i < cpu_counter; i++) {
-    struct cpu* new_cpu = &cpus[i];
-    int new_load = atomic_load (&new_cpu->proc_run_q_count);
-    if (new_load < load) {
-      load = new_load;
-      cpu = new_cpu;
+  for (int i = 1; i < (int)mp->cpu_count; i++) {
+    int idx = (start + i) % mp->cpu_count;
+    struct cpu* cpu = &cpus[idx];
+    int l = atomic_load (&cpu->proc_run_q_count);
+    if (l < best_load) {
+      best_load = l;
+      best_cpu = cpu;
     }
   }
 
-  return cpu;
+  return best_cpu;
 }
 
 /// Bootstrap code for non-BSP CPUs
 static void amd64_smp_bootstrap (struct limine_mp_info* mp_info) {
   amd64_load_kernel_cr3 ();
 
-  struct cpu* cpu = cpu_make (mp_info->lapic_id);
+  struct cpu* cpu = cpu_make (mp_info->lapic_id, mp_info->processor_id);
 
   amd64_init (cpu, true); /* gdt + idt */
   syscall_init ();
@@ -82,10 +82,11 @@ static void amd64_smp_bootstrap (struct limine_mp_info* mp_info) {
 
   DEBUG ("CPU %u is online!\n", thiscpu->id);
 
-  atomic_fetch_sub (&cpu_init_count, 1);
+  atomic_fetch_sub (&cpu_counter, 1);
 
   struct proc* spin_proc = proc_spawn_rd ("spin.exe");
-  proc_register (spin_proc, thiscpu);
+  struct cpu* spin_cpu = thiscpu;
+  proc_register (spin_proc, &spin_cpu);
 
   spin_lock_ctx_t ctxcpu;
   spin_lock (&spin_proc->cpu->lock, &ctxcpu);
@@ -98,7 +99,7 @@ void smp_init (void) {
 
   struct limine_mp_response* mp = limine_mp_request.response;
 
-  cpu_init_count = mp->cpu_count - 1; /* Don't include BSP */
+  cpu_counter = mp->cpu_count - 1;
 
   for (size_t i = 0; i < mp->cpu_count; i++) {
     if (mp->cpus[i]->lapic_id != thiscpu->lapic_id) {
@@ -106,8 +107,6 @@ void smp_init (void) {
     }
   }
 
-  while (atomic_load (&cpu_init_count) > 0)
+  while (atomic_load (&cpu_counter) > 0)
     ;
-
-  DEBUG ("All CPUs are online\n");
 }

kernel/amd64/smp.h

@@ -34,7 +34,7 @@ struct cpu {
   atomic_int proc_run_q_count;
 };
 
-struct cpu* cpu_make (uint64_t lapic_id);
+struct cpu* cpu_make (uint64_t lapic_id, uint64_t cpu_id);
 struct cpu* cpu_get (void);
 void cpu_request_sched (struct cpu* cpu);
 struct cpu* cpu_find_lightest (void);

kernel/amd64/syscall.c

@@ -21,6 +21,7 @@ uintptr_t amd64_syscall_dispatch (void* stack_ptr) {
 
   spin_lock (&thiscpu->lock, &ctxcpu);
   struct proc* caller = thiscpu->proc_current;
+  int caller_pid = caller->pid;
   spin_lock (&caller->lock, &ctxpr);
 
   memcpy (&caller->pdata.regs, regs, sizeof (struct saved_regs));
@@ -35,7 +36,24 @@ uintptr_t amd64_syscall_dispatch (void* stack_ptr) {
     return -ST_SYSCALL_NOT_FOUND;
   }
 
-  return func (caller, regs, regs->rdi, regs->rsi, regs->rdx, regs->r10, regs->r8, regs->r9);
+  bool reschedule = false;
+
+  struct cpu* reschedule_cpu = NULL;
+  uintptr_t r = func (caller, regs, &reschedule, &reschedule_cpu, regs->rdi, regs->rsi, regs->rdx,
+                      regs->r10, regs->r8, regs->r9);
+
+  caller = proc_find_pid (caller_pid);
+
+  if (caller != NULL) {
+    spin_lock (&caller->lock, &ctxpr);
+    caller->pdata.regs.rax = r;
+    spin_unlock (&caller->lock, &ctxpr);
+  }
+
+  if (reschedule)
+    cpu_request_sched (reschedule_cpu);
+
+  return r;
 }
 
 void syscall_init (void) {

kernel/proc/mutex.c

@@ -37,14 +37,15 @@ void proc_mutexes_cleanup (struct proc* proc) {
     if (resource->u.mutex.owner == proc && resource->u.mutex.locked) {
       spin_unlock (&resource->lock, &ctxrs);
 
-      proc_mutex_unlock (proc, &resource->u.mutex);
+      struct cpu* reschedule_cpu;
+      proc_mutex_unlock (proc, &resource->u.mutex, &reschedule_cpu);
     }
   }
 
   spin_unlock (&proc->procgroup->lock, &ctxpg);
 }
 
-bool proc_cleanup_resource_mutex (struct proc_resource* resource) {
+bool proc_cleanup_resource_mutex (struct proc_resource* resource, struct cpu** reschedule_cpu) {
   struct proc_mutex* mutex = &resource->u.mutex;
   spin_lock_ctx_t ctxmt, ctxsq;
 
@@ -62,7 +63,7 @@ bool proc_cleanup_resource_mutex (struct proc_resource* resource) {
     spin_unlock (&mutex->suspension_q.lock, &ctxsq);
     spin_unlock (&mutex->resource->lock, &ctxmt);
 
-    reschedule = reschedule || proc_sq_resume (suspended_proc, sq_entry);
+    reschedule = reschedule || proc_sq_resume (suspended_proc, sq_entry, reschedule_cpu);
 
     /* reacquire */
     spin_lock (&mutex->resource->lock, &ctxmt);
@@ -78,7 +79,7 @@ bool proc_cleanup_resource_mutex (struct proc_resource* resource) {
   return reschedule;
 }
 
-bool proc_mutex_lock (struct proc* proc, struct proc_mutex* mutex) {
+bool proc_mutex_lock (struct proc* proc, struct proc_mutex* mutex, struct cpu** reschedule_cpu) {
   spin_lock_ctx_t ctxmt;
 
   spin_lock (&mutex->resource->lock, &ctxmt);
@@ -90,10 +91,11 @@ bool proc_mutex_lock (struct proc* proc, struct proc_mutex* mutex) {
     return PROC_NO_RESCHEDULE;
   }
 
-  return proc_sq_suspend (proc, &mutex->suspension_q, &mutex->resource->lock, &ctxmt);
+  return proc_sq_suspend (proc, &mutex->suspension_q, &mutex->resource->lock, &ctxmt,
+                          reschedule_cpu);
 }
 
-bool proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex) {
+bool proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex, struct cpu** reschedule_cpu) {
   spin_lock_ctx_t ctxmt, ctxsq;
 
   spin_lock (&mutex->resource->lock, &ctxmt);
@@ -117,7 +119,7 @@ bool proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex) {
     spin_unlock (&mutex->suspension_q.lock, &ctxsq);
     spin_unlock (&mutex->resource->lock, &ctxmt);
 
-    return proc_sq_resume (resumed_proc, sq_entry);
+    return proc_sq_resume (resumed_proc, sq_entry, reschedule_cpu);
   }
 
   mutex->locked = false;
@@ -126,5 +128,5 @@ bool proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex) {
   spin_unlock (&mutex->suspension_q.lock, &ctxsq);
   spin_unlock (&mutex->resource->lock, &ctxmt);
 
-  return PROC_NEED_RESCHEDULE;
+  return PROC_NO_RESCHEDULE;
 }

kernel/proc/mutex.h

@@ -6,6 +6,7 @@
 
 struct proc;
 struct proc_resource;
+struct cpu;
 
 struct proc_mutex {
   struct proc_resource* resource;
@@ -15,9 +16,9 @@ struct proc_mutex {
   struct proc* owner;
 };
 
-bool proc_cleanup_resource_mutex (struct proc_resource* resource);
-bool proc_mutex_lock (struct proc* proc, struct proc_mutex* mutex);
-bool proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex);
+bool proc_cleanup_resource_mutex (struct proc_resource* resource, struct cpu** reschedule_cpu);
+bool proc_mutex_lock (struct proc* proc, struct proc_mutex* mutex, struct cpu** reschedule_cpu);
+bool proc_mutex_unlock (struct proc* proc, struct proc_mutex* mutex, struct cpu** reschedule_cpu);
 void proc_mutexes_cleanup (struct proc* proc);
 
 #endif // _KERNEL_PROC_MUTEX_H

kernel/proc/proc.c

@@ -122,15 +122,16 @@ struct proc* proc_find_pid (int pid) {
   return proc;
 }
 
-void proc_register (struct proc* proc, struct cpu* cpu1) {
-  spin_lock_ctx_t ctxcpu, ctxprtr;
+bool proc_register (struct proc* proc, struct cpu** reschedule_cpu) {
+  spin_lock_ctx_t ctxcpu, ctxprtr, ctxpr;
 
-  proc->cpu = cpu1 != NULL ? cpu1 : cpu_find_lightest ();
-
-  struct cpu* cpu = proc->cpu;
+  struct cpu* cpu = *reschedule_cpu != NULL ? *reschedule_cpu : cpu_find_lightest ();
 
   spin_lock (&proc_tree_lock, &ctxprtr);
   spin_lock (&cpu->lock, &ctxcpu);
+  spin_lock (&proc->lock, &ctxpr);
+
+  proc->cpu = cpu;
 
   rbtree_insert (struct proc, &proc_tree, &proc->proc_tree_link, proc_tree_link, pid);
 
@@ -139,8 +140,13 @@ void proc_register (struct proc* proc, struct cpu* cpu1) {
   if (cpu->proc_current == NULL)
     cpu->proc_current = proc;
 
-  spin_unlock (&proc_tree_lock, &ctxprtr);
+  spin_unlock (&proc->lock, &ctxpr);
   spin_unlock (&cpu->lock, &ctxcpu);
+  spin_unlock (&proc_tree_lock, &ctxprtr);
+
+  *reschedule_cpu = cpu;
+
+  return PROC_NEED_RESCHEDULE;
 }
 
 /* caller holds cpu->lock */
@@ -237,27 +243,32 @@ void proc_sched (void) {
   }
 }
 
-void proc_kill (struct proc* proc) {
+bool proc_kill (struct proc* proc, struct cpu** reschedule_cpu) {
   spin_lock_ctx_t ctxpr, ctxcpu;
-  struct cpu* cpu = proc->cpu;
 
   spin_lock (&proc->lock, &ctxpr);
-  atomic_store (&proc->state, PROC_DEAD);
-  proc->cpu = NULL;
+  struct cpu* cpu = proc->cpu;
   spin_unlock (&proc->lock, &ctxpr);
 
   spin_lock (&cpu->lock, &ctxcpu);
+  spin_lock (&proc->lock, &ctxpr);
+
+  atomic_store (&proc->state, PROC_DEAD);
+  proc->cpu = NULL;
 
   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;
 
+  spin_unlock (&proc->lock, &ctxpr);
   spin_unlock (&cpu->lock, &ctxcpu);
 
   DEBUG ("killed PID %d\n", proc->pid);
 
-  cpu_request_sched (cpu);
+  *reschedule_cpu = cpu;
+
+  return PROC_NEED_RESCHEDULE;
 }
 
 static void proc_irq_sched (void* arg, void* regs) {
@@ -272,10 +283,12 @@ void proc_init (void) {
 #endif
 
   struct proc* spin_proc = proc_spawn_rd ("spin.exe");
-  proc_register (spin_proc, thiscpu);
+  struct cpu* spin_cpu = thiscpu;
+  proc_register (spin_proc, &spin_cpu);
 
   struct proc* init = proc_spawn_rd ("init.exe");
-  proc_register (init, NULL);
+  struct cpu* init_cpu = thiscpu;
+  proc_register (init, &init_cpu);
 
   spin_lock_ctx_t ctxcpu;
   spin_lock (&spin_proc->cpu->lock, &ctxcpu);

kernel/proc/proc.h

@@ -47,9 +47,9 @@ struct proc {
 };
 
 void proc_sched (void);
-void proc_kill (struct proc* proc);
+bool proc_kill (struct proc* proc, struct cpu** reschedule_cpu);
 struct elf_aux proc_load_segments (struct proc* proc, uint8_t* elf);
-void proc_register (struct proc* proc, struct cpu* cpu);
+bool proc_register (struct proc* proc, struct cpu** reschedule_cpu);
 struct proc* proc_find_pid (int pid);
 struct proc* proc_spawn_rd (char* name);
 void proc_init (void);

kernel/proc/procgroup.c

@@ -197,7 +197,8 @@ void procgroup_detach (struct procgroup* procgroup, struct proc* proc) {
 
       rnode = next;
 
-      proc_delete_resource (resource);
+      struct cpu* reschedule_cpu;
+      proc_delete_resource (resource, &reschedule_cpu);
     }
 
     struct list_node_link *mapping_link, *mapping_link_tmp;

kernel/proc/resource.c

@@ -51,8 +51,8 @@ struct proc_resource* proc_create_resource_mutex (struct procgroup* procgroup, i
   return resource;
 }
 
-bool proc_delete_resource (struct proc_resource* resource) {
-  bool reschedule = resource->ops.cleanup (resource);
+bool proc_delete_resource (struct proc_resource* resource, struct cpu** reschedule_cpu) {
+  bool reschedule = resource->ops.cleanup (resource, reschedule_cpu);
   free (resource);
 
   return reschedule;

kernel/proc/resource.h

@@ -11,6 +11,7 @@
 
 struct proc;
 struct procgroup;
+struct cpu;
 
 struct proc_resource {
   int type;
@@ -21,12 +22,12 @@ struct proc_resource {
     struct proc_mutex mutex;
   } u;
   struct {
-    bool (*cleanup) (struct proc_resource* resource);
+    bool (*cleanup) (struct proc_resource* resource, struct cpu** reschedule_cpu);
   } ops;
 };
 
 struct proc_resource* proc_find_resource (struct procgroup* procgroup, int rid);
 struct proc_resource* proc_create_resource_mutex (struct procgroup* procgroup, int rid);
-bool proc_delete_resource (struct proc_resource* resource);
+bool proc_delete_resource (struct proc_resource* resource, struct cpu** reschedule_cpu);
 
 #endif // _KERNEL_PROC_RESOURCE_H

kernel/proc/suspension_q.c

@@ -9,7 +9,7 @@
 #include <sys/spin_lock.h>
 
 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* ctxrl, struct cpu** reschedule_cpu) {
   spin_lock_ctx_t ctxpr, ctxcpu, ctxsq;
   struct cpu* cpu = proc->cpu;
 
@@ -48,10 +48,13 @@ bool proc_sq_suspend (struct proc* proc, struct proc_suspension_q* sq, spin_lock
   spin_unlock (&proc->lock, &ctxpr);
   spin_unlock (&cpu->lock, &ctxcpu);
 
+  *reschedule_cpu = cpu;
+
   return PROC_NEED_RESCHEDULE;
 }
 
-bool proc_sq_resume (struct proc* proc, struct proc_sq_entry* sq_entry) {
+bool proc_sq_resume (struct proc* proc, struct proc_sq_entry* sq_entry,
+                     struct cpu** reschedule_cpu) {
   spin_lock_ctx_t ctxsq, ctxpr, ctxcpu;
   struct cpu* cpu = cpu_find_lightest ();
   struct proc_suspension_q* sq = sq_entry->sq;
@@ -80,6 +83,8 @@ bool proc_sq_resume (struct proc* proc, struct proc_sq_entry* sq_entry) {
 
   free (sq_entry);
 
+  *reschedule_cpu = cpu;
+
   return PROC_NEED_RESCHEDULE;
 }
 

kernel/proc/suspension_q.h

@@ -5,6 +5,7 @@
 #include <sync/spin_lock.h>
 
 struct proc;
+struct cpu;
 
 struct proc_suspension_q {
   struct list_node_link* proc_list;
@@ -20,7 +21,8 @@ struct proc_sq_entry {
 
 void proc_sqs_cleanup (struct proc* proc);
 bool proc_sq_suspend (struct proc* proc, struct proc_suspension_q* sq, spin_lock_t* resource_lock,
-                      spin_lock_ctx_t* ctxrl);
-bool proc_sq_resume (struct proc* proc, struct proc_sq_entry* sq_entry);
+                      spin_lock_ctx_t* ctxrl, struct cpu** reschedule_cpu);
+bool proc_sq_resume (struct proc* proc, struct proc_sq_entry* sq_entry,
+                     struct cpu** reschedule_cpu);
 
 #endif // _KERNEL_PROC_SUSPENTION_Q_H

kernel/syscall/syscall.c

@@ -16,9 +16,10 @@
 #include <syscall/syscall.h>
 
 #define DEFINE_SYSCALL(name)                                                                       \
-  uintptr_t name (struct proc* UNUSED proc, void* UNUSED regs, uintptr_t UNUSED a1,                \
-                  uintptr_t UNUSED a2, uintptr_t UNUSED a3, uintptr_t UNUSED a4,                   \
-                  uintptr_t UNUSED a5, uintptr_t UNUSED a6)
+  uintptr_t name (struct proc* UNUSED proc, void* UNUSED regs, bool* UNUSED reschedule,            \
+                  struct cpu** UNUSED reschedule_cpu, uintptr_t UNUSED a1, uintptr_t UNUSED a2,    \
+                  uintptr_t UNUSED a3, uintptr_t UNUSED a4, uintptr_t UNUSED a5,                   \
+                  uintptr_t UNUSED a6)
 
 #define SYSRESULT(x) ((uintptr_t)(x))
 
@@ -44,7 +45,9 @@ static void* sys_get_user_buffer (struct proc* proc, uintptr_t uvaddr, size_t si
 
 /* int quit (void) */
 DEFINE_SYSCALL (sys_quit) {
-  proc_kill (proc);
+  if (proc_kill (proc, reschedule_cpu) == PROC_NEED_RESCHEDULE)
+    *reschedule = true;
+
   return SYSRESULT (ST_OK);
 }
 
@@ -92,7 +95,8 @@ DEFINE_SYSCALL (sys_clone) {
 
   int pid = new->pid;
 
-  proc_register (new, NULL);
+  if (proc_register (new, reschedule_cpu) == PROC_NEED_RESCHEDULE)
+    *reschedule = true;
 
   return SYSRESULT (pid);
 }
@@ -102,7 +106,7 @@ DEFINE_SYSCALL (sys_argument_ptr) { return proc->uvaddr_argument; }
 
 /* int sched (void) */
 DEFINE_SYSCALL (sys_sched) {
-  proc_sched ();
+  *reschedule = true;
   return SYSRESULT (ST_OK);
 }
 
@@ -127,8 +131,8 @@ DEFINE_SYSCALL (sys_mutex_delete) {
   if (mutex_resource == NULL)
     return SYSRESULT (-ST_NOT_FOUND);
 
-  if (proc_delete_resource (mutex_resource) == PROC_NEED_RESCHEDULE)
-    proc_sched ();
+  if (proc_delete_resource (mutex_resource, reschedule_cpu) == PROC_NEED_RESCHEDULE)
+    *reschedule = true;
 
   return SYSRESULT (ST_OK);
 }
@@ -142,8 +146,8 @@ DEFINE_SYSCALL (sys_mutex_lock) {
   if (mutex_resource == NULL)
     return SYSRESULT (-ST_NOT_FOUND);
 
-  if (proc_mutex_lock (proc, &mutex_resource->u.mutex) == PROC_NEED_RESCHEDULE)
-    proc_sched ();
+  if (proc_mutex_lock (proc, &mutex_resource->u.mutex, reschedule_cpu) == PROC_NEED_RESCHEDULE)
+    *reschedule = true;
 
   return SYSRESULT (ST_OK);
 }
@@ -157,8 +161,8 @@ DEFINE_SYSCALL (sys_mutex_unlock) {
   if (mutex_resource == NULL)
     return SYSRESULT (-ST_NOT_FOUND);
 
-  if (proc_mutex_unlock (proc, &mutex_resource->u.mutex) == PROC_NEED_RESCHEDULE)
-    proc_sched ();
+  if (proc_mutex_unlock (proc, &mutex_resource->u.mutex, reschedule_cpu) == PROC_NEED_RESCHEDULE)
+    *reschedule = true;
 
   return SYSRESULT (ST_OK);
 }

kernel/syscall/syscall.h

@@ -4,7 +4,10 @@
 #include <libk/std.h>
 #include <proc/proc.h>
 
-typedef uintptr_t (*syscall_handler_func_t) (struct proc* proc, void* regs, uintptr_t a1,
+struct cpu;
+
+typedef uintptr_t (*syscall_handler_func_t) (struct proc* proc, void* regs, bool* reschedule,
+                                             struct cpu** reschedule_cpu, uintptr_t a1,
                                              uintptr_t a2, uintptr_t a3, uintptr_t a4, uintptr_t a5,
                                              uintptr_t a6);