Update work_pool.zig

1 files changed, 28 insertions, 21 deletions

diff --git a/src/work_pool.zig b/src/work_pool.zig
index 090f74216..a663e7186 100644
--- a/src/work_pool.zig
+++ b/src/work_pool.zig
@@ -1,32 +1,39 @@
 const ThreadPool = @import("thread_pool");
 const std = @import("std");
-var pool: ThreadPool = undefined;
-var loaded: bool = false;
 
 pub const Batch = ThreadPool.Batch;
 pub const Task = ThreadPool.Task;
 
-fn create() *ThreadPool {
-    @setCold(true);
+pub fn NewWorkPool(comptime max_threads: ?usize) type {
+    return struct {
+        var pool: ThreadPool = undefined;
+        var loaded: bool = false;
 
-    pool = ThreadPool.init(.{
-        .max_threads = @floatToInt(u32, @floor(@intToFloat(f32, @maximum(std.Thread.getCpuCount() catch 0, 2)) * 0.8)),
-        .stack_size = 2 * 1024 * 1024,
-    });
-    return &pool;
-}
-pub inline fn get() *ThreadPool {
-    // lil racy
-    if (loaded) return &pool;
-    loaded = true;
+        fn create() *ThreadPool {
+            @setCold(true);
 
-    return create();
-}
+            pool = ThreadPool.init(.{
+                .max_threads = max_threads orelse @floatToInt(u32, @floor(@intToFloat(f32, @maximum(std.Thread.getCpuCount() catch 0, 2)) * 0.8)),
+                .stack_size = 2 * 1024 * 1024,
+            });
+            return &pool;
+        }
+        pub inline fn get() *ThreadPool {
+            // lil racy
+            if (loaded) return &pool;
+            loaded = true;
 
-pub fn scheduleBatch(batch: ThreadPool.Batch) void {
-    get().schedule(batch);
-}
+            return create();
+        }
 
-pub fn schedule(task: *ThreadPool.Task) void {
-    get().schedule(ThreadPool.Batch.from(task));
+        pub fn scheduleBatch(batch: ThreadPool.Batch) void {
+            get().schedule(batch);
+        }
+
+        pub fn schedule(task: *ThreadPool.Task) void {
+            get().schedule(ThreadPool.Batch.from(task));
+        }
+    };
 }
+
+pub const WorkPool = NewWorkPool(null);