//// This module provides a distributed rate limiter that can be used to limit the //// number of requests or function calls per second for a given identifier. //// //// A single actor is used to assign one rate limiter actor per identifier. The //// rate limiter actor then uses a Token Bucket algorithm to determine if a //// request or function call should be allowed to proceed. A separate process is //// polling the rate limiters to remove full buckets to reduce unnecessary memory //// usage. //// //// The rate limits are configured using the following two options: //// //// - `per_second`: The rate of new available tokens per second. Think of this //// as the steady state rate limit. //// - `burst_limit`: The maximum number of available tokens. Think of this as //// the burst rate limit. The default value is the `per_second` rate limit. //// //// The rate limiter can be applied to a function or handler using the `apply` //// function, which returns a new function that checks the rate limit before //// calling the original function. //// //// # Example //// //// ```gleam //// import glimit //// //// let limiter = //// glimit.new() //// |> glimit.per_second(10) //// |> glimit.burst_limit(100) //// |> glimit.identifier(fn(request) { request.ip }) //// |> glimit.on_limit_exceeded(fn(_request) { "Rate limit reached" }) //// |> glimit.build() //// //// let handler = //// fn(_request) { "Hello, world!" } //// |> glimit.apply(limiter) //// ``` //// import gleam/option.{type Option, None, Some} import gleam/result import glimit/rate_limiter import glimit/registry.{type RateLimiterRegistryActor} /// A rate limiter. /// pub type RateLimiter(a, b, id) { RateLimiter( rate_limiter_registry: RateLimiterRegistryActor(id), on_limit_exceeded: fn(a) -> b, identifier: fn(a) -> id, ) } /// A builder for configuring the rate limiter. /// pub type RateLimiterBuilder(a, b, id) { RateLimiterBuilder( per_second: Option(Int), burst_limit: Option(Int), identifier: Option(fn(a) -> id), on_limit_exceeded: Option(fn(a) -> b), ) } /// Create a new rate limiter builder. /// pub fn new() -> RateLimiterBuilder(a, b, id) { RateLimiterBuilder( per_second: None, burst_limit: None, identifier: None, on_limit_exceeded: None, ) } /// Set the rate limit per second. /// /// The value is not only used for the rate at which tokens are added to the bucket, but /// also for the maximum number of available tokens. To set a different value fo the /// maximum number of available tokens, use the `burst_limit` function. /// pub fn per_second( limiter: RateLimiterBuilder(a, b, id), limit: Int, ) -> RateLimiterBuilder(a, b, id) { RateLimiterBuilder(..limiter, per_second: Some(limit)) } /// Set the maximum number of available tokens. /// /// The maximum number of available tokens is the maximum number of requests that can be /// made in a single second. The default value is the same as the rate limit per second. /// pub fn burst_limit( limiter: RateLimiterBuilder(a, b, id), burst_limit: Int, ) -> RateLimiterBuilder(a, b, id) { RateLimiterBuilder(..limiter, burst_limit: Some(burst_limit)) } /// Set the handler to be called when the rate limit is reached. /// pub fn on_limit_exceeded( limiter: RateLimiterBuilder(a, b, id), on_limit_exceeded: fn(a) -> b, ) -> RateLimiterBuilder(a, b, id) { RateLimiterBuilder(..limiter, on_limit_exceeded: Some(on_limit_exceeded)) } /// Set the identifier function to be used to identify the rate limit. /// pub fn identifier( limiter: RateLimiterBuilder(a, b, id), identifier: fn(a) -> id, ) -> RateLimiterBuilder(a, b, id) { RateLimiterBuilder(..limiter, identifier: Some(identifier)) } /// Build the rate limiter. /// /// Panics if the rate limiter registry cannot be started or if the `identifier` /// function or `on_limit_exceeded` function is missing. /// /// To handle errors instead of panicking, use `try_build`. /// pub fn build(config: RateLimiterBuilder(a, b, id)) -> RateLimiter(a, b, id) { case try_build(config) { Ok(limiter) -> limiter Error(message) -> panic as message } } /// Build the rate limiter, but return an error instead of panicking. /// pub fn try_build( config: RateLimiterBuilder(a, b, id), ) -> Result(RateLimiter(a, b, id), String) { use per_second <- result.try(case config.per_second { Some(per_second) -> Ok(per_second) None -> Error("`per_second` rate limit is required") }) let burst_limit = case config.burst_limit { Some(burst_limit) -> burst_limit None -> per_second } use rate_limiter_registry <- result.try( registry.new(per_second, burst_limit) |> result.map_error(fn(_) { "Failed to start rate limiter registry" }), ) use identifier <- result.try(case config.identifier { Some(identifier) -> Ok(identifier) None -> Error("`identifier` function is required") }) use on_limit_exceeded <- result.try(case config.on_limit_exceeded { Some(on_limit_exceeded) -> Ok(on_limit_exceeded) None -> Error("`on_limit_exceeded` function is required") }) Ok(RateLimiter( rate_limiter_registry: rate_limiter_registry, on_limit_exceeded: on_limit_exceeded, identifier: identifier, )) } /// Apply the rate limiter to a request handler or function. /// pub fn apply(func: fn(a) -> b, limiter: RateLimiter(a, b, id)) -> fn(a) -> b { fn(input: a) -> b { let identifier = limiter.identifier(input) case limiter.rate_limiter_registry |> registry.get_or_create(identifier) { Ok(rate_limiter) -> { case rate_limiter |> rate_limiter.hit { Ok(Nil) -> func(input) Error(Nil) -> limiter.on_limit_exceeded(input) } } Error(_) -> panic as "Failed to get rate limiter" } } }