feat(redis-client): declare ops into a RedisBatch and flush them as one pipeline

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

litellm-rust/crates/redis-client/AGENTS.md

@@ -38,7 +38,10 @@
 - Pipeline ops take one key each. A multi-key op would span slots inside one group
 - Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once
 
-## Tests
+## Batches
+- `RedisBatch` declares ops now and sends them later as one `Pipeline`, so it inherits every pipeline rule above. It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local
+- Each op returns a `Declared<T>` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once
+- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis
 - Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes
 - Tests that only need the public API live in `tests/<module>.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`)
 - Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection

litellm-rust/crates/redis-client/src/batch.rs

@@ -0,0 +1,152 @@
+use std::{
+    io,
+    sync::{Arc, Mutex, MutexGuard},
+    time::Duration,
+};
+
+use redis::{FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value};
+use tokio::sync::oneshot::{self, error::TryRecvError};
+
+use crate::{Error, Expiry, Pipeline, RedisClient, Script};
+
+#[derive(Clone)]
+pub struct RedisBatch(Arc<Shared>);
+
+struct Shared {
+    client: RedisClient,
+    pending: Mutex<Pending>,
+    flushing: tokio::sync::Mutex<()>,
+}
+
+struct Pending {
+    pipeline: Pipeline,
+    replies: Vec<oneshot::Sender<Result<Value, Error>>>,
+}
+
+#[must_use]
+pub struct Declared<T> {
+    batch: RedisBatch,
+    reply: oneshot::Receiver<Result<Value, Error>>,
+    decode: fn(Value) -> Result<T, Error>,
+}
+
+impl RedisBatch {
+    pub(crate) fn new(client: RedisClient) -> Self {
+        let pending = Pending::new(&client);
+        Self(Arc::new(Shared {
+            client,
+            pending: Mutex::new(pending),
+            flushing: tokio::sync::Mutex::new(()),
+        }))
+    }
+
+    pub fn get<T: FromRedisValue>(&self, key: &str) -> Declared<Option<T>> {
+        self.declare(|pipeline| pipeline.get(key), decode)
+    }
+
+    pub fn set(&self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Declared<()> {
+        self.declare(|pipeline| pipeline.set(key, value, ttl), decode)
+    }
+
+    pub fn del(&self, key: &str) -> Declared<usize> {
+        self.declare(|pipeline| pipeline.del(key), decode)
+    }
+
+    pub fn expire(&self, key: &str, ttl: Duration) -> Declared<bool> {
+        self.declare(|pipeline| pipeline.expire(key, ttl), decode)
+    }
+
+    pub fn incr_by(&self, key: &str, delta: i64, expiry: Expiry) -> Declared<i64> {
+        self.declare(|pipeline| pipeline.incr_by(key, delta, expiry), decode)
+    }
+
+    pub fn incr_by_float(&self, key: &str, delta: f64, expiry: Expiry) -> Declared<f64> {
+        self.declare(
+            |pipeline| pipeline.incr_by_float(key, delta, expiry),
+            decode,
+        )
+    }
+
+    pub fn rpush(&self, key: &str, values: impl ToRedisArgs, expiry: Expiry) -> Declared<usize> {
+        self.declare(|pipeline| pipeline.rpush(key, values, expiry), decode)
+    }
+
+    pub fn eval<T: FromRedisValue>(
+        &self,
+        script: &Script,
+        keys: &[&str],
+        args: impl ToRedisArgs,
+    ) -> Declared<T> {
+        self.declare(|pipeline| pipeline.eval(script, keys, args), decode)
+    }
+
+    pub async fn flush(&self) {
+        let _flushing = self.0.flushing.lock().await;
+        let Pending { pipeline, replies } =
+            std::mem::replace(&mut *self.pending(), Pending::new(&self.0.client));
+        if replies.is_empty() {
+            return;
+        }
+        let results = self.0.client.run(&pipeline).await;
+        for (reply, result) in replies.into_iter().zip(results) {
+            let _ = reply.send(result);
+        }
+    }
+
+    fn declare<T>(
+        &self,
+        add: impl FnOnce(Pipeline) -> Pipeline,
+        decode: fn(Value) -> Result<T, Error>,
+    ) -> Declared<T> {
+        let (sender, reply) = oneshot::channel();
+        let mut pending = self.pending();
+        let pipeline = std::mem::replace(&mut pending.pipeline, self.0.client.pipeline());
+        pending.pipeline = add(pipeline);
+        pending.replies.push(sender);
+        Declared {
+            batch: self.clone(),
+            reply,
+            decode,
+        }
+    }
+
+    fn pending(&self) -> MutexGuard<'_, Pending> {
+        self.0
+            .pending
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner())
+    }
+}
+
+impl Pending {
+    fn new(client: &RedisClient) -> Self {
+        Self {
+            pipeline: client.pipeline(),
+            replies: Vec::new(),
+        }
+    }
+}
+
+impl<T> Declared<T> {
+    pub async fn resolve(mut self) -> Result<T, Error> {
+        match self.reply.try_recv() {
+            Ok(reply) => return reply.and_then(self.decode),
+            Err(TryRecvError::Closed) => return Err(abandoned()),
+            Err(TryRecvError::Empty) => {}
+        }
+        self.batch.flush().await;
+        let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned()));
+        reply.and_then(self.decode)
+    }
+}
+
+fn decode<T: FromRedisValue>(value: Value) -> Result<T, Error> {
+    T::from_redis_value(value).map_err(|error| Error::Decode(error.into()))
+}
+
+fn abandoned() -> Error {
+    Error::Transport(RedisError::from(io::Error::new(
+        io::ErrorKind::Interrupted,
+        "the batch flush carrying this op was cancelled",
+    )))
+}

