objectstore_server/

rate_limits.rs

//! Admission-based rate limiting for throughput and bandwidth.
//!
//! This module provides [`RateLimiter`], which enforces configurable limits at three
//! levels of granularity — global, per-usecase, and per-scope — for both request
//! throughput (requests/s) and upload/download bandwidth (bytes/s).
//!
//! Throughput is enforced using token buckets; bandwidth is estimated with an
//! exponentially weighted moving average (EWMA) updated by a background task every
//! 50 ms. All rate-limit checks are synchronous and non-blocking.

use std::fmt;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::{Mutex, atomic::AtomicU64};
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

use bytes::Bytes;
use futures_util::Stream;
use objectstore_service::id::ObjectContext;
use objectstore_types::scope::Scopes;
use serde::{Deserialize, Serialize};

/// Identifies which rate limit triggered a rejection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RateLimitRejection {
    /// Global bandwidth limit exceeded.
    BandwidthGlobal,
    /// Per-usecase bandwidth limit exceeded.
    BandwidthUsecase,
    /// Per-scope bandwidth limit exceeded.
    BandwidthScope,
    /// Global throughput limit exceeded.
    ThroughputGlobal,
    /// Per-usecase throughput limit exceeded.
    ThroughputUsecase,
    /// Per-scope throughput limit exceeded.
    ThroughputScope,
    /// Per-rule throughput limit exceeded.
    ThroughputRule,
}

impl RateLimitRejection {
    /// Returns a static string identifier suitable for use as a metric tag.
    pub fn as_str(self) -> &'static str {
        match self {
            RateLimitRejection::BandwidthGlobal => "bandwidth_global",
            RateLimitRejection::BandwidthUsecase => "bandwidth_usecase",
            RateLimitRejection::BandwidthScope => "bandwidth_scope",
            RateLimitRejection::ThroughputGlobal => "throughput_global",
            RateLimitRejection::ThroughputUsecase => "throughput_usecase",
            RateLimitRejection::ThroughputScope => "throughput_scope",
            RateLimitRejection::ThroughputRule => "throughput_rule",
        }
    }
}

impl fmt::Display for RateLimitRejection {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

/// Rate limits for objectstore.
#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
pub struct RateLimits {
    /// Limits the number of requests per second per service instance.
    pub throughput: ThroughputLimits,
    /// Limits the concurrent bandwidth per service instance.
    pub bandwidth: BandwidthLimits,
}

/// Request throughput limits applied at global, per-usecase, and per-scope granularity.
///
/// All limits are optional. When a limit is `None`, that level of limiting is not enforced.
/// Per-usecase and per-scope limits are expressed as a percentage of the global limit.
#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
pub struct ThroughputLimits {
    /// The overall maximum number of requests per second per service instance.
    ///
    /// Defaults to `None`, meaning no global rate limit is enforced.
    pub global_rps: Option<u32>,

    /// The maximum burst for each rate limit.
    ///
    /// Defaults to `0`, meaning no bursting is allowed. If set to a value greater than `0`,
    /// short spikes above the rate limit are allowed up to the burst size.
    pub burst: u32,

    /// The maximum percentage of the global rate limit that can be used by any usecase.
    ///
    /// Value from `0` to `100`. Defaults to `None`, meaning no per-usecase limit is enforced.
    pub usecase_pct: Option<u8>,

    /// The maximum percentage of the global rate limit that can be used by any scope.
    ///
    /// This treats each full scope separately and applies across all use cases:
    ///  - Two requests with exact same scopes count against the same limit regardless of use case.
    ///  - Two requests that share the same top scope but differ in inner scopes count separately.
    ///
    /// Value from `0` to `100`. Defaults to `None`, meaning no per-scope limit is enforced.
    pub scope_pct: Option<u8>,

