cloudiful-rate-limiter 0.1.0

use crate::error::{RateLimitErrorKind, ValkeyError};
use crate::valkey_scripts;
use crate::{
    AcquireResult, PeekResult, RateLimitError, RateLimitPolicy, RateLimiter, TryAcquireResult,
};
use async_trait::async_trait;
use redis::{Client, ErrorKind, ServerErrorKind, aio::ConnectionManager};
use std::fmt::{self, Display, Formatter};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

const DEFAULT_KEY_PREFIX: &str = "rate-limiter:";

#[derive(Debug, Clone)]
pub struct ValkeyRateLimiter {
    connection: ConnectionManager,
    prefix: String,
    interval: Duration,
}

impl ValkeyRateLimiter {
    pub async fn new(
        url: impl AsRef<str>,
        prefix: impl Into<String>,
        policy: RateLimitPolicy,
    ) -> Result<Self, RateLimitError> {
        let interval = policy.interval()?;
        let client = Client::open(url.as_ref())
            .map_err(|error| RateLimitError::backend(classify_redis_error(&error), error))?;
        let connection = client
            .get_connection_manager()
            .await
            .map_err(|error| RateLimitError::backend(classify_redis_error(&error), error))?;
        Ok(Self {
            connection,
            prefix: normalize_prefix(prefix.into()),
            interval,
        })
    }

    pub async fn with_default_prefix(
        url: impl AsRef<str>,
        policy: RateLimitPolicy,
    ) -> Result<Self, RateLimitError> {
        Self::new(url, DEFAULT_KEY_PREFIX, policy).await
    }

    fn namespaced_key(&self, key: &str) -> String {
        format!("{}{}", self.prefix, key)
    }

    async fn run_script(&self, key: &str, consume: bool) -> Result<ScriptResult, RateLimitError> {
        let mut connection = self.connection.clone();
        let now_millis = now_millis()
            .map_err(|error| RateLimitError::backend(RateLimitErrorKind::Time, error))?;
        let interval_millis = u64::try_from(self.interval.as_millis())
            .map_err(ValkeyError::IntervalOutOfRange)
            .map_err(|error| RateLimitError::backend(RateLimitErrorKind::Backend, error))?;
        let value: (i64, i64) = valkey_scripts::script(valkey_scripts::ACQUIRE_OR_PEEK)
            .key(self.namespaced_key(key))
            .arg(now_millis)
            .arg(interval_millis)
            .arg(if consume { 1 } else { 0 })
            .invoke_async(&mut connection)
            .await
            .map_err(|error| RateLimitError::backend(classify_redis_error(&error), error))?;

        Ok(ScriptResult {
            acquired: value.0 == 1,
            wait_millis: value.1.max(0) as u64,
        })
    }
}

#[async_trait]
impl RateLimiter for ValkeyRateLimiter {
    async fn acquire(&self, key: &str) -> Result<AcquireResult, RateLimitError> {
        let mut total_waited = Duration::ZERO;
        loop {
            let result = self.run_script(key, true).await?;
            if result.acquired {
                return Ok(AcquireResult {
                    key: key.to_owned(),
                    waited: total_waited,
                });
            }

            let wait_duration = Duration::from_millis(result.wait_millis);
            total_waited += wait_duration;
            tokio::time::sleep(wait_duration).await;
        }
    }

    async fn try_acquire(&self, key: &str) -> Result<TryAcquireResult, RateLimitError> {
        let result = self.run_script(key, true).await?;
        if result.acquired {
            Ok(TryAcquireResult::Acquired)
        } else {
            Ok(TryAcquireResult::Limited {
                wait_duration: Duration::from_millis(result.wait_millis),
            })
        }
    }

    async fn peek(&self, key: &str) -> Result<PeekResult, RateLimitError> {
        let result = self.run_script(key, false).await?;
        Ok(PeekResult {
            allowed: result.acquired,
            wait_duration: Duration::from_millis(result.wait_millis),
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct ScriptResult {
    acquired: bool,
    wait_millis: u64,
}

#[derive(Debug)]
struct SystemClockError(std::time::SystemTimeError);

impl Display for SystemClockError {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl std::error::Error for SystemClockError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        Some(&self.0)
    }
}

fn now_millis() -> Result<u64, SystemClockError> {
    let duration = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_err(SystemClockError)?;
    Ok(duration.as_millis().min(u128::from(u64::MAX)) as u64)
}

fn normalize_prefix(prefix: String) -> String {
    if prefix.is_empty() {
        DEFAULT_KEY_PREFIX.to_owned()
    } else {
        prefix
    }
}

fn classify_redis_error(error: &redis::RedisError) -> RateLimitErrorKind {
    if error.is_connection_dropped()
        || error.is_connection_refusal()
        || error.is_timeout()
        || matches!(
            error.kind(),
            ErrorKind::Io
                | ErrorKind::ClusterConnectionNotFound
                | ErrorKind::Server(ServerErrorKind::BusyLoading)
                | ErrorKind::Server(ServerErrorKind::ClusterDown)
                | ErrorKind::Server(ServerErrorKind::MasterDown)
                | ErrorKind::Server(ServerErrorKind::TryAgain)
        )
    {
        RateLimitErrorKind::Connection
    } else {
        RateLimitErrorKind::Backend
    }
}