diff options
Diffstat (limited to 'src/thread_pool.cpp')
| -rw-r--r-- | src/thread_pool.cpp | 184 |
1 files changed, 184 insertions, 0 deletions
diff --git a/src/thread_pool.cpp b/src/thread_pool.cpp new file mode 100644 index 000000000..2467ba609 --- /dev/null +++ b/src/thread_pool.cpp @@ -0,0 +1,184 @@ +// worker_queue.cpp + +#define WORKER_TASK_PROC(name) isize name(void *data) +typedef WORKER_TASK_PROC(WorkerTaskProc); + +struct WorkerTask { + WorkerTaskProc *do_work; + void *data; + isize result; +}; + + +struct ThreadPool { + gbMutex mutex; + gbSemaphore sem_available; + gbAtomic32 processing_work_count; + bool is_running; + + gbAllocator allocator; + + WorkerTask *tasks; + isize volatile task_head; + isize volatile task_tail; + isize volatile task_capacity; + + gbThread *threads; + isize thread_count; + + char worker_prefix[10]; + i32 worker_prefix_len; +}; + +void thread_pool_init(ThreadPool *pool, gbAllocator const &a, isize thread_count, char const *worker_prefix = nullptr); +void thread_pool_destroy(ThreadPool *pool); +void thread_pool_start(ThreadPool *pool); +void thread_pool_join(ThreadPool *pool); +void thread_pool_add_task(ThreadPool *pool, WorkerTaskProc *proc, void *data); +void thread_pool_kick(ThreadPool *pool); +void thread_pool_kick_and_wait(ThreadPool *pool); +GB_THREAD_PROC(worker_thread_internal); + +void thread_pool_init(ThreadPool *pool, gbAllocator const &a, isize thread_count, char const *worker_prefix) { + pool->allocator = a; + pool->task_head = 0; + pool->task_tail = 0; + pool->task_capacity = 1024; + pool->tasks = gb_alloc_array(a, WorkerTask, pool->task_capacity); + pool->thread_count = gb_max(thread_count, 0); + pool->threads = gb_alloc_array(a, gbThread, pool->thread_count); + gb_mutex_init(&pool->mutex); + gb_semaphore_init(&pool->sem_available); + pool->is_running = true; + + pool->worker_prefix_len = 0; + if (worker_prefix) { + i32 worker_prefix_len = cast(i32)gb_strlen(worker_prefix); + worker_prefix_len = gb_min(worker_prefix_len, 10); + gb_memmove(pool->worker_prefix, worker_prefix, worker_prefix_len); + pool->worker_prefix_len = worker_prefix_len; + } + + for (isize i = 0; i < pool->thread_count; i++) { + gbThread *t = &pool->threads[i]; + gb_thread_init(t); + t->user_index = i; + #if 0 + // TODO(bill): Fix this on Linux as it causes a seg-fault + if (pool->worker_prefix_len > 0) { + char worker_name[16] = {}; + gb_snprintf(worker_name, gb_size_of(worker_name), "%.*s%u", pool->worker_prefix_len, pool->worker_prefix, cast(u16)i); + gb_thread_set_name(t, worker_name); + } + #endif + } +} + +void thread_pool_start(ThreadPool *pool) { + for (isize i = 0; i < pool->thread_count; i++) { + gbThread *t = &pool->threads[i]; + gb_thread_start(t, worker_thread_internal, pool); + } +} + +void thread_pool_join(ThreadPool *pool) { + pool->is_running = false; + + gb_semaphore_post(&pool->sem_available, cast(i32)pool->thread_count); + + gb_yield(); + + for (isize i = 0; i < pool->thread_count; i++) { + gbThread *t = &pool->threads[i]; + gb_thread_join(t); + } +} + + +void thread_pool_destroy(ThreadPool *pool) { + thread_pool_join(pool); + + gb_semaphore_destroy(&pool->sem_available); + gb_mutex_destroy(&pool->mutex); + gb_free(pool->allocator, pool->threads); + pool->thread_count = 0; + gb_free(pool->allocator, pool->tasks); + pool->task_head = 0; + pool->task_tail = 0; + pool->task_capacity = 0; +} + + +void thread_pool_add_task(ThreadPool *pool, WorkerTaskProc *proc, void *data) { + gb_mutex_lock(&pool->mutex); + + if (pool->task_tail == pool->task_capacity) { + isize new_cap = 2*pool->task_capacity + 8; + WorkerTask *new_tasks = gb_alloc_array(pool->allocator, WorkerTask, new_cap); + gb_memmove(new_tasks, pool->tasks, (pool->task_tail)*gb_size_of(WorkerTask)); + pool->tasks = new_tasks; + pool->task_capacity = new_cap; + } + WorkerTask task = {}; + task.do_work = proc; + task.data = data; + + pool->tasks[pool->task_tail++] = task; + gb_semaphore_post(&pool->sem_available, 1); + gb_mutex_unlock(&pool->mutex); +} + +bool thread_pool_try_and_pop_task(ThreadPool *pool, WorkerTask *task) { + bool got_task = false; + if (gb_mutex_try_lock(&pool->mutex)) { + if (pool->task_tail > pool->task_head) { + gb_atomic32_fetch_add(&pool->processing_work_count, +1); + *task = pool->tasks[pool->task_head++]; + got_task = true; + } + gb_mutex_unlock(&pool->mutex); + } + return got_task; +} +void thread_pool_do_work(ThreadPool *pool, WorkerTask *task) { + task->result = task->do_work(task->data); + gb_atomic32_fetch_add(&pool->processing_work_count, -1); +} + +void thread_pool_wait_to_process(ThreadPool *pool) { + while (pool->task_tail > pool->task_head || gb_atomic32_load(&pool->processing_work_count) != 0) { + WorkerTask task = {}; + if (thread_pool_try_and_pop_task(pool, &task)) { + thread_pool_do_work(pool, &task); + } + + // Safety-kick + if (pool->task_tail > pool->task_head && gb_atomic32_load(&pool->processing_work_count) == 0) { + gb_mutex_lock(&pool->mutex); + gb_semaphore_post(&pool->sem_available, cast(i32)(pool->task_tail-pool->task_head)); + gb_mutex_unlock(&pool->mutex); + } + + gb_yield(); + } + + thread_pool_join(pool); +} + + +GB_THREAD_PROC(worker_thread_internal) { + ThreadPool *pool = cast(ThreadPool *)thread->user_data; + while (pool->is_running) { + gb_semaphore_wait(&pool->sem_available); + + WorkerTask task = {}; + if (thread_pool_try_and_pop_task(pool, &task)) { + thread_pool_do_work(pool, &task); + } + } + // Cascade + gb_semaphore_release(&pool->sem_available); + + return 0; +} + |