//#![warn(missing_docs)] use bb8_redis::redis::cmd; pub use bb8_redis::redis::RedisError; pub use bb8_redis::{bb8, RedisConnectionManager}; use std::ops::Add; use std::time::{Duration, Instant}; pub type RedisClientPool = bb8::Pool; pub struct RedisCellClient { pool: RedisClientPool, key: String, default_max_burst: u32, default_count_per_period: u32, default_period: u32, } pub enum ThrottleResult { Allowed, RetryAt(Instant), } impl RedisCellClient { // todo: seems like this could be derived // TODO: take something generic for conn // TODO: use r2d2 for connection pooling? pub fn new( pool: RedisClientPool, app: &str, key: &str, default_max_burst: u32, default_count_per_period: u32, default_period: u32, ) -> Self { let key = format!("{}:redis-cell:{}", app, key); Self { pool, key, default_max_burst, default_count_per_period, default_period, } } #[inline] async fn _throttle( &self, key: &str, max_burst: Option, count_per_period: Option, period: Option, quantity: u32, ) -> anyhow::Result { let mut conn = self.pool.get().await?; let count_per_period = count_per_period.unwrap_or(self.default_count_per_period); if count_per_period == 0 { return Ok(ThrottleResult::Allowed); } let max_burst = max_burst.unwrap_or(self.default_max_burst); let period = period.unwrap_or(self.default_period); /* https://github.com/brandur/redis-cell#response CL.THROTTLE [] 0. Whether the action was limited: 0 indicates the action is allowed. 1 indicates that the action was limited/blocked. 1. The total limit of the key (max_burst + 1). This is equivalent to the common X-RateLimit-Limit HTTP header. 2. The remaining limit of the key. Equivalent to X-RateLimit-Remaining. 3. The number of seconds until the user should retry, and always -1 if the action was allowed. Equivalent to Retry-After. 4. The number of seconds until the limit will reset to its maximum capacity. Equivalent to X-RateLimit-Reset. */ // TODO: don't unwrap. maybe return anyhow::Result> // TODO: should we return more error info? let x: Vec = cmd("CL.THROTTLE") .arg(&(key, max_burst, count_per_period, period, quantity)) .query_async(&mut *conn) .await?; // TODO: trace log the result? if x.len() != 5 { return Err(anyhow::anyhow!("unexpected redis result")); } let retry_after = *x.get(3).expect("index exists above"); if retry_after == -1 { Ok(ThrottleResult::Allowed) } else { // don't return a duration, return an instant let retry_at = Instant::now().add(Duration::from_secs(retry_after as u64)); Ok(ThrottleResult::RetryAt(retry_at)) } } #[inline] pub async fn throttle(&self) -> anyhow::Result { self._throttle(&self.key, None, None, None, 1).await } #[inline] pub async fn throttle_key( &self, key: &str, max_burst: Option, count_per_period: Option, period: Option, ) -> anyhow::Result { let key = format!("{}:{}", self.key, key); self._throttle(key.as_ref(), max_burst, count_per_period, period, 1) .await } #[inline] pub async fn throttle_quantity(&self, quantity: u32) -> anyhow::Result { self._throttle(&self.key, None, None, None, quantity).await } }