From 7bff2e0016ecbba22da9039271e613ab5eceb6d7 Mon Sep 17 00:00:00 2001 From: ultd Date: Mon, 18 Sep 2023 14:10:13 -0500 Subject: [PATCH 01/15] initial prometheus metrics and thread pool fork --- src/lib.zig | 9 + src/prometheus/counter.zig | 115 +++ src/prometheus/gauge.zig | 206 ++++++ src/prometheus/histogram.zig | 271 +++++++ src/prometheus/metric.zig | 166 +++++ src/prometheus/registry.zig | 344 +++++++++ src/sync/thread_pool.zig | 1316 ++++++++++++++++++++++++++++++++++ 7 files changed, 2427 insertions(+) create mode 100644 src/prometheus/counter.zig create mode 100644 src/prometheus/gauge.zig create mode 100644 src/prometheus/histogram.zig create mode 100644 src/prometheus/metric.zig create mode 100644 src/prometheus/registry.zig create mode 100644 src/sync/thread_pool.zig diff --git a/src/lib.zig b/src/lib.zig index acb52257b..86cdf05bf 100644 --- a/src/lib.zig +++ b/src/lib.zig @@ -45,6 +45,7 @@ pub const sync = struct { pub usingnamespace @import("sync/mpmc.zig"); pub usingnamespace @import("sync/ref.zig"); pub usingnamespace @import("sync/mux.zig"); + pub usingnamespace @import("sync/thread_pool.zig"); }; pub const utils = struct { @@ -73,3 +74,11 @@ pub const net = struct { pub usingnamespace @import("net/net.zig"); pub usingnamespace @import("net/echo.zig"); }; + +pub const prometheus = struct { + pub usingnamespace @import("prometheus/counter.zig"); + pub usingnamespace @import("prometheus/gauge.zig"); + pub usingnamespace @import("prometheus/histogram.zig"); + pub usingnamespace @import("prometheus/metric.zig"); + pub usingnamespace @import("prometheus/registry.zig"); +}; diff --git a/src/prometheus/counter.zig b/src/prometheus/counter.zig new file mode 100644 index 000000000..0e7e2f2b7 --- /dev/null +++ b/src/prometheus/counter.zig @@ -0,0 +1,115 @@ +const std = @import("std"); +const mem = std.mem; +const testing = std.testing; +const Metric = @import("metric.zig").Metric; + +pub const Counter = struct { + metric: Metric = Metric{ + .getResultFn = getResult, + }, + value: std.atomic.Atomic(u64) = .{ .value = 0 }, + + const Self = @This(); + + pub fn init(allocator: mem.Allocator) !*Self { + const self = try allocator.create(Self); + + self.* = .{}; + + return self; + } + + pub fn inc(self: *Self) void { + _ = self.value.fetchAdd(1, .SeqCst); + } + + pub fn dec(self: *Self) void { + _ = self.value.fetchSub(1, .SeqCst); + } + + pub fn add(self: *Self, value: anytype) void { + if (!comptime std.meta.trait.isNumber(@TypeOf(value))) { + @compileError("can't add a non-number"); + } + + _ = self.value.fetchAdd(@intCast(value), .SeqCst); + } + + pub fn get(self: *const Self) u64 { + return self.value.load(.SeqCst); + } + + pub fn set(self: *Self, value: anytype) void { + if (!comptime std.meta.trait.isNumber(@TypeOf(value))) { + @compileError("can't set a non-number"); + } + + _ = self.value.store(@intCast(value), .SeqCst); + } + + fn getResult(metric: *Metric, _: mem.Allocator) Metric.Error!Metric.Result { + const self = @fieldParentPtr(Self, "metric", metric); + return Metric.Result{ .counter = self.get() }; + } +}; + +test "prometheus.counter: inc/add/dec/set/get" { + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + var counter = try Counter.init(testing.allocator); + defer testing.allocator.destroy(counter); + + try testing.expectEqual(@as(u64, 0), counter.get()); + + counter.inc(); + try testing.expectEqual(@as(u64, 1), counter.get()); + + counter.add(200); + try testing.expectEqual(@as(u64, 201), counter.get()); + + counter.dec(); + try testing.expectEqual(@as(u64, 200), counter.get()); + + counter.set(43); + try testing.expectEqual(@as(u64, 43), counter.get()); +} + +test "prometheus.counter: concurrent" { + var counter = try Counter.init(testing.allocator); + defer testing.allocator.destroy(counter); + + var threads: [4]std.Thread = undefined; + for (&threads) |*thread| { + thread.* = try std.Thread.spawn( + .{}, + struct { + fn run(c: *Counter) void { + var i: usize = 0; + while (i < 20) : (i += 1) { + c.inc(); + } + } + }.run, + .{counter}, + ); + } + + for (&threads) |*thread| thread.join(); + + try testing.expectEqual(@as(u64, 80), counter.get()); +} + +test "prometheus.counter: write" { + var counter = try Counter.init(testing.allocator); + defer testing.allocator.destroy(counter); + counter.set(340); + + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + var metric = &counter.metric; + try metric.write(testing.allocator, buffer.writer(), "mycounter"); + + try testing.expectEqualStrings("mycounter 340\n", buffer.items); +} diff --git a/src/prometheus/gauge.zig b/src/prometheus/gauge.zig new file mode 100644 index 000000000..cd1aef41f --- /dev/null +++ b/src/prometheus/gauge.zig @@ -0,0 +1,206 @@ +const std = @import("std"); +const mem = std.mem; +const testing = std.testing; + +const Metric = @import("metric.zig").Metric; + +pub fn GaugeCallFnType(comptime StateType: type, comptime Return: type) type { + const CallFnArgType = switch (@typeInfo(StateType)) { + .Pointer => StateType, + .Optional => |opt| opt.child, + .Void => void, + else => *StateType, + }; + + return *const fn (state: CallFnArgType) Return; +} + +pub fn Gauge(comptime StateType: type, comptime Return: type) type { + const CallFnType = GaugeCallFnType(StateType, Return); + + return struct { + const Self = @This(); + + metric: Metric = .{ + .getResultFn = getResult, + }, + callFn: CallFnType = undefined, + state: StateType = undefined, + + pub fn init(allocator: mem.Allocator, callFn: CallFnType, state: StateType) !*Self { + const self = try allocator.create(Self); + + self.* = .{}; + self.callFn = callFn; + self.state = state; + + return self; + } + + pub fn get(self: *Self) Return { + const TypeInfo = @typeInfo(StateType); + switch (TypeInfo) { + .Pointer, .Void => { + return self.callFn(self.state); + }, + .Optional => { + if (self.state) |state| { + return self.callFn(state); + } + return 0; + }, + else => { + return self.callFn(&self.state); + }, + } + } + + fn getResult(metric: *Metric, allocator: mem.Allocator) Metric.Error!Metric.Result { + _ = allocator; + + const self = @fieldParentPtr(Self, "metric", metric); + + return switch (Return) { + f64 => Metric.Result{ .gauge = self.get() }, + u64 => Metric.Result{ .gauge_int = self.get() }, + else => unreachable, // Gauge Return may only be 'f64' or 'u64' + }; + } + }; +} + +test "gauge: get" { + const TestCase = struct { + state_type: type, + typ: type, + }; + + const testCases = [_]TestCase{ + .{ + .state_type = struct { + value: f64, + }, + .typ = f64, + }, + }; + + inline for (testCases) |tc| { + const State = tc.state_type; + const InnerType = tc.typ; + + var state = State{ .value = 20 }; + + var gauge = try Gauge(*State, InnerType).init( + testing.allocator, + struct { + fn get(s: *State) InnerType { + return s.value + 1; + } + }.get, + &state, + ); + defer testing.allocator.destroy(gauge); + + try testing.expectEqual(@as(InnerType, 21), gauge.get()); + } +} + +test "gauge: optional state" { + const State = struct { + value: f64, + }; + var state = State{ .value = 20.0 }; + + var gauge = try Gauge(?*State, f64).init( + testing.allocator, + struct { + fn get(s: *State) f64 { + return s.value + 1.0; + } + }.get, + &state, + ); + defer testing.allocator.destroy(gauge); + + try testing.expectEqual(@as(f64, 21.0), gauge.get()); +} + +test "gauge: non-pointer state" { + var gauge = try Gauge(f64, f64).init( + testing.allocator, + struct { + fn get(s: *f64) f64 { + s.* += 1.0; + return s.*; + } + }.get, + 0.0, + ); + defer testing.allocator.destroy(gauge); + + try testing.expectEqual(@as(f64, 1.0), gauge.get()); +} + +test "gauge: shared state" { + const State = struct { + mutex: std.Thread.Mutex = .{}, + items: std.ArrayList(usize) = std.ArrayList(usize).init(testing.allocator), + }; + var shared_state = State{}; + defer shared_state.items.deinit(); + + var gauge = try Gauge(*State, f64).init( + testing.allocator, + struct { + fn get(state: *State) f64 { + return @floatFromInt(state.items.items.len); + } + }.get, + &shared_state, + ); + defer testing.allocator.destroy(gauge); + + var threads: [4]std.Thread = undefined; + for (&threads, 0..) |*thread, thread_index| { + thread.* = try std.Thread.spawn( + .{}, + struct { + fn run(thread_idx: usize, state: *State) !void { + var i: usize = 0; + while (i < 4) : (i += 1) { + state.mutex.lock(); + defer state.mutex.unlock(); + try state.items.append(thread_idx + i); + } + } + }.run, + .{ thread_index, &shared_state }, + ); + } + + for (&threads) |*thread| thread.join(); + + try testing.expectEqual(@as(usize, 16), @as(usize, @intFromFloat(gauge.get()))); +} + +test "gauge: write" { + var gauge = try Gauge(usize, f64).init( + testing.allocator, + struct { + fn get(state: *usize) f64 { + state.* += 340; + return @floatFromInt(state.*); + } + }.get, + @as(usize, 0), + ); + defer testing.allocator.destroy(gauge); + + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + var metric = &gauge.metric; + try metric.write(testing.allocator, buffer.writer(), "mygauge"); + + try testing.expectEqualStrings("mygauge 340.000000\n", buffer.items); +} diff --git a/src/prometheus/histogram.zig b/src/prometheus/histogram.zig new file mode 100644 index 000000000..5a93fc595 --- /dev/null +++ b/src/prometheus/histogram.zig @@ -0,0 +1,271 @@ +const std = @import("std"); +const fmt = std.fmt; +const math = std.math; +const mem = std.mem; +const testing = std.testing; + +const Metric = @import("metric.zig").Metric; +const HistogramResult = @import("metric.zig").HistogramResult; + +const e10_min = -9; +const e10_max = 18; +const buckets_per_decimal = 18; +const decimal_buckets_count = e10_max - e10_min; +const buckets_count = decimal_buckets_count * buckets_per_decimal; + +const lower_bucket_range = blk: { + var buf: [64]u8 = undefined; + break :blk fmt.bufPrint(&buf, "0...{e:.3}", .{math.pow(f64, 10, e10_min)}) catch unreachable; +}; +const upper_bucket_range = blk: { + var buf: [64]u8 = undefined; + break :blk fmt.bufPrint(&buf, "{e:.3}...+Inf", .{math.pow(f64, 10, e10_max)}) catch unreachable; +}; + +const bucket_ranges: [buckets_count][]const u8 = blk: { + const bucket_multiplier = math.pow(f64, 10.0, 1.0 / @as(f64, buckets_per_decimal)); + + var v = math.pow(f64, 10, e10_min); + + var start = blk2: { + var buf: [64]u8 = undefined; + break :blk2 fmt.bufPrint(&buf, "{e:.3}", .{v}) catch unreachable; + }; + + var result: [buckets_count][]const u8 = undefined; + for (&result) |*range| { + v *= bucket_multiplier; + + const end = blk3: { + var buf: [64]u8 = undefined; + break :blk3 fmt.bufPrint(&buf, "{e:.3}", .{v}) catch unreachable; + }; + + range.* = start ++ "..." ++ end; + + start = end; + } + + break :blk result; +}; + +test "bucket ranges" { + try testing.expectEqualStrings("0...1.000e-09", lower_bucket_range); + try testing.expectEqualStrings("1.000e+18...+Inf", upper_bucket_range); + + try testing.expectEqualStrings("1.000e-09...1.136e-09", bucket_ranges[0]); + try testing.expectEqualStrings("1.136e-09...1.292e-09", bucket_ranges[1]); + try testing.expectEqualStrings("8.799e-09...1.000e-08", bucket_ranges[buckets_per_decimal - 1]); + try testing.expectEqualStrings("1.000e-08...1.136e-08", bucket_ranges[buckets_per_decimal]); + try testing.expectEqualStrings("8.799e-01...1.000e+00", bucket_ranges[buckets_per_decimal * (-e10_min) - 1]); + try testing.expectEqualStrings("1.000e+00...1.136e+00", bucket_ranges[buckets_per_decimal * (-e10_min)]); + try testing.expectEqualStrings("8.799e+17...1.000e+18", bucket_ranges[buckets_per_decimal * (e10_max - e10_min) - 1]); +} + +/// Histogram based on https://github.com/VictoriaMetrics/metrics/blob/master/histogram.go. +pub const Histogram = struct { + const Self = @This(); + + metric: Metric = .{ + .getResultFn = getResult, + }, + + mutex: std.Thread.Mutex = .{}, + decimal_buckets: [decimal_buckets_count][buckets_per_decimal]u64 = undefined, + + lower: u64 = 0, + upper: u64 = 0, + + sum: f64 = 0.0, + + pub fn init(allocator: mem.Allocator) !*Self { + const self = try allocator.create(Self); + + self.* = .{}; + for (&self.decimal_buckets) |*bucket| { + @memset(bucket, 0); + } + + return self; + } + + pub fn update(self: *Self, value: f64) void { + if (math.isNan(value) or value < 0) { + return; + } + + const bucket_idx: f64 = (math.log10(value) - e10_min) * buckets_per_decimal; + + // Keep a lock while updating the histogram. + self.mutex.lock(); + defer self.mutex.unlock(); + + self.sum += value; + + if (bucket_idx < 0) { + self.lower += 1; + } else if (bucket_idx >= buckets_count) { + self.upper += 1; + } else { + const idx: usize = blk: { + const tmp: usize = @intFromFloat(bucket_idx); + + if (bucket_idx == @as(f64, @floatFromInt(tmp)) and tmp > 0) { + // Edge case for 10^n values, which must go to the lower bucket + // according to Prometheus logic for `le`-based histograms. + break :blk tmp - 1; + } else { + break :blk tmp; + } + }; + + const decimal_bucket_idx = idx / buckets_per_decimal; + const offset = idx % buckets_per_decimal; + + var bucket: []u64 = &self.decimal_buckets[decimal_bucket_idx]; + bucket[offset] += 1; + } + } + + pub fn get(self: *const Self) u64 { + _ = self; + return 0; + } + + fn isBucketAllZero(bucket: []const u64) bool { + for (bucket) |v| { + if (v != 0) return false; + } + return true; + } + + fn getResult(metric: *Metric, allocator: mem.Allocator) Metric.Error!Metric.Result { + const self = @fieldParentPtr(Histogram, "metric", metric); + + // Arbitrary maximum capacity + var buckets = try std.ArrayList(HistogramResult.Bucket).initCapacity(allocator, 16); + var count_total: u64 = 0; + + // Keep a lock while querying the histogram. + self.mutex.lock(); + defer self.mutex.unlock(); + + if (self.lower > 0) { + try buckets.append(.{ + .vmrange = lower_bucket_range, + .count = self.lower, + }); + count_total += self.lower; + } + + for (&self.decimal_buckets, 0..) |bucket, decimal_bucket_idx| { + if (isBucketAllZero(&bucket)) continue; + + for (bucket, 0..) |count, offset| { + if (count <= 0) continue; + + const bucket_idx = (decimal_bucket_idx * buckets_per_decimal) + offset; + const vmrange = bucket_ranges[bucket_idx]; + + try buckets.append(.{ + .vmrange = vmrange, + .count = count, + }); + count_total += count; + } + } + + if (self.upper > 0) { + try buckets.append(.{ + .vmrange = upper_bucket_range, + .count = self.upper, + }); + count_total += self.upper; + } + + return Metric.Result{ + .histogram = .{ + .buckets = try buckets.toOwnedSlice(), + .sum = .{ .value = self.sum }, + .count = count_total, + }, + }; + } +}; + +test "write empty" { + var histogram = try Histogram.init(testing.allocator); + defer testing.allocator.destroy(histogram); + + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + var metric = &histogram.metric; + try metric.write(testing.allocator, buffer.writer(), "myhistogram"); + + try testing.expectEqual(@as(usize, 0), buffer.items.len); +} + +test "update then write" { + var histogram = try Histogram.init(testing.allocator); + defer testing.allocator.destroy(histogram); + + var i: usize = 98; + while (i < 218) : (i += 1) { + histogram.update(@floatFromInt(i)); + } + + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + var metric = &histogram.metric; + try metric.write(testing.allocator, buffer.writer(), "myhistogram"); + + const exp = + \\myhistogram_bucket{vmrange="8.799e+01...1.000e+02"} 3 + \\myhistogram_bucket{vmrange="1.000e+02...1.136e+02"} 13 + \\myhistogram_bucket{vmrange="1.136e+02...1.292e+02"} 16 + \\myhistogram_bucket{vmrange="1.292e+02...1.468e+02"} 17 + \\myhistogram_bucket{vmrange="1.468e+02...1.668e+02"} 20 + \\myhistogram_bucket{vmrange="1.668e+02...1.896e+02"} 23 + \\myhistogram_bucket{vmrange="1.896e+02...2.154e+02"} 26 + \\myhistogram_bucket{vmrange="2.154e+02...2.448e+02"} 2 + \\myhistogram_sum 18900 + \\myhistogram_count 120 + \\ + ; + + try testing.expectEqualStrings(exp, buffer.items); +} + +test "update then write with labels" { + var histogram = try Histogram.init(testing.allocator); + defer testing.allocator.destroy(histogram); + + var i: usize = 98; + while (i < 218) : (i += 1) { + histogram.update(@floatFromInt(i)); + } + + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + var metric = &histogram.metric; + try metric.write(testing.allocator, buffer.writer(), "myhistogram{route=\"/api/v2/users\"}"); + + const exp = + \\myhistogram_bucket{route="/api/v2/users",vmrange="8.799e+01...1.000e+02"} 3 + \\myhistogram_bucket{route="/api/v2/users",vmrange="1.000e+02...1.136e+02"} 13 + \\myhistogram_bucket{route="/api/v2/users",vmrange="1.136e+02...1.292e+02"} 16 + \\myhistogram_bucket{route="/api/v2/users",vmrange="1.292e+02...1.468e+02"} 17 + \\myhistogram_bucket{route="/api/v2/users",vmrange="1.468e+02...1.668e+02"} 20 + \\myhistogram_bucket{route="/api/v2/users",vmrange="1.668e+02...1.896e+02"} 23 + \\myhistogram_bucket{route="/api/v2/users",vmrange="1.896e+02...2.154e+02"} 26 + \\myhistogram_bucket{route="/api/v2/users",vmrange="2.154e+02...2.448e+02"} 2 + \\myhistogram_sum{route="/api/v2/users"} 18900 + \\myhistogram_count{route="/api/v2/users"} 120 + \\ + ; + + try testing.expectEqualStrings(exp, buffer.items); +} diff --git a/src/prometheus/metric.zig b/src/prometheus/metric.zig new file mode 100644 index 000000000..82ba27397 --- /dev/null +++ b/src/prometheus/metric.zig @@ -0,0 +1,166 @@ +const std = @import("std"); +const fmt = std.fmt; +const mem = std.mem; +const testing = std.testing; + +pub const HistogramResult = struct { + pub const Bucket = struct { + vmrange: []const u8, + count: u64, + }; + + pub const SumValue = struct { + value: f64 = 0, + + pub fn format(self: @This(), comptime format_string: []const u8, options: fmt.FormatOptions, writer: anytype) !void { + _ = format_string; + + const as_int: u64 = @intFromFloat(self.value); + if (@as(f64, @floatFromInt(as_int)) == self.value) { + try fmt.formatInt(as_int, 10, .lower, options, writer); + } else { + try fmt.formatFloatDecimal(self.value, options, writer); + } + } + }; + + buckets: []Bucket, + sum: SumValue, + count: u64, +}; + +pub const Metric = struct { + pub const Error = error{OutOfMemory} || std.os.WriteError || std.http.Server.Response.Writer.Error; + + pub const Result = union(enum) { + const Self = @This(); + + counter: u64, + gauge: f64, + gauge_int: u64, + histogram: HistogramResult, + + pub fn deinit(self: Self, allocator: mem.Allocator) void { + switch (self) { + .histogram => |v| { + allocator.free(v.buckets); + }, + else => {}, + } + } + }; + + getResultFn: *const fn (self: *Metric, allocator: mem.Allocator) Error!Result, + + pub fn write(self: *Metric, allocator: mem.Allocator, writer: anytype, name: []const u8) Error!void { + const result = try self.getResultFn(self, allocator); + defer result.deinit(allocator); + + switch (result) { + .counter, .gauge_int => |v| { + return try writer.print("{s} {d}\n", .{ name, v }); + }, + .gauge => |v| { + return try writer.print("{s} {d:.6}\n", .{ name, v }); + }, + .histogram => |v| { + if (v.buckets.len <= 0) return; + + const name_and_labels = splitName(name); + + if (name_and_labels.labels.len > 0) { + for (v.buckets) |bucket| { + try writer.print("{s}_bucket{{{s},vmrange=\"{s}\"}} {d:.6}\n", .{ + name_and_labels.name, + name_and_labels.labels, + bucket.vmrange, + bucket.count, + }); + } + try writer.print("{s}_sum{{{s}}} {:.6}\n", .{ + name_and_labels.name, + name_and_labels.labels, + v.sum, + }); + try writer.print("{s}_count{{{s}}} {d}\n", .{ + name_and_labels.name, + name_and_labels.labels, + v.count, + }); + } else { + for (v.buckets) |bucket| { + try writer.print("{s}_bucket{{vmrange=\"{s}\"}} {d:.6}\n", .{ + name_and_labels.name, + bucket.vmrange, + bucket.count, + }); + } + try writer.print("{s}_sum {:.6}\n", .{ + name_and_labels.name, + v.sum, + }); + try writer.print("{s}_count {d}\n", .{ + name_and_labels.name, + v.count, + }); + } + }, + } + } +}; + +const NameAndLabels = struct { + name: []const u8, + labels: []const u8 = "", +}; + +fn splitName(name: []const u8) NameAndLabels { + const bracket_pos = mem.indexOfScalar(u8, name, '{'); + if (bracket_pos) |pos| { + return NameAndLabels{ + .name = name[0..pos], + .labels = name[pos + 1 .. name.len - 1], + }; + } else { + return NameAndLabels{ + .name = name, + }; + } +} + +test "prometheus.metric: ensure splitName works" { + const TestCase = struct { + input: []const u8, + exp: NameAndLabels, + }; + + const test_cases = &[_]TestCase{ + .{ + .input = "foobar", + .exp = .{ + .name = "foobar", + }, + }, + .{ + .input = "foobar{route=\"/home\"}", + .exp = .{ + .name = "foobar", + .labels = "route=\"/home\"", + }, + }, + .{ + .input = "foobar{route=\"/home\",status=\"500\"}", + .exp = .{ + .name = "foobar", + .labels = "route=\"/home\",status=\"500\"", + }, + }, + }; + + inline for (test_cases) |tc| { + const res = splitName(tc.input); + + try testing.expectEqualStrings(tc.exp.name, res.name); + try testing.expectEqualStrings(tc.exp.labels, res.labels); + } +} diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig new file mode 100644 index 000000000..350aa750d --- /dev/null +++ b/src/prometheus/registry.zig @@ -0,0 +1,344 @@ +const std = @import("std"); +const fmt = std.fmt; +const hash_map = std.hash_map; +const heap = std.heap; +const mem = std.mem; +const testing = std.testing; +const Metric = @import("metric.zig").Metric; +const Counter = @import("counter.zig").Counter; +const Gauge = @import("gauge.zig").Gauge; +const Histogram = @import("histogram.zig").Histogram; +const GaugeCallFnType = @import("gauge.zig").GaugeCallFnType; + +pub const GetMetricError = error{ + // Returned when trying to add a metric to an already full registry. + TooManyMetrics, + // Returned when the name of name is bigger than the configured max_name_len. + NameTooLong, + + OutOfMemory, +}; + +const RegistryOptions = struct { + max_metrics: comptime_int = 8192, + max_name_len: comptime_int = 1024, +}; + +var gpa = std.heap.GeneralPurposeAllocator(.{}){}; +const gpa_allocator = gpa.allocator(); +pub var registry: *Registry(.{}) = undefined; + +pub fn init() error{OutOfMemory}!void { + registry = try Registry(.{}).init(gpa_allocator); +} + +pub fn deinit() void { + registry.deinit(); +} + +pub fn Registry(comptime options: RegistryOptions) type { + return struct { + const Self = @This(); + const MetricMap = hash_map.StringHashMapUnmanaged(*Metric); + + root_allocator: mem.Allocator, + arena_state: heap.ArenaAllocator, + mutex: std.Thread.Mutex, + metrics: MetricMap, + + pub fn init(alloc: mem.Allocator) error{OutOfMemory}!*Self { + const self = try alloc.create(Self); + + self.* = .{ + .root_allocator = alloc, + .arena_state = heap.ArenaAllocator.init(alloc), + .mutex = .{}, + .metrics = MetricMap{}, + }; + + return self; + } + + pub fn deinit(self: *Self) void { + self.arena_state.deinit(); + self.root_allocator.destroy(self); + } + + fn nbMetrics(self: *const Self) usize { + return self.metrics.count(); + } + + pub fn getOrCreateCounter(self: *Self, name: []const u8) GetMetricError!*Counter { + if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; + if (name.len > options.max_name_len) return error.NameTooLong; + + var allocator = self.arena_state.allocator(); + + const duped_name = try allocator.dupe(u8, name); + + self.mutex.lock(); + defer self.mutex.unlock(); + + var gop = try self.metrics.getOrPut(allocator, duped_name); + if (!gop.found_existing) { + var real_metric = try Counter.init(allocator); + gop.value_ptr.* = &real_metric.metric; + } + + return @fieldParentPtr(Counter, "metric", gop.value_ptr.*); + } + + pub fn getOrCreateHistogram(self: *Self, name: []const u8) GetMetricError!*Histogram { + if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; + if (name.len > options.max_name_len) return error.NameTooLong; + + var allocator = self.arena_state.allocator(); + + const duped_name = try allocator.dupe(u8, name); + + self.mutex.lock(); + defer self.mutex.unlock(); + + var gop = try self.metrics.getOrPut(allocator, duped_name); + if (!gop.found_existing) { + var real_metric = try Histogram.init(allocator); + gop.value_ptr.* = &real_metric.metric; + } + + return @fieldParentPtr(Histogram, "metric", gop.value_ptr.*); + } + + pub fn getOrCreateGauge( + self: *Self, + name: []const u8, + state: anytype, + callFn: GaugeCallFnType(@TypeOf(state), f64), + ) GetMetricError!*Gauge(@TypeOf(state), f64) { + if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; + if (name.len > options.max_name_len) return error.NameTooLong; + + var allocator = self.arena_state.allocator(); + + const duped_name = try allocator.dupe(u8, name); + + self.mutex.lock(); + defer self.mutex.unlock(); + + var gop = try self.metrics.getOrPut(allocator, duped_name); + if (!gop.found_existing) { + var real_metric = try Gauge(@TypeOf(state), f64).init(allocator, callFn, state); + gop.value_ptr.* = &real_metric.metric; + } + + return @fieldParentPtr(Gauge(@TypeOf(state), f64), "metric", gop.value_ptr.*); + } + + pub fn getOrCreateGaugeInt( + self: *Self, + name: []const u8, + state: anytype, + callFn: GaugeCallFnType(@TypeOf(state), u64), + ) GetMetricError!*Gauge(@TypeOf(state), u64) { + if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; + if (name.len > options.max_name_len) return error.NameTooLong; + + var allocator = self.arena_state.allocator(); + + const duped_name = try allocator.dupe(u8, name); + + self.mutex.lock(); + defer self.mutex.unlock(); + + var gop = try self.metrics.getOrPut(allocator, duped_name); + if (!gop.found_existing) { + var real_metric = try Gauge(@TypeOf(state), u64).init(allocator, callFn, state); + gop.value_ptr.* = &real_metric.metric; + } + + return @fieldParentPtr(Gauge(@TypeOf(state), u64), "metric", gop.value_ptr.*); + } + + pub fn write(self: *Self, allocator: mem.Allocator, writer: anytype) !void { + var arena_state = heap.ArenaAllocator.init(allocator); + defer arena_state.deinit(); + + self.mutex.lock(); + defer self.mutex.unlock(); + + try writeMetrics(arena_state.allocator(), self.metrics, writer); + } + + fn writeMetrics(allocator: mem.Allocator, map: MetricMap, writer: anytype) !void { + // Get the keys, sorted + const keys = blk: { + var key_list = try std.ArrayList([]const u8).initCapacity(allocator, map.count()); + + var key_iter = map.keyIterator(); + while (key_iter.next()) |key| { + key_list.appendAssumeCapacity(key.*); + } + + break :blk key_list.items; + }; + defer allocator.free(keys); + + std.mem.sort([]const u8, keys, {}, stringLessThan); + + // Write each metric in key order + for (keys) |key| { + var metric = map.get(key) orelse unreachable; + try metric.write(allocator, writer, key); + } + } + }; +} + +fn stringLessThan(context: void, lhs: []const u8, rhs: []const u8) bool { + _ = context; + return mem.lessThan(u8, lhs, rhs); +} + +test "registry getOrCreateCounter" { + var reg = try Registry(.{}).create(testing.allocator); + defer reg.destroy(); + + const name = try fmt.allocPrint(testing.allocator, "http_requests{{status=\"{d}\"}}", .{500}); + defer testing.allocator.free(name); + + var i: usize = 0; + while (i < 10) : (i += 1) { + var counter = try reg.getOrCreateCounter(name); + counter.inc(); + } + + var counter = try reg.getOrCreateCounter(name); + try testing.expectEqual(@as(u64, 10), counter.get()); +} + +test "registry write" { + const TestCase = struct { + counter_name: []const u8, + gauge_name: []const u8, + histogram_name: []const u8, + exp: []const u8, + }; + + const exp1 = + \\http_conn_pool_size 4.000000 + \\http_request_size_bucket{vmrange="1.292e+02...1.468e+02"} 1 + \\http_request_size_bucket{vmrange="4.642e+02...5.275e+02"} 1 + \\http_request_size_bucket{vmrange="1.136e+03...1.292e+03"} 1 + \\http_request_size_sum 1870.360000 + \\http_request_size_count 3 + \\http_requests 2 + \\ + ; + + const exp2 = + \\http_conn_pool_size{route="/api/v2/users"} 4.000000 + \\http_request_size_bucket{route="/api/v2/users",vmrange="1.292e+02...1.468e+02"} 1 + \\http_request_size_bucket{route="/api/v2/users",vmrange="4.642e+02...5.275e+02"} 1 + \\http_request_size_bucket{route="/api/v2/users",vmrange="1.136e+03...1.292e+03"} 1 + \\http_request_size_sum{route="/api/v2/users"} 1870.360000 + \\http_request_size_count{route="/api/v2/users"} 3 + \\http_requests{route="/api/v2/users"} 2 + \\ + ; + + const test_cases = &[_]TestCase{ + .{ + .counter_name = "http_requests", + .gauge_name = "http_conn_pool_size", + .histogram_name = "http_request_size", + .exp = exp1, + }, + .{ + .counter_name = "http_requests{route=\"/api/v2/users\"}", + .gauge_name = "http_conn_pool_size{route=\"/api/v2/users\"}", + .histogram_name = "http_request_size{route=\"/api/v2/users\"}", + .exp = exp2, + }, + }; + + inline for (test_cases) |tc| { + var reg = try Registry(.{}).create(testing.allocator); + defer reg.destroy(); + + // Add some counters + { + var counter = try reg.getOrCreateCounter(tc.counter_name); + counter.set(2); + } + + // Add some gauges + { + _ = try reg.getOrCreateGauge( + tc.gauge_name, + @as(f64, 4.0), + struct { + fn get(s: *f64) f64 { + return s.*; + } + }.get, + ); + } + + // Add an histogram + { + var histogram = try reg.getOrCreateHistogram(tc.histogram_name); + + histogram.update(500.12); + histogram.update(1230.240); + histogram.update(140); + } + + // Write to a buffer + { + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + try reg.write(testing.allocator, buffer.writer()); + + try testing.expectEqualStrings(tc.exp, buffer.items); + } + + // Write to a file + { + const filename = "prometheus_metrics.txt"; + var file = try std.fs.cwd().createFile(filename, .{ .read = true }); + defer { + file.close(); + std.fs.cwd().deleteFile(filename) catch {}; + } + + try reg.write(testing.allocator, file.writer()); + + try file.seekTo(0); + const file_data = try file.readToEndAlloc(testing.allocator, std.math.maxInt(usize)); + defer testing.allocator.free(file_data); + + try testing.expectEqualStrings(tc.exp, file_data); + } + } +} + +test "registry options" { + var reg = try Registry(.{ .max_metrics = 1, .max_name_len = 4 }).create(testing.allocator); + defer reg.destroy(); + + { + try testing.expectError(error.NameTooLong, reg.getOrCreateCounter("hello")); + _ = try reg.getOrCreateCounter("foo"); + } + + { + try testing.expectError(error.TooManyMetrics, reg.getOrCreateCounter("bar")); + } +} + +test "prometheus.registry: test default registry" { + registry = try Registry(.{}).init(testing.allocator); + defer registry.deinit(); + var counter = try registry.getOrCreateCounter("hello"); + counter.inc(); +} diff --git a/src/sync/thread_pool.zig b/src/sync/thread_pool.zig new file mode 100644 index 000000000..5c5143f4f --- /dev/null +++ b/src/sync/thread_pool.zig @@ -0,0 +1,1316 @@ +// Thank you bun.sh: +// https://github.com/oven-sh/bun/blob/main/src/thread_pool.zig +// +// Thank you @kprotty: +// https://github.com/kprotty/zap/blob/blog/src/thread_pool.zig + +const std = @import("std"); +const builtin = @import("builtin"); +const Futex = std.Thread.Futex; +const assert = std.debug.assert; +const Atomic = std.atomic.Atomic; +pub const OnSpawnCallback = *const fn (ctx: ?*anyopaque) ?*anyopaque; + +pub const ThreadPool = struct { + sleep_on_idle_network_thread: bool = true, + /// executed on the thread + on_thread_spawn: ?OnSpawnCallback = null, + threadpool_context: ?*anyopaque = null, + stack_size: u32, + max_threads: u32, + sync: Atomic(u32) = Atomic(u32).init(@as(u32, @bitCast(Sync{}))), + idle_event: Event = .{}, + join_event: Event = .{}, + run_queue: Node.Queue = .{}, + threads: Atomic(?*Thread) = Atomic(?*Thread).init(null), + name: []const u8 = "", + spawned_thread_count: Atomic(u32) = Atomic(u32).init(0), + + const Sync = packed struct { + /// Tracks the number of threads not searching for Tasks + idle: u14 = 0, + /// Tracks the number of threads spawned + spawned: u14 = 0, + /// What you see is what you get + unused: bool = false, + /// Used to not miss notifications while state = waking + notified: bool = false, + /// The current state of the thread pool + state: enum(u2) { + /// A notification can be issued to wake up a sleeping as the "waking thread". + pending = 0, + /// The state was notified with a signal. A thread is woken up. + /// The first thread to transition to `waking` becomes the "waking thread". + signaled, + /// There is a "waking thread" among us. + /// No other thread should be woken up until the waking thread transitions the state. + waking, + /// The thread pool was terminated. Start decremented `spawned` so that it can be joined. + shutdown, + } = .pending, + }; + + /// Configuration options for the thread pool. + /// TODO: add CPU core affinity? + pub const Config = struct { + stack_size: u32 = (std.Thread.SpawnConfig{}).stack_size, + max_threads: u32, + }; + + /// Statically initialize the thread pool using the configuration. + pub fn init(config: Config) ThreadPool { + return .{ + .stack_size = @max(1, config.stack_size), + .max_threads = @max(1, config.max_threads), + }; + } + + pub fn wakeForIdleEvents(this: *ThreadPool) void { + // Wake all the threads to check for idle events. + this.idle_event.wake(Event.NOTIFIED, std.math.maxInt(u32)); + } + + /// Wait for a thread to call shutdown() on the thread pool and kill the worker threads. + pub fn deinit(self: *ThreadPool) void { + self.join(); + self.* = undefined; + } + + /// A Task represents the unit of Work / Job / Execution that the ThreadPool schedules. + /// The user provides a `callback` which is invoked when the *Task can run on a thread. + pub const Task = struct { + node: Node = .{}, + callback: *const (fn (*Task) void), + }; + + /// An unordered collection of Tasks which can be submitted for scheduling as a group. + pub const Batch = struct { + len: usize = 0, + head: ?*Task = null, + tail: ?*Task = null, + + pub fn pop(this: *Batch) ?*Task { + const len = @atomicLoad(usize, &this.len, .Monotonic); + if (len == 0) { + return null; + } + var task = this.head.?; + if (task.node.next) |node| { + this.head = @fieldParentPtr(Task, "node", node); + } else { + if (task != this.tail.?) unreachable; + this.tail = null; + this.head = null; + } + + this.len -= 1; + if (len == 0) { + this.tail = null; + } + return task; + } + + /// Create a batch from a single task. + pub fn from(task: *Task) Batch { + return Batch{ + .len = 1, + .head = task, + .tail = task, + }; + } + + /// Another batch into this one, taking ownership of its tasks. + pub fn push(self: *Batch, batch: Batch) void { + if (batch.len == 0) return; + if (self.len == 0) { + self.* = batch; + } else { + self.tail.?.node.next = if (batch.head) |h| &h.node else null; + self.tail = batch.tail; + self.len += batch.len; + } + } + }; + + pub const WaitGroup = struct { + mutex: std.Thread.Mutex = .{}, + counter: u32 = 0, + event: std.Thread.ResetEvent, + + pub fn init(self: *WaitGroup) void { + self.* = .{ + .mutex = .{}, + .counter = 0, + .event = undefined, + }; + } + + pub fn deinit(self: *WaitGroup) void { + self.event.reset(); + self.* = undefined; + } + + pub fn start(self: *WaitGroup) void { + self.mutex.lock(); + defer self.mutex.unlock(); + + self.counter += 1; + } + + pub fn isDone(this: *WaitGroup) bool { + return @atomicLoad(u32, &this.counter, .Monotonic) == 0; + } + + pub fn finish(self: *WaitGroup) void { + self.mutex.lock(); + defer self.mutex.unlock(); + + self.counter -= 1; + + if (self.counter == 0) { + self.event.set(); + } + } + + pub fn wait(self: *WaitGroup) void { + while (true) { + self.mutex.lock(); + + if (self.counter == 0) { + self.mutex.unlock(); + return; + } + + self.mutex.unlock(); + self.event.wait(); + } + } + + pub fn reset(self: *WaitGroup) void { + self.event.reset(); + } + }; + + pub fn ConcurrentFunction( + comptime Function: anytype, + ) type { + return struct { + const Fn = Function; + const Args = std.meta.ArgsTuple(@TypeOf(Fn)); + const Runner = @This(); + thread_pool: *ThreadPool, + states: []Routine = undefined, + batch: Batch = .{}, + allocator: std.mem.Allocator, + + pub fn init(allocator: std.mem.Allocator, thread_pool: *ThreadPool, count: usize) !Runner { + return Runner{ + .allocator = allocator, + .thread_pool = thread_pool, + .states = try allocator.alloc(Routine, count), + .batch = .{}, + }; + } + + pub fn call(this: *@This(), args: Args) void { + this.states[this.batch.len] = .{ + .args = args, + }; + this.batch.push(Batch.from(&this.states[this.batch.len].task)); + } + + pub fn run(this: *@This()) void { + this.thread_pool.schedule(this.batch); + } + + pub const Routine = struct { + args: Args, + task: Task = .{ .callback = callback }, + + pub fn callback(task: *Task) void { + var routine = @fieldParentPtr(@This(), "task", task); + @call(.always_inline, Fn, routine.args); + } + }; + + pub fn deinit(this: *@This()) void { + this.allocator.free(this.states); + } + }; + } + + pub fn runner( + this: *ThreadPool, + allocator: std.mem.Allocator, + comptime Function: anytype, + count: usize, + ) !ConcurrentFunction(Function) { + return try ConcurrentFunction(Function).init(allocator, this, count); + } + + /// Loop over an array of tasks and invoke `Run` on each one in a different thread + /// **Blocks the calling thread** until all tasks are completed. + pub fn do( + this: *ThreadPool, + allocator: std.mem.Allocator, + wg: ?*WaitGroup, + ctx: anytype, + comptime Run: anytype, + values: anytype, + ) !void { + return try Do(this, allocator, wg, @TypeOf(ctx), ctx, Run, @TypeOf(values), values, false); + } + + pub fn doPtr( + this: *ThreadPool, + allocator: std.mem.Allocator, + wg: ?*WaitGroup, + ctx: anytype, + comptime Run: anytype, + values: anytype, + ) !void { + return try Do(this, allocator, wg, @TypeOf(ctx), ctx, Run, @TypeOf(values), values, true); + } + + pub fn Do( + this: *ThreadPool, + allocator: std.mem.Allocator, + wg: ?*WaitGroup, + comptime Context: type, + ctx: Context, + comptime Function: anytype, + comptime ValuesType: type, + values: ValuesType, + comptime as_ptr: bool, + ) !void { + if (values.len == 0) + return; + var allocated_wait_group: ?*WaitGroup = null; + defer { + if (allocated_wait_group) |group| { + group.deinit(); + allocator.destroy(group); + } + } + + var wait_group = wg orelse brk: { + allocated_wait_group = try allocator.create(WaitGroup); + allocated_wait_group.?.init(); + break :brk allocated_wait_group.?; + }; + const WaitContext = struct { + wait_group: *WaitGroup = undefined, + ctx: Context, + values: ValuesType, + }; + + const RunnerTask = struct { + task: Task, + ctx: *WaitContext, + i: usize = 0, + + pub fn call(task: *Task) void { + var runner_task = @fieldParentPtr(@This(), "task", task); + const i = runner_task.i; + if (comptime as_ptr) { + Function(runner_task.ctx.ctx, &runner_task.ctx.values[i], i); + } else { + Function(runner_task.ctx.ctx, runner_task.ctx.values[i], i); + } + + runner_task.ctx.wait_group.finish(); + } + }; + var wait_context = allocator.create(WaitContext) catch unreachable; + wait_context.* = .{ + .ctx = ctx, + .wait_group = wait_group, + .values = values, + }; + defer allocator.destroy(wait_context); + var tasks = allocator.alloc(RunnerTask, values.len) catch unreachable; + defer allocator.free(tasks); + var batch: Batch = undefined; + var offset = tasks.len - 1; + + { + tasks[0] = .{ + .i = offset, + .task = .{ .callback = RunnerTask.call }, + .ctx = wait_context, + }; + batch = Batch.from(&tasks[0].task); + } + if (tasks.len > 1) { + for (tasks[1..]) |*runner_task| { + offset -= 1; + runner_task.* = .{ + .i = offset, + .task = .{ .callback = RunnerTask.call }, + .ctx = wait_context, + }; + batch.push(Batch.from(&runner_task.task)); + } + } + + wait_group.counter += @as(u32, @intCast(values.len)); + this.schedule(batch); + wait_group.wait(); + } + + /// Schedule a batch of tasks to be executed by some thread on the thread pool. + pub fn schedule(self: *ThreadPool, batch: Batch) void { + // Sanity check + if (batch.len == 0) { + return; + } + + // Extract out the Node's from the Tasks + var list = Node.List{ + .head = &batch.head.?.node, + .tail = &batch.tail.?.node, + }; + + // Push the task Nodes to the most appropriate queue + if (Thread.current) |thread| { + thread.run_buffer.push(&list) catch thread.run_queue.push(list); + } else { + self.run_queue.push(list); + } + + forceSpawn(self); + } + + pub fn forceSpawn(self: *ThreadPool) void { + // Try to notify a thread + const is_waking = false; + return self.notify(is_waking); + } + + inline fn notify(self: *ThreadPool, is_waking: bool) void { + // Fast path to check the Sync state to avoid calling into notifySlow(). + // If we're waking, then we need to update the state regardless + if (!is_waking) { + const sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + if (sync.notified) { + return; + } + } + + return self.notifySlow(is_waking); + } + + /// Warm the thread pool up to the given number of threads. + /// https://www.youtube.com/watch?v=ys3qcbO5KWw + pub fn warm(self: *ThreadPool, count: u14) void { + var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + if (sync.spawned >= count) + return; + + const to_spawn = @min(count - sync.spawned, @as(u14, @truncate(self.max_threads))); + while (sync.spawned < to_spawn) { + var new_sync = sync; + new_sync.spawned += 1; + sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( + @as(u32, @bitCast(sync)), + @as(u32, @bitCast(new_sync)), + .Release, + .Monotonic, + ) orelse break)); + const spawn_config = if (builtin.os.tag.isDarwin()) + // stack size must be a multiple of page_size + // macOS will fail to spawn a thread if the stack size is not a multiple of page_size + std.Thread.SpawnConfig{ .stack_size = ((std.Thread.SpawnConfig{}).stack_size + (std.mem.page_size / 2) / std.mem.page_size) * std.mem.page_size } + else + std.Thread.SpawnConfig{}; + + const thread = std.Thread.spawn(spawn_config, Thread.run, .{self}) catch return self.unregister(null); + thread.detach(); + } + } + + noinline fn notifySlow(self: *ThreadPool, is_waking: bool) void { + var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + while (sync.state != .shutdown) { + const can_wake = is_waking or (sync.state == .pending); + if (is_waking) { + assert(sync.state == .waking); + } + + var new_sync = sync; + new_sync.notified = true; + if (can_wake and sync.idle > 0) { // wake up an idle thread + new_sync.state = .signaled; + } else if (can_wake and sync.spawned < self.max_threads) { // spawn a new thread + new_sync.state = .signaled; + new_sync.spawned += 1; + } else if (is_waking) { // no other thread to pass on "waking" status + new_sync.state = .pending; + } else if (sync.notified) { // nothing to update + return; + } + + // Release barrier synchronizes with Acquire in wait() + // to ensure pushes to run queues happen before observing a posted notification. + sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( + @as(u32, @bitCast(sync)), + @as(u32, @bitCast(new_sync)), + .Release, + .Monotonic, + ) orelse { + // We signaled to notify an idle thread + if (can_wake and sync.idle > 0) { + return self.idle_event.notify(); + } + + // We signaled to spawn a new thread + if (can_wake and sync.spawned < self.max_threads) { + const spawn_config = if (builtin.os.tag.isDarwin()) + // stack size must be a multiple of page_size + // macOS will fail to spawn a thread if the stack size is not a multiple of page_size + std.Thread.SpawnConfig{ .stack_size = ((std.Thread.SpawnConfig{}).stack_size + (std.mem.page_size / 2) / std.mem.page_size) * std.mem.page_size } + else + std.Thread.SpawnConfig{}; + + const thread = std.Thread.spawn(spawn_config, Thread.run, .{self}) catch return self.unregister(null); + // if (self.name.len > 0) thread.setName(self.name) catch {}; + return thread.detach(); + } + + return; + })); + } + } + + noinline fn wait(self: *ThreadPool, _is_waking: bool) error{Shutdown}!bool { + var is_idle = false; + var is_waking = _is_waking; + var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + + while (true) { + if (sync.state == .shutdown) return error.Shutdown; + if (is_waking) assert(sync.state == .waking); + + // Consume a notification made by notify(). + if (sync.notified) { + var new_sync = sync; + new_sync.notified = false; + if (is_idle) + new_sync.idle -= 1; + if (sync.state == .signaled) + new_sync.state = .waking; + + // Acquire barrier synchronizes with notify() + // to ensure that pushes to run queue are observed after wait() returns. + sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( + @as(u32, @bitCast(sync)), + @as(u32, @bitCast(new_sync)), + .Acquire, + .Monotonic, + ) orelse { + return is_waking or (sync.state == .signaled); + })); + } else if (!is_idle) { + var new_sync = sync; + new_sync.idle += 1; + if (is_waking) + new_sync.state = .pending; + + sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( + @as(u32, @bitCast(sync)), + @as(u32, @bitCast(new_sync)), + .Monotonic, + .Monotonic, + ) orelse { + is_waking = false; + is_idle = true; + continue; + })); + } else { + if (Thread.current) |current| { + current.drainIdleEvents(); + } + + self.idle_event.wait(); + sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + } + } + } + + /// Marks the thread pool as shutdown + pub noinline fn shutdown(self: *ThreadPool) void { + var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + while (sync.state != .shutdown) { + var new_sync = sync; + new_sync.notified = true; + new_sync.state = .shutdown; + new_sync.idle = 0; + + // Full barrier to synchronize with both wait() and notify() + sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( + @as(u32, @bitCast(sync)), + @as(u32, @bitCast(new_sync)), + .AcqRel, + .Monotonic, + ) orelse { + // Wake up any threads sleeping on the idle_event. + // TODO: I/O polling notification here. + if (sync.idle > 0) self.idle_event.shutdown(); + return; + })); + } + } + + fn register(noalias self: *ThreadPool, noalias thread: *Thread) void { + // Push the thread onto the threads stack in a lock-free manner. + var threads = self.threads.load(.Monotonic); + while (true) { + thread.next = threads; + threads = self.threads.tryCompareAndSwap( + threads, + thread, + .Release, + .Monotonic, + ) orelse break; + } + } + + pub fn setThreadContext(noalias pool: *ThreadPool, ctx: ?*anyopaque) void { + pool.threadpool_context = ctx; + + var thread = pool.threads.load(.Monotonic) orelse return; + thread.ctx = pool.threadpool_context; + while (thread.next) |next| { + next.ctx = pool.threadpool_context; + thread = next; + } + } + + fn unregister(noalias self: *ThreadPool, noalias maybe_thread: ?*Thread) void { + // Un-spawn one thread, either due to a failed OS thread spawning or the thread is exiting. + const one_spawned = @as(u32, @bitCast(Sync{ .spawned = 1 })); + const sync = @as(Sync, @bitCast(self.sync.fetchSub(one_spawned, .Release))); + assert(sync.spawned > 0); + + // The last thread to exit must wake up the thread pool join()er + // who will start the chain to shutdown all the threads. + if (sync.state == .shutdown and sync.spawned == 1) { + self.join_event.notify(); + } + + // If this is a thread pool thread, wait for a shutdown signal by the thread pool join()er. + const thread = maybe_thread orelse return; + thread.join_event.wait(); + + // After receiving the shutdown signal, shutdown the next thread in the pool. + // We have to do that without touching the thread pool itself since it's memory is invalidated by now. + // So just follow our .next link. + const next_thread = thread.next orelse return; + next_thread.join_event.notify(); + } + + fn join(self: *ThreadPool) void { + // Wait for the thread pool to be shutdown() then for all threads to enter a joinable state + var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + if (!(sync.state == .shutdown and sync.spawned == 0)) { + self.join_event.wait(); + sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); + } + + assert(sync.state == .shutdown); + assert(sync.spawned == 0); + + // If there are threads, start off the chain sending it the shutdown signal. + // The thread receives the shutdown signal and sends it to the next thread, and the next.. + const thread = self.threads.load(.Acquire) orelse return; + thread.join_event.notify(); + } + + pub const Thread = struct { + next: ?*Thread = null, + target: ?*Thread = null, + join_event: Event = .{}, + run_queue: Node.Queue = .{}, + idle_queue: Node.Queue = .{}, + run_buffer: Node.Buffer = .{}, + ctx: ?*anyopaque = null, + + pub threadlocal var current: ?*Thread = null; + + pub fn pushIdleTask(self: *Thread, task: *Task) void { + const list = Node.List{ + .head = &task.node, + .tail = &task.node, + }; + self.idle_queue.push(list); + } + + /// Thread entry point which runs a worker for the ThreadPool + fn run(thread_pool: *ThreadPool) void { + var self_ = Thread{}; + var self = &self_; + current = self; + + if (thread_pool.on_thread_spawn) |spawn| { + current.?.ctx = spawn(thread_pool.threadpool_context); + } + + thread_pool.register(self); + + defer thread_pool.unregister(self); + + var is_waking = false; + while (true) { + is_waking = thread_pool.wait(is_waking) catch return; + + while (self.pop(thread_pool)) |result| { + if (result.pushed or is_waking) + thread_pool.notify(is_waking); + is_waking = false; + + const task = @fieldParentPtr(Task, "node", result.node); + (task.callback)(task); + } + + self.drainIdleEvents(); + } + } + + pub fn drainIdleEvents(noalias self: *Thread) void { + var consumer = self.idle_queue.tryAcquireConsumer() catch return; + defer self.idle_queue.releaseConsumer(consumer); + while (self.idle_queue.pop(&consumer)) |node| { + const task = @fieldParentPtr(Task, "node", node); + (task.callback)(task); + } + } + + /// Try to dequeue a Node/Task from the ThreadPool. + /// Spurious reports of dequeue() returning empty are allowed. + pub fn pop(noalias self: *Thread, noalias thread_pool: *ThreadPool) ?Node.Buffer.Stole { + // Check our local buffer first + if (self.run_buffer.pop()) |node| { + return Node.Buffer.Stole{ + .node = node, + .pushed = false, + }; + } + + // Then check our local queue + if (self.run_buffer.consume(&self.run_queue)) |stole| { + return stole; + } + + // Then the global queue + if (self.run_buffer.consume(&thread_pool.run_queue)) |stole| { + return stole; + } + + // Then try work stealing from other threads + var num_threads: u32 = @as(Sync, @bitCast(thread_pool.sync.load(.Monotonic))).spawned; + while (num_threads > 0) : (num_threads -= 1) { + // Traverse the stack of registered threads on the thread pool + const target = self.target orelse thread_pool.threads.load(.Acquire) orelse unreachable; + self.target = target.next; + + // Try to steal from their queue first to avoid contention (the target steal's from queue last). + if (self.run_buffer.consume(&target.run_queue)) |stole| { + return stole; + } + + // Skip stealing from the buffer if we're the target. + // We still steal from our own queue above given it may have just been locked the first time we tried. + if (target == self) { + continue; + } + + // Steal from the buffer of a remote thread as a last resort + if (self.run_buffer.steal(&target.run_buffer)) |stole| { + return stole; + } + } + + return null; + } + }; + + /// An event which stores 1 semaphore token and is multi-threaded safe. + /// The event can be shutdown(), waking up all wait()ing threads and + /// making subsequent wait()'s return immediately. + const Event = struct { + state: Atomic(u32) = Atomic(u32).init(EMPTY), + + const EMPTY = 0; + const WAITING = 1; + pub const NOTIFIED = 2; + const SHUTDOWN = 3; + + /// Wait for and consume a notification + /// or wait for the event to be shutdown entirely + noinline fn wait(self: *Event) void { + var acquire_with: u32 = EMPTY; + var state = self.state.load(.Monotonic); + + while (true) { + // If we're shutdown then exit early. + // Acquire barrier to ensure operations before the shutdown() are seen after the wait(). + // Shutdown is rare so it's better to have an Acquire barrier here instead of on CAS failure + load which are common. + if (state == SHUTDOWN) { + std.atomic.fence(.Acquire); + return; + } + + // Consume a notification when it pops up. + // Acquire barrier to ensure operations before the notify() appear after the wait(). + if (state == NOTIFIED) { + state = self.state.tryCompareAndSwap( + state, + acquire_with, + .Acquire, + .Monotonic, + ) orelse return; + continue; + } + + // There is no notification to consume, we should wait on the event by ensuring its WAITING. + if (state != WAITING) blk: { + state = self.state.tryCompareAndSwap( + state, + WAITING, + .Monotonic, + .Monotonic, + ) orelse break :blk; + continue; + } + + // Wait on the event until a notify() or shutdown(). + // If we wake up to a notification, we must acquire it with WAITING instead of EMPTY + // since there may be other threads sleeping on the Futex who haven't been woken up yet. + // + // Acquiring to WAITING will make the next notify() or shutdown() wake a sleeping futex thread + // who will either exit on SHUTDOWN or acquire with WAITING again, ensuring all threads are awoken. + // This unfortunately results in the last notify() or shutdown() doing an extra futex wake but that's fine. + Futex.wait(&self.state, WAITING); + state = self.state.load(.Monotonic); + acquire_with = WAITING; + } + } + + /// Wait for and consume a notification + /// or wait for the event to be shutdown entirely + noinline fn waitFor(self: *Event, timeout: usize) void { + _ = timeout; + var acquire_with: u32 = EMPTY; + var state = self.state.load(.Monotonic); + + while (true) { + // If we're shutdown then exit early. + // Acquire barrier to ensure operations before the shutdown() are seen after the wait(). + // Shutdown is rare so it's better to have an Acquire barrier here instead of on CAS failure + load which are common. + if (state == SHUTDOWN) { + std.atomic.fence(.Acquire); + return; + } + + // Consume a notification when it pops up. + // Acquire barrier to ensure operations before the notify() appear after the wait(). + if (state == NOTIFIED) { + state = self.state.tryCompareAndSwap( + state, + acquire_with, + .Acquire, + .Monotonic, + ) orelse return; + continue; + } + + // There is no notification to consume, we should wait on the event by ensuring its WAITING. + if (state != WAITING) blk: { + state = self.state.tryCompareAndSwap( + state, + WAITING, + .Monotonic, + .Monotonic, + ) orelse break :blk; + continue; + } + + // Wait on the event until a notify() or shutdown(). + // If we wake up to a notification, we must acquire it with WAITING instead of EMPTY + // since there may be other threads sleeping on the Futex who haven't been woken up yet. + // + // Acquiring to WAITING will make the next notify() or shutdown() wake a sleeping futex thread + // who will either exit on SHUTDOWN or acquire with WAITING again, ensuring all threads are awoken. + // This unfortunately results in the last notify() or shutdown() doing an extra futex wake but that's fine. + Futex.wait(&self.state, WAITING); + state = self.state.load(.Monotonic); + acquire_with = WAITING; + } + } + + /// Post a notification to the event if it doesn't have one already + /// then wake up a waiting thread if there is one as well. + fn notify(self: *Event) void { + return self.wake(NOTIFIED, 1); + } + + /// Marks the event as shutdown, making all future wait()'s return immediately. + /// Then wakes up any threads currently waiting on the Event. + fn shutdown(self: *Event) void { + return self.wake(SHUTDOWN, std.math.maxInt(u32)); + } + + fn wake(self: *Event, release_with: u32, wake_threads: u32) void { + // Update the Event to notify it with the new `release_with` state (either NOTIFIED or SHUTDOWN). + // Release barrier to ensure any operations before this are this to happen before the wait() in the other threads. + const state = self.state.swap(release_with, .Release); + + // Only wake threads sleeping in futex if the state is WAITING. + // Avoids unnecessary wake ups. + if (state == WAITING) { + Futex.wake(&self.state, wake_threads); + } + } + }; + + /// Linked list intrusive memory node and lock-free data structures to operate with it + pub const Node = struct { + next: ?*Node = null, + + /// A linked list of Nodes + const List = struct { + head: *Node, + tail: *Node, + }; + + /// An unbounded multi-producer-(non blocking)-multi-consumer queue of Node pointers. + const Queue = struct { + stack: Atomic(usize) = Atomic(usize).init(0), + cache: ?*Node = null, + + const HAS_CACHE: usize = 0b01; + const IS_CONSUMING: usize = 0b10; + const PTR_MASK: usize = ~(HAS_CACHE | IS_CONSUMING); + + comptime { + assert(@alignOf(Node) >= ((IS_CONSUMING | HAS_CACHE) + 1)); + } + + fn push(noalias self: *Queue, list: List) void { + var stack = self.stack.load(.Monotonic); + while (true) { + // Attach the list to the stack (pt. 1) + list.tail.next = @as(?*Node, @ptrFromInt(stack & PTR_MASK)); + + // Update the stack with the list (pt. 2). + // Don't change the HAS_CACHE and IS_CONSUMING bits of the consumer. + var new_stack = @intFromPtr(list.head); + assert(new_stack & ~PTR_MASK == 0); + new_stack |= (stack & ~PTR_MASK); + + // Push to the stack with a release barrier for the consumer to see the proper list links. + stack = self.stack.tryCompareAndSwap( + stack, + new_stack, + .Release, + .Monotonic, + ) orelse break; + } + } + + fn tryAcquireConsumer(self: *Queue) error{ Empty, Contended }!?*Node { + var stack = self.stack.load(.Monotonic); + while (true) { + if (stack & IS_CONSUMING != 0) + return error.Contended; // The queue already has a consumer. + if (stack & (HAS_CACHE | PTR_MASK) == 0) + return error.Empty; // The queue is empty when there's nothing cached and nothing in the stack. + + // When we acquire the consumer, also consume the pushed stack if the cache is empty. + var new_stack = stack | HAS_CACHE | IS_CONSUMING; + if (stack & HAS_CACHE == 0) { + assert(stack & PTR_MASK != 0); + new_stack &= ~PTR_MASK; + } + + // Acquire barrier on getting the consumer to see cache/Node updates done by previous consumers + // and to ensure our cache/Node updates in pop() happen after that of previous consumers. + stack = self.stack.tryCompareAndSwap( + stack, + new_stack, + .Acquire, + .Monotonic, + ) orelse return self.cache orelse @as(*Node, @ptrFromInt(stack & PTR_MASK)); + } + } + + fn releaseConsumer(noalias self: *Queue, noalias consumer: ?*Node) void { + // Stop consuming and remove the HAS_CACHE bit as well if the consumer's cache is empty. + // When HAS_CACHE bit is zeroed, the next consumer will acquire the pushed stack nodes. + var remove = IS_CONSUMING; + if (consumer == null) + remove |= HAS_CACHE; + + // Release the consumer with a release barrier to ensure cache/node accesses + // happen before the consumer was released and before the next consumer starts using the cache. + self.cache = consumer; + const stack = self.stack.fetchSub(remove, .Release); + assert(stack & remove != 0); + } + + fn pop(noalias self: *Queue, noalias consumer_ref: *?*Node) ?*Node { + // Check the consumer cache (fast path) + if (consumer_ref.*) |node| { + consumer_ref.* = node.next; + return node; + } + + // Load the stack to see if there was anything pushed that we could grab. + var stack = self.stack.load(.Monotonic); + assert(stack & IS_CONSUMING != 0); + if (stack & PTR_MASK == 0) { + return null; + } + + // Nodes have been pushed to the stack, grab then with an Acquire barrier to see the Node links. + stack = self.stack.swap(HAS_CACHE | IS_CONSUMING, .Acquire); + assert(stack & IS_CONSUMING != 0); + assert(stack & PTR_MASK != 0); + + const node = @as(*Node, @ptrFromInt(stack & PTR_MASK)); + consumer_ref.* = node.next; + return node; + } + }; + + /// A bounded single-producer, multi-consumer ring buffer for node pointers. + const Buffer = struct { + head: Atomic(Index) = Atomic(Index).init(0), + tail: Atomic(Index) = Atomic(Index).init(0), + array: [capacity]Atomic(*Node) = undefined, + + const Index = u32; + const capacity = 256; // Appears to be a pretty good trade-off in space vs contended throughput + comptime { + assert(std.math.maxInt(Index) >= capacity); + assert(std.math.isPowerOfTwo(capacity)); + } + + fn push(noalias self: *Buffer, noalias list: *List) error{Overflow}!void { + var head = self.head.load(.Monotonic); + var tail = self.tail.loadUnchecked(); // we're the only thread that can change this + + while (true) { + var size = tail -% head; + assert(size <= capacity); + + // Push nodes from the list to the buffer if it's not empty.. + if (size < capacity) { + var nodes: ?*Node = list.head; + while (size < capacity) : (size += 1) { + const node = nodes orelse break; + nodes = node.next; + + // Array written atomically with weakest ordering since it could be getting atomically read by steal(). + self.array[tail % capacity].store(node, .Unordered); + tail +%= 1; + } + + // Release barrier synchronizes with Acquire loads for steal()ers to see the array writes. + self.tail.store(tail, .Release); + + // Update the list with the nodes we pushed to the buffer and try again if there's more. + list.head = nodes orelse return; + std.atomic.spinLoopHint(); + head = self.head.load(.Monotonic); + continue; + } + + // Try to steal/overflow half of the tasks in the buffer to make room for future push()es. + // Migrating half amortizes the cost of stealing while requiring future pops to still use the buffer. + // Acquire barrier to ensure the linked list creation after the steal only happens after we successfully steal. + var migrate = size / 2; + head = self.head.tryCompareAndSwap( + head, + head +% migrate, + .Acquire, + .Monotonic, + ) orelse { + // Link the migrated Nodes together + const first = self.array[head % capacity].loadUnchecked(); + while (migrate > 0) : (migrate -= 1) { + const prev = self.array[head % capacity].loadUnchecked(); + head +%= 1; + prev.next = self.array[head % capacity].loadUnchecked(); + } + + // Append the list that was supposed to be pushed to the end of the migrated Nodes + const last = self.array[(head -% 1) % capacity].loadUnchecked(); + last.next = list.head; + list.tail.next = null; + + // Return the migrated nodes + the original list as overflowed + list.head = first; + return error.Overflow; + }; + } + } + + fn pop(self: *Buffer) ?*Node { + var head = self.head.load(.Monotonic); + var tail = self.tail.loadUnchecked(); // we're the only thread that can change this + + while (true) { + // Quick sanity check and return null when not empty + var size = tail -% head; + assert(size <= capacity); + if (size == 0) { + return null; + } + + // Dequeue with an acquire barrier to ensure any writes done to the Node + // only happens after we successfully claim it from the array. + head = self.head.tryCompareAndSwap( + head, + head +% 1, + .Acquire, + .Monotonic, + ) orelse return self.array[head % capacity].loadUnchecked(); + } + } + + const Stole = struct { + node: *Node, + pushed: bool, + }; + + fn consume(noalias self: *Buffer, noalias queue: *Queue) ?Stole { + var consumer = queue.tryAcquireConsumer() catch return null; + defer queue.releaseConsumer(consumer); + + const head = self.head.load(.Monotonic); + const tail = self.tail.loadUnchecked(); // we're the only thread that can change this + + const size = tail -% head; + assert(size <= capacity); + assert(size == 0); // we should only be consuming if our array is empty + + // Pop nodes from the queue and push them to our array. + // Atomic stores to the array as steal() threads may be atomically reading from it. + var pushed: Index = 0; + while (pushed < capacity) : (pushed += 1) { + const node = queue.pop(&consumer) orelse break; + self.array[(tail +% pushed) % capacity].store(node, .Unordered); + } + + // We will be returning one node that we stole from the queue. + // Get an extra, and if that's not possible, take one from our array. + const node = queue.pop(&consumer) orelse blk: { + if (pushed == 0) return null; + pushed -= 1; + break :blk self.array[(tail +% pushed) % capacity].loadUnchecked(); + }; + + // Update the array tail with the nodes we pushed to it. + // Release barrier to synchronize with Acquire barrier in steal()'s to see the written array Nodes. + if (pushed > 0) self.tail.store(tail +% pushed, .Release); + return Stole{ + .node = node, + .pushed = pushed > 0, + }; + } + + fn steal(noalias self: *Buffer, noalias buffer: *Buffer) ?Stole { + const head = self.head.load(.Monotonic); + const tail = self.tail.loadUnchecked(); // we're the only thread that can change this + + const size = tail -% head; + assert(size <= capacity); + assert(size == 0); // we should only be stealing if our array is empty + + while (true) : (std.atomic.spinLoopHint()) { + const buffer_head = buffer.head.load(.Acquire); + const buffer_tail = buffer.tail.load(.Acquire); + + // Overly large size indicates the the tail was updated a lot after the head was loaded. + // Reload both and try again. + const buffer_size = buffer_tail -% buffer_head; + if (buffer_size > capacity) { + continue; + } + + // Try to steal half (divCeil) to amortize the cost of stealing from other threads. + const steal_size = buffer_size - (buffer_size / 2); + if (steal_size == 0) { + return null; + } + + // Copy the nodes we will steal from the target's array to our own. + // Atomically load from the target buffer array as it may be pushing and atomically storing to it. + // Atomic store to our array as other steal() threads may be atomically loading from it as above. + var i: Index = 0; + while (i < steal_size) : (i += 1) { + const node = buffer.array[(buffer_head +% i) % capacity].load(.Unordered); + self.array[(tail +% i) % capacity].store(node, .Unordered); + } + + // Try to commit the steal from the target buffer using: + // - an Acquire barrier to ensure that we only interact with the stolen Nodes after the steal was committed. + // - a Release barrier to ensure that the Nodes are copied above prior to the committing of the steal + // because if they're copied after the steal, the could be getting rewritten by the target's push(). + _ = buffer.head.compareAndSwap( + buffer_head, + buffer_head +% steal_size, + .AcqRel, + .Monotonic, + ) orelse { + // Pop one from the nodes we stole as we'll be returning it + const pushed = steal_size - 1; + const node = self.array[(tail +% pushed) % capacity].loadUnchecked(); + + // Update the array tail with the nodes we pushed to it. + // Release barrier to synchronize with Acquire barrier in steal()'s to see the written array Nodes. + if (pushed > 0) self.tail.store(tail +% pushed, .Release); + return Stole{ + .node = node, + .pushed = pushed > 0, + }; + }; + } + } + }; + }; +}; + +test "parallel for loop" { + var thread_pool = ThreadPool.init(.{ .max_threads = 12 }); + var sleepy_time: u32 = 100; + var huge_array = &[_]u32{ + sleepy_time + std.rand.DefaultPrng.init(1).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(2).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(3).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(4).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(5).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(6).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(7).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(8).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(9).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(10).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(11).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(12).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(13).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(14).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(15).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(16).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(17).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(18).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(19).random().uintAtMost(u32, 20), + sleepy_time + std.rand.DefaultPrng.init(20).random().uintAtMost(u32, 20), + }; + const Runner = struct { + completed: usize = 0, + total: usize = 0, + pub fn run(ctx: *@This(), value: u32, _: usize) void { + std.time.sleep(value); + ctx.completed += 1; + std.debug.assert(ctx.completed <= ctx.total); + } + }; + var runny = try std.heap.page_allocator.create(Runner); + runny.* = .{ .total = huge_array.len }; + try thread_pool.doAndWait(std.heap.page_allocator, null, runny, Runner.run, std.mem.span(huge_array)); + try std.testing.expectEqual(huge_array.len, runny.completed); +} + +pub fn NewWorkPool(comptime max_threads: ?usize) type { + return struct { + var pool: ThreadPool = undefined; + var loaded: bool = false; + + fn create() *ThreadPool { + @setCold(true); + + pool = ThreadPool.init(.{ + .max_threads = max_threads orelse @max(@as(u32, @truncate(std.Thread.getCpuCount() catch 0)), 2), + .stack_size = 2 * 1024 * 1024, + }); + return &pool; + } + + pub fn deinit() void { + get().deinit(); + } + + pub inline fn get() *ThreadPool { + // lil racy + if (loaded) return &pool; + loaded = true; + + return create(); + } + + pub fn scheduleBatch(batch: ThreadPool.Batch) void { + get().schedule(batch); + } + + pub fn scheduleTask(task: *ThreadPool.Task) void { + get().schedule(ThreadPool.Batch.from(task)); + } + + pub fn go(allocator: std.mem.Allocator, comptime Context: type, context: Context, comptime function: *const fn (Context) void) !void { + const TaskType = struct { + task: ThreadPool.Task, + context: Context, + allocator: std.mem.Allocator, + + pub fn callback(task: *ThreadPool.Task) void { + var this_task = @fieldParentPtr(@This(), "task", task); + function(this_task.context); + this_task.allocator.destroy(this_task); + } + }; + + var task_ = try allocator.create(TaskType); + task_.* = .{ + .task = .{ .callback = TaskType.callback }, + .context = context, + .allocator = allocator, + }; + scheduleTask(&task_.task); + } + }; +} + +pub const WorkPool = NewWorkPool(null); +const testing = std.testing; + +const CrdsTableTrimContext = struct { + index: usize, + max_trim: usize, + self: *CrdsTable, +}; + +const CrdsTable = struct { + pub fn trim(context: CrdsTableTrimContext) void { + const self = context.self; + _ = self; + const max_trim = context.max_trim; + _ = max_trim; + const index = context.index; + _ = index; + + std.debug.print("I ran!\n\n", .{}); + // todo + + } +}; + +test "sync.thread_pool: workpool works" { + var crds: CrdsTable = CrdsTable{}; + var a = CrdsTableTrimContext{ .index = 1, .max_trim = 2, .self = &crds }; + defer WorkPool.deinit(); + try WorkPool.go(testing.allocator, CrdsTableTrimContext, a, CrdsTable.trim); + + std.time.sleep(std.time.ns_per_s * 1); + WorkPool.pool.shutdown(); +} From 8e3d9758bdae3b7964676e76052df41ddf58295b Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 08:23:35 -0500 Subject: [PATCH 02/15] merge: resolve conflict before merge: remove thread pool since there's already one in main --- src/lib.zig | 1 - src/sync/thread_pool.zig | 1316 -------------------------------------- 2 files changed, 1317 deletions(-) delete mode 100644 src/sync/thread_pool.zig diff --git a/src/lib.zig b/src/lib.zig index 86cdf05bf..c2972ca1f 100644 --- a/src/lib.zig +++ b/src/lib.zig @@ -45,7 +45,6 @@ pub const sync = struct { pub usingnamespace @import("sync/mpmc.zig"); pub usingnamespace @import("sync/ref.zig"); pub usingnamespace @import("sync/mux.zig"); - pub usingnamespace @import("sync/thread_pool.zig"); }; pub const utils = struct { diff --git a/src/sync/thread_pool.zig b/src/sync/thread_pool.zig deleted file mode 100644 index 5c5143f4f..000000000 --- a/src/sync/thread_pool.zig +++ /dev/null @@ -1,1316 +0,0 @@ -// Thank you bun.sh: -// https://github.com/oven-sh/bun/blob/main/src/thread_pool.zig -// -// Thank you @kprotty: -// https://github.com/kprotty/zap/blob/blog/src/thread_pool.zig - -const std = @import("std"); -const builtin = @import("builtin"); -const Futex = std.Thread.Futex; -const assert = std.debug.assert; -const Atomic = std.atomic.Atomic; -pub const OnSpawnCallback = *const fn (ctx: ?*anyopaque) ?*anyopaque; - -pub const ThreadPool = struct { - sleep_on_idle_network_thread: bool = true, - /// executed on the thread - on_thread_spawn: ?OnSpawnCallback = null, - threadpool_context: ?*anyopaque = null, - stack_size: u32, - max_threads: u32, - sync: Atomic(u32) = Atomic(u32).init(@as(u32, @bitCast(Sync{}))), - idle_event: Event = .{}, - join_event: Event = .{}, - run_queue: Node.Queue = .{}, - threads: Atomic(?*Thread) = Atomic(?*Thread).init(null), - name: []const u8 = "", - spawned_thread_count: Atomic(u32) = Atomic(u32).init(0), - - const Sync = packed struct { - /// Tracks the number of threads not searching for Tasks - idle: u14 = 0, - /// Tracks the number of threads spawned - spawned: u14 = 0, - /// What you see is what you get - unused: bool = false, - /// Used to not miss notifications while state = waking - notified: bool = false, - /// The current state of the thread pool - state: enum(u2) { - /// A notification can be issued to wake up a sleeping as the "waking thread". - pending = 0, - /// The state was notified with a signal. A thread is woken up. - /// The first thread to transition to `waking` becomes the "waking thread". - signaled, - /// There is a "waking thread" among us. - /// No other thread should be woken up until the waking thread transitions the state. - waking, - /// The thread pool was terminated. Start decremented `spawned` so that it can be joined. - shutdown, - } = .pending, - }; - - /// Configuration options for the thread pool. - /// TODO: add CPU core affinity? - pub const Config = struct { - stack_size: u32 = (std.Thread.SpawnConfig{}).stack_size, - max_threads: u32, - }; - - /// Statically initialize the thread pool using the configuration. - pub fn init(config: Config) ThreadPool { - return .{ - .stack_size = @max(1, config.stack_size), - .max_threads = @max(1, config.max_threads), - }; - } - - pub fn wakeForIdleEvents(this: *ThreadPool) void { - // Wake all the threads to check for idle events. - this.idle_event.wake(Event.NOTIFIED, std.math.maxInt(u32)); - } - - /// Wait for a thread to call shutdown() on the thread pool and kill the worker threads. - pub fn deinit(self: *ThreadPool) void { - self.join(); - self.* = undefined; - } - - /// A Task represents the unit of Work / Job / Execution that the ThreadPool schedules. - /// The user provides a `callback` which is invoked when the *Task can run on a thread. - pub const Task = struct { - node: Node = .{}, - callback: *const (fn (*Task) void), - }; - - /// An unordered collection of Tasks which can be submitted for scheduling as a group. - pub const Batch = struct { - len: usize = 0, - head: ?*Task = null, - tail: ?*Task = null, - - pub fn pop(this: *Batch) ?*Task { - const len = @atomicLoad(usize, &this.len, .Monotonic); - if (len == 0) { - return null; - } - var task = this.head.?; - if (task.node.next) |node| { - this.head = @fieldParentPtr(Task, "node", node); - } else { - if (task != this.tail.?) unreachable; - this.tail = null; - this.head = null; - } - - this.len -= 1; - if (len == 0) { - this.tail = null; - } - return task; - } - - /// Create a batch from a single task. - pub fn from(task: *Task) Batch { - return Batch{ - .len = 1, - .head = task, - .tail = task, - }; - } - - /// Another batch into this one, taking ownership of its tasks. - pub fn push(self: *Batch, batch: Batch) void { - if (batch.len == 0) return; - if (self.len == 0) { - self.* = batch; - } else { - self.tail.?.node.next = if (batch.head) |h| &h.node else null; - self.tail = batch.tail; - self.len += batch.len; - } - } - }; - - pub const WaitGroup = struct { - mutex: std.Thread.Mutex = .{}, - counter: u32 = 0, - event: std.Thread.ResetEvent, - - pub fn init(self: *WaitGroup) void { - self.* = .{ - .mutex = .{}, - .counter = 0, - .event = undefined, - }; - } - - pub fn deinit(self: *WaitGroup) void { - self.event.reset(); - self.* = undefined; - } - - pub fn start(self: *WaitGroup) void { - self.mutex.lock(); - defer self.mutex.unlock(); - - self.counter += 1; - } - - pub fn isDone(this: *WaitGroup) bool { - return @atomicLoad(u32, &this.counter, .Monotonic) == 0; - } - - pub fn finish(self: *WaitGroup) void { - self.mutex.lock(); - defer self.mutex.unlock(); - - self.counter -= 1; - - if (self.counter == 0) { - self.event.set(); - } - } - - pub fn wait(self: *WaitGroup) void { - while (true) { - self.mutex.lock(); - - if (self.counter == 0) { - self.mutex.unlock(); - return; - } - - self.mutex.unlock(); - self.event.wait(); - } - } - - pub fn reset(self: *WaitGroup) void { - self.event.reset(); - } - }; - - pub fn ConcurrentFunction( - comptime Function: anytype, - ) type { - return struct { - const Fn = Function; - const Args = std.meta.ArgsTuple(@TypeOf(Fn)); - const Runner = @This(); - thread_pool: *ThreadPool, - states: []Routine = undefined, - batch: Batch = .{}, - allocator: std.mem.Allocator, - - pub fn init(allocator: std.mem.Allocator, thread_pool: *ThreadPool, count: usize) !Runner { - return Runner{ - .allocator = allocator, - .thread_pool = thread_pool, - .states = try allocator.alloc(Routine, count), - .batch = .{}, - }; - } - - pub fn call(this: *@This(), args: Args) void { - this.states[this.batch.len] = .{ - .args = args, - }; - this.batch.push(Batch.from(&this.states[this.batch.len].task)); - } - - pub fn run(this: *@This()) void { - this.thread_pool.schedule(this.batch); - } - - pub const Routine = struct { - args: Args, - task: Task = .{ .callback = callback }, - - pub fn callback(task: *Task) void { - var routine = @fieldParentPtr(@This(), "task", task); - @call(.always_inline, Fn, routine.args); - } - }; - - pub fn deinit(this: *@This()) void { - this.allocator.free(this.states); - } - }; - } - - pub fn runner( - this: *ThreadPool, - allocator: std.mem.Allocator, - comptime Function: anytype, - count: usize, - ) !ConcurrentFunction(Function) { - return try ConcurrentFunction(Function).init(allocator, this, count); - } - - /// Loop over an array of tasks and invoke `Run` on each one in a different thread - /// **Blocks the calling thread** until all tasks are completed. - pub fn do( - this: *ThreadPool, - allocator: std.mem.Allocator, - wg: ?*WaitGroup, - ctx: anytype, - comptime Run: anytype, - values: anytype, - ) !void { - return try Do(this, allocator, wg, @TypeOf(ctx), ctx, Run, @TypeOf(values), values, false); - } - - pub fn doPtr( - this: *ThreadPool, - allocator: std.mem.Allocator, - wg: ?*WaitGroup, - ctx: anytype, - comptime Run: anytype, - values: anytype, - ) !void { - return try Do(this, allocator, wg, @TypeOf(ctx), ctx, Run, @TypeOf(values), values, true); - } - - pub fn Do( - this: *ThreadPool, - allocator: std.mem.Allocator, - wg: ?*WaitGroup, - comptime Context: type, - ctx: Context, - comptime Function: anytype, - comptime ValuesType: type, - values: ValuesType, - comptime as_ptr: bool, - ) !void { - if (values.len == 0) - return; - var allocated_wait_group: ?*WaitGroup = null; - defer { - if (allocated_wait_group) |group| { - group.deinit(); - allocator.destroy(group); - } - } - - var wait_group = wg orelse brk: { - allocated_wait_group = try allocator.create(WaitGroup); - allocated_wait_group.?.init(); - break :brk allocated_wait_group.?; - }; - const WaitContext = struct { - wait_group: *WaitGroup = undefined, - ctx: Context, - values: ValuesType, - }; - - const RunnerTask = struct { - task: Task, - ctx: *WaitContext, - i: usize = 0, - - pub fn call(task: *Task) void { - var runner_task = @fieldParentPtr(@This(), "task", task); - const i = runner_task.i; - if (comptime as_ptr) { - Function(runner_task.ctx.ctx, &runner_task.ctx.values[i], i); - } else { - Function(runner_task.ctx.ctx, runner_task.ctx.values[i], i); - } - - runner_task.ctx.wait_group.finish(); - } - }; - var wait_context = allocator.create(WaitContext) catch unreachable; - wait_context.* = .{ - .ctx = ctx, - .wait_group = wait_group, - .values = values, - }; - defer allocator.destroy(wait_context); - var tasks = allocator.alloc(RunnerTask, values.len) catch unreachable; - defer allocator.free(tasks); - var batch: Batch = undefined; - var offset = tasks.len - 1; - - { - tasks[0] = .{ - .i = offset, - .task = .{ .callback = RunnerTask.call }, - .ctx = wait_context, - }; - batch = Batch.from(&tasks[0].task); - } - if (tasks.len > 1) { - for (tasks[1..]) |*runner_task| { - offset -= 1; - runner_task.* = .{ - .i = offset, - .task = .{ .callback = RunnerTask.call }, - .ctx = wait_context, - }; - batch.push(Batch.from(&runner_task.task)); - } - } - - wait_group.counter += @as(u32, @intCast(values.len)); - this.schedule(batch); - wait_group.wait(); - } - - /// Schedule a batch of tasks to be executed by some thread on the thread pool. - pub fn schedule(self: *ThreadPool, batch: Batch) void { - // Sanity check - if (batch.len == 0) { - return; - } - - // Extract out the Node's from the Tasks - var list = Node.List{ - .head = &batch.head.?.node, - .tail = &batch.tail.?.node, - }; - - // Push the task Nodes to the most appropriate queue - if (Thread.current) |thread| { - thread.run_buffer.push(&list) catch thread.run_queue.push(list); - } else { - self.run_queue.push(list); - } - - forceSpawn(self); - } - - pub fn forceSpawn(self: *ThreadPool) void { - // Try to notify a thread - const is_waking = false; - return self.notify(is_waking); - } - - inline fn notify(self: *ThreadPool, is_waking: bool) void { - // Fast path to check the Sync state to avoid calling into notifySlow(). - // If we're waking, then we need to update the state regardless - if (!is_waking) { - const sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - if (sync.notified) { - return; - } - } - - return self.notifySlow(is_waking); - } - - /// Warm the thread pool up to the given number of threads. - /// https://www.youtube.com/watch?v=ys3qcbO5KWw - pub fn warm(self: *ThreadPool, count: u14) void { - var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - if (sync.spawned >= count) - return; - - const to_spawn = @min(count - sync.spawned, @as(u14, @truncate(self.max_threads))); - while (sync.spawned < to_spawn) { - var new_sync = sync; - new_sync.spawned += 1; - sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( - @as(u32, @bitCast(sync)), - @as(u32, @bitCast(new_sync)), - .Release, - .Monotonic, - ) orelse break)); - const spawn_config = if (builtin.os.tag.isDarwin()) - // stack size must be a multiple of page_size - // macOS will fail to spawn a thread if the stack size is not a multiple of page_size - std.Thread.SpawnConfig{ .stack_size = ((std.Thread.SpawnConfig{}).stack_size + (std.mem.page_size / 2) / std.mem.page_size) * std.mem.page_size } - else - std.Thread.SpawnConfig{}; - - const thread = std.Thread.spawn(spawn_config, Thread.run, .{self}) catch return self.unregister(null); - thread.detach(); - } - } - - noinline fn notifySlow(self: *ThreadPool, is_waking: bool) void { - var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - while (sync.state != .shutdown) { - const can_wake = is_waking or (sync.state == .pending); - if (is_waking) { - assert(sync.state == .waking); - } - - var new_sync = sync; - new_sync.notified = true; - if (can_wake and sync.idle > 0) { // wake up an idle thread - new_sync.state = .signaled; - } else if (can_wake and sync.spawned < self.max_threads) { // spawn a new thread - new_sync.state = .signaled; - new_sync.spawned += 1; - } else if (is_waking) { // no other thread to pass on "waking" status - new_sync.state = .pending; - } else if (sync.notified) { // nothing to update - return; - } - - // Release barrier synchronizes with Acquire in wait() - // to ensure pushes to run queues happen before observing a posted notification. - sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( - @as(u32, @bitCast(sync)), - @as(u32, @bitCast(new_sync)), - .Release, - .Monotonic, - ) orelse { - // We signaled to notify an idle thread - if (can_wake and sync.idle > 0) { - return self.idle_event.notify(); - } - - // We signaled to spawn a new thread - if (can_wake and sync.spawned < self.max_threads) { - const spawn_config = if (builtin.os.tag.isDarwin()) - // stack size must be a multiple of page_size - // macOS will fail to spawn a thread if the stack size is not a multiple of page_size - std.Thread.SpawnConfig{ .stack_size = ((std.Thread.SpawnConfig{}).stack_size + (std.mem.page_size / 2) / std.mem.page_size) * std.mem.page_size } - else - std.Thread.SpawnConfig{}; - - const thread = std.Thread.spawn(spawn_config, Thread.run, .{self}) catch return self.unregister(null); - // if (self.name.len > 0) thread.setName(self.name) catch {}; - return thread.detach(); - } - - return; - })); - } - } - - noinline fn wait(self: *ThreadPool, _is_waking: bool) error{Shutdown}!bool { - var is_idle = false; - var is_waking = _is_waking; - var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - - while (true) { - if (sync.state == .shutdown) return error.Shutdown; - if (is_waking) assert(sync.state == .waking); - - // Consume a notification made by notify(). - if (sync.notified) { - var new_sync = sync; - new_sync.notified = false; - if (is_idle) - new_sync.idle -= 1; - if (sync.state == .signaled) - new_sync.state = .waking; - - // Acquire barrier synchronizes with notify() - // to ensure that pushes to run queue are observed after wait() returns. - sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( - @as(u32, @bitCast(sync)), - @as(u32, @bitCast(new_sync)), - .Acquire, - .Monotonic, - ) orelse { - return is_waking or (sync.state == .signaled); - })); - } else if (!is_idle) { - var new_sync = sync; - new_sync.idle += 1; - if (is_waking) - new_sync.state = .pending; - - sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( - @as(u32, @bitCast(sync)), - @as(u32, @bitCast(new_sync)), - .Monotonic, - .Monotonic, - ) orelse { - is_waking = false; - is_idle = true; - continue; - })); - } else { - if (Thread.current) |current| { - current.drainIdleEvents(); - } - - self.idle_event.wait(); - sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - } - } - } - - /// Marks the thread pool as shutdown - pub noinline fn shutdown(self: *ThreadPool) void { - var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - while (sync.state != .shutdown) { - var new_sync = sync; - new_sync.notified = true; - new_sync.state = .shutdown; - new_sync.idle = 0; - - // Full barrier to synchronize with both wait() and notify() - sync = @as(Sync, @bitCast(self.sync.tryCompareAndSwap( - @as(u32, @bitCast(sync)), - @as(u32, @bitCast(new_sync)), - .AcqRel, - .Monotonic, - ) orelse { - // Wake up any threads sleeping on the idle_event. - // TODO: I/O polling notification here. - if (sync.idle > 0) self.idle_event.shutdown(); - return; - })); - } - } - - fn register(noalias self: *ThreadPool, noalias thread: *Thread) void { - // Push the thread onto the threads stack in a lock-free manner. - var threads = self.threads.load(.Monotonic); - while (true) { - thread.next = threads; - threads = self.threads.tryCompareAndSwap( - threads, - thread, - .Release, - .Monotonic, - ) orelse break; - } - } - - pub fn setThreadContext(noalias pool: *ThreadPool, ctx: ?*anyopaque) void { - pool.threadpool_context = ctx; - - var thread = pool.threads.load(.Monotonic) orelse return; - thread.ctx = pool.threadpool_context; - while (thread.next) |next| { - next.ctx = pool.threadpool_context; - thread = next; - } - } - - fn unregister(noalias self: *ThreadPool, noalias maybe_thread: ?*Thread) void { - // Un-spawn one thread, either due to a failed OS thread spawning or the thread is exiting. - const one_spawned = @as(u32, @bitCast(Sync{ .spawned = 1 })); - const sync = @as(Sync, @bitCast(self.sync.fetchSub(one_spawned, .Release))); - assert(sync.spawned > 0); - - // The last thread to exit must wake up the thread pool join()er - // who will start the chain to shutdown all the threads. - if (sync.state == .shutdown and sync.spawned == 1) { - self.join_event.notify(); - } - - // If this is a thread pool thread, wait for a shutdown signal by the thread pool join()er. - const thread = maybe_thread orelse return; - thread.join_event.wait(); - - // After receiving the shutdown signal, shutdown the next thread in the pool. - // We have to do that without touching the thread pool itself since it's memory is invalidated by now. - // So just follow our .next link. - const next_thread = thread.next orelse return; - next_thread.join_event.notify(); - } - - fn join(self: *ThreadPool) void { - // Wait for the thread pool to be shutdown() then for all threads to enter a joinable state - var sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - if (!(sync.state == .shutdown and sync.spawned == 0)) { - self.join_event.wait(); - sync = @as(Sync, @bitCast(self.sync.load(.Monotonic))); - } - - assert(sync.state == .shutdown); - assert(sync.spawned == 0); - - // If there are threads, start off the chain sending it the shutdown signal. - // The thread receives the shutdown signal and sends it to the next thread, and the next.. - const thread = self.threads.load(.Acquire) orelse return; - thread.join_event.notify(); - } - - pub const Thread = struct { - next: ?*Thread = null, - target: ?*Thread = null, - join_event: Event = .{}, - run_queue: Node.Queue = .{}, - idle_queue: Node.Queue = .{}, - run_buffer: Node.Buffer = .{}, - ctx: ?*anyopaque = null, - - pub threadlocal var current: ?*Thread = null; - - pub fn pushIdleTask(self: *Thread, task: *Task) void { - const list = Node.List{ - .head = &task.node, - .tail = &task.node, - }; - self.idle_queue.push(list); - } - - /// Thread entry point which runs a worker for the ThreadPool - fn run(thread_pool: *ThreadPool) void { - var self_ = Thread{}; - var self = &self_; - current = self; - - if (thread_pool.on_thread_spawn) |spawn| { - current.?.ctx = spawn(thread_pool.threadpool_context); - } - - thread_pool.register(self); - - defer thread_pool.unregister(self); - - var is_waking = false; - while (true) { - is_waking = thread_pool.wait(is_waking) catch return; - - while (self.pop(thread_pool)) |result| { - if (result.pushed or is_waking) - thread_pool.notify(is_waking); - is_waking = false; - - const task = @fieldParentPtr(Task, "node", result.node); - (task.callback)(task); - } - - self.drainIdleEvents(); - } - } - - pub fn drainIdleEvents(noalias self: *Thread) void { - var consumer = self.idle_queue.tryAcquireConsumer() catch return; - defer self.idle_queue.releaseConsumer(consumer); - while (self.idle_queue.pop(&consumer)) |node| { - const task = @fieldParentPtr(Task, "node", node); - (task.callback)(task); - } - } - - /// Try to dequeue a Node/Task from the ThreadPool. - /// Spurious reports of dequeue() returning empty are allowed. - pub fn pop(noalias self: *Thread, noalias thread_pool: *ThreadPool) ?Node.Buffer.Stole { - // Check our local buffer first - if (self.run_buffer.pop()) |node| { - return Node.Buffer.Stole{ - .node = node, - .pushed = false, - }; - } - - // Then check our local queue - if (self.run_buffer.consume(&self.run_queue)) |stole| { - return stole; - } - - // Then the global queue - if (self.run_buffer.consume(&thread_pool.run_queue)) |stole| { - return stole; - } - - // Then try work stealing from other threads - var num_threads: u32 = @as(Sync, @bitCast(thread_pool.sync.load(.Monotonic))).spawned; - while (num_threads > 0) : (num_threads -= 1) { - // Traverse the stack of registered threads on the thread pool - const target = self.target orelse thread_pool.threads.load(.Acquire) orelse unreachable; - self.target = target.next; - - // Try to steal from their queue first to avoid contention (the target steal's from queue last). - if (self.run_buffer.consume(&target.run_queue)) |stole| { - return stole; - } - - // Skip stealing from the buffer if we're the target. - // We still steal from our own queue above given it may have just been locked the first time we tried. - if (target == self) { - continue; - } - - // Steal from the buffer of a remote thread as a last resort - if (self.run_buffer.steal(&target.run_buffer)) |stole| { - return stole; - } - } - - return null; - } - }; - - /// An event which stores 1 semaphore token and is multi-threaded safe. - /// The event can be shutdown(), waking up all wait()ing threads and - /// making subsequent wait()'s return immediately. - const Event = struct { - state: Atomic(u32) = Atomic(u32).init(EMPTY), - - const EMPTY = 0; - const WAITING = 1; - pub const NOTIFIED = 2; - const SHUTDOWN = 3; - - /// Wait for and consume a notification - /// or wait for the event to be shutdown entirely - noinline fn wait(self: *Event) void { - var acquire_with: u32 = EMPTY; - var state = self.state.load(.Monotonic); - - while (true) { - // If we're shutdown then exit early. - // Acquire barrier to ensure operations before the shutdown() are seen after the wait(). - // Shutdown is rare so it's better to have an Acquire barrier here instead of on CAS failure + load which are common. - if (state == SHUTDOWN) { - std.atomic.fence(.Acquire); - return; - } - - // Consume a notification when it pops up. - // Acquire barrier to ensure operations before the notify() appear after the wait(). - if (state == NOTIFIED) { - state = self.state.tryCompareAndSwap( - state, - acquire_with, - .Acquire, - .Monotonic, - ) orelse return; - continue; - } - - // There is no notification to consume, we should wait on the event by ensuring its WAITING. - if (state != WAITING) blk: { - state = self.state.tryCompareAndSwap( - state, - WAITING, - .Monotonic, - .Monotonic, - ) orelse break :blk; - continue; - } - - // Wait on the event until a notify() or shutdown(). - // If we wake up to a notification, we must acquire it with WAITING instead of EMPTY - // since there may be other threads sleeping on the Futex who haven't been woken up yet. - // - // Acquiring to WAITING will make the next notify() or shutdown() wake a sleeping futex thread - // who will either exit on SHUTDOWN or acquire with WAITING again, ensuring all threads are awoken. - // This unfortunately results in the last notify() or shutdown() doing an extra futex wake but that's fine. - Futex.wait(&self.state, WAITING); - state = self.state.load(.Monotonic); - acquire_with = WAITING; - } - } - - /// Wait for and consume a notification - /// or wait for the event to be shutdown entirely - noinline fn waitFor(self: *Event, timeout: usize) void { - _ = timeout; - var acquire_with: u32 = EMPTY; - var state = self.state.load(.Monotonic); - - while (true) { - // If we're shutdown then exit early. - // Acquire barrier to ensure operations before the shutdown() are seen after the wait(). - // Shutdown is rare so it's better to have an Acquire barrier here instead of on CAS failure + load which are common. - if (state == SHUTDOWN) { - std.atomic.fence(.Acquire); - return; - } - - // Consume a notification when it pops up. - // Acquire barrier to ensure operations before the notify() appear after the wait(). - if (state == NOTIFIED) { - state = self.state.tryCompareAndSwap( - state, - acquire_with, - .Acquire, - .Monotonic, - ) orelse return; - continue; - } - - // There is no notification to consume, we should wait on the event by ensuring its WAITING. - if (state != WAITING) blk: { - state = self.state.tryCompareAndSwap( - state, - WAITING, - .Monotonic, - .Monotonic, - ) orelse break :blk; - continue; - } - - // Wait on the event until a notify() or shutdown(). - // If we wake up to a notification, we must acquire it with WAITING instead of EMPTY - // since there may be other threads sleeping on the Futex who haven't been woken up yet. - // - // Acquiring to WAITING will make the next notify() or shutdown() wake a sleeping futex thread - // who will either exit on SHUTDOWN or acquire with WAITING again, ensuring all threads are awoken. - // This unfortunately results in the last notify() or shutdown() doing an extra futex wake but that's fine. - Futex.wait(&self.state, WAITING); - state = self.state.load(.Monotonic); - acquire_with = WAITING; - } - } - - /// Post a notification to the event if it doesn't have one already - /// then wake up a waiting thread if there is one as well. - fn notify(self: *Event) void { - return self.wake(NOTIFIED, 1); - } - - /// Marks the event as shutdown, making all future wait()'s return immediately. - /// Then wakes up any threads currently waiting on the Event. - fn shutdown(self: *Event) void { - return self.wake(SHUTDOWN, std.math.maxInt(u32)); - } - - fn wake(self: *Event, release_with: u32, wake_threads: u32) void { - // Update the Event to notify it with the new `release_with` state (either NOTIFIED or SHUTDOWN). - // Release barrier to ensure any operations before this are this to happen before the wait() in the other threads. - const state = self.state.swap(release_with, .Release); - - // Only wake threads sleeping in futex if the state is WAITING. - // Avoids unnecessary wake ups. - if (state == WAITING) { - Futex.wake(&self.state, wake_threads); - } - } - }; - - /// Linked list intrusive memory node and lock-free data structures to operate with it - pub const Node = struct { - next: ?*Node = null, - - /// A linked list of Nodes - const List = struct { - head: *Node, - tail: *Node, - }; - - /// An unbounded multi-producer-(non blocking)-multi-consumer queue of Node pointers. - const Queue = struct { - stack: Atomic(usize) = Atomic(usize).init(0), - cache: ?*Node = null, - - const HAS_CACHE: usize = 0b01; - const IS_CONSUMING: usize = 0b10; - const PTR_MASK: usize = ~(HAS_CACHE | IS_CONSUMING); - - comptime { - assert(@alignOf(Node) >= ((IS_CONSUMING | HAS_CACHE) + 1)); - } - - fn push(noalias self: *Queue, list: List) void { - var stack = self.stack.load(.Monotonic); - while (true) { - // Attach the list to the stack (pt. 1) - list.tail.next = @as(?*Node, @ptrFromInt(stack & PTR_MASK)); - - // Update the stack with the list (pt. 2). - // Don't change the HAS_CACHE and IS_CONSUMING bits of the consumer. - var new_stack = @intFromPtr(list.head); - assert(new_stack & ~PTR_MASK == 0); - new_stack |= (stack & ~PTR_MASK); - - // Push to the stack with a release barrier for the consumer to see the proper list links. - stack = self.stack.tryCompareAndSwap( - stack, - new_stack, - .Release, - .Monotonic, - ) orelse break; - } - } - - fn tryAcquireConsumer(self: *Queue) error{ Empty, Contended }!?*Node { - var stack = self.stack.load(.Monotonic); - while (true) { - if (stack & IS_CONSUMING != 0) - return error.Contended; // The queue already has a consumer. - if (stack & (HAS_CACHE | PTR_MASK) == 0) - return error.Empty; // The queue is empty when there's nothing cached and nothing in the stack. - - // When we acquire the consumer, also consume the pushed stack if the cache is empty. - var new_stack = stack | HAS_CACHE | IS_CONSUMING; - if (stack & HAS_CACHE == 0) { - assert(stack & PTR_MASK != 0); - new_stack &= ~PTR_MASK; - } - - // Acquire barrier on getting the consumer to see cache/Node updates done by previous consumers - // and to ensure our cache/Node updates in pop() happen after that of previous consumers. - stack = self.stack.tryCompareAndSwap( - stack, - new_stack, - .Acquire, - .Monotonic, - ) orelse return self.cache orelse @as(*Node, @ptrFromInt(stack & PTR_MASK)); - } - } - - fn releaseConsumer(noalias self: *Queue, noalias consumer: ?*Node) void { - // Stop consuming and remove the HAS_CACHE bit as well if the consumer's cache is empty. - // When HAS_CACHE bit is zeroed, the next consumer will acquire the pushed stack nodes. - var remove = IS_CONSUMING; - if (consumer == null) - remove |= HAS_CACHE; - - // Release the consumer with a release barrier to ensure cache/node accesses - // happen before the consumer was released and before the next consumer starts using the cache. - self.cache = consumer; - const stack = self.stack.fetchSub(remove, .Release); - assert(stack & remove != 0); - } - - fn pop(noalias self: *Queue, noalias consumer_ref: *?*Node) ?*Node { - // Check the consumer cache (fast path) - if (consumer_ref.*) |node| { - consumer_ref.* = node.next; - return node; - } - - // Load the stack to see if there was anything pushed that we could grab. - var stack = self.stack.load(.Monotonic); - assert(stack & IS_CONSUMING != 0); - if (stack & PTR_MASK == 0) { - return null; - } - - // Nodes have been pushed to the stack, grab then with an Acquire barrier to see the Node links. - stack = self.stack.swap(HAS_CACHE | IS_CONSUMING, .Acquire); - assert(stack & IS_CONSUMING != 0); - assert(stack & PTR_MASK != 0); - - const node = @as(*Node, @ptrFromInt(stack & PTR_MASK)); - consumer_ref.* = node.next; - return node; - } - }; - - /// A bounded single-producer, multi-consumer ring buffer for node pointers. - const Buffer = struct { - head: Atomic(Index) = Atomic(Index).init(0), - tail: Atomic(Index) = Atomic(Index).init(0), - array: [capacity]Atomic(*Node) = undefined, - - const Index = u32; - const capacity = 256; // Appears to be a pretty good trade-off in space vs contended throughput - comptime { - assert(std.math.maxInt(Index) >= capacity); - assert(std.math.isPowerOfTwo(capacity)); - } - - fn push(noalias self: *Buffer, noalias list: *List) error{Overflow}!void { - var head = self.head.load(.Monotonic); - var tail = self.tail.loadUnchecked(); // we're the only thread that can change this - - while (true) { - var size = tail -% head; - assert(size <= capacity); - - // Push nodes from the list to the buffer if it's not empty.. - if (size < capacity) { - var nodes: ?*Node = list.head; - while (size < capacity) : (size += 1) { - const node = nodes orelse break; - nodes = node.next; - - // Array written atomically with weakest ordering since it could be getting atomically read by steal(). - self.array[tail % capacity].store(node, .Unordered); - tail +%= 1; - } - - // Release barrier synchronizes with Acquire loads for steal()ers to see the array writes. - self.tail.store(tail, .Release); - - // Update the list with the nodes we pushed to the buffer and try again if there's more. - list.head = nodes orelse return; - std.atomic.spinLoopHint(); - head = self.head.load(.Monotonic); - continue; - } - - // Try to steal/overflow half of the tasks in the buffer to make room for future push()es. - // Migrating half amortizes the cost of stealing while requiring future pops to still use the buffer. - // Acquire barrier to ensure the linked list creation after the steal only happens after we successfully steal. - var migrate = size / 2; - head = self.head.tryCompareAndSwap( - head, - head +% migrate, - .Acquire, - .Monotonic, - ) orelse { - // Link the migrated Nodes together - const first = self.array[head % capacity].loadUnchecked(); - while (migrate > 0) : (migrate -= 1) { - const prev = self.array[head % capacity].loadUnchecked(); - head +%= 1; - prev.next = self.array[head % capacity].loadUnchecked(); - } - - // Append the list that was supposed to be pushed to the end of the migrated Nodes - const last = self.array[(head -% 1) % capacity].loadUnchecked(); - last.next = list.head; - list.tail.next = null; - - // Return the migrated nodes + the original list as overflowed - list.head = first; - return error.Overflow; - }; - } - } - - fn pop(self: *Buffer) ?*Node { - var head = self.head.load(.Monotonic); - var tail = self.tail.loadUnchecked(); // we're the only thread that can change this - - while (true) { - // Quick sanity check and return null when not empty - var size = tail -% head; - assert(size <= capacity); - if (size == 0) { - return null; - } - - // Dequeue with an acquire barrier to ensure any writes done to the Node - // only happens after we successfully claim it from the array. - head = self.head.tryCompareAndSwap( - head, - head +% 1, - .Acquire, - .Monotonic, - ) orelse return self.array[head % capacity].loadUnchecked(); - } - } - - const Stole = struct { - node: *Node, - pushed: bool, - }; - - fn consume(noalias self: *Buffer, noalias queue: *Queue) ?Stole { - var consumer = queue.tryAcquireConsumer() catch return null; - defer queue.releaseConsumer(consumer); - - const head = self.head.load(.Monotonic); - const tail = self.tail.loadUnchecked(); // we're the only thread that can change this - - const size = tail -% head; - assert(size <= capacity); - assert(size == 0); // we should only be consuming if our array is empty - - // Pop nodes from the queue and push them to our array. - // Atomic stores to the array as steal() threads may be atomically reading from it. - var pushed: Index = 0; - while (pushed < capacity) : (pushed += 1) { - const node = queue.pop(&consumer) orelse break; - self.array[(tail +% pushed) % capacity].store(node, .Unordered); - } - - // We will be returning one node that we stole from the queue. - // Get an extra, and if that's not possible, take one from our array. - const node = queue.pop(&consumer) orelse blk: { - if (pushed == 0) return null; - pushed -= 1; - break :blk self.array[(tail +% pushed) % capacity].loadUnchecked(); - }; - - // Update the array tail with the nodes we pushed to it. - // Release barrier to synchronize with Acquire barrier in steal()'s to see the written array Nodes. - if (pushed > 0) self.tail.store(tail +% pushed, .Release); - return Stole{ - .node = node, - .pushed = pushed > 0, - }; - } - - fn steal(noalias self: *Buffer, noalias buffer: *Buffer) ?Stole { - const head = self.head.load(.Monotonic); - const tail = self.tail.loadUnchecked(); // we're the only thread that can change this - - const size = tail -% head; - assert(size <= capacity); - assert(size == 0); // we should only be stealing if our array is empty - - while (true) : (std.atomic.spinLoopHint()) { - const buffer_head = buffer.head.load(.Acquire); - const buffer_tail = buffer.tail.load(.Acquire); - - // Overly large size indicates the the tail was updated a lot after the head was loaded. - // Reload both and try again. - const buffer_size = buffer_tail -% buffer_head; - if (buffer_size > capacity) { - continue; - } - - // Try to steal half (divCeil) to amortize the cost of stealing from other threads. - const steal_size = buffer_size - (buffer_size / 2); - if (steal_size == 0) { - return null; - } - - // Copy the nodes we will steal from the target's array to our own. - // Atomically load from the target buffer array as it may be pushing and atomically storing to it. - // Atomic store to our array as other steal() threads may be atomically loading from it as above. - var i: Index = 0; - while (i < steal_size) : (i += 1) { - const node = buffer.array[(buffer_head +% i) % capacity].load(.Unordered); - self.array[(tail +% i) % capacity].store(node, .Unordered); - } - - // Try to commit the steal from the target buffer using: - // - an Acquire barrier to ensure that we only interact with the stolen Nodes after the steal was committed. - // - a Release barrier to ensure that the Nodes are copied above prior to the committing of the steal - // because if they're copied after the steal, the could be getting rewritten by the target's push(). - _ = buffer.head.compareAndSwap( - buffer_head, - buffer_head +% steal_size, - .AcqRel, - .Monotonic, - ) orelse { - // Pop one from the nodes we stole as we'll be returning it - const pushed = steal_size - 1; - const node = self.array[(tail +% pushed) % capacity].loadUnchecked(); - - // Update the array tail with the nodes we pushed to it. - // Release barrier to synchronize with Acquire barrier in steal()'s to see the written array Nodes. - if (pushed > 0) self.tail.store(tail +% pushed, .Release); - return Stole{ - .node = node, - .pushed = pushed > 0, - }; - }; - } - } - }; - }; -}; - -test "parallel for loop" { - var thread_pool = ThreadPool.init(.{ .max_threads = 12 }); - var sleepy_time: u32 = 100; - var huge_array = &[_]u32{ - sleepy_time + std.rand.DefaultPrng.init(1).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(2).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(3).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(4).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(5).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(6).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(7).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(8).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(9).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(10).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(11).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(12).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(13).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(14).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(15).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(16).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(17).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(18).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(19).random().uintAtMost(u32, 20), - sleepy_time + std.rand.DefaultPrng.init(20).random().uintAtMost(u32, 20), - }; - const Runner = struct { - completed: usize = 0, - total: usize = 0, - pub fn run(ctx: *@This(), value: u32, _: usize) void { - std.time.sleep(value); - ctx.completed += 1; - std.debug.assert(ctx.completed <= ctx.total); - } - }; - var runny = try std.heap.page_allocator.create(Runner); - runny.* = .{ .total = huge_array.len }; - try thread_pool.doAndWait(std.heap.page_allocator, null, runny, Runner.run, std.mem.span(huge_array)); - try std.testing.expectEqual(huge_array.len, runny.completed); -} - -pub fn NewWorkPool(comptime max_threads: ?usize) type { - return struct { - var pool: ThreadPool = undefined; - var loaded: bool = false; - - fn create() *ThreadPool { - @setCold(true); - - pool = ThreadPool.init(.{ - .max_threads = max_threads orelse @max(@as(u32, @truncate(std.Thread.getCpuCount() catch 0)), 2), - .stack_size = 2 * 1024 * 1024, - }); - return &pool; - } - - pub fn deinit() void { - get().deinit(); - } - - pub inline fn get() *ThreadPool { - // lil racy - if (loaded) return &pool; - loaded = true; - - return create(); - } - - pub fn scheduleBatch(batch: ThreadPool.Batch) void { - get().schedule(batch); - } - - pub fn scheduleTask(task: *ThreadPool.Task) void { - get().schedule(ThreadPool.Batch.from(task)); - } - - pub fn go(allocator: std.mem.Allocator, comptime Context: type, context: Context, comptime function: *const fn (Context) void) !void { - const TaskType = struct { - task: ThreadPool.Task, - context: Context, - allocator: std.mem.Allocator, - - pub fn callback(task: *ThreadPool.Task) void { - var this_task = @fieldParentPtr(@This(), "task", task); - function(this_task.context); - this_task.allocator.destroy(this_task); - } - }; - - var task_ = try allocator.create(TaskType); - task_.* = .{ - .task = .{ .callback = TaskType.callback }, - .context = context, - .allocator = allocator, - }; - scheduleTask(&task_.task); - } - }; -} - -pub const WorkPool = NewWorkPool(null); -const testing = std.testing; - -const CrdsTableTrimContext = struct { - index: usize, - max_trim: usize, - self: *CrdsTable, -}; - -const CrdsTable = struct { - pub fn trim(context: CrdsTableTrimContext) void { - const self = context.self; - _ = self; - const max_trim = context.max_trim; - _ = max_trim; - const index = context.index; - _ = index; - - std.debug.print("I ran!\n\n", .{}); - // todo - - } -}; - -test "sync.thread_pool: workpool works" { - var crds: CrdsTable = CrdsTable{}; - var a = CrdsTableTrimContext{ .index = 1, .max_trim = 2, .self = &crds }; - defer WorkPool.deinit(); - try WorkPool.go(testing.allocator, CrdsTableTrimContext, a, CrdsTable.trim); - - std.time.sleep(std.time.ns_per_s * 1); - WorkPool.pool.shutdown(); -} From 0ad8c7b2d5eba3067d1092b41f9311c739fc12ab Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 09:59:45 -0500 Subject: [PATCH 03/15] feat(prometheus): histogram, gauge, counter implement histogram and gauge. make existing counter conform to standard. rename old "gauge" to "gauge_fn" since it is not a true prometheus gauge but an equivalent of the golang GaugeFunc --- src/lib.zig | 2 + src/prometheus/counter.zig | 82 +++---- src/prometheus/gauge.zig | 210 +++------------- src/prometheus/gauge_fn.zig | 204 ++++++++++++++++ src/prometheus/histogram.zig | 452 ++++++++++++++++++----------------- src/prometheus/http.zig | 60 +++++ src/prometheus/metric.zig | 62 +++-- src/prometheus/registry.zig | 179 ++++++-------- 8 files changed, 662 insertions(+), 589 deletions(-) create mode 100644 src/prometheus/gauge_fn.zig create mode 100644 src/prometheus/http.zig diff --git a/src/lib.zig b/src/lib.zig index d8089a943..50af87cf0 100644 --- a/src/lib.zig +++ b/src/lib.zig @@ -79,6 +79,8 @@ pub const net = struct { pub const prometheus = struct { pub usingnamespace @import("prometheus/counter.zig"); pub usingnamespace @import("prometheus/gauge.zig"); + pub usingnamespace @import("prometheus/gauge_fn.zig"); + pub usingnamespace @import("prometheus/http.zig"); pub usingnamespace @import("prometheus/histogram.zig"); pub usingnamespace @import("prometheus/metric.zig"); pub usingnamespace @import("prometheus/registry.zig"); diff --git a/src/prometheus/counter.zig b/src/prometheus/counter.zig index 0e7e2f2b7..55cd239f5 100644 --- a/src/prometheus/counter.zig +++ b/src/prometheus/counter.zig @@ -1,63 +1,53 @@ const std = @import("std"); const mem = std.mem; const testing = std.testing; -const Metric = @import("metric.zig").Metric; -pub const Counter = struct { - metric: Metric = Metric{ - .getResultFn = getResult, - }, - value: std.atomic.Atomic(u64) = .{ .value = 0 }, +const Metric = @import("metric.zig").Metric; - const Self = @This(); +const Self = @This(); - pub fn init(allocator: mem.Allocator) !*Self { - const self = try allocator.create(Self); +metric: Metric = Metric{ .getResultFn = getResult }, +value: std.atomic.Atomic(u64) = std.atomic.Atomic(u64).init(0), - self.* = .{}; +pub fn init(allocator: mem.Allocator) !*Self { + const self = try allocator.create(Self); - return self; - } - - pub fn inc(self: *Self) void { - _ = self.value.fetchAdd(1, .SeqCst); - } + self.* = .{}; - pub fn dec(self: *Self) void { - _ = self.value.fetchSub(1, .SeqCst); - } + return self; +} - pub fn add(self: *Self, value: anytype) void { - if (!comptime std.meta.trait.isNumber(@TypeOf(value))) { - @compileError("can't add a non-number"); - } +pub fn inc(self: *Self) void { + _ = self.value.fetchAdd(1, .SeqCst); +} - _ = self.value.fetchAdd(@intCast(value), .SeqCst); +pub fn add(self: *Self, value: anytype) void { + switch (@typeInfo(@TypeOf(value))) { + .Int, .Float, .ComptimeInt, .ComptimeFloat => {}, + else => @compileError("can't add a non-number"), } - pub fn get(self: *const Self) u64 { - return self.value.load(.SeqCst); - } + _ = self.value.fetchAdd(@intCast(value), .SeqCst); +} - pub fn set(self: *Self, value: anytype) void { - if (!comptime std.meta.trait.isNumber(@TypeOf(value))) { - @compileError("can't set a non-number"); - } +pub fn get(self: *const Self) u64 { + return self.value.load(.SeqCst); +} - _ = self.value.store(@intCast(value), .SeqCst); - } +pub fn reset(self: *Self) void { + _ = self.value.store(0, .SeqCst); +} - fn getResult(metric: *Metric, _: mem.Allocator) Metric.Error!Metric.Result { - const self = @fieldParentPtr(Self, "metric", metric); - return Metric.Result{ .counter = self.get() }; - } -}; +fn getResult(metric: *Metric, _: mem.Allocator) Metric.Error!Metric.Result { + const self = @fieldParentPtr(Self, "metric", metric); + return Metric.Result{ .counter = self.get() }; +} test "prometheus.counter: inc/add/dec/set/get" { var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); - var counter = try Counter.init(testing.allocator); + var counter = try Self.init(testing.allocator); defer testing.allocator.destroy(counter); try testing.expectEqual(@as(u64, 0), counter.get()); @@ -67,16 +57,10 @@ test "prometheus.counter: inc/add/dec/set/get" { counter.add(200); try testing.expectEqual(@as(u64, 201), counter.get()); - - counter.dec(); - try testing.expectEqual(@as(u64, 200), counter.get()); - - counter.set(43); - try testing.expectEqual(@as(u64, 43), counter.get()); } test "prometheus.counter: concurrent" { - var counter = try Counter.init(testing.allocator); + var counter = try Self.init(testing.allocator); defer testing.allocator.destroy(counter); var threads: [4]std.Thread = undefined; @@ -84,7 +68,7 @@ test "prometheus.counter: concurrent" { thread.* = try std.Thread.spawn( .{}, struct { - fn run(c: *Counter) void { + fn run(c: *Self) void { var i: usize = 0; while (i < 20) : (i += 1) { c.inc(); @@ -101,9 +85,9 @@ test "prometheus.counter: concurrent" { } test "prometheus.counter: write" { - var counter = try Counter.init(testing.allocator); + var counter = try Self.init(testing.allocator); defer testing.allocator.destroy(counter); - counter.set(340); + counter.* = .{ .value = .{ .value = 340 } }; var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); diff --git a/src/prometheus/gauge.zig b/src/prometheus/gauge.zig index cd1aef41f..992226e6e 100644 --- a/src/prometheus/gauge.zig +++ b/src/prometheus/gauge.zig @@ -1,66 +1,48 @@ const std = @import("std"); -const mem = std.mem; -const testing = std.testing; const Metric = @import("metric.zig").Metric; -pub fn GaugeCallFnType(comptime StateType: type, comptime Return: type) type { - const CallFnArgType = switch (@typeInfo(StateType)) { - .Pointer => StateType, - .Optional => |opt| opt.child, - .Void => void, - else => *StateType, - }; - - return *const fn (state: CallFnArgType) Return; -} +/// A gauge that stores the value it reports. +/// Read and write operations are atomic and unordered. +pub fn Gauge(comptime T: type) type { + return struct { + value: std.atomic.Atomic(T) = .{ .value = 0 }, + metric: Metric = .{ .getResultFn = getResult }, -pub fn Gauge(comptime StateType: type, comptime Return: type) type { - const CallFnType = GaugeCallFnType(StateType, Return); + pub fn init(_: anytype) @This() { + return .{}; + } - return struct { - const Self = @This(); + pub fn inc(self: *@This()) void { + self.value.fetchAdd(1, .Unordered); + } - metric: Metric = .{ - .getResultFn = getResult, - }, - callFn: CallFnType = undefined, - state: StateType = undefined, + pub fn add(self: *@This(), v: T) void { + self.value.fetchAdd(v, .Unordered); + } - pub fn init(allocator: mem.Allocator, callFn: CallFnType, state: StateType) !*Self { - const self = try allocator.create(Self); + pub fn dec(self: *@This()) void { + self.value.fetchSub(1, .Unordered); + } - self.* = .{}; - self.callFn = callFn; - self.state = state; + pub fn sub(self: *@This(), v: T) void { + self.value.fetchAdd(v, .Unordered); + } - return self; + pub fn set(self: *@This(), v: T) void { + self.value.store(v, .Unordered); } - pub fn get(self: *Self) Return { - const TypeInfo = @typeInfo(StateType); - switch (TypeInfo) { - .Pointer, .Void => { - return self.callFn(self.state); - }, - .Optional => { - if (self.state) |state| { - return self.callFn(state); - } - return 0; - }, - else => { - return self.callFn(&self.state); - }, - } + pub fn get(self: *@This()) T { + return self.value.load(.Unordered); } - fn getResult(metric: *Metric, allocator: mem.Allocator) Metric.Error!Metric.Result { + fn getResult(metric: *Metric, allocator: std.mem.Allocator) Metric.Error!Metric.Result { _ = allocator; - const self = @fieldParentPtr(Self, "metric", metric); + const self = @fieldParentPtr(@This(), "metric", metric); - return switch (Return) { + return switch (T) { f64 => Metric.Result{ .gauge = self.get() }, u64 => Metric.Result{ .gauge_int = self.get() }, else => unreachable, // Gauge Return may only be 'f64' or 'u64' @@ -68,139 +50,3 @@ pub fn Gauge(comptime StateType: type, comptime Return: type) type { } }; } - -test "gauge: get" { - const TestCase = struct { - state_type: type, - typ: type, - }; - - const testCases = [_]TestCase{ - .{ - .state_type = struct { - value: f64, - }, - .typ = f64, - }, - }; - - inline for (testCases) |tc| { - const State = tc.state_type; - const InnerType = tc.typ; - - var state = State{ .value = 20 }; - - var gauge = try Gauge(*State, InnerType).init( - testing.allocator, - struct { - fn get(s: *State) InnerType { - return s.value + 1; - } - }.get, - &state, - ); - defer testing.allocator.destroy(gauge); - - try testing.expectEqual(@as(InnerType, 21), gauge.get()); - } -} - -test "gauge: optional state" { - const State = struct { - value: f64, - }; - var state = State{ .value = 20.0 }; - - var gauge = try Gauge(?*State, f64).init( - testing.allocator, - struct { - fn get(s: *State) f64 { - return s.value + 1.0; - } - }.get, - &state, - ); - defer testing.allocator.destroy(gauge); - - try testing.expectEqual(@as(f64, 21.0), gauge.get()); -} - -test "gauge: non-pointer state" { - var gauge = try Gauge(f64, f64).init( - testing.allocator, - struct { - fn get(s: *f64) f64 { - s.* += 1.0; - return s.*; - } - }.get, - 0.0, - ); - defer testing.allocator.destroy(gauge); - - try testing.expectEqual(@as(f64, 1.0), gauge.get()); -} - -test "gauge: shared state" { - const State = struct { - mutex: std.Thread.Mutex = .{}, - items: std.ArrayList(usize) = std.ArrayList(usize).init(testing.allocator), - }; - var shared_state = State{}; - defer shared_state.items.deinit(); - - var gauge = try Gauge(*State, f64).init( - testing.allocator, - struct { - fn get(state: *State) f64 { - return @floatFromInt(state.items.items.len); - } - }.get, - &shared_state, - ); - defer testing.allocator.destroy(gauge); - - var threads: [4]std.Thread = undefined; - for (&threads, 0..) |*thread, thread_index| { - thread.* = try std.Thread.spawn( - .{}, - struct { - fn run(thread_idx: usize, state: *State) !void { - var i: usize = 0; - while (i < 4) : (i += 1) { - state.mutex.lock(); - defer state.mutex.unlock(); - try state.items.append(thread_idx + i); - } - } - }.run, - .{ thread_index, &shared_state }, - ); - } - - for (&threads) |*thread| thread.join(); - - try testing.expectEqual(@as(usize, 16), @as(usize, @intFromFloat(gauge.get()))); -} - -test "gauge: write" { - var gauge = try Gauge(usize, f64).init( - testing.allocator, - struct { - fn get(state: *usize) f64 { - state.* += 340; - return @floatFromInt(state.*); - } - }.get, - @as(usize, 0), - ); - defer testing.allocator.destroy(gauge); - - var buffer = std.ArrayList(u8).init(testing.allocator); - defer buffer.deinit(); - - var metric = &gauge.metric; - try metric.write(testing.allocator, buffer.writer(), "mygauge"); - - try testing.expectEqualStrings("mygauge 340.000000\n", buffer.items); -} diff --git a/src/prometheus/gauge_fn.zig b/src/prometheus/gauge_fn.zig new file mode 100644 index 000000000..6096fcf58 --- /dev/null +++ b/src/prometheus/gauge_fn.zig @@ -0,0 +1,204 @@ +const std = @import("std"); +const mem = std.mem; +const testing = std.testing; + +const Metric = @import("metric.zig").Metric; + +pub fn GaugeCallFnType(comptime StateType: type, comptime Return: type) type { + const CallFnArgType = switch (@typeInfo(StateType)) { + .Pointer => StateType, + .Optional => |opt| opt.child, + .Void => void, + else => *StateType, + }; + + return *const fn (state: CallFnArgType) Return; +} + +pub fn GaugeFn(comptime StateType: type, comptime Return: type) type { + const CallFnType = GaugeCallFnType(StateType, Return); + + return struct { + const Self = @This(); + + metric: Metric = .{ .getResultFn = getResult }, + callFn: CallFnType = undefined, + state: StateType = undefined, + + pub fn init(allocator: mem.Allocator, callFn: CallFnType, state: StateType) !*Self { + const self = try allocator.create(Self); + + self.* = .{}; + self.callFn = callFn; + self.state = state; + + return self; + } + + pub fn get(self: *Self) Return { + const TypeInfo = @typeInfo(StateType); + switch (TypeInfo) { + .Pointer, .Void => { + return self.callFn(self.state); + }, + .Optional => { + if (self.state) |state| { + return self.callFn(state); + } + return 0; + }, + else => { + return self.callFn(&self.state); + }, + } + } + + fn getResult(metric: *Metric, allocator: mem.Allocator) Metric.Error!Metric.Result { + _ = allocator; + + const self = @fieldParentPtr(Self, "metric", metric); + + return switch (Return) { + f64 => Metric.Result{ .gauge = self.get() }, + u64 => Metric.Result{ .gauge_int = self.get() }, + else => unreachable, // Gauge Return may only be 'f64' or 'u64' + }; + } + }; +} + +test "prometheus.gauge_fn: get" { + const TestCase = struct { + state_type: type, + typ: type, + }; + + const testCases = [_]TestCase{ + .{ + .state_type = struct { + value: f64, + }, + .typ = f64, + }, + }; + + inline for (testCases) |tc| { + const State = tc.state_type; + const InnerType = tc.typ; + + var state = State{ .value = 20 }; + + var gauge = try GaugeFn(*State, InnerType).init( + testing.allocator, + struct { + fn get(s: *State) InnerType { + return s.value + 1; + } + }.get, + &state, + ); + defer testing.allocator.destroy(gauge); + + try testing.expectEqual(@as(InnerType, 21), gauge.get()); + } +} + +test "prometheus.gauge_fn: optional state" { + const State = struct { + value: f64, + }; + var state = State{ .value = 20.0 }; + + var gauge = try GaugeFn(?*State, f64).init( + testing.allocator, + struct { + fn get(s: *State) f64 { + return s.value + 1.0; + } + }.get, + &state, + ); + defer testing.allocator.destroy(gauge); + + try testing.expectEqual(@as(f64, 21.0), gauge.get()); +} + +test "prometheus.gauge_fn: non-pointer state" { + var gauge = try GaugeFn(f64, f64).init( + testing.allocator, + struct { + fn get(s: *f64) f64 { + s.* += 1.0; + return s.*; + } + }.get, + 0.0, + ); + defer testing.allocator.destroy(gauge); + + try testing.expectEqual(@as(f64, 1.0), gauge.get()); +} + +test "prometheus.gauge_fn: shared state" { + const State = struct { + mutex: std.Thread.Mutex = .{}, + items: std.ArrayList(usize) = std.ArrayList(usize).init(testing.allocator), + }; + var shared_state = State{}; + defer shared_state.items.deinit(); + + var gauge = try GaugeFn(*State, f64).init( + testing.allocator, + struct { + fn get(state: *State) f64 { + return @floatFromInt(state.items.items.len); + } + }.get, + &shared_state, + ); + defer testing.allocator.destroy(gauge); + + var threads: [4]std.Thread = undefined; + for (&threads, 0..) |*thread, thread_index| { + thread.* = try std.Thread.spawn( + .{}, + struct { + fn run(thread_idx: usize, state: *State) !void { + var i: usize = 0; + while (i < 4) : (i += 1) { + state.mutex.lock(); + defer state.mutex.unlock(); + try state.items.append(thread_idx + i); + } + } + }.run, + .{ thread_index, &shared_state }, + ); + } + + for (&threads) |*thread| thread.join(); + + try testing.expectEqual(@as(usize, 16), @as(usize, @intFromFloat(gauge.get()))); +} + +test "prometheus.gauge_fn: write" { + var gauge = try GaugeFn(usize, f64).init( + testing.allocator, + struct { + fn get(state: *usize) f64 { + state.* += 340; + return @floatFromInt(state.*); + } + }.get, + @as(usize, 0), + ); + defer testing.allocator.destroy(gauge); + + var buffer = std.ArrayList(u8).init(testing.allocator); + defer buffer.deinit(); + + var metric = &gauge.metric; + try metric.write(testing.allocator, buffer.writer(), "mygauge"); + + try testing.expectEqualStrings("mygauge 340.000000\n", buffer.items); +} diff --git a/src/prometheus/histogram.zig b/src/prometheus/histogram.zig index 5a93fc595..5b0a010b3 100644 --- a/src/prometheus/histogram.zig +++ b/src/prometheus/histogram.zig @@ -1,271 +1,289 @@ const std = @import("std"); -const fmt = std.fmt; -const math = std.math; -const mem = std.mem; -const testing = std.testing; +const Allocator = std.mem.Allocator; +const ArrayList = std.ArrayList; +const Atomic = std.atomic.Atomic; +const Ordering = std.atomic.Ordering; const Metric = @import("metric.zig").Metric; -const HistogramResult = @import("metric.zig").HistogramResult; -const e10_min = -9; -const e10_max = 18; -const buckets_per_decimal = 18; -const decimal_buckets_count = e10_max - e10_min; -const buckets_count = decimal_buckets_count * buckets_per_decimal; +const default_buckets: [11]f64 = .{ 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0 }; -const lower_bucket_range = blk: { - var buf: [64]u8 = undefined; - break :blk fmt.bufPrint(&buf, "0...{e:.3}", .{math.pow(f64, 10, e10_min)}) catch unreachable; -}; -const upper_bucket_range = blk: { - var buf: [64]u8 = undefined; - break :blk fmt.bufPrint(&buf, "{e:.3}...+Inf", .{math.pow(f64, 10, e10_max)}) catch unreachable; -}; - -const bucket_ranges: [buckets_count][]const u8 = blk: { - const bucket_multiplier = math.pow(f64, 10.0, 1.0 / @as(f64, buckets_per_decimal)); - - var v = math.pow(f64, 10, e10_min); - - var start = blk2: { - var buf: [64]u8 = undefined; - break :blk2 fmt.bufPrint(&buf, "{e:.3}", .{v}) catch unreachable; - }; - - var result: [buckets_count][]const u8 = undefined; - for (&result) |*range| { - v *= bucket_multiplier; - - const end = blk3: { - var buf: [64]u8 = undefined; - break :blk3 fmt.bufPrint(&buf, "{e:.3}", .{v}) catch unreachable; - }; - - range.* = start ++ "..." ++ end; - - start = end; +pub fn defaultBuckets(allocator: Allocator) !ArrayList(f64) { + var l = try ArrayList(f64).initCapacity(allocator, default_buckets.len); + for (default_buckets) |b| { + try l.append(b); } - - break :blk result; -}; - -test "bucket ranges" { - try testing.expectEqualStrings("0...1.000e-09", lower_bucket_range); - try testing.expectEqualStrings("1.000e+18...+Inf", upper_bucket_range); - - try testing.expectEqualStrings("1.000e-09...1.136e-09", bucket_ranges[0]); - try testing.expectEqualStrings("1.136e-09...1.292e-09", bucket_ranges[1]); - try testing.expectEqualStrings("8.799e-09...1.000e-08", bucket_ranges[buckets_per_decimal - 1]); - try testing.expectEqualStrings("1.000e-08...1.136e-08", bucket_ranges[buckets_per_decimal]); - try testing.expectEqualStrings("8.799e-01...1.000e+00", bucket_ranges[buckets_per_decimal * (-e10_min) - 1]); - try testing.expectEqualStrings("1.000e+00...1.136e+00", bucket_ranges[buckets_per_decimal * (-e10_min)]); - try testing.expectEqualStrings("8.799e+17...1.000e+18", bucket_ranges[buckets_per_decimal * (e10_max - e10_min) - 1]); + return l; } -/// Histogram based on https://github.com/VictoriaMetrics/metrics/blob/master/histogram.go. +/// Histogram optimized for fast writes. +/// Reads and writes are thread-safe if you use the public methods. +/// Writes are lock-free. Reads are locked with a mutex because they occupy a shard. +/// +/// The histogram state is represented in a shard. There are two shards, hot and cold. +/// Writes incremenent the hot shard. +/// Reads flip a switch to change which shard is considered hot for writes, +/// then wait for the previous hot shard to cool down before reading it. pub const Histogram = struct { - const Self = @This(); - - metric: Metric = .{ - .getResultFn = getResult, + allocator: Allocator, + + /// The highest value to include in each bucket. + upper_bounds: ArrayList(f64), + + /// One hot shard for writing, one cold shard for reading. + shards: [2]struct { + /// Total of all observed values. + sum: Atomic(f64) = Atomic(f64).init(0.0), + /// Total number of observations that have finished being recorded to this shard. + count: Atomic(u64) = Atomic(u64).init(0), + /// Cumulative counts for each upper bound. + buckets: ArrayList(Atomic(u64)), }, - mutex: std.Thread.Mutex = .{}, - decimal_buckets: [decimal_buckets_count][buckets_per_decimal]u64 = undefined, + /// Used to ensure reads and writes occur on separate shards. + /// Atomic representation of `ShardSync`. + shard_sync: Atomic(u64), - lower: u64 = 0, - upper: u64 = 0, + /// Prevents more than one reader at a time, since read operations actually + /// execute an internal write by swapping the hot and cold shards. + read_mutex: std.Thread.Mutex = .{}, - sum: f64 = 0.0, + /// Used by registry to report the histogram + metric: Metric = .{ .getResultFn = getResult }, - pub fn init(allocator: mem.Allocator) !*Self { - const self = try allocator.create(Self); + const ShardSync = packed struct { + /// The total count of events that have started to be recorded (including those that finished). + /// If this is larger than the shard count, it means a write is in progress. + count: u63 = 0, + /// Index of the shard currently being used for writes. + shard: u1 = 0, + }; - self.* = .{}; - for (&self.decimal_buckets) |*bucket| { - @memset(bucket, 0); - } + const Self = @This(); + + pub fn init(allocator: Allocator, buckets: ArrayList(f64)) !@This() { + return .{ + .allocator = allocator, + .upper_bounds = buckets, + .shards = .{ + .{ .buckets = try shardBuckets(allocator, buckets.items.len) }, + .{ .buckets = try shardBuckets(allocator, buckets.items.len) }, + }, + .shard_sync = Atomic(u64).init(0), + }; + } - return self; + pub fn deinit(self: *Self) void { + self.shards[0].buckets.deinit(); + self.shards[1].buckets.deinit(); + self.upper_bounds.deinit(); } - pub fn update(self: *Self, value: f64) void { - if (math.isNan(value) or value < 0) { - return; + fn shardBuckets(allocator: Allocator, size: usize) !ArrayList(Atomic(u64)) { + var shard_buckets = try ArrayList(Atomic(u64)).initCapacity(allocator, size); + for (0..size) |_| { + shard_buckets.appendAssumeCapacity(Atomic(u64).init(0)); } + return shard_buckets; + } - const bucket_idx: f64 = (math.log10(value) - e10_min) * buckets_per_decimal; + /// Writes a value into the histogram. + pub fn observe(self: *Self, item: f64) void { + const shard_sync = self.incrementCount(.Acquire); // acquires lock. must be first step. + const shard = &self.shards[shard_sync.shard]; + for (0.., self.upper_bounds.items) |i, bound| { + if (item <= bound) { + _ = shard.buckets.items[i].fetchAdd(1, .Monotonic); + break; + } + } + _ = shard.count.fetchAdd(1, .Release); // releases lock. must be last step. + } - // Keep a lock while updating the histogram. - self.mutex.lock(); - defer self.mutex.unlock(); + /// Reads the current state of the histogram. + pub fn getSnapshot(self: *Self, allocator: ?Allocator) !HistogramSnapshot { + var alloc = self.allocator; + if (allocator) |a| alloc = a; + + // Acquire the lock so no one else executes this function at the same time. + self.read_mutex.lock(); + defer self.read_mutex.unlock(); + + // Make the hot shard cold. Some writers may still be writing to it, + // but no more will start after this. + const shard_sync = self.flipShard(.Monotonic); + const cold_shard = &self.shards[shard_sync.shard]; + const hot_shard = &self.shards[shard_sync.shard +% 1]; + + // Wait until all writers are done writing to the cold shard + // TODO: switch to a condvar. see: `std.Thread.Condition` + while (cold_shard.count.tryCompareAndSwap(shard_sync.count, 0, .Acquire, .Monotonic)) |_| { + // Acquire on success: keeps shard usage after. + } - self.sum += value; + // Now the cold shard is totally cold and unused by other threads. + // - read the cold shard's data + // - zero out the cold shard. + // - write the cold shard's data into the hot shard. + const cold_shard_sum = cold_shard.sum.swap(0.0, .Monotonic); + var buckets = try ArrayList(Bucket).initCapacity(alloc, self.upper_bounds.items.len); + var cumulative_count: u64 = 0; + for (0.., self.upper_bounds.items) |i, upper_bound| { + const count = cold_shard.buckets.items[i].swap(0, .Monotonic); + cumulative_count += count; + buckets.appendAssumeCapacity(.{ + .cumulative_count = cumulative_count, + .upper_bound = upper_bound, + }); + _ = hot_shard.buckets.items[i].fetchAdd(count, .Monotonic); + } + _ = hot_shard.sum.fetchAdd(cold_shard_sum, .Monotonic); + _ = hot_shard.count.fetchAdd(shard_sync.count, .Monotonic); - if (bucket_idx < 0) { - self.lower += 1; - } else if (bucket_idx >= buckets_count) { - self.upper += 1; - } else { - const idx: usize = blk: { - const tmp: usize = @intFromFloat(bucket_idx); + return HistogramSnapshot.init(cold_shard_sum, shard_sync.count, buckets); + } - if (bucket_idx == @as(f64, @floatFromInt(tmp)) and tmp > 0) { - // Edge case for 10^n values, which must go to the lower bucket - // according to Prometheus logic for `le`-based histograms. - break :blk tmp - 1; - } else { - break :blk tmp; - } - }; + fn getResult(metric: *Metric, allocator: Allocator) Metric.Error!Metric.Result { + const self = @fieldParentPtr(Self, "metric", metric); + const snapshot = try self.getSnapshot(allocator); + return Metric.Result{ .histogram = snapshot }; + } - const decimal_bucket_idx = idx / buckets_per_decimal; - const offset = idx % buckets_per_decimal; + /// Increases the global count (used for synchronization), not a count within a shard. + /// Returns the state from before this operation, which was replaced by this operation. + fn incrementCount(self: *@This(), comptime ordering: Ordering) ShardSync { + return @bitCast(self.shard_sync.fetchAdd(1, ordering)); + } - var bucket: []u64 = &self.decimal_buckets[decimal_bucket_idx]; - bucket[offset] += 1; - } + /// Makes the hot shard cold and vice versa. + /// Returns the state from before this operation, which was replaced by this operation. + fn flipShard(self: *@This(), comptime ordering: Ordering) ShardSync { + const data = self.shard_sync.fetchAdd(@bitCast(ShardSync{ .shard = 1 }), ordering); + return @bitCast(data); } +}; - pub fn get(self: *const Self) u64 { - _ = self; - return 0; +/// A snapshot of the histogram state from a point in time. +pub const HistogramSnapshot = struct { + /// Sum of all values observed by the histogram. + sum: f64, + /// Total number of events observed by the histogram. + count: u64, + /// Cumulative histogram counts. + /// + /// The len *must* be the same as the amount of memory that was + /// allocated for this slice, or else the memory will leak. + buckets: []Bucket, + /// Allocator that was used to allocate the buckets. + allocator: Allocator, + + pub fn init(sum: f64, count: u64, buckets: ArrayList(Bucket)) @This() { + std.debug.assert(buckets.capacity == buckets.items.len); + return .{ + .sum = sum, + .count = count, + .buckets = buckets.items, + .allocator = buckets.allocator, + }; } - fn isBucketAllZero(bucket: []const u64) bool { - for (bucket) |v| { - if (v != 0) return false; - } - return true; + pub fn deinit(self: *@This()) void { + self.allocator.free(self.buckets); } +}; - fn getResult(metric: *Metric, allocator: mem.Allocator) Metric.Error!Metric.Result { - const self = @fieldParentPtr(Histogram, "metric", metric); +pub const Bucket = struct { + cumulative_count: u64 = 0, + upper_bound: f64 = 0, +}; - // Arbitrary maximum capacity - var buckets = try std.ArrayList(HistogramResult.Bucket).initCapacity(allocator, 16); - var count_total: u64 = 0; +test "prometheus.histogram: empty" { + const allocator = std.testing.allocator; + var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + defer hist.deinit(); - // Keep a lock while querying the histogram. - self.mutex.lock(); - defer self.mutex.unlock(); + var snapshot = try hist.getSnapshot(null); + defer snapshot.deinit(); - if (self.lower > 0) { - try buckets.append(.{ - .vmrange = lower_bucket_range, - .count = self.lower, - }); - count_total += self.lower; - } + try expectSnapshot(0, &default_buckets, &(.{0} ** 11), snapshot); +} - for (&self.decimal_buckets, 0..) |bucket, decimal_bucket_idx| { - if (isBucketAllZero(&bucket)) continue; +test "prometheus.histogram: data goes in correct buckets" { + const allocator = std.testing.allocator; + var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + defer hist.deinit(); - for (bucket, 0..) |count, offset| { - if (count <= 0) continue; + const expected_buckets = observeVarious(&hist); - const bucket_idx = (decimal_bucket_idx * buckets_per_decimal) + offset; - const vmrange = bucket_ranges[bucket_idx]; + var snapshot = try hist.getSnapshot(null); + defer snapshot.deinit(); - try buckets.append(.{ - .vmrange = vmrange, - .count = count, - }); - count_total += count; - } - } + try expectSnapshot(7, &default_buckets, &expected_buckets, snapshot); +} - if (self.upper > 0) { - try buckets.append(.{ - .vmrange = upper_bucket_range, - .count = self.upper, - }); - count_total += self.upper; - } +test "prometheus.histogram: repeated snapshots measure the same thing" { + const allocator = std.testing.allocator; + var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + defer hist.deinit(); - return Metric.Result{ - .histogram = .{ - .buckets = try buckets.toOwnedSlice(), - .sum = .{ .value = self.sum }, - .count = count_total, - }, - }; - } -}; + const expected_buckets = observeVarious(&hist); -test "write empty" { - var histogram = try Histogram.init(testing.allocator); - defer testing.allocator.destroy(histogram); + var snapshot1 = try hist.getSnapshot(null); + snapshot1.deinit(); + var snapshot = try hist.getSnapshot(null); + defer snapshot.deinit(); - var buffer = std.ArrayList(u8).init(testing.allocator); - defer buffer.deinit(); + try expectSnapshot(7, &default_buckets, &expected_buckets, snapshot); +} - var metric = &histogram.metric; - try metric.write(testing.allocator, buffer.writer(), "myhistogram"); +test "prometheus.histogram: values accumulate across snapshots" { + const allocator = std.testing.allocator; + var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + defer hist.deinit(); - try testing.expectEqual(@as(usize, 0), buffer.items.len); -} + _ = observeVarious(&hist); -test "update then write" { - var histogram = try Histogram.init(testing.allocator); - defer testing.allocator.destroy(histogram); + var snapshot1 = try hist.getSnapshot(null); + snapshot1.deinit(); - var i: usize = 98; - while (i < 218) : (i += 1) { - histogram.update(@floatFromInt(i)); - } + hist.observe(1.0); + + var snapshot = try hist.getSnapshot(null); + defer snapshot.deinit(); - var buffer = std.ArrayList(u8).init(testing.allocator); - defer buffer.deinit(); - - var metric = &histogram.metric; - try metric.write(testing.allocator, buffer.writer(), "myhistogram"); - - const exp = - \\myhistogram_bucket{vmrange="8.799e+01...1.000e+02"} 3 - \\myhistogram_bucket{vmrange="1.000e+02...1.136e+02"} 13 - \\myhistogram_bucket{vmrange="1.136e+02...1.292e+02"} 16 - \\myhistogram_bucket{vmrange="1.292e+02...1.468e+02"} 17 - \\myhistogram_bucket{vmrange="1.468e+02...1.668e+02"} 20 - \\myhistogram_bucket{vmrange="1.668e+02...1.896e+02"} 23 - \\myhistogram_bucket{vmrange="1.896e+02...2.154e+02"} 26 - \\myhistogram_bucket{vmrange="2.154e+02...2.448e+02"} 2 - \\myhistogram_sum 18900 - \\myhistogram_count 120 - \\ - ; - - try testing.expectEqualStrings(exp, buffer.items); + const expected_buckets: [11]u64 = .{ 1, 1, 1, 1, 4, 4, 4, 6, 7, 7, 7 }; + try expectSnapshot(8, &default_buckets, &expected_buckets, snapshot); } -test "update then write with labels" { - var histogram = try Histogram.init(testing.allocator); - defer testing.allocator.destroy(histogram); +fn observeVarious(hist: *Histogram) [11]u64 { + hist.observe(1.0); + hist.observe(0.1); + hist.observe(2.0); + hist.observe(0.1); + hist.observe(0.0000000001); + hist.observe(0.1); + hist.observe(100.0); + return .{ 1, 1, 1, 1, 4, 4, 4, 5, 6, 6, 6 }; +} - var i: usize = 98; - while (i < 218) : (i += 1) { - histogram.update(@floatFromInt(i)); +fn expectSnapshot( + expected_total: u64, + expected_bounds: []const f64, + expected_buckets: []const u64, + snapshot: anytype, +) !void { + try std.testing.expectEqual(expected_total, snapshot.count); + try std.testing.expectEqual(default_buckets.len, snapshot.buckets.len); + for (0.., snapshot.buckets) |i, bucket| { + try expectEqual(expected_buckets[i], bucket.cumulative_count, "value in bucket {}\n", .{i}); + try expectEqual(expected_bounds[i], bucket.upper_bound, "bound for bucket {}\n", .{i}); } +} - var buffer = std.ArrayList(u8).init(testing.allocator); - defer buffer.deinit(); - - var metric = &histogram.metric; - try metric.write(testing.allocator, buffer.writer(), "myhistogram{route=\"/api/v2/users\"}"); - - const exp = - \\myhistogram_bucket{route="/api/v2/users",vmrange="8.799e+01...1.000e+02"} 3 - \\myhistogram_bucket{route="/api/v2/users",vmrange="1.000e+02...1.136e+02"} 13 - \\myhistogram_bucket{route="/api/v2/users",vmrange="1.136e+02...1.292e+02"} 16 - \\myhistogram_bucket{route="/api/v2/users",vmrange="1.292e+02...1.468e+02"} 17 - \\myhistogram_bucket{route="/api/v2/users",vmrange="1.468e+02...1.668e+02"} 20 - \\myhistogram_bucket{route="/api/v2/users",vmrange="1.668e+02...1.896e+02"} 23 - \\myhistogram_bucket{route="/api/v2/users",vmrange="1.896e+02...2.154e+02"} 26 - \\myhistogram_bucket{route="/api/v2/users",vmrange="2.154e+02...2.448e+02"} 2 - \\myhistogram_sum{route="/api/v2/users"} 18900 - \\myhistogram_count{route="/api/v2/users"} 120 - \\ - ; - - try testing.expectEqualStrings(exp, buffer.items); +fn expectEqual(expected: anytype, actual: anytype, comptime fmt: anytype, args: anytype) !void { + std.testing.expectEqual(expected, actual) catch |e| { + std.debug.print(fmt, args); + return e; + }; + return; } diff --git a/src/prometheus/http.zig b/src/prometheus/http.zig new file mode 100644 index 000000000..bf532a001 --- /dev/null +++ b/src/prometheus/http.zig @@ -0,0 +1,60 @@ +//! Basic HTTP adapter for prometheus using the http server in std. + +const std = @import("std"); +const Registry = @import("registry.zig").Registry; + +// TODO: log with correct logger + +pub fn servePrometheus( + allocator: std.mem.Allocator, + registry: *Registry(.{}), + listen_addr: std.net.Address, +) !void { + var server = std.http.Server.init(allocator, .{}); + defer server.deinit(); + try server.listen(listen_addr); + + outer: while (true) { + var response = try server.accept(.{ + .allocator = allocator, + }); + defer response.deinit(); + + while (response.reset() != .closing) { + response.wait() catch |err| switch (err) { + error.HttpHeadersInvalid => continue :outer, + error.EndOfStream => continue, + else => return err, + }; + handleRequest(allocator, &response, registry) catch |e| { + std.debug.print("Failed while handling prometheus http request. {}", .{e}); + }; + } + } +} + +fn handleRequest( + allocator: std.mem.Allocator, + response: *std.http.Server.Response, + registry: *Registry(.{}), +) !void { + std.debug.print("{s} {s} {s}\n", .{ + @tagName(response.request.method), + @tagName(response.request.version), + response.request.target, + }); + + if (response.request.method == .GET and + std.mem.startsWith(u8, response.request.target, "/metrics")) + { + response.transfer_encoding = .chunked; + try response.headers.append("content-type", "text/plain"); + try response.do(); + try registry.write(allocator, response.writer()); + try response.finish(); + } else { + response.status = .not_found; + try response.do(); + try response.finish(); + } +} diff --git a/src/prometheus/metric.zig b/src/prometheus/metric.zig index 82ba27397..849475f1f 100644 --- a/src/prometheus/metric.zig +++ b/src/prometheus/metric.zig @@ -3,31 +3,7 @@ const fmt = std.fmt; const mem = std.mem; const testing = std.testing; -pub const HistogramResult = struct { - pub const Bucket = struct { - vmrange: []const u8, - count: u64, - }; - - pub const SumValue = struct { - value: f64 = 0, - - pub fn format(self: @This(), comptime format_string: []const u8, options: fmt.FormatOptions, writer: anytype) !void { - _ = format_string; - - const as_int: u64 = @intFromFloat(self.value); - if (@as(f64, @floatFromInt(as_int)) == self.value) { - try fmt.formatInt(as_int, 10, .lower, options, writer); - } else { - try fmt.formatFloatDecimal(self.value, options, writer); - } - } - }; - - buckets: []Bucket, - sum: SumValue, - count: u64, -}; +const HistogramSnapshot = @import("histogram.zig").HistogramSnapshot; pub const Metric = struct { pub const Error = error{OutOfMemory} || std.os.WriteError || std.http.Server.Response.Writer.Error; @@ -38,7 +14,7 @@ pub const Metric = struct { counter: u64, gauge: f64, gauge_int: u64, - histogram: HistogramResult, + histogram: HistogramSnapshot, pub fn deinit(self: Self, allocator: mem.Allocator) void { switch (self) { @@ -70,17 +46,17 @@ pub const Metric = struct { if (name_and_labels.labels.len > 0) { for (v.buckets) |bucket| { - try writer.print("{s}_bucket{{{s},vmrange=\"{s}\"}} {d:.6}\n", .{ + try writer.print("{s}_bucket{{{s},le=\"{s}\"}} {d:.6}\n", .{ name_and_labels.name, name_and_labels.labels, - bucket.vmrange, - bucket.count, + floatMetric(bucket.upper_bound), + bucket.cumulative_count, }); } try writer.print("{s}_sum{{{s}}} {:.6}\n", .{ name_and_labels.name, name_and_labels.labels, - v.sum, + floatMetric(v.sum), }); try writer.print("{s}_count{{{s}}} {d}\n", .{ name_and_labels.name, @@ -89,15 +65,15 @@ pub const Metric = struct { }); } else { for (v.buckets) |bucket| { - try writer.print("{s}_bucket{{vmrange=\"{s}\"}} {d:.6}\n", .{ + try writer.print("{s}_bucket{{le=\"{s}\"}} {d:.6}\n", .{ name_and_labels.name, - bucket.vmrange, - bucket.count, + floatMetric(bucket.upper_bound), + bucket.cumulative_count, }); } try writer.print("{s}_sum {:.6}\n", .{ name_and_labels.name, - v.sum, + floatMetric(v.sum), }); try writer.print("{s}_count {d}\n", .{ name_and_labels.name, @@ -109,6 +85,24 @@ pub const Metric = struct { } }; +/// Converts a float into an anonymous type that can be formatted properly for prometheus. +pub fn floatMetric(value: anytype) struct { + value: @TypeOf(value), + + pub fn format(self: @This(), comptime format_string: []const u8, options: fmt.FormatOptions, writer: anytype) !void { + _ = format_string; + + const as_int: u64 = @intFromFloat(self.value); + if (@as(f64, @floatFromInt(as_int)) == self.value) { + try fmt.formatInt(as_int, 10, .lower, options, writer); + } else { + try fmt.formatFloatDecimal(self.value, options, writer); + } + } +} { + return .{ .value = value }; +} + const NameAndLabels = struct { name: []const u8, labels: []const u8 = "", diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig index 350aa750d..0ae93db1f 100644 --- a/src/prometheus/registry.zig +++ b/src/prometheus/registry.zig @@ -4,11 +4,14 @@ const hash_map = std.hash_map; const heap = std.heap; const mem = std.mem; const testing = std.testing; + const Metric = @import("metric.zig").Metric; -const Counter = @import("counter.zig").Counter; +const Counter = @import("counter.zig"); const Gauge = @import("gauge.zig").Gauge; +const GaugeFn = @import("gauge_fn.zig").GaugeFn; +const GaugeCallFnType = @import("gauge_fn.zig").GaugeCallFnType; const Histogram = @import("histogram.zig").Histogram; -const GaugeCallFnType = @import("gauge.zig").GaugeCallFnType; +const defaultBuckets = @import("histogram.zig").defaultBuckets; pub const GetMetricError = error{ // Returned when trying to add a metric to an already full registry. @@ -24,21 +27,10 @@ const RegistryOptions = struct { max_name_len: comptime_int = 1024, }; -var gpa = std.heap.GeneralPurposeAllocator(.{}){}; -const gpa_allocator = gpa.allocator(); -pub var registry: *Registry(.{}) = undefined; - -pub fn init() error{OutOfMemory}!void { - registry = try Registry(.{}).init(gpa_allocator); -} - -pub fn deinit() void { - registry.deinit(); -} - pub fn Registry(comptime options: RegistryOptions) type { return struct { const Self = @This(); + const MetricMap = hash_map.StringHashMapUnmanaged(*Metric); root_allocator: mem.Allocator, @@ -46,12 +38,12 @@ pub fn Registry(comptime options: RegistryOptions) type { mutex: std.Thread.Mutex, metrics: MetricMap, - pub fn init(alloc: mem.Allocator) error{OutOfMemory}!*Self { - const self = try alloc.create(Self); + pub fn init(allocator: mem.Allocator) !*Self { + const self = try allocator.create(Self); self.* = .{ - .root_allocator = alloc, - .arena_state = heap.ArenaAllocator.init(alloc), + .root_allocator = allocator, + .arena_state = heap.ArenaAllocator.init(allocator), .mutex = .{}, .metrics = MetricMap{}, }; @@ -69,76 +61,40 @@ pub fn Registry(comptime options: RegistryOptions) type { } pub fn getOrCreateCounter(self: *Self, name: []const u8) GetMetricError!*Counter { - if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; - if (name.len > options.max_name_len) return error.NameTooLong; - - var allocator = self.arena_state.allocator(); - - const duped_name = try allocator.dupe(u8, name); - - self.mutex.lock(); - defer self.mutex.unlock(); - - var gop = try self.metrics.getOrPut(allocator, duped_name); - if (!gop.found_existing) { - var real_metric = try Counter.init(allocator); - gop.value_ptr.* = &real_metric.metric; - } - - return @fieldParentPtr(Counter, "metric", gop.value_ptr.*); + return self.getOrCreateMetric(name, Counter, .{}); } - pub fn getOrCreateHistogram(self: *Self, name: []const u8) GetMetricError!*Histogram { - if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; - if (name.len > options.max_name_len) return error.NameTooLong; - - var allocator = self.arena_state.allocator(); - - const duped_name = try allocator.dupe(u8, name); - - self.mutex.lock(); - defer self.mutex.unlock(); - - var gop = try self.metrics.getOrPut(allocator, duped_name); - if (!gop.found_existing) { - var real_metric = try Histogram.init(allocator); - gop.value_ptr.* = &real_metric.metric; - } - - return @fieldParentPtr(Histogram, "metric", gop.value_ptr.*); + pub fn getOrCreateGauge(self: *Self, name: []const u8) GetMetricError!*Gauge { + return self.getOrCreateMetric(name, Gauge, .{}); } - pub fn getOrCreateGauge( + pub fn getOrCreateGaugeFn( self: *Self, name: []const u8, state: anytype, callFn: GaugeCallFnType(@TypeOf(state), f64), - ) GetMetricError!*Gauge(@TypeOf(state), f64) { - if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; - if (name.len > options.max_name_len) return error.NameTooLong; - - var allocator = self.arena_state.allocator(); - - const duped_name = try allocator.dupe(u8, name); - - self.mutex.lock(); - defer self.mutex.unlock(); - - var gop = try self.metrics.getOrPut(allocator, duped_name); - if (!gop.found_existing) { - var real_metric = try Gauge(@TypeOf(state), f64).init(allocator, callFn, state); - gop.value_ptr.* = &real_metric.metric; - } + ) GetMetricError!*GaugeFn(@TypeOf(state), Return(@TypeOf(callFn))) { + return self.getOrCreateMetric( + name, + GaugeFn(@TypeOf(state), Return(@TypeOf(callFn))), + .{ callFn, state }, + ); + } - return @fieldParentPtr(Gauge(@TypeOf(state), f64), "metric", gop.value_ptr.*); + pub fn getOrCreateHistogram( + self: *Self, + name: []const u8, + buckets: std.ArrayList(f64), + ) GetMetricError!*Histogram { + return self.getOrCreateMetric(name, Histogram, .{buckets}); } - pub fn getOrCreateGaugeInt( + fn getOrCreateMetric( self: *Self, name: []const u8, - state: anytype, - callFn: GaugeCallFnType(@TypeOf(state), u64), - ) GetMetricError!*Gauge(@TypeOf(state), u64) { + comptime MetricType: type, + args: anytype, + ) GetMetricError!*MetricType { if (self.nbMetrics() >= options.max_metrics) return error.TooManyMetrics; if (name.len > options.max_name_len) return error.NameTooLong; @@ -149,13 +105,13 @@ pub fn Registry(comptime options: RegistryOptions) type { self.mutex.lock(); defer self.mutex.unlock(); - var gop = try self.metrics.getOrPut(allocator, duped_name); + const gop = try self.metrics.getOrPut(allocator, duped_name); if (!gop.found_existing) { - var real_metric = try Gauge(@TypeOf(state), u64).init(allocator, callFn, state); + var real_metric = try @call(.auto, MetricType.init, .{allocator} ++ args); gop.value_ptr.* = &real_metric.metric; } - return @fieldParentPtr(Gauge(@TypeOf(state), u64), "metric", gop.value_ptr.*); + return @fieldParentPtr(MetricType, "metric", gop.value_ptr.*); } pub fn write(self: *Self, allocator: mem.Allocator, writer: anytype) !void { @@ -193,29 +149,38 @@ pub fn Registry(comptime options: RegistryOptions) type { }; } +/// Gets the return type of a function or function pointer +fn Return(comptime FnPtr: type) type { + return switch (@typeInfo(FnPtr)) { + .Fn => |fun| fun.return_type.?, + .Pointer => |ptr| @typeInfo(ptr.child).Fn.return_type.?, + else => @compileError("not a function or function pointer"), + }; +} + fn stringLessThan(context: void, lhs: []const u8, rhs: []const u8) bool { _ = context; return mem.lessThan(u8, lhs, rhs); } -test "registry getOrCreateCounter" { - var reg = try Registry(.{}).create(testing.allocator); - defer reg.destroy(); +test "prometheus.registry: getOrCreateCounter" { + var registry = try Registry(.{}).init(testing.allocator); + defer registry.deinit(); const name = try fmt.allocPrint(testing.allocator, "http_requests{{status=\"{d}\"}}", .{500}); defer testing.allocator.free(name); var i: usize = 0; while (i < 10) : (i += 1) { - var counter = try reg.getOrCreateCounter(name); + var counter = try registry.getOrCreateCounter(name); counter.inc(); } - var counter = try reg.getOrCreateCounter(name); + var counter = try registry.getOrCreateCounter(name); try testing.expectEqual(@as(u64, 10), counter.get()); } -test "registry write" { +test "prometheus.registry: write" { const TestCase = struct { counter_name: []const u8, gauge_name: []const u8, @@ -261,18 +226,18 @@ test "registry write" { }; inline for (test_cases) |tc| { - var reg = try Registry(.{}).create(testing.allocator); - defer reg.destroy(); + var registry = try Registry(.{}).init(testing.allocator); + defer registry.deinit(); // Add some counters { - var counter = try reg.getOrCreateCounter(tc.counter_name); - counter.set(2); + var counter = try registry.getOrCreateCounter(tc.counter_name); + counter.* = .{ .value = .{ .value = 2 } }; } // Add some gauges { - _ = try reg.getOrCreateGauge( + _ = try registry.getOrCreateGaugeFn( tc.gauge_name, @as(f64, 4.0), struct { @@ -283,13 +248,16 @@ test "registry write" { ); } - // Add an histogram + // TODO: redesign buckets code so it uses the registry's allocator + const buckets = try defaultBuckets(testing.allocator); + defer buckets.deinit(); + // Add a histogram { - var histogram = try reg.getOrCreateHistogram(tc.histogram_name); + var histogram = try registry.getOrCreateHistogram(tc.histogram_name, buckets); - histogram.update(500.12); - histogram.update(1230.240); - histogram.update(140); + histogram.observe(500.12); + histogram.observe(1230.240); + histogram.observe(140); } // Write to a buffer @@ -297,7 +265,7 @@ test "registry write" { var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); - try reg.write(testing.allocator, buffer.writer()); + try registry.write(testing.allocator, buffer.writer()); try testing.expectEqualStrings(tc.exp, buffer.items); } @@ -311,7 +279,7 @@ test "registry write" { std.fs.cwd().deleteFile(filename) catch {}; } - try reg.write(testing.allocator, file.writer()); + try registry.write(testing.allocator, file.writer()); try file.seekTo(0); const file_data = try file.readToEndAlloc(testing.allocator, std.math.maxInt(usize)); @@ -322,23 +290,20 @@ test "registry write" { } } -test "registry options" { - var reg = try Registry(.{ .max_metrics = 1, .max_name_len = 4 }).create(testing.allocator); - defer reg.destroy(); +test "prometheus.registry: options" { + var registry = try Registry(.{ .max_metrics = 1, .max_name_len = 4 }).init(testing.allocator); + defer registry.deinit(); { - try testing.expectError(error.NameTooLong, reg.getOrCreateCounter("hello")); - _ = try reg.getOrCreateCounter("foo"); + try testing.expectError(error.NameTooLong, registry.getOrCreateCounter("hello")); + _ = try registry.getOrCreateCounter("foo"); } { - try testing.expectError(error.TooManyMetrics, reg.getOrCreateCounter("bar")); + try testing.expectError(error.TooManyMetrics, registry.getOrCreateCounter("bar")); } } -test "prometheus.registry: test default registry" { - registry = try Registry(.{}).init(testing.allocator); - defer registry.deinit(); - var counter = try registry.getOrCreateCounter("hello"); - counter.inc(); +test { + testing.refAllDecls(@This()); } From 426e76e5e73bb3c3720f0e2d74887c390f8e75c7 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 13:50:15 -0500 Subject: [PATCH 04/15] fix(prometheus): segfault due to improper histogram init --- src/prometheus/gauge.zig | 6 ++++-- src/prometheus/histogram.zig | 20 ++++++++++++-------- src/prometheus/registry.zig | 28 ++++++++++++++++++++++------ 3 files changed, 38 insertions(+), 16 deletions(-) diff --git a/src/prometheus/gauge.zig b/src/prometheus/gauge.zig index 992226e6e..60430fbf2 100644 --- a/src/prometheus/gauge.zig +++ b/src/prometheus/gauge.zig @@ -9,8 +9,10 @@ pub fn Gauge(comptime T: type) type { value: std.atomic.Atomic(T) = .{ .value = 0 }, metric: Metric = .{ .getResultFn = getResult }, - pub fn init(_: anytype) @This() { - return .{}; + pub fn init(allocator: std.mem.Allocator) @This() { + const self = try allocator.create(@This()); + self.* = .{}; + return self; } pub fn inc(self: *@This()) void { diff --git a/src/prometheus/histogram.zig b/src/prometheus/histogram.zig index 5b0a010b3..d8ded7c50 100644 --- a/src/prometheus/histogram.zig +++ b/src/prometheus/histogram.zig @@ -16,7 +16,7 @@ pub fn defaultBuckets(allocator: Allocator) !ArrayList(f64) { return l; } -/// Histogram optimized for fast writes. +/// Histogram optimized for fast concurrent writes. /// Reads and writes are thread-safe if you use the public methods. /// Writes are lock-free. Reads are locked with a mutex because they occupy a shard. /// @@ -61,8 +61,9 @@ pub const Histogram = struct { const Self = @This(); - pub fn init(allocator: Allocator, buckets: ArrayList(f64)) !@This() { - return .{ + pub fn init(allocator: Allocator, buckets: ArrayList(f64)) !*@This() { + const self = try allocator.create(Self); + self.* = .{ .allocator = allocator, .upper_bounds = buckets, .shards = .{ @@ -71,12 +72,14 @@ pub const Histogram = struct { }, .shard_sync = Atomic(u64).init(0), }; + return self; } pub fn deinit(self: *Self) void { self.shards[0].buckets.deinit(); self.shards[1].buckets.deinit(); self.upper_bounds.deinit(); + self.allocator.destroy(self); } fn shardBuckets(allocator: Allocator, size: usize) !ArrayList(Atomic(u64)) { @@ -88,15 +91,16 @@ pub const Histogram = struct { } /// Writes a value into the histogram. - pub fn observe(self: *Self, item: f64) void { + pub fn observe(self: *Self, value: f64) void { const shard_sync = self.incrementCount(.Acquire); // acquires lock. must be first step. const shard = &self.shards[shard_sync.shard]; for (0.., self.upper_bounds.items) |i, bound| { - if (item <= bound) { + if (value <= bound) { _ = shard.buckets.items[i].fetchAdd(1, .Monotonic); break; } } + _ = shard.sum.fetchAdd(value, .Release); // releases lock. must be last step. _ = shard.count.fetchAdd(1, .Release); // releases lock. must be last step. } @@ -213,7 +217,7 @@ test "prometheus.histogram: data goes in correct buckets" { var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); defer hist.deinit(); - const expected_buckets = observeVarious(&hist); + const expected_buckets = observeVarious(hist); var snapshot = try hist.getSnapshot(null); defer snapshot.deinit(); @@ -226,7 +230,7 @@ test "prometheus.histogram: repeated snapshots measure the same thing" { var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); defer hist.deinit(); - const expected_buckets = observeVarious(&hist); + const expected_buckets = observeVarious(hist); var snapshot1 = try hist.getSnapshot(null); snapshot1.deinit(); @@ -241,7 +245,7 @@ test "prometheus.histogram: values accumulate across snapshots" { var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); defer hist.deinit(); - _ = observeVarious(&hist); + _ = observeVarious(hist); var snapshot1 = try hist.getSnapshot(null); snapshot1.deinit(); diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig index 0ae93db1f..3a8a5f718 100644 --- a/src/prometheus/registry.zig +++ b/src/prometheus/registry.zig @@ -190,9 +190,17 @@ test "prometheus.registry: write" { const exp1 = \\http_conn_pool_size 4.000000 - \\http_request_size_bucket{vmrange="1.292e+02...1.468e+02"} 1 - \\http_request_size_bucket{vmrange="4.642e+02...5.275e+02"} 1 - \\http_request_size_bucket{vmrange="1.136e+03...1.292e+03"} 1 + \\http_request_size_bucket{le="0.005"} 0 + \\http_request_size_bucket{le="0.01"} 0 + \\http_request_size_bucket{le="0.025"} 0 + \\http_request_size_bucket{le="0.05"} 0 + \\http_request_size_bucket{le="0.1"} 0 + \\http_request_size_bucket{le="0.25"} 0 + \\http_request_size_bucket{le="0.5"} 0 + \\http_request_size_bucket{le="1"} 0 + \\http_request_size_bucket{le="2.5"} 0 + \\http_request_size_bucket{le="5"} 0 + \\http_request_size_bucket{le="10"} 0 \\http_request_size_sum 1870.360000 \\http_request_size_count 3 \\http_requests 2 @@ -201,9 +209,17 @@ test "prometheus.registry: write" { const exp2 = \\http_conn_pool_size{route="/api/v2/users"} 4.000000 - \\http_request_size_bucket{route="/api/v2/users",vmrange="1.292e+02...1.468e+02"} 1 - \\http_request_size_bucket{route="/api/v2/users",vmrange="4.642e+02...5.275e+02"} 1 - \\http_request_size_bucket{route="/api/v2/users",vmrange="1.136e+03...1.292e+03"} 1 + \\http_request_size_bucket{route="/api/v2/users",le="0.005"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="0.01"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="0.025"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="0.05"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="0.1"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="0.25"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="0.5"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="1"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="2.5"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="5"} 0 + \\http_request_size_bucket{route="/api/v2/users",le="10"} 0 \\http_request_size_sum{route="/api/v2/users"} 1870.360000 \\http_request_size_count{route="/api/v2/users"} 3 \\http_requests{route="/api/v2/users"} 2 From 9b05339b2e5eca2bd38d76fa1a3685130a9a4cc1 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 14:40:30 -0500 Subject: [PATCH 05/15] refactor(prometheus): more straightforward data ownership metrics structs can use normal initialization so they aren's so esoteric. the registry is the context where we'd like to allocate everything in an arena and manage the memory centrally. so, for clarity, that scope can also be in charge of allocating the structs. that way, the structs can behave as normal structs on their own, without requiring special memory management --- src/prometheus/counter.zig | 20 ++++--------------- src/prometheus/gauge.zig | 20 ++++++++----------- src/prometheus/gauge_fn.zig | 33 ++++++++++--------------------- src/prometheus/histogram.zig | 13 +++++------- src/prometheus/registry.zig | 38 ++++++++++++++++++++++++------------ 5 files changed, 53 insertions(+), 71 deletions(-) diff --git a/src/prometheus/counter.zig b/src/prometheus/counter.zig index 55cd239f5..c89e99567 100644 --- a/src/prometheus/counter.zig +++ b/src/prometheus/counter.zig @@ -9,14 +9,6 @@ const Self = @This(); metric: Metric = Metric{ .getResultFn = getResult }, value: std.atomic.Atomic(u64) = std.atomic.Atomic(u64).init(0), -pub fn init(allocator: mem.Allocator) !*Self { - const self = try allocator.create(Self); - - self.* = .{}; - - return self; -} - pub fn inc(self: *Self) void { _ = self.value.fetchAdd(1, .SeqCst); } @@ -47,8 +39,7 @@ test "prometheus.counter: inc/add/dec/set/get" { var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); - var counter = try Self.init(testing.allocator); - defer testing.allocator.destroy(counter); + var counter = Self{}; try testing.expectEqual(@as(u64, 0), counter.get()); @@ -60,8 +51,7 @@ test "prometheus.counter: inc/add/dec/set/get" { } test "prometheus.counter: concurrent" { - var counter = try Self.init(testing.allocator); - defer testing.allocator.destroy(counter); + var counter = Self{}; var threads: [4]std.Thread = undefined; for (&threads) |*thread| { @@ -75,7 +65,7 @@ test "prometheus.counter: concurrent" { } } }.run, - .{counter}, + .{&counter}, ); } @@ -85,9 +75,7 @@ test "prometheus.counter: concurrent" { } test "prometheus.counter: write" { - var counter = try Self.init(testing.allocator); - defer testing.allocator.destroy(counter); - counter.* = .{ .value = .{ .value = 340 } }; + var counter = Self{ .value = .{ .value = 340 } }; var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); diff --git a/src/prometheus/gauge.zig b/src/prometheus/gauge.zig index 60430fbf2..4cb9b72ab 100644 --- a/src/prometheus/gauge.zig +++ b/src/prometheus/gauge.zig @@ -9,40 +9,36 @@ pub fn Gauge(comptime T: type) type { value: std.atomic.Atomic(T) = .{ .value = 0 }, metric: Metric = .{ .getResultFn = getResult }, - pub fn init(allocator: std.mem.Allocator) @This() { - const self = try allocator.create(@This()); - self.* = .{}; - return self; - } + const Self = @This(); - pub fn inc(self: *@This()) void { + pub fn inc(self: *Self) void { self.value.fetchAdd(1, .Unordered); } - pub fn add(self: *@This(), v: T) void { + pub fn add(self: *Self, v: T) void { self.value.fetchAdd(v, .Unordered); } - pub fn dec(self: *@This()) void { + pub fn dec(self: *Self) void { self.value.fetchSub(1, .Unordered); } - pub fn sub(self: *@This(), v: T) void { + pub fn sub(self: *Self, v: T) void { self.value.fetchAdd(v, .Unordered); } - pub fn set(self: *@This(), v: T) void { + pub fn set(self: *Self, v: T) void { self.value.store(v, .Unordered); } - pub fn get(self: *@This()) T { + pub fn get(self: *Self) T { return self.value.load(.Unordered); } fn getResult(metric: *Metric, allocator: std.mem.Allocator) Metric.Error!Metric.Result { _ = allocator; - const self = @fieldParentPtr(@This(), "metric", metric); + const self = @fieldParentPtr(Self, "metric", metric); return switch (T) { f64 => Metric.Result{ .gauge = self.get() }, diff --git a/src/prometheus/gauge_fn.zig b/src/prometheus/gauge_fn.zig index 6096fcf58..0f07c1e7d 100644 --- a/src/prometheus/gauge_fn.zig +++ b/src/prometheus/gauge_fn.zig @@ -25,14 +25,11 @@ pub fn GaugeFn(comptime StateType: type, comptime Return: type) type { callFn: CallFnType = undefined, state: StateType = undefined, - pub fn init(allocator: mem.Allocator, callFn: CallFnType, state: StateType) !*Self { - const self = try allocator.create(Self); - - self.* = .{}; - self.callFn = callFn; - self.state = state; - - return self; + pub fn init(callFn: CallFnType, state: StateType) Self { + return .{ + .callFn = callFn, + .state = state, + }; } pub fn get(self: *Self) Return { @@ -88,8 +85,7 @@ test "prometheus.gauge_fn: get" { var state = State{ .value = 20 }; - var gauge = try GaugeFn(*State, InnerType).init( - testing.allocator, + var gauge = GaugeFn(*State, InnerType).init( struct { fn get(s: *State) InnerType { return s.value + 1; @@ -97,7 +93,6 @@ test "prometheus.gauge_fn: get" { }.get, &state, ); - defer testing.allocator.destroy(gauge); try testing.expectEqual(@as(InnerType, 21), gauge.get()); } @@ -109,8 +104,7 @@ test "prometheus.gauge_fn: optional state" { }; var state = State{ .value = 20.0 }; - var gauge = try GaugeFn(?*State, f64).init( - testing.allocator, + var gauge = GaugeFn(?*State, f64).init( struct { fn get(s: *State) f64 { return s.value + 1.0; @@ -118,14 +112,12 @@ test "prometheus.gauge_fn: optional state" { }.get, &state, ); - defer testing.allocator.destroy(gauge); try testing.expectEqual(@as(f64, 21.0), gauge.get()); } test "prometheus.gauge_fn: non-pointer state" { - var gauge = try GaugeFn(f64, f64).init( - testing.allocator, + var gauge = GaugeFn(f64, f64).init( struct { fn get(s: *f64) f64 { s.* += 1.0; @@ -134,7 +126,6 @@ test "prometheus.gauge_fn: non-pointer state" { }.get, 0.0, ); - defer testing.allocator.destroy(gauge); try testing.expectEqual(@as(f64, 1.0), gauge.get()); } @@ -147,8 +138,7 @@ test "prometheus.gauge_fn: shared state" { var shared_state = State{}; defer shared_state.items.deinit(); - var gauge = try GaugeFn(*State, f64).init( - testing.allocator, + var gauge = GaugeFn(*State, f64).init( struct { fn get(state: *State) f64 { return @floatFromInt(state.items.items.len); @@ -156,7 +146,6 @@ test "prometheus.gauge_fn: shared state" { }.get, &shared_state, ); - defer testing.allocator.destroy(gauge); var threads: [4]std.Thread = undefined; for (&threads, 0..) |*thread, thread_index| { @@ -182,8 +171,7 @@ test "prometheus.gauge_fn: shared state" { } test "prometheus.gauge_fn: write" { - var gauge = try GaugeFn(usize, f64).init( - testing.allocator, + var gauge = GaugeFn(usize, f64).init( struct { fn get(state: *usize) f64 { state.* += 340; @@ -192,7 +180,6 @@ test "prometheus.gauge_fn: write" { }.get, @as(usize, 0), ); - defer testing.allocator.destroy(gauge); var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); diff --git a/src/prometheus/histogram.zig b/src/prometheus/histogram.zig index d8ded7c50..c9efbfd68 100644 --- a/src/prometheus/histogram.zig +++ b/src/prometheus/histogram.zig @@ -61,9 +61,8 @@ pub const Histogram = struct { const Self = @This(); - pub fn init(allocator: Allocator, buckets: ArrayList(f64)) !*@This() { - const self = try allocator.create(Self); - self.* = .{ + pub fn init(allocator: Allocator, buckets: ArrayList(f64)) !Self { + return Self{ .allocator = allocator, .upper_bounds = buckets, .shards = .{ @@ -72,14 +71,12 @@ pub const Histogram = struct { }, .shard_sync = Atomic(u64).init(0), }; - return self; } pub fn deinit(self: *Self) void { self.shards[0].buckets.deinit(); self.shards[1].buckets.deinit(); self.upper_bounds.deinit(); - self.allocator.destroy(self); } fn shardBuckets(allocator: Allocator, size: usize) !ArrayList(Atomic(u64)) { @@ -217,7 +214,7 @@ test "prometheus.histogram: data goes in correct buckets" { var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); defer hist.deinit(); - const expected_buckets = observeVarious(hist); + const expected_buckets = observeVarious(&hist); var snapshot = try hist.getSnapshot(null); defer snapshot.deinit(); @@ -230,7 +227,7 @@ test "prometheus.histogram: repeated snapshots measure the same thing" { var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); defer hist.deinit(); - const expected_buckets = observeVarious(hist); + const expected_buckets = observeVarious(&hist); var snapshot1 = try hist.getSnapshot(null); snapshot1.deinit(); @@ -245,7 +242,7 @@ test "prometheus.histogram: values accumulate across snapshots" { var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); defer hist.deinit(); - _ = observeVarious(hist); + _ = observeVarious(&hist); var snapshot1 = try hist.getSnapshot(null); snapshot1.deinit(); diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig index 3a8a5f718..b6f596d5c 100644 --- a/src/prometheus/registry.zig +++ b/src/prometheus/registry.zig @@ -89,6 +89,10 @@ pub fn Registry(comptime options: RegistryOptions) type { return self.getOrCreateMetric(name, Histogram, .{buckets}); } + /// MetricType must be initializable in one of these ways: + /// - try MetricType.init(allocator, ...args) + /// - MetricType.init(...args) + /// - as args struct (only if no init method is defined) fn getOrCreateMetric( self: *Self, name: []const u8, @@ -107,7 +111,17 @@ pub fn Registry(comptime options: RegistryOptions) type { const gop = try self.metrics.getOrPut(allocator, duped_name); if (!gop.found_existing) { - var real_metric = try @call(.auto, MetricType.init, .{allocator} ++ args); + var real_metric = try allocator.create(MetricType); + if (@hasDecl(MetricType, "init")) { + const params = @typeInfo(@TypeOf(MetricType.init)).Fn.params; + if (params.len != 0 and params[0].type.? == mem.Allocator) { + real_metric.* = try @call(.auto, MetricType.init, .{allocator} ++ args); + } else { + real_metric.* = @call(.auto, MetricType.init, args); + } + } else { + real_metric.* = args; + } gop.value_ptr.* = &real_metric.metric; } @@ -198,10 +212,10 @@ test "prometheus.registry: write" { \\http_request_size_bucket{le="0.25"} 0 \\http_request_size_bucket{le="0.5"} 0 \\http_request_size_bucket{le="1"} 0 - \\http_request_size_bucket{le="2.5"} 0 - \\http_request_size_bucket{le="5"} 0 - \\http_request_size_bucket{le="10"} 0 - \\http_request_size_sum 1870.360000 + \\http_request_size_bucket{le="2.5"} 1 + \\http_request_size_bucket{le="5"} 1 + \\http_request_size_bucket{le="10"} 2 + \\http_request_size_sum 18.703600 \\http_request_size_count 3 \\http_requests 2 \\ @@ -217,10 +231,10 @@ test "prometheus.registry: write" { \\http_request_size_bucket{route="/api/v2/users",le="0.25"} 0 \\http_request_size_bucket{route="/api/v2/users",le="0.5"} 0 \\http_request_size_bucket{route="/api/v2/users",le="1"} 0 - \\http_request_size_bucket{route="/api/v2/users",le="2.5"} 0 - \\http_request_size_bucket{route="/api/v2/users",le="5"} 0 - \\http_request_size_bucket{route="/api/v2/users",le="10"} 0 - \\http_request_size_sum{route="/api/v2/users"} 1870.360000 + \\http_request_size_bucket{route="/api/v2/users",le="2.5"} 1 + \\http_request_size_bucket{route="/api/v2/users",le="5"} 1 + \\http_request_size_bucket{route="/api/v2/users",le="10"} 2 + \\http_request_size_sum{route="/api/v2/users"} 18.703600 \\http_request_size_count{route="/api/v2/users"} 3 \\http_requests{route="/api/v2/users"} 2 \\ @@ -271,9 +285,9 @@ test "prometheus.registry: write" { { var histogram = try registry.getOrCreateHistogram(tc.histogram_name, buckets); - histogram.observe(500.12); - histogram.observe(1230.240); - histogram.observe(140); + histogram.observe(5.0012); + histogram.observe(12.30240); + histogram.observe(1.40); } // Write to a buffer From 9bc31a62b030f4894d9e1bc7f0cf60a249652383 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 16:25:28 -0500 Subject: [PATCH 06/15] refactor(prometheus): monotonic for non-sync. clarify default histogram buckets ownership --- src/prometheus/counter.zig | 8 ++++---- src/prometheus/gauge.zig | 14 ++++++------- src/prometheus/histogram.zig | 39 ++++++++++++++---------------------- src/prometheus/registry.zig | 9 +++------ 4 files changed, 29 insertions(+), 41 deletions(-) diff --git a/src/prometheus/counter.zig b/src/prometheus/counter.zig index c89e99567..53ad9c5b4 100644 --- a/src/prometheus/counter.zig +++ b/src/prometheus/counter.zig @@ -10,7 +10,7 @@ metric: Metric = Metric{ .getResultFn = getResult }, value: std.atomic.Atomic(u64) = std.atomic.Atomic(u64).init(0), pub fn inc(self: *Self) void { - _ = self.value.fetchAdd(1, .SeqCst); + _ = self.value.fetchAdd(1, .Monotonic); } pub fn add(self: *Self, value: anytype) void { @@ -19,15 +19,15 @@ pub fn add(self: *Self, value: anytype) void { else => @compileError("can't add a non-number"), } - _ = self.value.fetchAdd(@intCast(value), .SeqCst); + _ = self.value.fetchAdd(@intCast(value), .Monotonic); } pub fn get(self: *const Self) u64 { - return self.value.load(.SeqCst); + return self.value.load(.Monotonic); } pub fn reset(self: *Self) void { - _ = self.value.store(0, .SeqCst); + _ = self.value.store(0, .Monotonic); } fn getResult(metric: *Metric, _: mem.Allocator) Metric.Error!Metric.Result { diff --git a/src/prometheus/gauge.zig b/src/prometheus/gauge.zig index 4cb9b72ab..92cd108f6 100644 --- a/src/prometheus/gauge.zig +++ b/src/prometheus/gauge.zig @@ -3,7 +3,7 @@ const std = @import("std"); const Metric = @import("metric.zig").Metric; /// A gauge that stores the value it reports. -/// Read and write operations are atomic and unordered. +/// Read and write operations are atomic and monotonic. pub fn Gauge(comptime T: type) type { return struct { value: std.atomic.Atomic(T) = .{ .value = 0 }, @@ -12,27 +12,27 @@ pub fn Gauge(comptime T: type) type { const Self = @This(); pub fn inc(self: *Self) void { - self.value.fetchAdd(1, .Unordered); + self.value.fetchAdd(1, .Monotonic); } pub fn add(self: *Self, v: T) void { - self.value.fetchAdd(v, .Unordered); + self.value.fetchAdd(v, .Monotonic); } pub fn dec(self: *Self) void { - self.value.fetchSub(1, .Unordered); + self.value.fetchSub(1, .Monotonic); } pub fn sub(self: *Self, v: T) void { - self.value.fetchAdd(v, .Unordered); + self.value.fetchAdd(v, .Monotonic); } pub fn set(self: *Self, v: T) void { - self.value.store(v, .Unordered); + self.value.store(v, .Monotonic); } pub fn get(self: *Self) T { - return self.value.load(.Unordered); + return self.value.load(.Monotonic); } fn getResult(metric: *Metric, allocator: std.mem.Allocator) Metric.Error!Metric.Result { diff --git a/src/prometheus/histogram.zig b/src/prometheus/histogram.zig index c9efbfd68..48b252ab3 100644 --- a/src/prometheus/histogram.zig +++ b/src/prometheus/histogram.zig @@ -6,15 +6,7 @@ const Ordering = std.atomic.Ordering; const Metric = @import("metric.zig").Metric; -const default_buckets: [11]f64 = .{ 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0 }; - -pub fn defaultBuckets(allocator: Allocator) !ArrayList(f64) { - var l = try ArrayList(f64).initCapacity(allocator, default_buckets.len); - for (default_buckets) |b| { - try l.append(b); - } - return l; -} +pub const default_buckets: [11]f64 = .{ 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0 }; /// Histogram optimized for fast concurrent writes. /// Reads and writes are thread-safe if you use the public methods. @@ -42,7 +34,7 @@ pub const Histogram = struct { /// Used to ensure reads and writes occur on separate shards. /// Atomic representation of `ShardSync`. - shard_sync: Atomic(u64), + shard_sync: Atomic(u64) = Atomic(u64).init(0), /// Prevents more than one reader at a time, since read operations actually /// execute an internal write by swapping the hot and cold shards. @@ -61,15 +53,16 @@ pub const Histogram = struct { const Self = @This(); - pub fn init(allocator: Allocator, buckets: ArrayList(f64)) !Self { + pub fn init(allocator: Allocator, buckets: []const f64) !Self { + var upper_bounds = try ArrayList(f64).initCapacity(allocator, buckets.len); + upper_bounds.appendSliceAssumeCapacity(buckets); return Self{ .allocator = allocator, - .upper_bounds = buckets, + .upper_bounds = upper_bounds, .shards = .{ - .{ .buckets = try shardBuckets(allocator, buckets.items.len) }, - .{ .buckets = try shardBuckets(allocator, buckets.items.len) }, + .{ .buckets = try shardBuckets(allocator, buckets.len) }, + .{ .buckets = try shardBuckets(allocator, buckets.len) }, }, - .shard_sync = Atomic(u64).init(0), }; } @@ -80,11 +73,9 @@ pub const Histogram = struct { } fn shardBuckets(allocator: Allocator, size: usize) !ArrayList(Atomic(u64)) { - var shard_buckets = try ArrayList(Atomic(u64)).initCapacity(allocator, size); - for (0..size) |_| { - shard_buckets.appendAssumeCapacity(Atomic(u64).init(0)); - } - return shard_buckets; + var slice = try allocator.alloc(u64, size); + @memset(slice, 0); + return ArrayList(Atomic(u64)).fromOwnedSlice(allocator, @ptrCast(slice)); } /// Writes a value into the histogram. @@ -200,7 +191,7 @@ pub const Bucket = struct { test "prometheus.histogram: empty" { const allocator = std.testing.allocator; - var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + var hist = try Histogram.init(allocator, &default_buckets); defer hist.deinit(); var snapshot = try hist.getSnapshot(null); @@ -211,7 +202,7 @@ test "prometheus.histogram: empty" { test "prometheus.histogram: data goes in correct buckets" { const allocator = std.testing.allocator; - var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + var hist = try Histogram.init(allocator, &default_buckets); defer hist.deinit(); const expected_buckets = observeVarious(&hist); @@ -224,7 +215,7 @@ test "prometheus.histogram: data goes in correct buckets" { test "prometheus.histogram: repeated snapshots measure the same thing" { const allocator = std.testing.allocator; - var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + var hist = try Histogram.init(allocator, &default_buckets); defer hist.deinit(); const expected_buckets = observeVarious(&hist); @@ -239,7 +230,7 @@ test "prometheus.histogram: repeated snapshots measure the same thing" { test "prometheus.histogram: values accumulate across snapshots" { const allocator = std.testing.allocator; - var hist = try Histogram.init(allocator, try defaultBuckets(allocator)); + var hist = try Histogram.init(allocator, &default_buckets); defer hist.deinit(); _ = observeVarious(&hist); diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig index b6f596d5c..b9f3453e7 100644 --- a/src/prometheus/registry.zig +++ b/src/prometheus/registry.zig @@ -11,7 +11,7 @@ const Gauge = @import("gauge.zig").Gauge; const GaugeFn = @import("gauge_fn.zig").GaugeFn; const GaugeCallFnType = @import("gauge_fn.zig").GaugeCallFnType; const Histogram = @import("histogram.zig").Histogram; -const defaultBuckets = @import("histogram.zig").defaultBuckets; +const default_buckets = @import("histogram.zig").default_buckets; pub const GetMetricError = error{ // Returned when trying to add a metric to an already full registry. @@ -84,7 +84,7 @@ pub fn Registry(comptime options: RegistryOptions) type { pub fn getOrCreateHistogram( self: *Self, name: []const u8, - buckets: std.ArrayList(f64), + buckets: []const f64, ) GetMetricError!*Histogram { return self.getOrCreateMetric(name, Histogram, .{buckets}); } @@ -278,12 +278,9 @@ test "prometheus.registry: write" { ); } - // TODO: redesign buckets code so it uses the registry's allocator - const buckets = try defaultBuckets(testing.allocator); - defer buckets.deinit(); // Add a histogram { - var histogram = try registry.getOrCreateHistogram(tc.histogram_name, buckets); + var histogram = try registry.getOrCreateHistogram(tc.histogram_name, &default_buckets); histogram.observe(5.0012); histogram.observe(12.30240); From a90fac5f3a1c32905b9517bcedb0f93d694d6b68 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 16:53:23 -0500 Subject: [PATCH 07/15] fix(prometheus): specify gauge type and add gauge to registry tests --- src/prometheus/registry.zig | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig index b9f3453e7..98b8f5940 100644 --- a/src/prometheus/registry.zig +++ b/src/prometheus/registry.zig @@ -64,8 +64,8 @@ pub fn Registry(comptime options: RegistryOptions) type { return self.getOrCreateMetric(name, Counter, .{}); } - pub fn getOrCreateGauge(self: *Self, name: []const u8) GetMetricError!*Gauge { - return self.getOrCreateMetric(name, Gauge, .{}); + pub fn getOrCreateGauge(self: *Self, name: []const u8, comptime T: type) GetMetricError!*Gauge(T) { + return self.getOrCreateMetric(name, Gauge(T), .{}); } pub fn getOrCreateGaugeFn( @@ -124,7 +124,7 @@ pub fn Registry(comptime options: RegistryOptions) type { } gop.value_ptr.* = &real_metric.metric; } - + // std.debug.assert(ok); return @fieldParentPtr(MetricType, "metric", gop.value_ptr.*); } @@ -198,12 +198,14 @@ test "prometheus.registry: write" { const TestCase = struct { counter_name: []const u8, gauge_name: []const u8, + gauge_fn_name: []const u8, histogram_name: []const u8, exp: []const u8, }; const exp1 = \\http_conn_pool_size 4.000000 + \\http_gauge 13 \\http_request_size_bucket{le="0.005"} 0 \\http_request_size_bucket{le="0.01"} 0 \\http_request_size_bucket{le="0.025"} 0 @@ -223,6 +225,7 @@ test "prometheus.registry: write" { const exp2 = \\http_conn_pool_size{route="/api/v2/users"} 4.000000 + \\http_gauge{route="/api/v2/users"} 13 \\http_request_size_bucket{route="/api/v2/users",le="0.005"} 0 \\http_request_size_bucket{route="/api/v2/users",le="0.01"} 0 \\http_request_size_bucket{route="/api/v2/users",le="0.025"} 0 @@ -243,13 +246,15 @@ test "prometheus.registry: write" { const test_cases = &[_]TestCase{ .{ .counter_name = "http_requests", - .gauge_name = "http_conn_pool_size", + .gauge_name = "http_gauge", + .gauge_fn_name = "http_conn_pool_size", .histogram_name = "http_request_size", .exp = exp1, }, .{ .counter_name = "http_requests{route=\"/api/v2/users\"}", - .gauge_name = "http_conn_pool_size{route=\"/api/v2/users\"}", + .gauge_name = "http_gauge{route=\"/api/v2/users\"}", + .gauge_fn_name = "http_conn_pool_size{route=\"/api/v2/users\"}", .histogram_name = "http_request_size{route=\"/api/v2/users\"}", .exp = exp2, }, @@ -266,9 +271,15 @@ test "prometheus.registry: write" { } // Add some gauges + { + var counter = try registry.getOrCreateGauge(tc.gauge_name, u64); + counter.* = .{ .value = .{ .value = 13 } }; + } + + // Add some gauge_fns { _ = try registry.getOrCreateGaugeFn( - tc.gauge_name, + tc.gauge_fn_name, @as(f64, 4.0), struct { fn get(s: *f64) f64 { From 5a8306df2dee9131ad68732a32797ffe5d501d9a Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 20:01:58 -0500 Subject: [PATCH 08/15] fix(prometheus): UB bug allowed casting a metric as the wrong type --- src/prometheus/registry.zig | 38 +++++++++++++++++++++++++++++-------- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig index 98b8f5940..d96c1664f 100644 --- a/src/prometheus/registry.zig +++ b/src/prometheus/registry.zig @@ -14,12 +14,14 @@ const Histogram = @import("histogram.zig").Histogram; const default_buckets = @import("histogram.zig").default_buckets; pub const GetMetricError = error{ - // Returned when trying to add a metric to an already full registry. + /// Returned when trying to add a metric to an already full registry. TooManyMetrics, - // Returned when the name of name is bigger than the configured max_name_len. + /// Returned when the name of name is bigger than the configured max_name_len. NameTooLong, OutOfMemory, + /// Attempted to get a metric of the wrong type. + InvalidType, }; const RegistryOptions = struct { @@ -31,7 +33,11 @@ pub fn Registry(comptime options: RegistryOptions) type { return struct { const Self = @This(); - const MetricMap = hash_map.StringHashMapUnmanaged(*Metric); + const MetricMap = hash_map.StringHashMapUnmanaged(struct { + /// Used to validate the pointer is cast into a valid  type. + type_name: []const u8, + metric: *Metric, + }); root_allocator: mem.Allocator, arena_state: heap.ArenaAllocator, @@ -122,10 +128,15 @@ pub fn Registry(comptime options: RegistryOptions) type { } else { real_metric.* = args; } - gop.value_ptr.* = &real_metric.metric; + gop.value_ptr.* = .{ + .type_name = @typeName(MetricType), + .metric = &real_metric.metric, + }; + } else if (!std.mem.eql(u8, gop.value_ptr.*.type_name, @typeName(MetricType))) { + return GetMetricError.InvalidType; } - // std.debug.assert(ok); - return @fieldParentPtr(MetricType, "metric", gop.value_ptr.*); + + return @fieldParentPtr(MetricType, "metric", gop.value_ptr.*.metric); } pub fn write(self: *Self, allocator: mem.Allocator, writer: anytype) !void { @@ -156,8 +167,8 @@ pub fn Registry(comptime options: RegistryOptions) type { // Write each metric in key order for (keys) |key| { - var metric = map.get(key) orelse unreachable; - try metric.write(allocator, writer, key); + var value = map.get(key) orelse unreachable; + try value.metric.write(allocator, writer, key); } } }; @@ -194,6 +205,17 @@ test "prometheus.registry: getOrCreateCounter" { try testing.expectEqual(@as(u64, 10), counter.get()); } +test "prometheus.registry: getOrCreateX requires the same type" { + var registry = try Registry(.{}).init(testing.allocator); + defer registry.deinit(); + + const name = try fmt.allocPrint(testing.allocator, "http_requests{{status=\"{d}\"}}", .{500}); + defer testing.allocator.free(name); + + _ = try registry.getOrCreateCounter(name); + if (registry.getOrCreateGauge(name, u64)) |_| try testing.expect(false) else |_| {} +} + test "prometheus.registry: write" { const TestCase = struct { counter_name: []const u8, From 817b462f2c6c09d914696fc434326c8548861fa4 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 20:29:38 -0500 Subject: [PATCH 09/15] feat(prometheus): multithreading test for histogram --- src/prometheus/histogram.zig | 41 ++++++++++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/src/prometheus/histogram.zig b/src/prometheus/histogram.zig index 48b252ab3..407cc0774 100644 --- a/src/prometheus/histogram.zig +++ b/src/prometheus/histogram.zig @@ -179,7 +179,7 @@ pub const HistogramSnapshot = struct { }; } - pub fn deinit(self: *@This()) void { + pub fn deinit(self: *const @This()) void { self.allocator.free(self.buckets); } }; @@ -247,6 +247,41 @@ test "prometheus.histogram: values accumulate across snapshots" { try expectSnapshot(8, &default_buckets, &expected_buckets, snapshot); } +test "prometheus.histogram: totals add up after concurrent reads and writes" { + const allocator = std.testing.allocator; + var hist = try Histogram.init(allocator, &default_buckets); + defer hist.deinit(); + + var threads: [4]std.Thread = undefined; + for (&threads) |*thread| { + thread.* = try std.Thread.spawn( + .{}, + struct { + fn run(h: *Histogram) void { + for (0..1000) |i| { + _ = observeVarious(h); + if (i % 10 == 0) { + (h.getSnapshot(null) catch @panic("snapshot")).deinit(); + } + } + } + }.run, + .{&hist}, + ); + } + for (&threads) |*thread| thread.join(); + + const snapshot = try hist.getSnapshot(allocator); + defer snapshot.deinit(); + + var expected = ArrayList(u64).init(allocator); + defer expected.deinit(); + for (result) |r| { + try expected.append(4000 * r); + } + try expectSnapshot(28000, &default_buckets, expected.items, snapshot); +} + fn observeVarious(hist: *Histogram) [11]u64 { hist.observe(1.0); hist.observe(0.1); @@ -255,9 +290,11 @@ fn observeVarious(hist: *Histogram) [11]u64 { hist.observe(0.0000000001); hist.observe(0.1); hist.observe(100.0); - return .{ 1, 1, 1, 1, 4, 4, 4, 5, 6, 6, 6 }; + return result; } +const result: [11]u64 = .{ 1, 1, 1, 1, 4, 4, 4, 5, 6, 6, 6 }; + fn expectSnapshot( expected_total: u64, expected_bounds: []const f64, From 9fbc74f1ddf1de6c9496e1f25f19cf8db7b056f8 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Mon, 18 Dec 2023 20:35:21 -0500 Subject: [PATCH 10/15] feat(prometheus): better http logging --- src/prometheus/http.zig | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/prometheus/http.zig b/src/prometheus/http.zig index bf532a001..66c396559 100644 --- a/src/prometheus/http.zig +++ b/src/prometheus/http.zig @@ -1,23 +1,22 @@ //! Basic HTTP adapter for prometheus using the http server in std. const std = @import("std"); -const Registry = @import("registry.zig").Registry; -// TODO: log with correct logger +const Registry = @import("registry.zig").Registry; +const Logger = @import("../trace/log.zig").Logger; pub fn servePrometheus( allocator: std.mem.Allocator, registry: *Registry(.{}), listen_addr: std.net.Address, + logger: Logger, ) !void { var server = std.http.Server.init(allocator, .{}); defer server.deinit(); try server.listen(listen_addr); outer: while (true) { - var response = try server.accept(.{ - .allocator = allocator, - }); + var response = try server.accept(.{ .allocator = allocator }); defer response.deinit(); while (response.reset() != .closing) { @@ -26,8 +25,8 @@ pub fn servePrometheus( error.EndOfStream => continue, else => return err, }; - handleRequest(allocator, &response, registry) catch |e| { - std.debug.print("Failed while handling prometheus http request. {}", .{e}); + handleRequest(allocator, &response, registry, logger) catch |e| { + logger.err("prometheus http: Failed to handle request. {}", .{e}); }; } } @@ -37,8 +36,9 @@ fn handleRequest( allocator: std.mem.Allocator, response: *std.http.Server.Response, registry: *Registry(.{}), + logger: Logger, ) !void { - std.debug.print("{s} {s} {s}\n", .{ + logger.debug("prometheus http: {s} {s} {s}\n", .{ @tagName(response.request.method), @tagName(response.request.version), response.request.target, From e5813bd0bee3ad51d577c5ba2ae8203cde0415bb Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Tue, 19 Dec 2023 09:55:32 -0500 Subject: [PATCH 11/15] add test endpoint for prometheus --- build.zig | 17 ++++++++++++++++- src/prometheus/http.zig | 37 +++++++++++++++++++++++++++++++++++-- 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/build.zig b/build.zig index 25c3865ba..135797c0c 100644 --- a/build.zig +++ b/build.zig @@ -163,6 +163,21 @@ pub fn build(b: *std.Build) void { if (b.args) |args| { benchmark_cmd.addArgs(args); } - b.step("benchmark", "benchmark gossip").dependOn(&benchmark_cmd.step); + + // test prometheus http endpoint + const test_prometheus = b.addExecutable(.{ + .name = "test-prometheus", + .root_source_file = .{ .path = "src/prometheus/http.zig" }, + .target = target, + .optimize = optimize, + .main_pkg_path = .{ .path = "src" }, + }); + b.installArtifact(test_prometheus); + const test_prometheus_cmd = b.addRunArtifact(test_prometheus); + if (b.args) |args| { + test_prometheus_cmd.addArgs(args); + } + b.step("test-prometheus", "run test prometheus endpoint with dummy data") + .dependOn(&test_prometheus_cmd.step); } diff --git a/src/prometheus/http.zig b/src/prometheus/http.zig index 66c396559..1e674c059 100644 --- a/src/prometheus/http.zig +++ b/src/prometheus/http.zig @@ -3,7 +3,9 @@ const std = @import("std"); const Registry = @import("registry.zig").Registry; +const default_buckets = @import("histogram.zig").default_buckets; const Logger = @import("../trace/log.zig").Logger; +const Level = @import("../trace/level.zig").Level; pub fn servePrometheus( allocator: std.mem.Allocator, @@ -26,7 +28,7 @@ pub fn servePrometheus( else => return err, }; handleRequest(allocator, &response, registry, logger) catch |e| { - logger.err("prometheus http: Failed to handle request. {}", .{e}); + logger.errf("prometheus http: Failed to handle request. {}", .{e}); }; } } @@ -38,7 +40,7 @@ fn handleRequest( registry: *Registry(.{}), logger: Logger, ) !void { - logger.debug("prometheus http: {s} {s} {s}\n", .{ + logger.debugf("prometheus http: {s} {s} {s}\n", .{ @tagName(response.request.method), @tagName(response.request.version), response.request.target, @@ -58,3 +60,34 @@ fn handleRequest( try response.finish(); } } + +/// Runs a test prometheus endpoint with dummy data. +pub fn main() !void { + const a = std.heap.page_allocator; + var registry = try Registry(.{}).init(a); + _ = try std.Thread.spawn( + .{}, + struct { + fn run(r: *Registry(.{})) !void { + var secs_counter = try r.getOrCreateCounter("seconds_since_start"); + var gauge = try r.getOrCreateGauge("seconds_hand", u64); + var hist = try r.getOrCreateHistogram("hist", &default_buckets); + while (true) { + std.time.sleep(1_000_000_000); + secs_counter.inc(); + gauge.set(@as(u64, @intCast(std.time.timestamp())) % @as(u64, 60)); + hist.observe(1.1); + hist.observe(0.02); + } + } + }.run, + .{registry}, + ); + const logger = Logger.init(a, Level.debug); + try servePrometheus( + a, + registry, + try std.net.Address.parseIp4("0.0.0.0", 1234), + logger, + ); +} From 7eca0e45b30901f5f4d8cde28edb6731a68f1f24 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Thu, 21 Dec 2023 22:38:16 -0500 Subject: [PATCH 12/15] fix(prometheus): pr feedback - initialize prometheus registry and http adapter with sig main() - add global registry singleton. - implement OnceCell to safely support global singletons that need to be initialized at runtime - make Counter a normal looking struct (not a file struct) - switch unnecessary .Release operation in histogram to .Monotonic - switch from std.http to httpz --- build.zig | 29 +++++------ build.zig.zon | 4 ++ src/cmd/cmd.zig | 14 ++++++ src/lib.zig | 1 + src/prometheus/counter.zig | 56 +++++++++++---------- src/prometheus/histogram.zig | 2 +- src/prometheus/http.zig | 96 ++++++++++++++--------------------- src/prometheus/registry.zig | 27 +++++----- src/sync/once_cell.zig | 98 ++++++++++++++++++++++++++++++++++++ 9 files changed, 210 insertions(+), 117 deletions(-) create mode 100644 src/sync/once_cell.zig diff --git a/build.zig b/build.zig index 135797c0c..6cfea8ba4 100644 --- a/build.zig +++ b/build.zig @@ -23,6 +23,7 @@ pub fn build(b: *std.Build) void { const zig_network_module = b.dependency("zig-network", opts).module("network"); const zig_cli_module = b.dependency("zig-cli", opts).module("zig-cli"); const getty_mod = b.dependency("getty", opts).module("getty"); + const httpz_mod = b.dependency("httpz", opts).module("httpz"); const lib = b.addStaticLibrary(.{ .name = "sig", @@ -53,6 +54,10 @@ pub fn build(b: *std.Build) void { .name = "getty", .module = getty_mod, }, + .{ + .name = "httpz", + .module = httpz_mod, + }, }, }); @@ -60,6 +65,7 @@ pub fn build(b: *std.Build) void { lib.addModule("zig-network", zig_network_module); lib.addModule("zig-cli", zig_cli_module); lib.addModule("getty", getty_mod); + lib.addModule("httpz", httpz_mod); // This declares intent for the library to be installed into the standard // location when the user invokes the "install" step (the default step when @@ -77,6 +83,8 @@ pub fn build(b: *std.Build) void { tests.addModule("base58-zig", base58_module); tests.addModule("zig-cli", zig_cli_module); tests.addModule("getty", getty_mod); + tests.addModule("httpz", httpz_mod); + const run_tests = b.addRunArtifact(tests); const test_step = b.step("test", "Run library tests"); test_step.dependOn(&lib.step); @@ -94,6 +102,7 @@ pub fn build(b: *std.Build) void { exe.addModule("zig-network", zig_network_module); exe.addModule("zig-cli", zig_cli_module); exe.addModule("getty", getty_mod); + exe.addModule("httpz", httpz_mod); // This declares intent for the executable to be installed into the // standard location when the user invokes the "install" step (the default @@ -137,6 +146,8 @@ pub fn build(b: *std.Build) void { fuzz_exe.addModule("zig-network", zig_network_module); fuzz_exe.addModule("zig-cli", zig_cli_module); fuzz_exe.addModule("getty", getty_mod); + fuzz_exe.addModule("httpz", httpz_mod); + b.installArtifact(fuzz_exe); const fuzz_cmd = b.addRunArtifact(fuzz_exe); if (b.args) |args| { @@ -158,26 +169,12 @@ pub fn build(b: *std.Build) void { benchmark_exe.addModule("zig-network", zig_network_module); benchmark_exe.addModule("zig-cli", zig_cli_module); benchmark_exe.addModule("getty", getty_mod); + benchmark_exe.addModule("httpz", httpz_mod); + b.installArtifact(benchmark_exe); const benchmark_cmd = b.addRunArtifact(benchmark_exe); if (b.args) |args| { benchmark_cmd.addArgs(args); } b.step("benchmark", "benchmark gossip").dependOn(&benchmark_cmd.step); - - // test prometheus http endpoint - const test_prometheus = b.addExecutable(.{ - .name = "test-prometheus", - .root_source_file = .{ .path = "src/prometheus/http.zig" }, - .target = target, - .optimize = optimize, - .main_pkg_path = .{ .path = "src" }, - }); - b.installArtifact(test_prometheus); - const test_prometheus_cmd = b.addRunArtifact(test_prometheus); - if (b.args) |args| { - test_prometheus_cmd.addArgs(args); - } - b.step("test-prometheus", "run test prometheus endpoint with dummy data") - .dependOn(&test_prometheus_cmd.step); } diff --git a/build.zig.zon b/build.zig.zon index c97036005..bd2c5fc23 100644 --- a/build.zig.zon +++ b/build.zig.zon @@ -18,5 +18,9 @@ .url = "https://github.com/getty-zig/getty/archive/5b0e750d92ee4ef8e46ad743bb8ced63723acd00.tar.gz", .hash = "12209398657d260abcd6dae946d8da4cd3057b8c7990608476a9f8011aae570d2ebb", }, + .httpz = .{ + .url = "https://github.com/karlseguin/http.zig/archive/7a751549a751d9b45952037abdb127b3225b2ac1.tar.gz", + .hash = "122004f74adf46001fe9129d8cec54bd4a98895ce89f0897790e13b60fa99e527b99", + }, }, } diff --git a/src/cmd/cmd.zig b/src/cmd/cmd.zig index ff73d8674..fd653ae86 100644 --- a/src/cmd/cmd.zig +++ b/src/cmd/cmd.zig @@ -8,6 +8,9 @@ const io = std.io; const Pubkey = @import("../core/pubkey.zig").Pubkey; const SocketAddr = @import("../net/net.zig").SocketAddr; const GossipService = @import("../gossip/gossip_service.zig").GossipService; +const servePrometheus = @import("../prometheus/http.zig").servePrometheus; +const global_registry = @import("../prometheus/registry.zig").global_registry; +const Registry = @import("../prometheus/registry.zig").Registry; var gpa = std.heap.GeneralPurposeAllocator(.{}){}; const gpa_allocator = gpa.allocator(); @@ -76,6 +79,8 @@ fn gossip(_: []const []const u8) !void { // var logger: Logger = .noop; + _ = try spawnMetrics(gpa_allocator, 12345); + var my_keypair = try getOrInitIdentity(gpa_allocator, logger); var gossip_port: u16 = @intCast(gossip_port_option.value.int.?); @@ -121,6 +126,15 @@ fn gossip(_: []const []const u8) !void { handle.join(); } +/// Initializes the global registry. +/// Spawns a thread to serve the metrics over http. +/// Returns error if registry was already initialized. +/// Uses same allocator for both registry and http adapter. +fn spawnMetrics(allocator: std.mem.Allocator, port: u16) !std.Thread { + const registry = try global_registry.initialize(Registry(.{}).init, .{allocator}); + return try std.Thread.spawn(.{}, servePrometheus, .{ allocator, registry, port }); +} + pub fn run() !void { return cli.run(app, gpa_allocator); } diff --git a/src/lib.zig b/src/lib.zig index 50af87cf0..826e685e9 100644 --- a/src/lib.zig +++ b/src/lib.zig @@ -46,6 +46,7 @@ pub const sync = struct { pub usingnamespace @import("sync/mpmc.zig"); pub usingnamespace @import("sync/ref.zig"); pub usingnamespace @import("sync/mux.zig"); + pub usingnamespace @import("sync/once_cell.zig"); pub usingnamespace @import("sync/thread_pool.zig"); }; diff --git a/src/prometheus/counter.zig b/src/prometheus/counter.zig index 53ad9c5b4..9044f6af4 100644 --- a/src/prometheus/counter.zig +++ b/src/prometheus/counter.zig @@ -4,42 +4,44 @@ const testing = std.testing; const Metric = @import("metric.zig").Metric; -const Self = @This(); +pub const Counter = struct { + const Self = @This(); -metric: Metric = Metric{ .getResultFn = getResult }, -value: std.atomic.Atomic(u64) = std.atomic.Atomic(u64).init(0), + metric: Metric = Metric{ .getResultFn = getResult }, + value: std.atomic.Atomic(u64) = std.atomic.Atomic(u64).init(0), -pub fn inc(self: *Self) void { - _ = self.value.fetchAdd(1, .Monotonic); -} - -pub fn add(self: *Self, value: anytype) void { - switch (@typeInfo(@TypeOf(value))) { - .Int, .Float, .ComptimeInt, .ComptimeFloat => {}, - else => @compileError("can't add a non-number"), + pub fn inc(self: *Self) void { + _ = self.value.fetchAdd(1, .Monotonic); } - _ = self.value.fetchAdd(@intCast(value), .Monotonic); -} + pub fn add(self: *Self, value: anytype) void { + switch (@typeInfo(@TypeOf(value))) { + .Int, .Float, .ComptimeInt, .ComptimeFloat => {}, + else => @compileError("can't add a non-number"), + } -pub fn get(self: *const Self) u64 { - return self.value.load(.Monotonic); -} + _ = self.value.fetchAdd(@intCast(value), .Monotonic); + } -pub fn reset(self: *Self) void { - _ = self.value.store(0, .Monotonic); -} + pub fn get(self: *const Self) u64 { + return self.value.load(.Monotonic); + } -fn getResult(metric: *Metric, _: mem.Allocator) Metric.Error!Metric.Result { - const self = @fieldParentPtr(Self, "metric", metric); - return Metric.Result{ .counter = self.get() }; -} + pub fn reset(self: *Self) void { + _ = self.value.store(0, .Monotonic); + } + + fn getResult(metric: *Metric, _: mem.Allocator) Metric.Error!Metric.Result { + const self = @fieldParentPtr(Self, "metric", metric); + return Metric.Result{ .counter = self.get() }; + } +}; test "prometheus.counter: inc/add/dec/set/get" { var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); - var counter = Self{}; + var counter = Counter{}; try testing.expectEqual(@as(u64, 0), counter.get()); @@ -51,14 +53,14 @@ test "prometheus.counter: inc/add/dec/set/get" { } test "prometheus.counter: concurrent" { - var counter = Self{}; + var counter = Counter{}; var threads: [4]std.Thread = undefined; for (&threads) |*thread| { thread.* = try std.Thread.spawn( .{}, struct { - fn run(c: *Self) void { + fn run(c: *Counter) void { var i: usize = 0; while (i < 20) : (i += 1) { c.inc(); @@ -75,7 +77,7 @@ test "prometheus.counter: concurrent" { } test "prometheus.counter: write" { - var counter = Self{ .value = .{ .value = 340 } }; + var counter = Counter{ .value = .{ .value = 340 } }; var buffer = std.ArrayList(u8).init(testing.allocator); defer buffer.deinit(); diff --git a/src/prometheus/histogram.zig b/src/prometheus/histogram.zig index 407cc0774..b85eeeb15 100644 --- a/src/prometheus/histogram.zig +++ b/src/prometheus/histogram.zig @@ -88,7 +88,7 @@ pub const Histogram = struct { break; } } - _ = shard.sum.fetchAdd(value, .Release); // releases lock. must be last step. + _ = shard.sum.fetchAdd(value, .Monotonic); _ = shard.count.fetchAdd(1, .Release); // releases lock. must be last step. } diff --git a/src/prometheus/http.zig b/src/prometheus/http.zig index 1e674c059..201321b8c 100644 --- a/src/prometheus/http.zig +++ b/src/prometheus/http.zig @@ -1,77 +1,57 @@ -//! Basic HTTP adapter for prometheus using the http server in std. - const std = @import("std"); +const httpz = @import("httpz"); + +const Level = @import("../trace/level.zig").Level; const Registry = @import("registry.zig").Registry; +const global_registry = @import("registry.zig").global_registry; const default_buckets = @import("histogram.zig").default_buckets; -const Logger = @import("../trace/log.zig").Logger; -const Level = @import("../trace/level.zig").Level; pub fn servePrometheus( allocator: std.mem.Allocator, registry: *Registry(.{}), - listen_addr: std.net.Address, - logger: Logger, + port: u16, ) !void { - var server = std.http.Server.init(allocator, .{}); - defer server.deinit(); - try server.listen(listen_addr); - - outer: while (true) { - var response = try server.accept(.{ .allocator = allocator }); - defer response.deinit(); - - while (response.reset() != .closing) { - response.wait() catch |err| switch (err) { - error.HttpHeadersInvalid => continue :outer, - error.EndOfStream => continue, - else => return err, - }; - handleRequest(allocator, &response, registry, logger) catch |e| { - logger.errf("prometheus http: Failed to handle request. {}", .{e}); - }; - } - } + const endpoint = MetricsEndpoint{ + .allocator = allocator, + .registry = registry, + }; + var server = try httpz.ServerCtx(*const MetricsEndpoint, *const MetricsEndpoint).init( + allocator, + .{ .port = port }, + &endpoint, + ); + var router = server.router(); + router.get("/metrics", getMetrics); + return server.listen(); } -fn handleRequest( +const MetricsEndpoint = struct { allocator: std.mem.Allocator, - response: *std.http.Server.Response, registry: *Registry(.{}), - logger: Logger, -) !void { - logger.debugf("prometheus http: {s} {s} {s}\n", .{ - @tagName(response.request.method), - @tagName(response.request.version), - response.request.target, - }); +}; - if (response.request.method == .GET and - std.mem.startsWith(u8, response.request.target, "/metrics")) - { - response.transfer_encoding = .chunked; - try response.headers.append("content-type", "text/plain"); - try response.do(); - try registry.write(allocator, response.writer()); - try response.finish(); - } else { - response.status = .not_found; - try response.do(); - try response.finish(); - } +pub fn getMetrics( + self: *const MetricsEndpoint, + _: *httpz.Request, + response: *httpz.Response, +) !void { + try self.registry.write(self.allocator, response.writer()); } /// Runs a test prometheus endpoint with dummy data. pub fn main() !void { - const a = std.heap.page_allocator; - var registry = try Registry(.{}).init(a); + const alloc = std.heap.page_allocator; + _ = try global_registry.initialize(Registry(.{}).init, .{alloc}); + _ = try std.Thread.spawn( .{}, struct { - fn run(r: *Registry(.{})) !void { - var secs_counter = try r.getOrCreateCounter("seconds_since_start"); - var gauge = try r.getOrCreateGauge("seconds_hand", u64); - var hist = try r.getOrCreateHistogram("hist", &default_buckets); + fn run() !void { + const reg = try global_registry.get(); + var secs_counter = try reg.getOrCreateCounter("seconds_since_start"); + var gauge = try reg.getOrCreateGauge("seconds_hand", u64); + var hist = try reg.getOrCreateHistogram("hist", &default_buckets); while (true) { std.time.sleep(1_000_000_000); secs_counter.inc(); @@ -81,13 +61,11 @@ pub fn main() !void { } } }.run, - .{registry}, + .{}, ); - const logger = Logger.init(a, Level.debug); try servePrometheus( - a, - registry, - try std.net.Address.parseIp4("0.0.0.0", 1234), - logger, + alloc, + try global_registry.get(), + 12345, ); } diff --git a/src/prometheus/registry.zig b/src/prometheus/registry.zig index d96c1664f..4b4eb3ff2 100644 --- a/src/prometheus/registry.zig +++ b/src/prometheus/registry.zig @@ -5,8 +5,10 @@ const heap = std.heap; const mem = std.mem; const testing = std.testing; +const OnceCell = @import("../sync/once_cell.zig").OnceCell; + const Metric = @import("metric.zig").Metric; -const Counter = @import("counter.zig"); +const Counter = @import("counter.zig").Counter; const Gauge = @import("gauge.zig").Gauge; const GaugeFn = @import("gauge_fn.zig").GaugeFn; const GaugeCallFnType = @import("gauge_fn.zig").GaugeCallFnType; @@ -24,6 +26,10 @@ pub const GetMetricError = error{ InvalidType, }; +/// Global registry singleton for convenience. +pub const global_registry: *OnceCell(Registry(.{})) = &global_registry_owned; +var global_registry_owned: OnceCell(Registry(.{})) = .{}; + const RegistryOptions = struct { max_metrics: comptime_int = 8192, max_name_len: comptime_int = 1024, @@ -39,27 +45,20 @@ pub fn Registry(comptime options: RegistryOptions) type { metric: *Metric, }); - root_allocator: mem.Allocator, arena_state: heap.ArenaAllocator, mutex: std.Thread.Mutex, metrics: MetricMap, - pub fn init(allocator: mem.Allocator) !*Self { - const self = try allocator.create(Self); - - self.* = .{ - .root_allocator = allocator, + pub fn init(allocator: mem.Allocator) Self { + return .{ .arena_state = heap.ArenaAllocator.init(allocator), .mutex = .{}, .metrics = MetricMap{}, }; - - return self; } pub fn deinit(self: *Self) void { self.arena_state.deinit(); - self.root_allocator.destroy(self); } fn nbMetrics(self: *const Self) usize { @@ -189,7 +188,7 @@ fn stringLessThan(context: void, lhs: []const u8, rhs: []const u8) bool { } test "prometheus.registry: getOrCreateCounter" { - var registry = try Registry(.{}).init(testing.allocator); + var registry = Registry(.{}).init(testing.allocator); defer registry.deinit(); const name = try fmt.allocPrint(testing.allocator, "http_requests{{status=\"{d}\"}}", .{500}); @@ -206,7 +205,7 @@ test "prometheus.registry: getOrCreateCounter" { } test "prometheus.registry: getOrCreateX requires the same type" { - var registry = try Registry(.{}).init(testing.allocator); + var registry = Registry(.{}).init(testing.allocator); defer registry.deinit(); const name = try fmt.allocPrint(testing.allocator, "http_requests{{status=\"{d}\"}}", .{500}); @@ -283,7 +282,7 @@ test "prometheus.registry: write" { }; inline for (test_cases) |tc| { - var registry = try Registry(.{}).init(testing.allocator); + var registry = Registry(.{}).init(testing.allocator); defer registry.deinit(); // Add some counters @@ -351,7 +350,7 @@ test "prometheus.registry: write" { } test "prometheus.registry: options" { - var registry = try Registry(.{ .max_metrics = 1, .max_name_len = 4 }).init(testing.allocator); + var registry = Registry(.{ .max_metrics = 1, .max_name_len = 4 }).init(testing.allocator); defer registry.deinit(); { diff --git a/src/sync/once_cell.zig b/src/sync/once_cell.zig new file mode 100644 index 000000000..e8e5f9ac3 --- /dev/null +++ b/src/sync/once_cell.zig @@ -0,0 +1,98 @@ +const std = @import("std"); + +/// Thread-safe data structure that can only be written to once. +/// WARNING: This does not make the inner type thread-safe. +/// +/// All fields are private. Direct access leads to undefined behavior. +/// +/// 1. When this struct is initialized, the contained type is missing. +/// 2. Call one of the init methods to initialize the contained type. +/// 3. After initialization: +/// - get methods will return the initialized value. +/// - value may not be re-initialized. +pub fn OnceCell(comptime T: type) type { + return struct { + value: T = undefined, + started: std.atomic.Atomic(bool) = std.atomic.Atomic(bool).init(false), + finished: std.atomic.Atomic(bool) = std.atomic.Atomic(bool).init(false), + + const Self = @This(); + + pub fn init() Self { + return .{}; + } + + /// Initializes the inner value and returns pointer to it. + /// Returns error if it was already initialized. + /// Blocks while other threads are in the process of initialization. + pub fn initialize(self: *Self, initLogic: anytype, init_args: anytype) error{AlreadyInitialized}!*T { + if (!self.acquire()) return error.AlreadyInitialized; + self.value = @call(.auto, initLogic, init_args); + self.finished.store(true, .Release); + return &self.value; + } + + /// Tries to initialize the inner value and returns pointer to it, or return error if it fails. + /// Returns error if it was already initialized. + /// Blocks while other threads are in the process of initialization. + pub fn tryInitialize(self: *Self, initLogic: anytype, init_args: anytype) !*T { + if (!self.acquire()) return error.AlreadyInitialized; + errdefer self.started.store(false, .Release); + self.value = try @call(.auto, initLogic, init_args); + self.finished.store(true, .Release); + return &self.value; + } + + /// Returns pointer to inner value if already initialized. + /// Otherwise initializes the value and returns it. + /// Blocks while other threads are in the process of initialization. + pub fn getOrInit(self: *Self, initLogic: anytype, init_args: anytype) *T { + if (self.acquire()) { + self.value = @call(.auto, initLogic, init_args); + self.finished.store(true, .Release); + } + return &self.value; + } + + /// Returns pointer to inner value if already initialized. + /// Otherwise tries to initialize the value and returns it, or return error if it fails. + /// Blocks while other threads are in the process of initialization. + pub fn getOrTryInit(self: *Self, initLogic: anytype, init_args: anytype) !*T { + if (self.acquire()) { + errdefer self.started.store(false, .Release); + self.value = try @call(.auto, initLogic, init_args); + self.finished.store(true, .Release); + } + return &self.value; + } + + /// Tries to acquire the write lock. + /// returns: + /// - true if write lock is acquired. + /// - false if write lock is not acquirable because a write was already completed. + /// - waits if another thread has a write in progress. if the other thread fails, this may acquire the lock. + fn acquire(self: *Self) bool { + while (self.started.compareAndSwap(false, true, .Acquire, .Monotonic)) |_| { + if (self.finished.load(.Acquire)) { + return false; + } + } + return true; + } + + /// Returns the value if initialized. + /// Returns error if not initialized. + /// Blocks while other threads are in the process of initialization. + pub fn get(self: *Self) error{NotInitialized}!*T { + if (self.finished.load(.Acquire)) { + return &self.value; + } + while (self.started.load(.Monotonic)) { + if (self.finished.load(.Acquire)) { + return &self.value; + } + } + return error.NotInitialized; + } + }; +} From 4609f9b890bdf602cf2941c6c64dec3be1ca8a53 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Fri, 22 Dec 2023 10:03:17 -0500 Subject: [PATCH 13/15] feat(prometheus): expose metrics port config to cli --- src/cmd/cmd.zig | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/src/cmd/cmd.zig b/src/cmd/cmd.zig index fd653ae86..8b147766b 100644 --- a/src/cmd/cmd.zig +++ b/src/cmd/cmd.zig @@ -34,11 +34,21 @@ var gossip_entrypoints_option = cli.Option{ .value_name = "Entrypoints", }; +var metrics_port_option = cli.Option{ + .long_name = "metrics-port", + .help = "port to expose prometheus metrics via http", + .short_alias = 'm', + .value = cli.OptionValue{ .int = 12345 }, + .required = false, + .value_name = "port_number", +}; + var app = &cli.App{ .name = "sig", .description = "Sig is a Solana client implementation written in Zig.\nThis is still a WIP, PRs welcome.", .version = "0.1.1", .author = "Syndica & Contributors", + .options = &.{&metrics_port_option}, .subcommands = &.{ &cli.Command{ .name = "identity", @@ -79,7 +89,7 @@ fn gossip(_: []const []const u8) !void { // var logger: Logger = .noop; - _ = try spawnMetrics(gpa_allocator, 12345); + _ = try spawnMetrics(gpa_allocator); var my_keypair = try getOrInitIdentity(gpa_allocator, logger); @@ -126,13 +136,13 @@ fn gossip(_: []const []const u8) !void { handle.join(); } -/// Initializes the global registry. -/// Spawns a thread to serve the metrics over http. -/// Returns error if registry was already initialized. +/// Initializes the global registry. Returns error if registry was already initialized. +/// Spawns a thread to serve the metrics over http on the CLI configured port. /// Uses same allocator for both registry and http adapter. -fn spawnMetrics(allocator: std.mem.Allocator, port: u16) !std.Thread { +fn spawnMetrics(allocator: std.mem.Allocator) !std.Thread { + var metrics_port: u16 = @intCast(metrics_port_option.value.int.?); const registry = try global_registry.initialize(Registry(.{}).init, .{allocator}); - return try std.Thread.spawn(.{}, servePrometheus, .{ allocator, registry, port }); + return try std.Thread.spawn(.{}, servePrometheus, .{ allocator, registry, metrics_port }); } pub fn run() !void { From 9cd4be047219d4b7364d502ebc3828953ab2d933 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Fri, 22 Dec 2023 11:05:27 -0500 Subject: [PATCH 14/15] test(sync): add unit tests for OnceCell --- src/sync/once_cell.zig | 104 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 103 insertions(+), 1 deletion(-) diff --git a/src/sync/once_cell.zig b/src/sync/once_cell.zig index e8e5f9ac3..f8610fa7c 100644 --- a/src/sync/once_cell.zig +++ b/src/sync/once_cell.zig @@ -35,7 +35,7 @@ pub fn OnceCell(comptime T: type) type { /// Tries to initialize the inner value and returns pointer to it, or return error if it fails. /// Returns error if it was already initialized. /// Blocks while other threads are in the process of initialization. - pub fn tryInitialize(self: *Self, initLogic: anytype, init_args: anytype) !*T { + pub fn tryInit(self: *Self, initLogic: anytype, init_args: anytype) !*T { if (!self.acquire()) return error.AlreadyInitialized; errdefer self.started.store(false, .Release); self.value = try @call(.auto, initLogic, init_args); @@ -96,3 +96,105 @@ pub fn OnceCell(comptime T: type) type { } }; } + +test "sync.once_cell: init returns correctly" { + var oc = OnceCell(u64).init(); + const x = try oc.initialize(returns(10), .{}); + try std.testing.expect(10 == x.*); +} + +test "sync.once_cell: cannot get uninitialized" { + var oc = OnceCell(u64).init(); + if (oc.get()) |_| { + try std.testing.expect(false); + } else |_| {} +} + +test "sync.once_cell: can get initialized" { + var oc = OnceCell(u64).init(); + _ = try oc.initialize(returns(10), .{}); + const x = try oc.get(); + try std.testing.expect(10 == x.*); +} + +test "sync.once_cell: tryInit returns error on failure" { + var oc = OnceCell(u64).init(); + const err = oc.tryInit(returnErr, .{}); + try std.testing.expectError(error.TestErr, err); +} + +test "sync.once_cell: tryInit works on success" { + var oc = OnceCell(u64).init(); + const x1 = try oc.tryInit(returnNotErr(10), .{}); + const x2 = try oc.get(); + try std.testing.expect(10 == x1.*); + try std.testing.expect(10 == x2.*); +} + +test "sync.once_cell: tryInit returns error if initialized" { + var oc = OnceCell(u64).init(); + const x1 = try oc.tryInit(returnNotErr(10), .{}); + const err = oc.tryInit(returnNotErr(11), .{}); + const x2 = try oc.get(); + try std.testing.expect(10 == x1.*); + try std.testing.expectError(error.AlreadyInitialized, err); + try std.testing.expect(10 == x2.*); +} + +test "sync.once_cell: getOrInit can initialize when needed" { + var oc = OnceCell(u64).init(); + const x1 = oc.getOrInit(returns(10), .{}); + const x2 = try oc.get(); + try std.testing.expect(10 == x1.*); + try std.testing.expect(10 == x2.*); +} + +test "sync.once_cell: getOrInit uses already initialized value" { + var oc = OnceCell(u64).init(); + const x1 = oc.getOrInit(returns(10), .{}); + const x2 = oc.getOrInit(returns(11), .{}); + try std.testing.expect(10 == x1.*); + try std.testing.expect(10 == x2.*); +} + +test "sync.once_cell: getOrTryInit returns error on failure" { + var oc = OnceCell(u64).init(); + const err = oc.getOrTryInit(returnErr, .{}); + try std.testing.expectError(error.TestErr, err); +} + +test "sync.once_cell: getOrTryInit works on success" { + var oc = OnceCell(u64).init(); + const x1 = try oc.getOrTryInit(returnNotErr(10), .{}); + const x2 = try oc.get(); + try std.testing.expect(10 == x1.*); + try std.testing.expect(10 == x2.*); +} + +test "sync.once_cell: getOrTryInit uses already initialized value" { + var oc = OnceCell(u64).init(); + const x1 = try oc.getOrTryInit(returnNotErr(10), .{}); + const x2 = try oc.getOrTryInit(returnNotErr(11), .{}); + try std.testing.expect(10 == x1.*); + try std.testing.expect(10 == x2.*); +} + +fn returns(comptime x: u64) fn () u64 { + return struct { + fn get() u64 { + return x; + } + }.get; +} + +fn returnNotErr(comptime x: u64) fn () error{}!u64 { + return struct { + fn get() !u64 { + return x; + } + }.get; +} + +fn returnErr() !u64 { + return error.TestErr; +} From 47dce1a63c9b2407298cf7eb2af91b46e19abfb8 Mon Sep 17 00:00:00 2001 From: Drew Nutter Date: Fri, 22 Dec 2023 16:52:19 -0500 Subject: [PATCH 15/15] fix(prometheus): join prometheus server thread on exit --- src/cmd/cmd.zig | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cmd/cmd.zig b/src/cmd/cmd.zig index 8b147766b..bc29717d5 100644 --- a/src/cmd/cmd.zig +++ b/src/cmd/cmd.zig @@ -89,7 +89,7 @@ fn gossip(_: []const []const u8) !void { // var logger: Logger = .noop; - _ = try spawnMetrics(gpa_allocator); + const metrics_thread = try spawnMetrics(gpa_allocator); var my_keypair = try getOrInitIdentity(gpa_allocator, logger); @@ -134,6 +134,7 @@ fn gossip(_: []const []const u8) !void { ); handle.join(); + metrics_thread.detach(); } /// Initializes the global registry. Returns error if registry was already initialized.