litellm-rust/crates/redis-client/src/client.rs

@@ -9,7 +9,7 @@ use redis::{Cmd, FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Valu
 
 use crate::{
     CallEvent, ConnectError, Error, Expiry, FailureKind, Namespace, Operation, Outcome, Pipeline,
-    RedisConfig, RedisEvents, Script,
+    RedisBatch, RedisConfig, RedisEvents, Script,
     commands::{self, ScriptCall},
     connection::Connection,
 };
@@ -141,6 +141,10 @@ impl RedisClient {
         Pipeline::new(self.namespace.clone())
     }
 
+    pub fn batch(&self) -> RedisBatch {
+        RedisBatch::new(self.clone())
+    }
+
     pub async fn run(&self, pipeline: &Pipeline) -> Vec<Result<Value, Error>> {
         let operation = Operation::Pipeline {
             ops: pipeline.len(),

litellm-rust/crates/redis-client/src/lib.rs

@@ -1,3 +1,4 @@
+mod batch;
 mod client;
 mod commands;
 mod config;
@@ -8,6 +9,7 @@ mod namespace;
 mod pipeline;
 mod script;
 
+pub use batch::{Declared, RedisBatch};
 pub use client::{RedisClient, Ttl};
 pub use commands::Expiry;
 pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts};

litellm-rust/crates/redis-client/tests/batch.rs

@@ -0,0 +1,147 @@
+mod support;
+
+use std::{
+    sync::{Arc, Mutex},
+    time::Duration,
+};
+
+use litellm_circuit_breaker::{State, Transition};
+use litellm_redis_client::{
+    CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Namespace, Operation, RedisClient, RedisConfig,
+    RedisEvents, Script, Timeouts,
+};
+use litellm_redis_testing::{RedisDeployment, Topology};
+use rstest::rstest;
+use support::{KEEP, connect, endpoint};
+
+#[derive(Default)]
+struct RoundTrips(Mutex<Vec<Operation>>);
+
+impl RoundTrips {
+    fn seen(&self) -> Vec<Operation> {
+        self.0.lock().unwrap().clone()
+    }
+}
+
+impl RedisEvents for RoundTrips {
+    fn call(&self, event: &CallEvent) {
+        self.0.lock().unwrap().push(event.operation);
+    }
+
+    fn breaker(&self, _transition: &Transition<State>) {}
+}
+
+fn observed(server: &RedisDeployment) -> (RedisClient, Arc<RoundTrips>) {
+    let trips = Arc::new(RoundTrips::default());
+    let client = RedisClient::new(RedisConfig {
+        endpoint: endpoint(server),
+        credentials: Credentials::NONE,
+        namespace: Namespace::new("ns"),
+        timeouts: Timeouts::default(),
+        breaker: None,
+        events: Some(trips.clone()),
+    })
+    .unwrap();
+    (client, trips)
+}
+
+#[rstest]
+#[tokio::test]
+async fn declared_ops_wait_for_the_flush_and_share_one_round_trip(
+    #[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
+) {
+    let server = RedisDeployment::start(topology).await.unwrap();
+    let (client, trips) = observed(&server);
+    let batch = client.batch();
+    let counter = batch.incr_by("counter", 3, KEEP);
+    let total = batch.incr_by_float("total", 1.5, KEEP);
+    let written = batch.set("value", "v", Duration::from_secs(60));
+    let read = batch.get::<String>("value");
+
+    let outside = connect(&server, "ns", DEFAULT_TIMEOUT);
+    assert_eq!(outside.get::<i64>("counter").await.unwrap(), None);
+    assert!(trips.seen().is_empty());
+
+    batch.flush().await;
+
+    assert_eq!(trips.seen(), [Operation::Pipeline { ops: 4 }]);
+    assert_eq!(counter.resolve().await.unwrap(), 3);
+    assert_eq!(total.resolve().await.unwrap(), 1.5);
+    written.resolve().await.unwrap();
+    assert_eq!(read.resolve().await.unwrap().as_deref(), Some("v"));
+    assert_eq!(trips.seen().len(), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn resolving_before_a_flush_sends_everything_declared_so_far(
+    #[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
+) {
+    let server = RedisDeployment::start(topology).await.unwrap();
+    let (client, trips) = observed(&server);
+    let batch = client.batch();
+    let first = batch.incr_by("a", 1, KEEP);
+    let second = batch.incr_by("b", 2, KEEP);
+
+    assert_eq!(second.resolve().await.unwrap(), 2);
+    let late = batch.incr_by("a", 10, KEEP);
+    assert_eq!(first.resolve().await.unwrap(), 1);
+    assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
+
+    assert_eq!(late.resolve().await.unwrap(), 11);
+    assert_eq!(
+        trips.seen(),
+        [
+            Operation::Pipeline { ops: 2 },
+            Operation::Pipeline { ops: 1 }
+        ]
+    );
+    batch.flush().await;
+    assert_eq!(trips.seen().len(), 2);
+}
+
+#[rstest]
+#[tokio::test]
+async fn each_declared_op_gets_its_own_typed_result(
+    #[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
+) {
+    let server = RedisDeployment::start(topology).await.unwrap();
+    let client = connect(&server, "", DEFAULT_TIMEOUT);
+    client
+        .set("text", "not a number", Duration::from_secs(60))
+        .await
+        .unwrap();
+    let batch = client.batch();
+    let rejected = batch.incr_by("text", 1, KEEP);
+    let undecodable = batch.get::<i64>("text");
+    let missing = batch.get::<String>("missing");
+    let script =
+        batch.eval::<Vec<String>>(&Script::new("return {KEYS[1], ARGV[1]}"), &["k"], "arg");
+    let removed = batch.del("text");
+
+    batch.flush().await;
+
+    assert!(matches!(rejected.resolve().await, Err(Error::Server(_))));
+    assert!(matches!(undecodable.resolve().await, Err(Error::Decode(_))));
+    assert_eq!(missing.resolve().await.unwrap(), None);
+    assert_eq!(script.resolve().await.unwrap(), ["k", "arg"]);
+    assert_eq!(removed.resolve().await.unwrap(), 1);
+}
+
+#[rstest]
+#[tokio::test]
+async fn an_op_whose_flush_was_cancelled_resolves_as_transport(
+    #[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
+) {
+    let server = RedisDeployment::start(topology).await.unwrap();
+    let client = connect(&server, "", Duration::from_secs(5));
+    let batch = client.batch();
+    let counter = batch.incr_by("counter", 1, KEEP);
+
+    server.freeze().await.unwrap();
+    let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await;
+    server.thaw().await.unwrap();
+
+    assert!(cancelled.is_err());
+    assert!(matches!(counter.resolve().await, Err(Error::Transport(_))));
+}