diff --git a/src/metrics.zig b/src/metrics.zig index 6507b0d..7524539 100644 --- a/src/metrics.zig +++ b/src/metrics.zig @@ -160,3 +160,162 @@ pub fn exponentialBuckets(allocator: std.mem.Allocator, start: f64, factor: f64, 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 }); + } +