3 files changed, 148 insertions, 15 deletions

diff --git a/src/threadpool.c b/src/threadpool.c
index 3b1168c..953d0ee 100644
--- a/src/threadpool.c
+++ b/src/threadpool.c
@@ -1,10 +1,134 @@
 #include "common.h"
 #include "threadpool.h"
+#include "log.h"
 #include "smrt_arena.h"
 #include "deque.h"
 
+#include <bits/pthreadtypes.h>
+#include <iso646.h>
+#include <pthread.h>
+#include <stdatomic.h>
+#include <string.h>
+
+
 typedef struct {
-    tp_job_proc proc;
-    void *args;
+    tp_job_proc  proc;
+           void *args;
 } tp_job_t;
 
+typedef struct {
+          u64  id;
+thread_pool_t *tp;
+} worker_args;
+
+void *worker_proc(void *args) {
+    worker_args *wargs = (worker_args*)args;
+    thread_pool_t *tp = wargs->tp;
+
+    pthread_mutex_lock(tp->count_mtx); {
+        tp->num_alive++;
+    } pthread_mutex_unlock(tp->count_mtx);
+
+    for (;;) {
+        if (!wargs->tp->is_running) {
+            break;
+        }
+
+        tp_job_t *j = (tp_job_t*)ts_deque_pop(wargs->tp->jobs);
+
+        pthread_mutex_lock(tp->count_mtx);
+        tp->num_active++;
+        pthread_mutex_unlock(tp->count_mtx);
+
+        j->proc(j->args);
+
+        pthread_mutex_lock(tp->count_mtx);
+        tp->num_active--;
+        // If this is the last active thread finishing, tell someone about it
+        if (!tp->num_active) {
+            pthread_cond_broadcast(tp->done_signal);
+        }
+        pthread_mutex_unlock(tp->count_mtx);
+    }
+
+    pthread_mutex_lock(tp->count_mtx); {
+        tp->num_alive--;
+    } pthread_mutex_unlock(tp->count_mtx);
+
+    return NULL;
+}
+
+thread_pool_t *tp_create(smrt_arena_t *arena, u64 max_jobs, u64 num_threads) {
+    thread_pool_t *tp = smrt_arena_push(arena, sizeof(thread_pool_t), true);
+
+    ez_deque_t *q = ez_deque_create(arena, sizeof(tp_job_t), max_jobs);
+    ts_deque_t tsq = ts_deque_create(q);
+
+    pthread_t *threads = SMRTA_ALLOC_ARRAY(arena, pthread_t, num_threads);
+    worker_args * args = SMRTA_ALLOC_ARRAY(arena, worker_args, num_threads);
+    for (u64 i = 0; i < num_threads; i++) {
+        args[i] = (worker_args){
+            .id = i,
+            .tp = tp,
+        };
+        pthread_create(&threads[i], NULL, worker_proc, (void*)&args[i]);
+    }
+
+    pthread_mutex_t *count_mtx = smrt_arena_push(arena, sizeof(pthread_cond_t), true);
+    pthread_mutex_init(count_mtx, NULL);
+    pthread_cond_t *all_done = smrt_arena_push(arena, sizeof(pthread_cond_t), true);
+    pthread_cond_init(all_done, NULL);
+
+    *tp = (thread_pool_t){
+        .num_threads=num_threads,
+        .threads=threads,
+
+        .jobs=tsq,
+        .is_running=true,
+
+        .count_mtx=count_mtx,
+        .done_signal=all_done,
+        .num_alive=0,
+        .num_active=0,
+    };
+
+    return tp;
+}
+
+void no_op(void *nothing) { (void)nothing; }
+
+i32 tp_destroy(thread_pool_t *tp) {
+    tp->is_running = false;
+    for (;;) {
+        // TODO: add timeout to skip this if a thread is hung
+        u64 na = 0;
+        pthread_mutex_lock(tp->count_mtx);
+        if ((na = tp->num_alive) == 0) break;
+        pthread_mutex_unlock(tp->count_mtx);
+
+        for (u64 i = 0; i < na; i++) {
+            tp_push_job(tp, no_op, NULL);
+        }
+    }
+    ts_deque_destroy(tp->jobs);
+    pthread_mutex_destroy(tp->count_mtx);
+    pthread_cond_destroy(tp->done_signal);
+    return 0;
+}
+
+i32 tp_push_job(thread_pool_t *tp, tp_job_proc job, void *args) {
+    tp_job_t j = {
+        .proc = job,
+        .args = args,
+    };
+
+    return ts_deque_enqueue(tp->jobs, &j);
+}
+
+void tp_wait(thread_pool_t *tp) {
+    pthread_mutex_lock(tp->count_mtx);
+    while (tp->jobs.queue->occupied || tp->num_active) {
+        pthread_cond_wait(tp->done_signal, tp->count_mtx);
+    }
+    pthread_mutex_unlock(tp->count_mtx);
+}
diff --git a/src/threadpool.h b/src/threadpool.h
index 121d803..ae97583 100644
--- a/src/threadpool.h
+++ b/src/threadpool.h
@@ -10,23 +10,31 @@
 typedef void (*tp_job_proc)(void *args);
 
 typedef struct {
-    u64 num_threads;
-    pthread_t *threads;
+            u64  num_threads;
+      pthread_t *    threads;
 
-    sem_t threads_running;
-    ts_deque_t jobs;
+            b32   is_running;
+
+pthread_mutex_t *  count_mtx;
+ pthread_cond_t *done_signal;
+
+            u64   num_active;
+            u64    num_alive;
+
+     ts_deque_t         jobs;
 } thread_pool_t;
 
-// Job queue will be allocated on the arena.
-thread_pool_t tp_create(smrt_arena_t *      arena ,
-                                  u64    max_jobs ,
-                                  u64 num_threads);
+// Thread pool will be allocated on the arena.
+thread_pool_t *tp_create(smrt_arena_t *    arena ,
+                                 u64    max_jobs ,
+                                 u64 num_threads);
 
-i32 tp_destroy(thread_pool_t tp);
+// This function waits for threads to finish, frees relevant allocations, and then frees *tp.
+i32 tp_destroy(thread_pool_t *tp);
 
-i32 tp_push_job(thread_pool_t     tp ,
-                      tp_job_proc   job ,
-                          void *args);
+i32 tp_push_job(thread_pool_t *  tp ,
+                  tp_job_proc   job ,
+                         void *args);
 
 // Wait for all threads to finish
-void tp_wait(thread_pool_t tp);
+void tp_wait(thread_pool_t *tp);
diff --git a/src/unity.c b/src/unity.c
index 69dd9d1..6fe4973 100644
--- a/src/unity.c
+++ b/src/unity.c
@@ -28,4 +28,5 @@
 
 // Misc
 #include "log.c"
+#include "threadpool.c"
 #include "version.c"