diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index 8be4af65eab..78c9593aa17 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/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` 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/.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 diff --git a/litellm-rust/crates/redis-client/src/batch.rs b/litellm-rust/crates/redis-client/src/batch.rs new file mode 100644 index 00000000000..37955e1350c --- /dev/null +++ b/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); + +struct Shared { + client: RedisClient, + pending: Mutex, + flushing: tokio::sync::Mutex<()>, +} + +struct Pending { + pipeline: Pipeline, + replies: Vec>>, +} + +#[must_use] +pub struct Declared { + batch: RedisBatch, + reply: oneshot::Receiver>, + decode: fn(Value) -> Result, +} + +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(&self, key: &str) -> Declared> { + 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 { + self.declare(|pipeline| pipeline.del(key), decode) + } + + pub fn expire(&self, key: &str, ttl: Duration) -> Declared { + self.declare(|pipeline| pipeline.expire(key, ttl), decode) + } + + pub fn incr_by(&self, key: &str, delta: i64, expiry: Expiry) -> Declared { + self.declare(|pipeline| pipeline.incr_by(key, delta, expiry), decode) + } + + pub fn incr_by_float(&self, key: &str, delta: f64, expiry: Expiry) -> Declared { + self.declare( + |pipeline| pipeline.incr_by_float(key, delta, expiry), + decode, + ) + } + + pub fn rpush(&self, key: &str, values: impl ToRedisArgs, expiry: Expiry) -> Declared { + self.declare(|pipeline| pipeline.rpush(key, values, expiry), decode) + } + + pub fn eval( + &self, + script: &Script, + keys: &[&str], + args: impl ToRedisArgs, + ) -> Declared { + 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( + &self, + add: impl FnOnce(Pipeline) -> Pipeline, + decode: fn(Value) -> Result, + ) -> Declared { + 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 Declared { + pub async fn resolve(mut self) -> Result { + 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(value: Value) -> Result { + 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", + ))) +} diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index d7fbdf5da21..0e3b37a8c67 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/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> { let operation = Operation::Pipeline { ops: pipeline.len(), diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 91e6530cec9..f46720030e6 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/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}; diff --git a/litellm-rust/crates/redis-client/tests/batch.rs b/litellm-rust/crates/redis-client/tests/batch.rs new file mode 100644 index 00000000000..4bf5f6fc911 --- /dev/null +++ b/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>); + +impl RoundTrips { + fn seen(&self) -> Vec { + 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) {} +} + +fn observed(server: &RedisDeployment) -> (RedisClient, Arc) { + 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::("value"); + + let outside = connect(&server, "ns", DEFAULT_TIMEOUT); + assert_eq!(outside.get::("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::("text"); + let missing = batch.get::("missing"); + let script = + batch.eval::>(&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(_)))); +}