Optionally use redis for throttles

Cargo.lock

@@ -3096,6 +3096,8 @@ dependencies = [
 name = "throttle"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
+ "mod-redis",
  "once_cell",
  "redis-cell",
  "serde",

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)| {

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;

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"]}

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<Mutex<MemoryStore>> = OnceCell::new();
+static REDIS: OnceCell<RedisConnection> = 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<u64>,
+) -> Result<ThrottleResult, Error> {
+    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 = <Vec<i64> 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<u64>,
 ) -> Result<ThrottleResult, Error> {
-    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)]

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)
+```

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