    /// Overrides for specific usecases and scopes.
    pub rules: Vec<ThroughputRule>,
}

/// An override rule that applies a specific throughput limit to matching request contexts.
///
/// A rule matches when all specified fields match the request context. Fields not set match
/// any value. When multiple rules match, each is enforced independently via its own token
/// bucket — all matching rules must admit the request.
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct ThroughputRule {
    /// Optional usecase to match.
    ///
    /// If `None`, matches any usecase.
    pub usecase: Option<String>,

    /// Scopes to match.
    ///
    /// If empty, matches any scopes. Additional scopes in the context are ignored, so a rule
    /// matches if all of the specified scopes are present in the request with matching values.
    pub scopes: Vec<(String, String)>,

    /// The rate limit to apply when this rule matches.
    ///
    /// If both a rate and pct are specified, the more restrictive limit applies.
    /// Should be greater than `0`. To block traffic entirely, use killswitches instead.
    pub rps: Option<u32>,

    /// The percentage of the global rate limit to apply when this rule matches.
    ///
    /// If both a rate and pct are specified, the more restrictive limit applies.
    /// Should be greater than `0`. To block traffic entirely, use killswitches instead.
    pub pct: Option<u8>,
}

impl ThroughputRule {
    /// Returns `true` if this rule matches the given context.
    pub fn matches(&self, context: &ObjectContext) -> bool {
        if let Some(ref rule_usecase) = self.usecase
            && rule_usecase != &context.usecase
        {
            return false;
        }

        for (scope_name, scope_value) in &self.scopes {
            match context.scopes.get_value(scope_name) {
                Some(value) if value == scope_value => (),
                _ => return false,
            }
        }

        true
    }
}

/// Bandwidth limits applied at global, per-usecase, and per-scope granularity.
///
/// Bandwidth is measured as bytes transferred per second (upload + download combined)
/// and estimated using an EWMA updated every 50 ms. All limits are optional.
#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq)]
pub struct BandwidthLimits {
    /// The overall maximum bandwidth (in bytes per second) per service instance.
    ///
    /// Defaults to `None`, meaning no global bandwidth limit is enforced.
    pub global_bps: Option<u64>,

    /// The maximum percentage of the global bandwidth limit that can be used by any usecase.
    ///
    /// Value from `0` to `100`. Defaults to `None`, meaning no per-usecase bandwidth limit is enforced.
    pub usecase_pct: Option<u8>,

    /// The maximum percentage of the global bandwidth limit that can be used by any scope.
    ///
    /// Value from `0` to `100`. Defaults to `None`, meaning no per-scope bandwidth limit is enforced.
    pub scope_pct: Option<u8>,
}

/// Combined rate limiter that enforces both bandwidth and throughput limits.
///
/// Checks are synchronous and non-blocking. Bandwidth is checked before
/// throughput so that rejected requests are never counted toward the admitted
/// throughput EWMA. See [`check`](RateLimiter::check) for details.
///
/// Call [`start`](RateLimiter::start) after construction to launch the background
/// estimation tasks. Without it, bandwidth EWMAs remain at zero and bandwidth
/// limits are never triggered.
#[derive(Debug)]
pub struct RateLimiter {
    bandwidth: BandwidthRateLimiter,
    throughput: ThroughputRateLimiter,
}

impl RateLimiter {
    /// Creates a new rate limiter from the given configuration.
    ///
    /// Background estimation tasks are not started until [`start`](RateLimiter::start) is called.
    pub fn new(config: RateLimits) -> Self {
        Self {
            bandwidth: BandwidthRateLimiter::new(config.bandwidth),
            throughput: ThroughputRateLimiter::new(config.throughput),
        }
    }

    /// Starts background tasks for rate limit estimation and monitoring.
    ///
    /// Must be called from within a Tokio runtime.
    pub fn start(&self) {
        self.bandwidth.start();
        self.throughput.start();
    }

