Something went wrong. Try again.
fork of gitlab.com/devnw/zig/thread-pool for zig 0.17
Something went wrong. Try again.
6.2 kB · 160 lines
Zig
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161const std = @import("std");const thread_pool = @import("thread_pool");
pub const Counter = struct { value: std.atomic.Value(u64) = std.atomic.Value(u64).init(0),};
/// COMPTIME-ELISION: returns the build-pinned timestamp mode, or `.full`/// (the ambient default) when the build is unpinned. Tests that exercise/// `.disabled`-specific behavior must check this and skip when it returns/// `.full`; tests that call this function pick up the pinned value so/// `Core.init` does not reject the runtime config.pub const default_timestamp_mode: thread_pool.types.TimestampMode = blk: { const build_options = @import("thread_pool").build_options; break :blk switch (build_options.timestamp_mode) { .any => .full, .full => .full, .disabled => .disabled, };};
/// COMPTIME-ELISION: returns `true` iff the current build pinned the/// scheduler variant to something other than `variant`. Tests that/// explicitly exercise one variant via `defaultConfig(variant)` should/// call `try skipIfSchedulerVariantIncompatible(variant)` at the top so/// pinned-variant builds cleanly skip incompatible scenarios rather than/// fail with `error.BuildOptionMismatch`.pub fn skipIfSchedulerVariantIncompatible(variant: thread_pool.SchedulerVariant) !void { const build_options = @import("thread_pool").build_options; const incompatible = switch (build_options.scheduler_variant) { .any => false, .global_queue => variant != .global_queue, .work_stealing => variant != .work_stealing, .lane_extension => variant != .lane_extension, }; if (incompatible) return error.SkipZigTest;}
pub fn defaultConfig(variant: thread_pool.SchedulerVariant) thread_pool.Config { return .{ .min_threads = 2, .max_threads = 2, .queue_capacity = 512, .local_queue_capacity = 64, .inject_batch_size = 8, .worker_scratch_cap_bytes = 4096, .idle_park_after_ns = 100_000, .idle_purge_after_ns = 500_000, .scheduler_variant = variant, .timestamp_mode = default_timestamp_mode, };}
pub fn counterTask(raw_ctx: *anyopaque, runtime: *thread_pool.TaskContext) void { _ = runtime; const counter: *Counter = @ptrCast(@alignCast(raw_ctx)); _ = counter.value.fetchAdd(1, .monotonic);}
pub const ScratchTaskContext = struct { counter: *Counter, bytes: usize,};
pub fn scratchTask(raw_ctx: *anyopaque, runtime: *thread_pool.TaskContext) void { const ctx: *ScratchTaskContext = @ptrCast(@alignCast(raw_ctx)); _ = runtime.allocScratch(u8, ctx.bytes) catch {}; _ = ctx.counter.value.fetchAdd(1, .monotonic);}
pub const SleepTaskContext = struct { counter: *Counter, sleep_ns: u64,};
pub fn sleepTask(raw_ctx: *anyopaque, runtime: *thread_pool.TaskContext) void { _ = runtime; const ctx: *SleepTaskContext = @ptrCast(@alignCast(raw_ctx)); thread_pool.clock.sleep(ctx.sleep_ns); _ = ctx.counter.value.fetchAdd(1, .monotonic);}
pub const GateTaskContext = struct { counter: *Counter, started: std.atomic.Value(bool) = std.atomic.Value(bool).init(false), release: std.atomic.Value(bool) = std.atomic.Value(bool).init(false),};
pub fn gateTask(raw_ctx: *anyopaque, runtime: *thread_pool.TaskContext) void { _ = runtime; const ctx: *GateTaskContext = @ptrCast(@alignCast(raw_ctx)); ctx.started.store(true, .release); while (!ctx.release.load(.acquire)) { thread_pool.clock.sleep(100_000); } _ = ctx.counter.value.fetchAdd(1, .monotonic);}
/// Waits until `wg` has no pending tasks by polling `pendingCount()`, and/// fails with `error.WaitGroupDeadline` after `timeout_ns`. Unlike/// `wg.wait()`, a task the pool loses fails the test fast instead of/// hanging it. The poll sleeps rather than spins so it never competes for/// CPU with the workers it waits on.////// `pendingCount()` is an atomic load that bypasses `WaitGroup.mutex`/// entirely (unlike `wg.wait()`, which takes the mutex and so cannot/// return before `done()` has released it). That means this function can/// observe `pending == 0` and return while a worker thread is still/// inside `WaitGroup.done()`, between its `pending` decrement and its/// `mutex.unlock()`. Callers therefore must not let `wg` go out of scope/// (or be reused for another round) until the pool itself is/// deinit-joined (`pool.deinit()`, which joins every worker thread) — at/// that point no worker can still be touching `wg`. This is true at every/// current call site.pub fn waitPendingZero(wg: *const thread_pool.WaitGroup, timeout_ns: i96) !void { const start_ns = thread_pool.clock.nanoTimestamp(); while (wg.pendingCount() != 0) { if (thread_pool.clock.nanoTimestamp() - start_ns > timeout_ns) { std.debug.print("wait group still has {d} pending task(s) after {d} ms\n", .{ wg.pendingCount(), @divTrunc(timeout_ns, std.time.ns_per_ms), }); return error.WaitGroupDeadline; } thread_pool.clock.sleep(50 * std.time.ns_per_us); }}
pub const RecursiveContext = struct { pool: *thread_pool.ThreadPool, counter: *Counter, wait_group: *thread_pool.WaitGroup, child_count: usize,};
pub fn recursiveTask(raw_ctx: *anyopaque, runtime: *thread_pool.TaskContext) void { _ = runtime; const ctx: *RecursiveContext = @ptrCast(@alignCast(raw_ctx)); _ = ctx.counter.value.fetchAdd(1, .monotonic); for (0..ctx.child_count) |_| { _ = ctx.pool.submit(counterTask, @ptrCast(ctx.counter), .{ .wait_group = ctx.wait_group }) catch {}; }}
pub const ProducerContext = struct { pool: *thread_pool.ThreadPool, counter: *Counter, wait_group: *thread_pool.WaitGroup, task_count: usize,};
pub fn producerMain(ctx: *ProducerContext) void { var i: usize = 0; while (i < ctx.task_count) : (i += 1) { _ = ctx.pool.submit(counterTask, @ptrCast(ctx.counter), .{ .wait_group = ctx.wait_group }) catch |err| switch (err) { error.ShutdownInProgress => return, else => return, }; }}