zfin/src/net/RateLimiter.zig

//! Token-bucket rate limiter.
//!
//! Enforces a maximum number of requests per time window using the
//! token bucket algorithm. Tokens refill continuously; each request
//! consumes one token. When the bucket is empty, callers can either
//! poll with `tryAcquire` or block with `acquire`.
//!
//! wall-clock required: a rate limiter is, by definition, a clock
//! consumer. Every refill computation needs the actual elapsed time
//! since the last refill. Threading a caller-captured timestamp would
//! collapse every `acquire` in the same frame to the same "now," which
//! would under-refill the bucket across a series of rate-limited calls.

const std = @import("std");

io: std.Io,
/// Maximum tokens (requests) in the bucket
max_tokens: u32,
/// Current available tokens
tokens: f64,
/// Tokens added per nanosecond
refill_rate_per_ns: f64,
/// Last time tokens were refilled (nanoseconds from clock.real)
last_refill: i128,

const RateLimiter = @This();

/// Create a rate limiter.
/// `max_per_window` is the max requests allowed in `window_ns` nanoseconds.
pub fn init(io: std.Io, max_per_window: u32, window_ns: u64) RateLimiter {
    return .{
        .io = io,
        .max_tokens = max_per_window,
        .tokens = @floatFromInt(max_per_window),
        .refill_rate_per_ns = @as(f64, @floatFromInt(max_per_window)) / @as(f64, @floatFromInt(window_ns)),
        .last_refill = @intCast(std.Io.Timestamp.now(io, .real).nanoseconds),
    };
}

/// Convenience: N requests per minute.
/// Starts with 1 token (no burst) to stay within provider sliding-window limits.
pub fn perMinute(io: std.Io, n: u32) RateLimiter {
    var rl = init(io, n, std.time.ns_per_min);
    rl.tokens = 1.0;
    return rl;
}

/// Convenience: N requests per day
pub fn perDay(io: std.Io, n: u32) RateLimiter {
    return init(io, n, std.time.ns_per_day);
}

/// Try to acquire a token. Returns true if granted, false if rate-limited.
/// Caller should sleep and retry if false.
pub fn tryAcquire(self: *RateLimiter) bool {
    self.refill();
    if (self.tokens >= 1.0) {
        self.tokens -= 1.0;
        return true;
    }
    return false;
}

/// Acquire a token, blocking (sleeping) until one is available.
pub fn acquire(self: *RateLimiter) void {
    while (!self.tryAcquire()) {
        // Sleep for the time needed to generate 1 token
        const wait_ns: u64 = @intFromFloat(1.0 / self.refill_rate_per_ns);
        std.Io.sleep(self.io, .{ .nanoseconds = @intCast(wait_ns) }, .awake) catch {};
    }
}

/// Sleep until a token is likely available, with a minimum 2-second floor.
/// Use after receiving a server-side 429 to wait before retrying.
pub fn backoff(self: *RateLimiter) void {
    const wait_ns: u64 = @max(self.estimateWaitNs(), 2 * std.time.ns_per_s);
    std.Io.sleep(self.io, .{ .nanoseconds = @intCast(wait_ns) }, .awake) catch {};
}

/// Returns estimated wait time in nanoseconds until a token is available.
/// Returns 0 if a token is available now.
pub fn estimateWaitNs(self: *RateLimiter) u64 {
    self.refill();
    if (self.tokens >= 1.0) return 0;
    const deficit = 1.0 - self.tokens;
    return @intFromFloat(deficit / self.refill_rate_per_ns);
}

fn refill(self: *RateLimiter) void {
    const now: i128 = @intCast(std.Io.Timestamp.now(self.io, .real).nanoseconds);
    const elapsed = now - self.last_refill;
    if (elapsed <= 0) return;

    const new_tokens = @as(f64, @floatFromInt(elapsed)) * self.refill_rate_per_ns;
    self.tokens = @min(self.tokens + new_tokens, @as(f64, @floatFromInt(self.max_tokens)));
    self.last_refill = now;
}

test "rate limiter basic" {
    var rl = RateLimiter.perMinute(std.testing.io, 60);
    // perMinute starts with 1 token (no burst)
    try std.testing.expect(rl.tryAcquire());
    // Second call should be rate-limited immediately
    try std.testing.expect(!rl.tryAcquire());
}

test "rate limiter perDay keeps full burst" {
    var rl = RateLimiter.perDay(std.testing.io, 25);
    // perDay starts with full bucket
    for (0..25) |_| {
        try std.testing.expect(rl.tryAcquire());
    }
    try std.testing.expect(!rl.tryAcquire());
}

test "rate limiter exhaustion" {
    var rl = RateLimiter.init(std.testing.io, 2, std.time.ns_per_s);
    try std.testing.expect(rl.tryAcquire());
    try std.testing.expect(rl.tryAcquire());
    // Bucket should be empty now
    try std.testing.expect(!rl.tryAcquire());
}