feat: add TokenBucket and refillLocked to metrics

This commit is contained in:
dualfroz committed 2026-09-07 17:58:20 +02:00
1 parent 95c1bed183
commit 0bc2663e45
1 file changed
+159
+159
View File
@@ -160,3 +160,162 @@ pub fn exponentialBuckets(allocator: std.mem.Allocator, start: f64, factor: f64,
return out; return out;
} }
/// Classic token-bucket rate limiter. Time is passed in explicitly as
/// nanoseconds so the algorithm is deterministic and testable; `tryAcquireNow`
/// wraps it with the wall clock for real use.
pub const TokenBucket = struct {
capacity: f64,
refillPerSecond: f64,
tokens: f64,
lastRefillNanos: u64,
mutex: std.Thread.Mutex = .{},
pub fn init(capacity: f64, refillPerSecond: f64, nowNanos: u64) TokenBucket {
std.debug.assert(capacity > 0);
std.debug.assert(refillPerSecond > 0);
return .{
.capacity = capacity,
.refillPerSecond = refillPerSecond,
.tokens = capacity,
.lastRefillNanos = nowNanos,
};
}
fn refillLocked(self: *TokenBucket, nowNanos: u64) void {
if (nowNanos <= self.lastRefillNanos) return;
const elapsedNanos = nowNanos - self.lastRefillNanos;
const elapsedSeconds = @as(f64, @floatFromInt(elapsedNanos)) / std.time.ns_per_s;
self.tokens = @min(self.capacity, self.tokens + elapsedSeconds * self.refillPerSecond);
self.lastRefillNanos = nowNanos;
}
/// Attempts to consume `n` tokens at time `nowNanos`. All-or-nothing:
/// never partially consumes on failure.
pub fn tryAcquireAt(self: *TokenBucket, n: f64, nowNanos: u64) bool {
self.mutex.lock();
defer self.mutex.unlock();
self.refillLocked(nowNanos);
if (self.tokens >= n) {
self.tokens -= n;
return true;
}
return false;
}
pub fn tryAcquireNow(self: *TokenBucket, n: f64) bool {
return self.tryAcquireAt(n, @intCast(std.time.nanoTimestamp()));
}
pub fn availableAt(self: *TokenBucket, nowNanos: u64) f64 {
self.mutex.lock();
defer self.mutex.unlock();
self.refillLocked(nowNanos);
return self.tokens;
}
};
/// Sliding-window-counter rate limiter (the "Cloudflare" approximation):
/// tracks a count for the current fixed window and the previous one, and
/// weights the previous window's count by how much of it still overlaps the
/// trailing `windowNanos` interval. O(1) memory unlike a sliding log, more
/// accurate at window boundaries than a plain fixed window.
pub const SlidingWindowLimiter = struct {
limit: u64,
windowNanos: u64,
currentWindowStart: u64,
currentCount: u64 = 0,
previousCount: u64 = 0,
mutex: std.Thread.Mutex = .{},
pub fn init(limit: u64, windowNanos: u64, nowNanos: u64) SlidingWindowLimiter {
std.debug.assert(limit > 0);
std.debug.assert(windowNanos > 0);
return .{
.limit = limit,
.windowNanos = windowNanos,
.currentWindowStart = alignedWindowStart(nowNanos, windowNanos),
};
}
fn alignedWindowStart(nowNanos: u64, windowNanos: u64) u64 {
return (nowNanos / windowNanos) * windowNanos;
}
fn rotateLocked(self: *SlidingWindowLimiter, nowNanos: u64) void {
const windowStart = alignedWindowStart(nowNanos, self.windowNanos);
if (windowStart == self.currentWindowStart) return;
const windowsElapsed = (windowStart - self.currentWindowStart) / self.windowNanos;
if (windowsElapsed == 1) {
self.previousCount = self.currentCount;
} else {
self.previousCount = 0;
}
self.currentCount = 0;
self.currentWindowStart = windowStart;
}
pub fn allowAt(self: *SlidingWindowLimiter, nowNanos: u64) bool {
self.mutex.lock();
defer self.mutex.unlock();
self.rotateLocked(nowNanos);
const elapsedInWindow = nowNanos - self.currentWindowStart;
const overlap = 1.0 - @as(f64, @floatFromInt(elapsedInWindow)) / @as(f64, @floatFromInt(self.windowNanos));
const clampedOverlap = std.math.clamp(overlap, 0, 1);
const weighted = @as(f64, @floatFromInt(self.previousCount)) * clampedOverlap + @as(f64, @floatFromInt(self.currentCount));
if (weighted >= @as(f64, @floatFromInt(self.limit))) return false;
self.currentCount += 1;
return true;
}
pub fn allowNow(self: *SlidingWindowLimiter) bool {
return self.allowAt(@intCast(std.time.nanoTimestamp()));
}
};
const MetricKind = enum { counter, gauge, histogram };
const MetricRef = union(MetricKind) {
counter: *const MetricCounter,
gauge: *const Gauge,
histogram: *Histogram,
};
/// Named-metric registry with a Prometheus-style text exporter. Stores only
/// borrowed pointers (owned name copies) - callers keep the metrics alive.
pub const Registry = struct {
allocator: std.mem.Allocator,
entries: std.StringHashMapUnmanaged(MetricRef) = .{},
pub fn init(allocator: std.mem.Allocator) Registry {
return .{ .allocator = allocator };
}
pub fn deinit(self: *Registry) void {
var it = self.entries.keyIterator();
while (it.next()) |key| self.allocator.free(key.*);
self.entries.deinit(self.allocator);
self.* = undefined;
}
fn put(self: *Registry, name: []const u8, ref: MetricRef) !void {
if (self.entries.getPtr(name)) |existing| {
existing.* = ref;
return;
}
const owned = try self.allocator.dupe(u8, name);
errdefer self.allocator.free(owned);
try self.entries.put(self.allocator, owned, ref);
}
pub fn registerCounter(self: *Registry, name: []const u8, c: *const MetricCounter) !void {
try self.put(name, .{ .counter = c });
}
pub fn registerGauge(self: *Registry, name: []const u8, g: *const Gauge) !void {
try self.put(name, .{ .gauge = g });
}
pub fn registerHistogram(self: *Registry, name: []const u8, h: *Histogram) !void {
try self.put(name, .{ .histogram = h });
}