    /// Checks if the given context is within the rate limits.
    ///
    /// Returns `true` if the request is admitted, `false` if it was rejected. On rejection, emits a
    /// `server.request.rate_limited` metric counter and a `warn!` log. Bandwidth is checked before
    /// throughput so that rejected requests are never counted toward admitted traffic.
    pub fn check(&self, context: &ObjectContext) -> bool {
        // Bandwidth is checked first because it is a pure read (no token consumption).
        // Throughput increments the EWMA accumulator only on success, so checking it
        // second ensures rejected requests are never counted toward admitted traffic.
        let rejection = self
            .bandwidth
            .check(context)
            .or_else(|| self.throughput.check(context));

        let Some(rejection) = rejection else {
            return true;
        };

        objectstore_metrics::count!("server.request.rate_limited", reason = rejection.as_str());
        objectstore_log::warn!(
            reason = rejection.as_str(),
            "Request rejected: rate limit exceeded"
        );
        false
    }

    /// Returns all bandwidth accumulators (global + per-usecase + per-scope) for the given context.
    ///
    /// Creates entries in the per-usecase/per-scope maps if they don't exist yet.
    pub fn bytes_accumulators(&self, context: &ObjectContext) -> Vec<Arc<AtomicU64>> {
        self.bandwidth.accumulators(context)
    }

    /// Records bandwidth usage across all accumulators for the given context.
    ///
    /// This is used for cases where bytes are known upfront (e.g. batch INSERT) rather than
    /// streamed through a `MeteredPayloadStream`.
    pub fn record_bandwidth(&self, context: &ObjectContext, bytes: u64) {
        for acc in self.bandwidth.accumulators(context) {
            acc.fetch_add(bytes, std::sync::atomic::Ordering::Relaxed);
        }
    }

