Module rate_limit 

Per-token QPS + burst rate limiting (token-bucket).

This module implements PATH-TO-V1 v0.4 exit-criterion Rate limiting per token. It layers behind crate::middleware::bearer_auth: every request that survives auth carries an AuthContext in its extensions, and the rate_limit middleware uses the contained TokenId to credit a per-token TokenBucket.

Design

Token-bucket variant

We use a refill-on-take bucket (sometimes called lazy refill) rather than a background-tick refill. On every request:

1. Compute elapsed nanos since the bucket’s last_refill.
2. Add elapsed * qps / 1_000_000_000 permits to tokens, clamped at burst.
3. If tokens >= 1.0 deduct one and admit; otherwise reject with 429 Too Many Requests and the appropriate Retry-After value.

Tradeoff: refill-on-take has zero background CPU cost and zero coordination overhead (no ticker task), at the price of cold buckets sitting in the DashMap until the process restarts. Since the allowlist of bearer tokens is bounded by configuration size (TENSOR_WASM_API_TOKENS is a finite comma-separated list), the cardinality is small and bounded — a future TTL eviction sweep is a non-goal for v0.4.0. A background-tick refiller would have been the wrong choice: it requires a per-bucket schedule, awakens for idle tokens, and either holds a global lock on the wake task or fragments scheduling per shard. The refill-on-take math is two adds and a clamp; the lock is held for microseconds.

Sharding

Buckets live in Arc<DashMap<TokenId, Mutex<BucketState>>>. DashMap provides shard-level read/write locks; the inner std::sync::Mutex serialises refill arithmetic for a single bucket. We use std::sync::Mutex rather than parking_lot::Mutex to avoid pulling a new dependency into tensor-wasm-api; the critical section is a handful of integer ops with no await points, so OS-mutex contention behaviour is acceptable.

Clock injection

Unit tests need deterministic refill behaviour without tokio::time::sleep (slow + flaky). The Clock trait abstracts “now”. Production uses RealClock; tests inject ManualClock and advance it explicitly.

Wiring

crate::server::build_router layers rate_limit after bearer_auth and before any route handler. If RateLimitConfig::is_disabled returns true (qps == 0 or burst == 0) the layer is installed but short-circuits to a pass-through — equivalent to no rate limiting.

Structs

AuthContext

Per-request authentication context inserted into axum::http::Extensions by crate::middleware::bearer_auth after a successful auth check.

ManualClock

Test clock: holds an explicit Instant that the test advances.

PerTenantRateLimitConfig

Per-tenant (secondary) rate-limit configuration.

RateLimitConfig

Static configuration for the rate limiter.

RateLimiter

In-process two-layer rate limiter.

RealClock

Production clock: delegates to Instant::now.

TokenId

Stable identifier for a bearer token within a single process lifetime.

Enums

AdmitResult

Outcome of an attempt to claim a permit from the bucket.

Traits

Clock

Abstract monotonic clock. Implemented by RealClock (production) and ManualClock (tests).

Functions

rate_limit

Axum middleware that enforces the per-token rate limit.