pdk_rate_limit_lib/

lib.rs

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

//! PDK Rate Limit Library
//!
//! Provides rate limiting functionality with support for both local and distributed
//! storage backends.
//!
//! It allows to control the rate at which operations can be performed by
//! implementing token bucket algorithms with configurable limits and time windows.
//!
//! **Local mode** handles all rate limiting quota in-memory, in a single node.
//!
//! **Clustered mode** enables distributed rate limiting across multiple nodes, coordinating quota and statistics.
//!
//! ## Highlights
//!
//! - Local and distributed rate limiting
//! - Configurable quota limits and time windows
//! - Support for grouped rate limits with selectors
//! - Asynchronous API for high-performance applications
//!
//! ## Primary types
//!
//! - [`RateLimit`]: trait which checks if a request should be allowed based on rules
//! - [`RateLimitResult`]: result indicating if request is allowed or rejected
//! - [`RateLimitStatistics`]: metadata about remaining quota and limits
//! - [`RateLimitBuilder`]: instances builder with configurable options (e.g. clustering, custom buckets, distributed storage)
//! - [`RateLimitError`]: error type for rate limiting operations

use data_storage_lib::ll::distributed::DistributedStorageError;
use thiserror::Error;

use data_storage_lib::ll::local::LocalStorageError;

mod bucket;
mod builder;
mod distribution_formula;
mod implementation;
mod key_manager;

pub use builder::{RateLimitBuilder, RateLimitBuilderInstance};
pub use implementation::RateLimitInstance;

use crate::bucket::Bucket;

/// Metadata about the remaining quota for a particular group and key pair.
///
/// This struct provides information about the current state of a rate limit bucket,
/// including how much quota remains, the total limit, and when the quota will reset.
pub struct RateLimitStatistics {
    /// The amount of quota still available.
    ///
    /// For cluster scenarios, to improve performance, this value is a best effort
    /// approximation and may not be perfectly accurate across all nodes.
    pub remaining: u64,

    /// The maximum amount of quota that can be consumed in a given time frame.
    ///
    /// This represents the total capacity of the rate limit bucket.
    pub limit: u64,

    /// The amount of time in milliseconds until more quota is refreshed.
    ///
    /// This indicates when the rate limit window will reset and quota will be replenished.
    pub reset: u64,
}

/// This enum represents the result of a rate limit operation, indicating whether a
/// request should be allowed or rejected based on the current rate limit state,
/// along with statistics about the rate limit bucket.
pub enum RateLimitResult {
    /// The request is allowed to proceed.
    ///
    /// Contains statistics about the current state of the rate limit bucket.
    Allowed(RateLimitStatistics),

    /// The request should be rejected due to rate limit exceeded.
    ///
    /// Contains statistics about the current state of the rate limit bucket,
    /// including when the limit will reset.
    TooManyRequests(RateLimitStatistics),
}

/// Errors that can occur during rate limiting operations.
///
/// This enum represents all possible error conditions that can arise when
/// performing rate limit checks or managing rate limit state.
#[non_exhaustive]
#[derive(Debug, Error)]
pub enum RateLimitError {
    /// An unexpected error occurred during rate limiting operations.
    ///
    /// This wraps any underlying errors from storage operations or other
    /// unexpected conditions.
    #[error("Unexpected error {0}.")]
    Unexpected(Box<dyn std::error::Error>),

    /// The maximum allowed number of hops was reached in distributed rate limiting.
    ///
    /// This error occurs when the distributed rate limiting algorithm
    /// has reached its maximum number of outgoing requests while trying to find available quota.
    #[error("Max Hops.")]
    MaxHops,

    /// A serialization error occurred.
    ///
    /// This error occurs when there is an error serializing or deserializing data.
    #[error("Serialization error: {0}.")]
    Serialization(#[from] bincode::Error),
}

/// RateLimit trait defines the core interface for rate limiting operations.
/// Implementations can use different storage backends (local, distributed)
/// and different algorithms while providing a consistent API.
#[allow(async_fn_in_trait)]
pub trait RateLimit {
    /// Check if a request should be allowed based on rate limiting rules.
    ///
    /// # Arguments
    ///
    /// * `group_selector` - A string identifying the rate limit group. This allows
    ///   different rate limits to be applied to different categories of requests.
    /// * `bucket_selector` - A string identifying the specific bucket within the group.
    ///   This is typically used to identify individual users, IPs, or other entities.
    /// * `increment` - The amount of quota to consume for this request. Typically 1,
    ///   but can be higher for operations that should consume more quota.
    ///
    /// # Returns
    ///
    /// Returns a [`RateLimitResult`] indicating whether the request should be allowed
    /// or rejected, along with statistics about the current rate limit state.
    ///
    /// # Errors
    ///
    /// Returns a [`RateLimitError`] if there was an error checking the rate limit,
    /// such as storage errors or network issues in distributed scenarios.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use rate_limit_lib::{RateLimit, RateLimitResult};
    /// # async fn example(rate_limiter: impl RateLimit) -> Result<(), rate_limit_lib::RateLimitError> {
    /// let result = rate_limiter.is_allowed("sla1", "550e8400-e29b-41d4-a716-446655440000", 1).await?;
    /// match result {
    ///     RateLimitResult::Allowed(stats) => {
    ///         println!("Request allowed. Remaining: {}", stats.remaining);
    ///         // Process the request
    ///     }
    ///     RateLimitResult::TooManyRequests(stats) => {
    ///         println!("Rate limit exceeded. Reset in: {}ms", stats.reset);
    ///         // Reject the request
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    async fn is_allowed(
        &self,
        group_selector: &str,
        bucket_selector: &str,
        increment: usize,
    ) -> Result<RateLimitResult, RateLimitError>;
}

impl From<LocalStorageError> for RateLimitError {
    fn from(value: LocalStorageError) -> Self {
        RateLimitError::Unexpected(value.into())
    }
}

impl From<DistributedStorageError> for RateLimitError {
    fn from(value: DistributedStorageError) -> Self {
        RateLimitError::Unexpected(value.into())
    }
}

impl RateLimitStatistics {
    fn from(bucket: &Bucket, now: u128) -> Self {
        if let Some(status) = bucket.status() {
            Self {
                remaining: status.remaining(),
                limit: status.limit(),
                reset: status.reset(now) as u64,
            }
        } else {
            Self {
                remaining: 0,
                limit: 0,
                reset: 0,
            }
        }
    }
}