From 95b97118de09c8d25e621cef009a76a5341fad45 Mon Sep 17 00:00:00 2001
From: "Michael D. Lowis" <mike@mdlowis.com>
Date: Tue, 13 Sep 2022 20:18:03 -0400
Subject: [PATCH] cleaned up some multithreading issues based on sanitizer
 results

---
 Kernel.c | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/Kernel.c b/Kernel.c
index cd30388..f513923 100644
--- a/Kernel.c
+++ b/Kernel.c
@@ -19,6 +19,7 @@
 #define WORD_SIZE           sizeof(void*)
 #define MAX_TASKS           1024
 #define TASK_NONE           ((TaskID_T)-1)
+#define DEFAULT_STACK_SIZE  (32768)
 
 #define WORD_ALIGN(x) \
     (((x) + WORD_SIZE) & ~(WORD_SIZE))
@@ -34,7 +35,7 @@ typedef enum {
 } TaskState_T;
 
 typedef struct Task_T {
-    long* stack_top;
+    long* _Atomic stack_top;
     long* heap_top;
     long* memory;
     struct Task_T* next;
@@ -91,10 +92,7 @@ static void WaitForTaskFree(TaskID_T prev, TaskID_T next)
 {
     if (next != TASK_NONE && prev != next)
     {
-        Task_T* ntask = &Tasks[next];
-        assert(ntask);
-
-        while (!atomic_load(&ntask->stack_top))
+        while (!atomic_load(&Tasks[next].stack_top))
         {
         }
     }
@@ -249,10 +247,14 @@ void PickNewTask(bool dead)
 
     WaitForTaskFree(prev, Running[CpuID].task);
     SwapTask(prev_task, next_task);
+
     if (Running[CpuID].dead != TASK_NONE)
     {
-        WaitForTaskFree(prev, Running[CpuID].dead);
+        WaitForTaskFree(TASK_NONE, Running[CpuID].dead);
+        AcquireLock();
         free(Tasks[Running[CpuID].dead].memory);
+        Tasks[Running[CpuID].dead].memory = NULL;
+        ReleaseLock();
         Running[CpuID].dead = TASK_NONE;
     }
 }
@@ -269,7 +271,6 @@ void Kernel_Exit(void)
 
 static TaskID_T AllocateTaskID(void)
 {
-    AcquireLock();
     TaskID_T tid = TASK_NONE;
     if (TaskCount < MAX_TASKS)
     {
@@ -281,7 +282,6 @@ static TaskID_T AllocateTaskID(void)
         /* TODO: iterate over table to find free slot... */
         assert(!"out of task IDs");
     }
-    ReleaseLock();
     return tid;
 }
 
@@ -290,7 +290,7 @@ static void InitializeTask(Task_T* task, void (*task_fn)(void*), void* arg, long
     /* allocate a new task. default memory size is used if 0 provided */
     if (memsz == 0)
     {
-        memsz = 32768;
+        memsz = DEFAULT_STACK_SIZE;
     }
     memsz = WORD_ALIGN(memsz);
     task->memory    = calloc(1, WORD_ALIGN(memsz));
@@ -318,14 +318,14 @@ static void InitializeTask(Task_T* task, void (*task_fn)(void*), void* arg, long
 
 TaskID_T Kernel_Spawn(void (*task_fn)(void*), void* arg, long int argsz, long int memsz)
 {
+    AcquireLock();
     TaskID_T tid = AllocateTaskID();
     if (tid != TASK_NONE)
     {
         InitializeTask(&Tasks[tid], task_fn, arg, argsz, memsz);
-        AcquireLock();
         Enter(tid);
-        ReleaseLock();
     }
+    ReleaseLock();
     return tid;
 }
 
@@ -420,6 +420,7 @@ void task1(void* arg)
     int* val = arg;
     for (int i = 0; i < 1000; i++)
     {
+        (void)val;
         printf("%d %d\n", *val, CpuID);
         Kernel_Yield();
     }
-- 
2.54.0