pdk-rate-limit-lib 1.7.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! Rate limiting buckets
//!
//! Module which provides the bucket implementations for the rate limiting functionality.
//! It supports both local and distributed (cluster) rate limiting with configurable
//! tiers and quota management. This module exposes:
//!
//! - [`Bucket`]: enum representing either local or cluster rate limiting buckets
//! - [`BucketFactory`]: factory to create bucket instances based on configuration
//! - [`RequestAllowed`]: result of rate limit checks (`Allowed`, `OutOfLocalQuota`, `OutOfQuota`)
//! - [`LimitStats`]: metadata about current quota state (remaining, limit, reset time)
//! - [`QuotaInfo`]: quota information for distributed rate limiting coordination

use cluster::ClusterBucket;
use local::LocalBucket;
use pdk_core::log::error;
use pdk_core::policy_context::api::Tier;
use serde::{Deserialize, Serialize};
use Bucket::{Cluster, Local};

use super::distribution_formula::DistributionFormula;

pub mod cluster;
pub mod local;

/// Factory to create bucket instances based on configuration.
/// Cluster mode true represents distributed rate limiting.
/// Tiers configure the rate limiting rules.
#[derive(Debug)]
pub struct BucketFactory {
    cluster: bool,
    tiers: Vec<(String, Vec<Tier>)>,
}

impl BucketFactory {
    /// Create a new factory to create buckets with configuration received.
    pub fn new(cluster: bool, tiers: Vec<(String, Vec<Tier>)>) -> Self {
        BucketFactory { cluster, tiers }
    }

    /// Create a new bucket with configuration received.
    pub fn create(&self, now: u128, group_selector: &str) -> Bucket {
        let tiers = self
            .tiers
            .iter()
            .find(|item| item.0.eq(group_selector))
            .map(|item| item.1.as_slice())
            .unwrap_or(&[]);

        match self.cluster {
            true => Cluster(ClusterBucket::new(now, tiers)),
            false => Local(LocalBucket::new(now, tiers)),
        }
    }
}

/// Enum representing either local or cluster rate limiting buckets.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum Bucket {
    Cluster(ClusterBucket),
    Local(LocalBucket),
}

impl Bucket {
    /// To use in case of the bucket being used in cluster mode
    pub fn request_allowed(&mut self, now: u128, amount: usize) -> RequestAllowed {
        match self {
            Cluster(bucket) => bucket.request_allowed(now, amount),
            Local(bucket) => bucket.request_allowed(now, amount),
        }
    }

    /// Get current status of the bucket.
    pub fn status(&self) -> Option<LimitStats> {
        match self {
            Cluster(bucket) => bucket.status(),
            Local(bucket) => bucket.status(),
        }
    }

    /// To use in the primary bucket to get quota.
    pub fn get_quota(
        &mut self,
        now: u128,
        bucket: &Bucket,
        formula: &DistributionFormula,
        amount: usize,
    ) -> Vec<QuotaInfo> {
        match (self, bucket) {
            (Cluster(self_bucket), Cluster(bucket)) => {
                self_bucket.get_quota(now, bucket, formula, amount)
            }
            _ => {
                //This should be unreachable
                error!("Unexpected error: Quota request for non cluster buckets");
                vec![]
            }
        }
    }

    /// To use in the secondary bucket to get quota.
    pub fn update_quota(&mut self, quota: &[QuotaInfo]) {
        match self {
            Cluster(bucket) => bucket.update_quota(quota),
            Local(_) => error!("Unexpected error: Quota update for non cluster buckets"),
        }
    }
}

/// Enum representing the result of a rate limit check.
pub enum RequestAllowed {
    Allowed,
    OutOfLocalQuota,
    OutOfQuota,
}

/// Metadata about the current quota state.
#[derive(Clone)]
pub struct LimitStats {
    left: u64,
    max: u64,
    reset: u128,
}

impl LimitStats {
    fn new(left: u64, max: u64, reset: u128) -> Self {
        Self { left, max, reset }
    }

    fn most_restrictive(stat1: LimitStats, stat2: LimitStats) -> LimitStats {
        match (
            (stat1.left == 0, stat1.reset),
            (stat2.left == 0, stat2.reset),
        ) {
            ((true, _), (false, _)) => stat1,
            ((false, _), (true, _)) => stat2,
            ((_, reset1), (_, reset2)) if reset1 < reset2 => stat1,
            _ => stat2,
        }
    }

    /// Get the remaining quota.
    pub fn remaining(&self) -> u64 {
        self.left
    }

    /// Get the limit.
    pub fn limit(&self) -> u64 {
        self.max
    }

    /// Get the time until the next reset.
    pub fn reset(&self, now: u128) -> u128 {
        if now >= self.reset {
            1u128
        } else {
            self.reset - now
        }
    }
}

/// Quota information for distributed rate limiting coordination.
#[derive(Debug)]
pub struct QuotaInfo {
    given: u64,
    remaining: u64,
    reset: u128,
}

impl QuotaInfo {
    /// Create a new quota info.
    pub fn new(given: u64, remaining: u64, reset: u128) -> Self {
        QuotaInfo {
            given,
            remaining,
            reset,
        }
    }

    /// Check if quota is there's any quota available.
    pub fn is_some(&self) -> bool {
        self.given > 0
    }
}