    /// Returns the current global bandwidth EWMA in bytes per second.
    pub fn bandwidth_ewma(&self) -> u64 {
        self.bandwidth
            .global
            .estimate
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Returns the configured global bandwidth limit in bytes/s, if set.
    pub fn bandwidth_limit(&self) -> Option<u64> {
        self.bandwidth.config.global_bps
    }

    /// Returns the current estimated throughput in requests per second.
    pub fn throughput_rps(&self) -> u64 {
        self.throughput
            .global_estimator
            .estimate
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Returns the configured global throughput limit in requests/s, if set.
    pub fn throughput_limit(&self) -> Option<u32> {
        self.throughput.config.global_rps
    }

    /// Returns total bytes transferred since startup.
    pub fn bandwidth_total_bytes(&self) -> u64 {
        self.bandwidth
            .total_bytes
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Returns total admitted requests since startup.
    pub fn throughput_total_admitted(&self) -> u64 {
        self.throughput
            .total_admitted
            .load(std::sync::atomic::Ordering::Relaxed)
    }
}

/// Shared EWMA estimator state.
///
/// The accumulator is incremented as events occur, and the estimate is updated
/// periodically by a background task using an exponentially weighted moving average.
#[derive(Debug)]
struct EwmaEstimator {
    accumulator: Arc<AtomicU64>,
    estimate: Arc<AtomicU64>,
}

impl EwmaEstimator {
    fn new() -> Self {
        Self {
            accumulator: Arc::new(AtomicU64::new(0)),
            estimate: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Updates the EWMA from the accumulator.
    ///
    /// Swaps the accumulator to zero, converts the count to a per-second rate using
    /// `to_rate` (i.e., `1.0 / tick_duration.as_secs_f64()`), then applies the
    /// EWMA smoothing and stores the floored result in `estimate`.
    fn update_ewma(&self, ewma: &mut f64, to_rate: f64) {
        const ALPHA: f64 = 0.2;
        let current = self
            .accumulator
            .swap(0, std::sync::atomic::Ordering::Relaxed);
        let rate = (current as f64) * to_rate;
        *ewma = ALPHA * rate + (1.0 - ALPHA) * *ewma;
        self.estimate
            .store(ewma.floor() as u64, std::sync::atomic::Ordering::Relaxed);
    }
}

#[derive(Debug)]
struct BandwidthRateLimiter {
    config: BandwidthLimits,
    /// Global accumulator/estimator pair.
    global: Arc<EwmaEstimator>,
    /// Cumulative bytes transferred since startup. Never reset.
    total_bytes: Arc<AtomicU64>,
    // NB: These maps grow unbounded but we accept this as we expect an overall limited
    // number of usecases and scopes. We emit gauge metrics to monitor their size.
    usecases: Arc<papaya::HashMap<String, Arc<EwmaEstimator>>>,
    scopes: Arc<papaya::HashMap<Scopes, Arc<EwmaEstimator>>>,
}

impl BandwidthRateLimiter {
    fn new(config: BandwidthLimits) -> Self {
        Self {
            config,
            global: Arc::new(EwmaEstimator::new()),
            total_bytes: Arc::new(AtomicU64::new(0)),
            usecases: Arc::new(papaya::HashMap::new()),
            scopes: Arc::new(papaya::HashMap::new()),
        }
    }

    fn start(&self) {
        let global = Arc::clone(&self.global);
        let usecases = Arc::clone(&self.usecases);
        let scopes = Arc::clone(&self.scopes);
        let global_limit = self.config.global_bps;
        // NB: This task has no shutdown mechanism — the rate limiter is only created once.
        // The task is aborted when the Tokio runtime is dropped on process exit.
        tokio::task::spawn(async move {
            Self::estimator(global, usecases, scopes, global_limit).await;
        });
    }

    /// Estimates the current bandwidth utilization using an exponentially weighted moving average.
    ///
    /// Iterates over the global estimator as well as all per-usecase and per-scope estimators
    /// on each tick, updating their EWMAs.
    async fn estimator(
        global: Arc<EwmaEstimator>,
        usecases: Arc<papaya::HashMap<String, Arc<EwmaEstimator>>>,
        scopes: Arc<papaya::HashMap<Scopes, Arc<EwmaEstimator>>>,
        global_limit: Option<u64>,
    ) {
        const TICK: Duration = Duration::from_millis(50); // Recompute EWMA on every TICK

        let mut interval = tokio::time::interval(TICK);
        let to_bps = 1.0 / TICK.as_secs_f64(); // Conversion factor from bytes to bps
        let mut global_ewma: f64 = 0.0;
        // Shadow EWMAs for per-usecase/per-scope entries, keyed the same way as the maps.
        let mut usecase_ewmas: std::collections::HashMap<String, f64> =
            std::collections::HashMap::new();
        let mut scope_ewmas: std::collections::HashMap<Scopes, f64> =
            std::collections::HashMap::new();

        loop {
            interval.tick().await;

            // Global
            global.update_ewma(&mut global_ewma, to_bps);
            objectstore_metrics::gauge!("server.bandwidth.ewma" = global_ewma.floor() as u64);
            if let Some(limit) = global_limit {
                objectstore_metrics::gauge!("server.bandwidth.limit" = limit);
            }

            // Per-usecase
            {
                let guard = usecases.pin();
                for (key, estimator) in guard.iter() {
                    let ewma = usecase_ewmas.entry(key.clone()).or_insert(0.0);
                    estimator.update_ewma(ewma, to_bps);
                }
            }

            // Per-scope
            {
                let guard = scopes.pin();
                for (key, estimator) in guard.iter() {
                    let ewma = scope_ewmas.entry(key.clone()).or_insert(0.0);
                    estimator.update_ewma(ewma, to_bps);
                }
            }

            objectstore_metrics::gauge!(
                "server.rate_limiter.bandwidth.scope_map_size" = scopes.len()
            );
            objectstore_metrics::gauge!(
                "server.rate_limiter.bandwidth.usecase_map_size" = usecases.len()
            );
        }
    }

    fn check(&self, context: &ObjectContext) -> Option<RateLimitRejection> {
        let global_bps = self.config.global_bps?;

        // Global check
        if self
            .global
            .estimate
            .load(std::sync::atomic::Ordering::Relaxed)
            > global_bps
        {
            return Some(RateLimitRejection::BandwidthGlobal);
        }

        // Per-usecase check
        if let Some(usecase_bps) = self.usecase_bps() {
            let guard = self.usecases.pin();
            if let Some(estimator) = guard.get(&context.usecase)
                && estimator
                    .estimate
                    .load(std::sync::atomic::Ordering::Relaxed)
                    > usecase_bps
            {
                return Some(RateLimitRejection::BandwidthUsecase);
            }
        }

        // Per-scope check
        if let Some(scope_bps) = self.scope_bps() {
            let guard = self.scopes.pin();
            if let Some(estimator) = guard.get(&context.scopes)
                && estimator
                    .estimate
                    .load(std::sync::atomic::Ordering::Relaxed)
                    > scope_bps
            {
                return Some(RateLimitRejection::BandwidthScope);
            }
        }

        None
    }

    /// Returns all accumulators (global + per-usecase + per-scope) for the given context.
    ///
    /// Creates entries in the per-usecase/per-scope maps if they don't exist yet.
    /// Always includes `total_bytes` (cumulative, never reset) as the first entry.
    fn accumulators(&self, context: &ObjectContext) -> Vec<Arc<AtomicU64>> {
        let mut accs = vec![
            Arc::clone(&self.total_bytes),
            Arc::clone(&self.global.accumulator),
        ];

        if self.usecase_bps().is_some() {
            let guard = self.usecases.pin();
            let estimator = guard
                .get_or_insert_with(context.usecase.clone(), || Arc::new(EwmaEstimator::new()));
            accs.push(Arc::clone(&estimator.accumulator));
        }

        if self.scope_bps().is_some() {
            let guard = self.scopes.pin();
            let estimator =
                guard.get_or_insert_with(context.scopes.clone(), || Arc::new(EwmaEstimator::new()));
            accs.push(Arc::clone(&estimator.accumulator));
        }

        accs
    }

    /// Returns the effective BPS for per-usecase limiting, if configured.
    fn usecase_bps(&self) -> Option<u64> {
        let global_bps = self.config.global_bps?;
        let pct = self.config.usecase_pct?;
        Some(((global_bps as f64) * (pct as f64 / 100.0)) as u64)
    }

    /// Returns the effective BPS for per-scope limiting, if configured.
    fn scope_bps(&self) -> Option<u64> {
        let global_bps = self.config.global_bps?;
        let pct = self.config.scope_pct?;
        Some(((global_bps as f64) * (pct as f64 / 100.0)) as u64)
    }
}

#[derive(Debug)]
struct ThroughputRateLimiter {
    config: ThroughputLimits,
    global: Option<Mutex<TokenBucket>>,
    /// Global EWMA estimator for admitted request rate.
    global_estimator: Arc<EwmaEstimator>,
    /// Cumulative admitted requests since startup. Never reset.
    total_admitted: Arc<AtomicU64>,
    // NB: These maps grow unbounded but we accept this as we expect an overall limited
    // number of usecases and scopes. We emit gauge metrics to monitor their size.
    usecases: Arc<papaya::HashMap<String, Mutex<TokenBucket>>>,
    scopes: Arc<papaya::HashMap<Scopes, Mutex<TokenBucket>>>,
    rules: papaya::HashMap<usize, Mutex<TokenBucket>>,
}

impl ThroughputRateLimiter {
    fn new(config: ThroughputLimits) -> Self {
        let global = config
            .global_rps
            .map(|rps| Mutex::new(TokenBucket::new(rps, config.burst)));

        Self {
            config,
            global,
            global_estimator: Arc::new(EwmaEstimator::new()),
            total_admitted: Arc::new(AtomicU64::new(0)),
            usecases: Arc::new(papaya::HashMap::new()),
            scopes: Arc::new(papaya::HashMap::new()),
            rules: papaya::HashMap::new(),
        }
    }

    fn start(&self) {
        let usecases = Arc::clone(&self.usecases);
        let scopes = Arc::clone(&self.scopes);
        let global_estimator = Arc::clone(&self.global_estimator);
        let global_limit = self.config.global_rps;
        // NB: This task has no shutdown mechanism — the rate limiter is only created once.
        // The task is aborted when the Tokio runtime is dropped on process exit.
        tokio::task::spawn(async move {
            const TICK: Duration = Duration::from_millis(50);
            let mut interval = tokio::time::interval(TICK);
            let to_rps = 1.0 / TICK.as_secs_f64();
            let mut global_ewma: f64 = 0.0;
            loop {
                interval.tick().await;
                global_estimator.update_ewma(&mut global_ewma, to_rps);
                objectstore_metrics::gauge!("server.throughput.ewma" = global_ewma.floor() as u64);
                if let Some(limit) = global_limit {
                    objectstore_metrics::gauge!(
                        "server.rate_limiter.throughput.limit" = u64::from(limit)
                    );
                }
                objectstore_metrics::gauge!(
                    "server.rate_limiter.throughput.scope_map_size" = scopes.len()
                );
                objectstore_metrics::gauge!(
                    "server.rate_limiter.throughput.usecase_map_size" = usecases.len()
                );
            }
        });
    }

    fn check(&self, context: &ObjectContext) -> Option<RateLimitRejection> {
        // NB: We intentionally use unwrap and crash the server if the mutexes are poisoned.

        // Global check
        if let Some(ref global) = self.global {
            let acquired = global.lock().unwrap().try_acquire();
            if !acquired {
                return Some(RateLimitRejection::ThroughputGlobal);
            }
        }

        // Usecase check - only if both global_rps and usecase_pct are set
        if let Some(usecase_rps) = self.usecase_rps() {
            let guard = self.usecases.pin();
            let bucket = guard
                .get_or_insert_with(context.usecase.clone(), || self.create_bucket(usecase_rps));
            if !bucket.lock().unwrap().try_acquire() {
                return Some(RateLimitRejection::ThroughputUsecase);
            }
        }

        // Scope check - only if both global_rps and scope_pct are set
        if let Some(scope_rps) = self.scope_rps() {
            let guard = self.scopes.pin();
            let bucket =
                guard.get_or_insert_with(context.scopes.clone(), || self.create_bucket(scope_rps));
            if !bucket.lock().unwrap().try_acquire() {
                return Some(RateLimitRejection::ThroughputScope);
            }
        }

        // Rule checks - each matching rule has its own dedicated bucket
        for (idx, rule) in self.config.rules.iter().enumerate() {
            if !rule.matches(context) {
                continue;
            }
            let Some(rule_rps) = self.rule_rps(rule) else {
                continue;
            };
            let guard = self.rules.pin();
            let bucket = guard.get_or_insert_with(idx, || self.create_bucket(rule_rps));
            if !bucket.lock().unwrap().try_acquire() {
                return Some(RateLimitRejection::ThroughputRule);
            }
        }

        // Count this admitted request in the EWMA accumulator and the cumulative counter.
        // NB: u64 wrapping is not a practical concern — at 1M rps it takes ~585k years.
        // Prometheus irate() also handles counter resets gracefully should it ever occur.
        self.global_estimator
            .accumulator
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        self.total_admitted
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        None
    }

    fn create_bucket(&self, rps: u32) -> Mutex<TokenBucket> {
        Mutex::new(TokenBucket::new(rps, self.config.burst))
    }

    /// Returns the effective RPS for per-usecase limiting, if configured.
    fn usecase_rps(&self) -> Option<u32> {
        let global_rps = self.config.global_rps?;
        let pct = self.config.usecase_pct?;
        Some(((global_rps as f64) * (pct as f64 / 100.0)) as u32)
    }

    /// Returns the effective RPS for per-scope limiting, if configured.
    fn scope_rps(&self) -> Option<u32> {
        let global_rps = self.config.global_rps?;
        let pct = self.config.scope_pct?;
        Some(((global_rps as f64) * (pct as f64 / 100.0)) as u32)
    }

    /// Returns the effective RPS for a rule, if it has a valid limit.
    fn rule_rps(&self, rule: &ThroughputRule) -> Option<u32> {
        let pct_limit = rule.pct.and_then(|p| {
            self.config
                .global_rps
                .map(|g| ((g as f64) * (p as f64 / 100.0)) as u32)
        });

        match (rule.rps, pct_limit) {
            (Some(r), Some(p)) => Some(r.min(p)),
            (Some(r), None) => Some(r),
            (None, Some(p)) => Some(p),
            (None, None) => None,
        }
    }
}

/// A token bucket rate limiter.
///
/// Tokens refill at a constant rate up to capacity. Each request consumes one token.
/// When no tokens are available, requests are rejected.
///
/// This implementation is not thread-safe on its own. Wrap in a `Mutex` for concurrent access.
#[derive(Debug)]
struct TokenBucket {
    refill_rate: f64,
    capacity: f64,
    tokens: f64,
    last_update: Instant,
}

impl TokenBucket {
    /// Creates a new, full token bucket with the specified rate limit and burst capacity.
    ///
    /// - `rps`: tokens refilled per second (sustained rate limit)
    /// - `burst`: initial tokens and burst allowance above sustained rate
    pub fn new(rps: u32, burst: u32) -> Self {
        Self {
            refill_rate: rps as f64,
            capacity: (rps + burst) as f64,
            tokens: (rps + burst) as f64,
            last_update: Instant::now(),
        }
    }

    /// Attempts to acquire a token from the bucket.
    ///
    /// Returns `true` if a token was acquired, `false` if no tokens available.
    pub fn try_acquire(&mut self) -> bool {
        let now = Instant::now();
        let refill = now.duration_since(self.last_update).as_secs_f64() * self.refill_rate;
        let refilled = (self.tokens + refill).min(self.capacity);

        // Only apply refill if we'd gain at least 1 whole token
        if refilled.floor() > self.tokens.floor() {
            self.last_update = now;
            self.tokens = refilled;
        }

        // Try to consume one token
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }
}

/// A wrapper around a byte stream that measures bandwidth usage.
///
/// Every time a chunk is polled successfully, all accumulators are incremented
/// by its size. Generic over both the stream type `S` and its error type.
pub(crate) struct MeteredPayloadStream<S> {
    inner: S,
    accumulators: Vec<Arc<AtomicU64>>,
}

impl<S> MeteredPayloadStream<S> {
    pub fn new(inner: S, accumulators: Vec<Arc<AtomicU64>>) -> Self {
        Self {
            inner,
            accumulators,
        }
    }
}

impl<S> fmt::Debug for MeteredPayloadStream<S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("MeteredPayloadStream")
            .field("accumulators", &self.accumulators)
            .finish()
    }
}

impl<S, E> Stream for MeteredPayloadStream<S>
where
    S: Stream<Item = Result<Bytes, E>> + Unpin,
{
    type Item = Result<Bytes, E>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        let res = Pin::new(&mut this.inner).poll_next(cx);
        if let Poll::Ready(Some(Ok(ref bytes))) = res {
            let len = bytes.len() as u64;
            for acc in &this.accumulators {
                acc.fetch_add(len, std::sync::atomic::Ordering::Relaxed);
            }
        }
        res
    }
}

#[cfg(test)]
mod tests {
    use objectstore_service::id::ObjectContext;
    use objectstore_types::scope::{Scope, Scopes};

    use super::*;

    fn make_context() -> ObjectContext {
        ObjectContext {
            usecase: "testing".into(),
            scopes: Scopes::from_iter([Scope::create("org", "1").unwrap()]),
        }
    }

    #[test]
    fn ewma_estimator_update_applies_alpha() {
        let estimator = EwmaEstimator::new();
        const TICK: f64 = 0.05; // 50ms
        let to_rate = 1.0 / TICK;
        let mut ewma: f64 = 0.0;

        // Simulate 10 events in one 50ms tick → 200 /s raw rate.
        // After one step: 0.2 * 200 + 0.8 * 0 = 40.
        estimator
            .accumulator
            .store(10, std::sync::atomic::Ordering::Relaxed);
        estimator.update_ewma(&mut ewma, to_rate);
        assert_eq!(
            estimator
                .estimate
                .load(std::sync::atomic::Ordering::Relaxed),
            40
        );

        // Accumulator must have been zeroed.
        assert_eq!(
            estimator
                .accumulator
                .load(std::sync::atomic::Ordering::Relaxed),
            0
        );
    }

    #[test]
    fn throughput_check_increments_accumulator() {
        let limiter = ThroughputRateLimiter::new(ThroughputLimits {
            global_rps: Some(1000),
            ..Default::default()
        });

        assert_eq!(
            limiter
                .global_estimator
                .accumulator
                .load(std::sync::atomic::Ordering::Relaxed),
            0
        );

        let context = make_context();
        assert!(limiter.check(&context).is_none());
        assert!(limiter.check(&context).is_none());

        assert_eq!(
            limiter
                .global_estimator
                .accumulator
                .load(std::sync::atomic::Ordering::Relaxed),
            2
        );
    }

    #[test]
    fn throughput_rejected_does_not_increment_accumulator() {
        let limiter = ThroughputRateLimiter::new(ThroughputLimits {
            global_rps: Some(1),
            burst: 0,
            ..Default::default()
        });

        let context = make_context();
        // First call admitted (consumes the one token), second rejected.
        assert!(limiter.check(&context).is_none());
        assert!(limiter.check(&context).is_some());

        assert_eq!(
            limiter
                .global_estimator
                .accumulator
                .load(std::sync::atomic::Ordering::Relaxed),
            1 // only the admitted request
        );
    }

    #[test]
    fn bandwidth_rejection_does_not_increment_throughput_accumulator() {
        // global_bps of 1 means the estimate (0 initially) is not > 1, so the first
        // call passes the bandwidth check. Use 0 to guarantee an immediate reject.
        // BandwidthRateLimiter::check rejects when estimate > global_bps, so set
        // global_bps = 0 to make the bandwidth check always reject.
        let limiter = RateLimiter::new(RateLimits {
            throughput: ThroughputLimits {
                global_rps: Some(1000),
                ..Default::default()
            },
            bandwidth: BandwidthLimits {
                global_bps: Some(0),
                ..Default::default()
            },
        });

        // Prime the bandwidth EWMA so it exceeds the limit.
        limiter
            .bandwidth
            .global
            .estimate
            .store(1, std::sync::atomic::Ordering::Relaxed);

        let context = make_context();
        assert!(!limiter.check(&context));

        // The throughput accumulator must still be 0.
        assert_eq!(
            limiter
                .throughput
                .global_estimator
                .accumulator
                .load(std::sync::atomic::Ordering::Relaxed),
            0
        );
    }

    #[test]
    fn rate_limiter_accessors_with_config() {
        let rate_limiter = RateLimiter::new(RateLimits {
            throughput: ThroughputLimits {
                global_rps: Some(500),
                ..Default::default()
            },
            bandwidth: BandwidthLimits {
                global_bps: Some(1_000_000),
                ..Default::default()
            },
        });

        assert_eq!(rate_limiter.bandwidth_limit(), Some(1_000_000));
        assert_eq!(rate_limiter.throughput_limit(), Some(500));
        // Estimates start at 0 before any background tick.
        assert_eq!(rate_limiter.bandwidth_ewma(), 0);
        assert_eq!(rate_limiter.throughput_rps(), 0);
    }

    #[test]
    fn rate_limiter_accessors_no_limits() {
        let rate_limiter = RateLimiter::new(RateLimits::default());

        assert_eq!(rate_limiter.bandwidth_limit(), None);
        assert_eq!(rate_limiter.throughput_limit(), None);
        assert_eq!(rate_limiter.bandwidth_ewma(), 0);
        assert_eq!(rate_limiter.throughput_rps(), 0);
    }
}