From 0ca585bd5aee4cd99ef51db86032873bd2534fda Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Sun, 5 Mar 2023 10:39:56 -0700 Subject: [PATCH] Optionally use redis for throttles --- Cargo.lock | 2 + crates/kumod/src/mod_kumo.rs | 12 ++++- crates/mod-redis/src/lib.rs | 6 +-- crates/throttle/Cargo.toml | 2 + crates/throttle/src/lib.rs | 49 ++++++++++++++++++- .../kumo/configure_redis_throttles.md | 19 +++++++ simple_policy.lua | 3 ++ 7 files changed, 87 insertions(+), 6 deletions(-) create mode 100644 docs/reference/kumo/configure_redis_throttles.md diff --git a/Cargo.lock b/Cargo.lock index c55c3431..b99d6c43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3096,6 +3096,8 @@ dependencies = [ name = "throttle" version = "0.1.0" dependencies = [ + "anyhow", + "mod-redis", "once_cell", "redis-cell", "serde", diff --git a/crates/kumod/src/mod_kumo.rs b/crates/kumod/src/mod_kumo.rs index b283b45e..f11d1318 100644 --- a/crates/kumod/src/mod_kumo.rs +++ b/crates/kumod/src/mod_kumo.rs @@ -5,8 +5,9 @@ use crate::lifecycle::LifeCycle; use crate::logging::{ClassifierParams, LogFileParams}; use crate::queue::QueueConfig; use crate::smtp_server::{EsmtpListenerParams, RejectError}; -use config::get_or_create_module; +use config::{any_err, get_or_create_module}; use mlua::{Function, Lua, LuaSerdeExt, Value}; +use mod_redis::RedisConnKey; use serde::Deserialize; use std::path::PathBuf; @@ -81,6 +82,15 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { })?, )?; + kumo_mod.set( + "configure_redis_throttles", + lua.create_async_function(|lua, params: Value| async move { + let key: RedisConnKey = lua.from_value(params)?; + let conn = key.open().await.map_err(any_err)?; + throttle::use_redis(conn).map_err(any_err) + })?, + )?; + kumo_mod.set( "reject", lua.create_function(move |_lua, (code, message): (u16, String)| { diff --git a/crates/mod-redis/src/lib.rs b/crates/mod-redis/src/lib.rs index fdcb485e..e38e6580 100644 --- a/crates/mod-redis/src/lib.rs +++ b/crates/mod-redis/src/lib.rs @@ -4,10 +4,8 @@ use mlua::{Lua, LuaSerdeExt, MultiValue, UserData, UserDataMethods, Value}; use once_cell::sync::Lazy; use r2d2::{ManageConnection, Pool, PooledConnection}; use redis::cluster::{ClusterClient, ClusterConnection}; -use redis::{ - Client, Cmd, Connection, ConnectionLike, RedisError, RedisWrite, ToRedisArgs, - Value as RedisValue, -}; +use redis::{Client, Connection, ConnectionLike, RedisWrite, ToRedisArgs}; +pub use redis::{Cmd, FromRedisValue, RedisError, Value as RedisValue}; use serde::Deserialize; use serde_json::Value as JsonValue; use std::collections::HashMap; diff --git a/crates/throttle/Cargo.toml b/crates/throttle/Cargo.toml index d25e893e..2ae18d8c 100644 --- a/crates/throttle/Cargo.toml +++ b/crates/throttle/Cargo.toml @@ -6,6 +6,8 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +anyhow = "1.0" +mod-redis = {path="../mod-redis"} once_cell = "1.17" redis-cell = { git = "https://github.com/brandur/redis-cell.git" } serde = {version="1.0", features=["derive"]} diff --git a/crates/throttle/src/lib.rs b/crates/throttle/src/lib.rs index da541ae6..73baf498 100644 --- a/crates/throttle/src/lib.rs +++ b/crates/throttle/src/lib.rs @@ -2,6 +2,7 @@ //! The implementation uses an in-memory store, but can be adjusted in the future //! to support using a redis-cell equipped redis server to share the throttles //! among multiple machines. +use mod_redis::{Cmd, FromRedisValue, RedisConnection, RedisError}; use once_cell::sync::OnceCell; use redis_cell::cell::store::MemoryStore; use redis_cell::cell::{Rate, RateLimiter, RateQuota}; @@ -12,11 +13,16 @@ use std::time::Duration; use thiserror::Error; static MEMORY: OnceCell> = OnceCell::new(); +static REDIS: OnceCell = OnceCell::new(); #[derive(Error, Debug)] pub enum Error { #[error("{0}")] Generic(String), + #[error("{0}")] + AnyHow(#[from] anyhow::Error), + #[error("{0}")] + Redis(#[from] RedisError), } #[derive(Debug, Eq, PartialEq, Clone, Copy, Serialize, Deserialize)] @@ -151,6 +157,36 @@ fn local_throttle( }) } +async fn redis_throttle( + conn: RedisConnection, + key: &str, + limit: u64, + period: Duration, + max_burst: u64, + quantity: Option, +) -> Result { + let mut cmd = Cmd::new(); + cmd.arg("CL.THROTTLE") + .arg(key) + .arg(max_burst) + .arg(limit) + .arg(period.as_secs()) + .arg(quantity.unwrap_or(1)); + let result = conn.query(cmd).await?; + let result = as FromRedisValue>::from_redis_value(&result)?; + + Ok(ThrottleResult { + throttled: result[0] != 0, + limit: result[1] as u64, + remaining: result[2] as u64, + reset_after: Duration::from_secs(result[3].max(0) as u64), + retry_after: match result[4] { + n if n < 0 => None, + n => Some(Duration::from_secs(n as u64)), + }, + }) +} + /// It is very important for `key` to be used with the same `limit`, /// `period` and `max_burst` values in order to produce meaningful /// results. @@ -175,7 +211,18 @@ pub async fn throttle( max_burst: u64, quantity: Option, ) -> Result { - local_throttle(key, limit, period, max_burst, quantity) + if let Some(redis) = REDIS.get().cloned() { + redis_throttle(redis, key, limit, period, max_burst, quantity).await + } else { + local_throttle(key, limit, period, max_burst, quantity) + } +} + +pub fn use_redis(conn: RedisConnection) -> Result<(), Error> { + REDIS + .set(conn) + .map_err(|_| Error::Generic("redis already configured for throttles".to_string()))?; + Ok(()) } #[cfg(test)] diff --git a/docs/reference/kumo/configure_redis_throttles.md b/docs/reference/kumo/configure_redis_throttles.md new file mode 100644 index 00000000..3c264fa2 --- /dev/null +++ b/docs/reference/kumo/configure_redis_throttles.md @@ -0,0 +1,19 @@ +# `kumo.configure_redis_throttles { PARAMS }` + +Configure the throttle layer to use a [Redis](https://redis.io/) data store to +manage throttling across multiple MTA nodes. + +The redis server must have [redis-cell](https://github.com/brandur/redis-cell) +installed for throttles to work in this way. + +*PARAMS* behaves exactly as described in [redis.open](../redis/open.md). + +This function should be called only from inside your [init](../events/init.md) +event handler. + +```lua +kumo.on('init', function() + -- Use shared throttles rather than in-process throttles + kumo.configure_redis_throttles { node = 'redis://my-redis-host/' } +end) +``` diff --git a/simple_policy.lua b/simple_policy.lua index 2bc40b9b..0826faa6 100644 --- a/simple_policy.lua +++ b/simple_policy.lua @@ -76,6 +76,9 @@ kumo.on('init', function() kind = 'RocksDB', } + -- Use shared throttles rather than in-process throttles + -- kumo.configure_redis_throttles { node = 'redis://127.0.0.1/' } + -- Create some example sources local entries = {} for i = 1, 10 do