From af0d29030e4db5f63d6b5125816fb7d45b040688 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Mon, 26 Oct 2026 18:41:46 -0700 Subject: [PATCH] feat(redis-client): run Lua scripts and pipelines with per-op errors (#48273) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/redis-client/AGENTS.md | 8 + litellm-rust/crates/redis-client/Cargo.toml | 1 + .../crates/redis-client/src/client.rs | 247 +++++---------- .../crates/redis-client/src/commands.rs | 58 ++++ .../crates/redis-client/src/connection.rs | 80 +++++ litellm-rust/crates/redis-client/src/lib.rs | 9 + .../crates/redis-client/src/pipeline.rs | 288 ++++++++++++++++++ .../crates/redis-client/src/script.rs | 88 ++++++ .../crates/redis-client/src/test_support.rs | 101 ++++++ 10 files changed, 707 insertions(+), 174 deletions(-) create mode 100644 litellm-rust/crates/redis-client/src/commands.rs create mode 100644 litellm-rust/crates/redis-client/src/connection.rs create mode 100644 litellm-rust/crates/redis-client/src/pipeline.rs create mode 100644 litellm-rust/crates/redis-client/src/script.rs create mode 100644 litellm-rust/crates/redis-client/src/test_support.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f31c7e9cf34..d3bb1ee3813 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4299,6 +4299,7 @@ dependencies = [ name = "litellm-redis-client" version = "0.1.0" dependencies = [ + "futures-util", "litellm-redis-testing", "redis", "rstest", diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index 1e6f53bff7e..e0f1a2da392 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -21,6 +21,14 @@ - On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test - TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`) +## Scripts and pipelines +- A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production +- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async` +- A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip +- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order +- 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 - Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes - Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning) diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index 59e4cbe55b8..e15bbb6e6a5 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +futures-util.workspace = true redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } thiserror.workspace = true diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 7df258f1733..0de40a058a2 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,29 +1,18 @@ use std::time::Duration; -use redis::{ - Cmd, ConnectionInfo, FromRedisValue, IntoConnectionInfo, RedisResult, ToRedisArgs, - ToSingleRedisArg, - aio::{ConnectionManager, ConnectionManagerConfig}, - cluster::ClusterClientBuilder, - cluster_async::ClusterConnection, +use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value}; + +use crate::{ + ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands, + connection::Connection, }; -use crate::{ClusterEndpoint, ConnectError, Endpoint, Error, Namespace, RedisConfig}; - -const RETRIES: u32 = 3; - #[derive(Clone)] pub struct RedisClient { - connection: Connection, + pub(crate) connection: Connection, namespace: Namespace, } -#[derive(Clone)] -enum Connection { - Standalone(ConnectionManager), - Cluster(ClusterConnection), -} - #[derive(Debug, PartialEq, Eq)] pub enum Ttl { Missing, @@ -35,10 +24,12 @@ impl RedisClient { pub async fn connect(config: RedisConfig) -> Result { let connection = match &config.endpoint { Endpoint::Url(url) => { - standalone(url.as_str().into_connection_info()?, config.timeout).await? + Connection::standalone(url.as_str().into_connection_info()?, config.timeout).await? } - Endpoint::Host(host) => standalone(host.connection_info()?, config.timeout).await?, - Endpoint::Cluster(cluster) => self::cluster(cluster, config.timeout).await?, + Endpoint::Host(host) => { + Connection::standalone(host.connection_info()?, config.timeout).await? + } + Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout).await?, }; Ok(Self { connection, @@ -47,14 +38,14 @@ impl RedisClient { } pub async fn get(&self, key: &str) -> Result, Error> { - self.query(redis::cmd("GET").arg(self.key(key))).await + self.query(&commands::get(self.key(key))).await } pub async fn mget(&self, keys: &[&str]) -> Result>, Error> { if keys.is_empty() { return Ok(Vec::new()); } - self.query(redis::cmd("MGET").arg(self.keys(keys))).await + self.query(&commands::mget(self.keys(keys))).await } pub async fn set( @@ -63,14 +54,7 @@ impl RedisClient { value: impl ToSingleRedisArg, ttl: Duration, ) -> Result<(), Error> { - self.query( - redis::cmd("SET") - .arg(self.key(key)) - .arg(value) - .arg("PX") - .arg(millis(ttl)), - ) - .await + self.query(&commands::set(self.key(key), value, ttl)).await } pub async fn set_nx( @@ -80,14 +64,7 @@ impl RedisClient { ttl: Duration, ) -> Result { let reply: Option = self - .query( - redis::cmd("SET") - .arg(self.key(key)) - .arg(value) - .arg("NX") - .arg("PX") - .arg(millis(ttl)), - ) + .query(&commands::set_nx(self.key(key), value, ttl)) .await?; Ok(reply.is_some()) } @@ -96,16 +73,15 @@ impl RedisClient { if keys.is_empty() { return Ok(0); } - self.query(redis::cmd("DEL").arg(self.keys(keys))).await + self.query(&commands::del(self.keys(keys))).await } pub async fn expire(&self, key: &str, ttl: Duration) -> Result { - self.query(redis::cmd("PEXPIRE").arg(self.key(key)).arg(millis(ttl))) - .await + self.query(&commands::expire(self.key(key), ttl)).await } pub async fn ttl(&self, key: &str) -> Result { - let reply: i64 = self.query(redis::cmd("PTTL").arg(self.key(key))).await?; + let reply: i64 = self.query(&commands::ttl(self.key(key))).await?; Ok(match reply { -2 => Ttl::Missing, -1 => Ttl::Persistent, @@ -114,18 +90,38 @@ impl RedisClient { } pub async fn incr_by(&self, key: &str, delta: i64) -> Result { - self.query(redis::cmd("INCRBY").arg(self.key(key)).arg(delta)) - .await + self.query(&commands::incr_by(self.key(key), delta)).await } pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result { - self.query(redis::cmd("INCRBYFLOAT").arg(self.key(key)).arg(delta)) + self.query(&commands::incr_by_float(self.key(key), delta)) .await } pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { - self.query(redis::cmd("RPUSH").arg(self.key(key)).arg(values)) - .await + self.query(&commands::rpush(self.key(key), values)).await + } + + pub async fn eval( + &self, + script: &Script, + keys: &[&str], + args: impl ToRedisArgs, + ) -> Result { + let keys = self.keys(keys); + let args = args.to_redis_args(); + Ok(script + .invocation(&keys, &args) + .invoke_async(&mut self.connection.clone()) + .await?) + } + + pub fn pipeline(&self) -> Pipeline { + Pipeline::new(self.namespace.clone()) + } + + pub async fn run(&self, pipeline: &Pipeline) -> Vec> { + pipeline.run(&self.connection).await } fn key(&self, key: &str) -> String { @@ -137,145 +133,23 @@ impl RedisClient { } async fn query(&self, command: &Cmd) -> Result { - let reply = match self.connection.clone() { - Connection::Standalone(mut connection) => command.query_async(&mut connection).await, - Connection::Cluster(mut connection) => command.query_async(&mut connection).await, - }; - Ok(reply?) + Ok(command.query_async(&mut self.connection.clone()).await?) } } -async fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult { - let config = ConnectionManagerConfig::new() - .set_connection_timeout(Some(timeout)) - .set_response_timeout(Some(timeout)) - .set_number_of_retries(RETRIES as usize); - let manager = redis::Client::open(info)? - .get_connection_manager_with_config(config) - .await?; - Ok(Connection::Standalone(manager)) -} - -async fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult { - let builder = ClusterClientBuilder::new(endpoint.addrs()) - .connection_timeout(timeout) - .response_timeout(timeout) - .retries(RETRIES); - let builder = match &endpoint.username { - Some(username) => builder.username(username), - None => builder, - }; - let builder = match &endpoint.password { - Some(password) => builder.password(password), - None => builder, - }; - Ok(Connection::Cluster( - builder.build()?.get_async_connection().await?, - )) -} - -fn millis(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} - #[cfg(test)] mod tests { use std::time::{Duration, Instant}; - use litellm_redis_testing::{RedisCluster, RedisServer}; - use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; + use litellm_redis_testing::RedisServer; use rstest::rstest; - use super::{Connection, RedisClient, Ttl}; + use super::{RedisClient, Ttl}; use crate::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, - RedisConfig, + DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig, + test_support::{Server, Topology, client_id}, }; - #[derive(Clone, Copy, Debug)] - enum Topology { - Standalone, - Cluster, - } - - enum Server { - Standalone(RedisServer), - Cluster(RedisCluster), - } - - impl Server { - async fn start(topology: Topology) -> Self { - match topology { - Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()), - Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()), - } - } - - fn endpoint(&self) -> Endpoint { - match self { - Self::Standalone(server) => Endpoint::Url(server.url().to_owned()), - Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { - nodes: cluster - .nodes() - .map(|(host, port)| ClusterNode { - host: host.to_owned(), - port, - }) - .collect(), - username: None, - password: None, - tls: false, - }), - } - } - - async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient { - RedisClient::connect(RedisConfig { - endpoint: self.endpoint(), - namespace: Namespace::new(namespace), - timeout, - }) - .await - .unwrap() - } - - async fn restart(&self) { - match self { - Self::Standalone(server) => server.restart().await.unwrap(), - Self::Cluster(cluster) => cluster.restart().await.unwrap(), - } - } - - async fn freeze(&self) { - match self { - Self::Standalone(server) => server.freeze().await.unwrap(), - Self::Cluster(cluster) => cluster.freeze().await.unwrap(), - } - } - - async fn thaw(&self) { - match self { - Self::Standalone(server) => server.thaw().await.unwrap(), - Self::Cluster(cluster) => cluster.thaw().await.unwrap(), - } - } - } - - async fn client_id(client: &RedisClient, key: &str) -> i64 { - let command = redis::cmd("CLIENT").arg("ID").to_owned(); - match client.connection.clone() { - Connection::Standalone(mut connection) => { - command.query_async(&mut connection).await.unwrap() - } - Connection::Cluster(mut connection) => { - let node = Route::with_key(key, SlotAddr::Master); - let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node)); - redis::from_redis_value(connection.route_command(command, routing).await.unwrap()) - .unwrap() - } - } - } - #[rstest] #[tokio::test] async fn namespaces_isolate_keys_on_the_server( @@ -500,6 +374,31 @@ mod tests { assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75); } + #[rstest] + #[tokio::test] + async fn concurrent_commands_on_one_cloned_client_all_apply( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + + let tasks: Vec<_> = (0..100) + .map(|_| { + let client = client.clone(); + tokio::spawn(async move { + for _ in 0..10 { + client.incr_by("shared", 1).await.unwrap(); + } + }) + }) + .collect(); + for task in tasks { + task.await.unwrap(); + } + + assert_eq!(client.get::("shared").await.unwrap(), Some(1000)); + } + #[rstest] #[tokio::test] async fn a_host_endpoint_connects_to_its_database() { diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs new file mode 100644 index 00000000000..34ec7432c98 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -0,0 +1,58 @@ +use std::time::Duration; + +use redis::{Cmd, ToRedisArgs, ToSingleRedisArg}; + +pub(crate) fn get(key: String) -> Cmd { + redis::cmd("GET").arg(key).to_owned() +} + +pub(crate) fn mget(keys: Vec) -> Cmd { + redis::cmd("MGET").arg(keys).to_owned() +} + +pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { + redis::cmd("SET") + .arg(key) + .arg(value) + .arg("PX") + .arg(millis(ttl)) + .to_owned() +} + +pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { + redis::cmd("SET") + .arg(key) + .arg(value) + .arg("NX") + .arg("PX") + .arg(millis(ttl)) + .to_owned() +} + +pub(crate) fn del(keys: Vec) -> Cmd { + redis::cmd("DEL").arg(keys).to_owned() +} + +pub(crate) fn expire(key: String, ttl: Duration) -> Cmd { + redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned() +} + +pub(crate) fn ttl(key: String) -> Cmd { + redis::cmd("PTTL").arg(key).to_owned() +} + +pub(crate) fn incr_by(key: String, delta: i64) -> Cmd { + redis::cmd("INCRBY").arg(key).arg(delta).to_owned() +} + +pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd { + redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned() +} + +pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd { + redis::cmd("RPUSH").arg(key).arg(values).to_owned() +} + +fn millis(duration: Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs new file mode 100644 index 00000000000..f70449bd43e --- /dev/null +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -0,0 +1,80 @@ +use std::time::Duration; + +use redis::{ + Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value, + aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig}, + cluster::ClusterClientBuilder, + cluster_async::ClusterConnection, +}; + +use crate::ClusterEndpoint; + +const RETRIES: u32 = 3; + +#[derive(Clone)] +pub(crate) enum Connection { + Standalone(ConnectionManager), + Cluster(ClusterConnection), +} + +impl Connection { + pub(crate) async fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult { + let config = ConnectionManagerConfig::new() + .set_connection_timeout(Some(timeout)) + .set_response_timeout(Some(timeout)) + .set_number_of_retries(RETRIES as usize); + let manager = redis::Client::open(info)? + .get_connection_manager_with_config(config) + .await?; + Ok(Self::Standalone(manager)) + } + + pub(crate) async fn cluster( + endpoint: &ClusterEndpoint, + timeout: Duration, + ) -> RedisResult { + let builder = ClusterClientBuilder::new(endpoint.addrs()) + .connection_timeout(timeout) + .response_timeout(timeout) + .retries(RETRIES); + let builder = match &endpoint.username { + Some(username) => builder.username(username), + None => builder, + }; + let builder = match &endpoint.password { + Some(password) => builder.password(password), + None => builder, + }; + Ok(Self::Cluster( + builder.build()?.get_async_connection().await?, + )) + } +} + +impl ConnectionLike for Connection { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + match self { + Self::Standalone(connection) => connection.req_packed_command(cmd), + Self::Cluster(connection) => connection.req_packed_command(cmd), + } + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + match self { + Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count), + Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count), + } + } + + fn get_db(&self) -> i64 { + match self { + Self::Standalone(connection) => connection.get_db(), + Self::Cluster(connection) => connection.get_db(), + } + } +} diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 32c4354b803..47ff8677e06 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -1,7 +1,13 @@ mod client; +mod commands; mod config; +mod connection; mod error; mod namespace; +mod pipeline; +mod script; +#[cfg(test)] +mod test_support; pub use client::{RedisClient, Ttl}; pub use config::{ @@ -9,3 +15,6 @@ pub use config::{ }; pub use error::{ConnectError, Error}; pub use namespace::Namespace; +pub use pipeline::Pipeline; +pub use redis::Value; +pub use script::Script; diff --git a/litellm-rust/crates/redis-client/src/pipeline.rs b/litellm-rust/crates/redis-client/src/pipeline.rs new file mode 100644 index 00000000000..74d659170e9 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/pipeline.rs @@ -0,0 +1,288 @@ +use std::{collections::HashMap, time::Duration}; + +use futures_util::future::join_all; +use redis::{ + Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, + cluster_routing::Slot, +}; + +use crate::{Error, Namespace, Script, commands, connection::Connection}; + +pub struct Pipeline { + namespace: Namespace, + ops: Vec, +} + +struct Op { + command: Cmd, + key: String, + script: Option, +} + +struct ScriptCall { + script: Script, + keys: Vec, + args: Vec>, +} + +type Settled = (usize, Result); + +impl Pipeline { + pub(crate) fn new(namespace: Namespace) -> Self { + Self { + namespace, + ops: Vec::new(), + } + } + + pub fn get(self, key: &str) -> Self { + let key = self.namespace.apply(key); + self.command(commands::get(key.clone()), key) + } + + pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self { + let key = self.namespace.apply(key); + self.command(commands::set(key.clone(), value, ttl), key) + } + + pub fn del(self, key: &str) -> Self { + let key = self.namespace.apply(key); + self.command(commands::del(vec![key.clone()]), key) + } + + pub fn expire(self, key: &str, ttl: Duration) -> Self { + let key = self.namespace.apply(key); + self.command(commands::expire(key.clone(), ttl), key) + } + + pub fn incr_by(self, key: &str, delta: i64) -> Self { + let key = self.namespace.apply(key); + self.command(commands::incr_by(key.clone(), delta), key) + } + + pub fn incr_by_float(self, key: &str, delta: f64) -> Self { + let key = self.namespace.apply(key); + self.command(commands::incr_by_float(key.clone(), delta), key) + } + + pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self { + let keys: Vec = keys.iter().map(|key| self.namespace.apply(key)).collect(); + let args = args.to_redis_args(); + let op = Op { + command: script.evalsha(&keys, &args), + key: keys.first().cloned().unwrap_or_default(), + script: Some(ScriptCall { + script: script.clone(), + keys, + args, + }), + }; + self.push(op) + } + + fn command(self, command: Cmd, key: String) -> Self { + self.push(Op { + command, + key, + script: None, + }) + } + + fn push(mut self, op: Op) -> Self { + self.ops.push(op); + self + } + + pub(crate) async fn run(&self, connection: &Connection) -> Vec> { + let groups = self.groups(connection); + let settled = join_all( + groups + .iter() + .map(|group| self.run_group(connection.clone(), group)), + ) + .await; + let mut settled: Vec = settled.into_iter().flatten().collect(); + settled.sort_by_key(|(index, _)| *index); + settled.into_iter().map(|(_, result)| result).collect() + } + + fn groups(&self, connection: &Connection) -> Vec> { + match connection { + _ if self.ops.is_empty() => Vec::new(), + Connection::Standalone(_) => vec![(0..self.ops.len()).collect()], + Connection::Cluster(_) => { + let mut by_slot: HashMap> = HashMap::new(); + for (index, op) in self.ops.iter().enumerate() { + by_slot + .entry(Slot::for_key(&op.key)) + .or_default() + .push(index); + } + by_slot.into_values().collect() + } + } + } + + async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec { + let mut pipeline = redis::pipe(); + pipeline.ignore_errors(); + for &index in indices { + pipeline.add_command(self.ops[index].command.clone()); + } + let replies: RedisResult>> = + pipeline.query_async(&mut connection).await; + let replies = match replies { + Ok(replies) => replies, + Err(error) => { + return indices + .iter() + .map(|&index| (index, Err(Error::from(error.clone())))) + .collect(); + } + }; + let mut settled = Vec::with_capacity(indices.len()); + for (&index, reply) in indices.iter().zip(replies) { + settled.push((index, self.settle(&mut connection, index, reply).await)); + } + settled + } + + async fn settle( + &self, + connection: &mut Connection, + index: usize, + reply: RedisResult, + ) -> Result { + match (reply, &self.ops[index].script) { + (Err(error), Some(call)) + if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) => + { + Ok(call + .script + .invocation(&call.keys, &call.args) + .invoke_async(connection) + .await?) + } + (reply, _) => Ok(reply?), + } + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use redis::Value; + use rstest::rstest; + + use crate::{ + DEFAULT_TIMEOUT, Error, Script, + test_support::{Server, Topology, flush_scripts}, + }; + + #[rstest] + #[tokio::test] + async fn each_op_settles_on_its_own_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + client + .set("text", "not a number", Duration::from_secs(60)) + .await + .unwrap(); + let pipeline = client + .pipeline() + .incr_by("a", 1) + .incr_by("text", 1) + .incr_by("c", 5) + .set("b", "2", Duration::from_secs(60)) + .get("b") + .incr_by("a", 1) + .get("missing"); + + let results = client.run(&pipeline).await; + + assert_eq!(results.len(), 7); + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1)); + assert!( + matches!(results[1], Err(Error::Server(_))), + "{:?}", + results[1] + ); + assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5)); + assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); + assert_eq!( + results[4].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2)); + assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); + } + + #[rstest] + #[tokio::test] + async fn a_flushed_script_in_a_pipeline_reruns_alone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("ns", DEFAULT_TIMEOUT).await; + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.eval::(&incr, &["counter"], 1).await.unwrap(); + flush_scripts(&client).await; + let pipeline = client + .pipeline() + .eval(&incr, &["counter"], 2) + .incr_by("other", 7); + + let results = client.run(&pipeline).await; + + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3)); + assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7)); + assert_eq!(client.get::("counter").await.unwrap(), Some(3)); + } + + #[rstest] + #[tokio::test] + async fn an_unreachable_server_fails_every_op_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", Duration::from_millis(500)).await; + let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); + + server.freeze().await; + let results = client.run(&pipeline).await; + server.thaw().await; + + assert_eq!(results.len(), 2); + assert!( + matches!(results[0], Err(Error::Transport(_))), + "{:?}", + results[0] + ); + assert!( + matches!(results[1], Err(Error::Transport(_))), + "{:?}", + results[1] + ); + } + + #[rstest] + #[tokio::test] + async fn an_empty_pipeline_returns_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", Duration::from_secs(2)).await; + + server.freeze().await; + let started = Instant::now(); + let results = client.run(&client.pipeline()).await; + let elapsed = started.elapsed(); + server.thaw().await; + + assert!(results.is_empty()); + assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); + } +} diff --git a/litellm-rust/crates/redis-client/src/script.rs b/litellm-rust/crates/redis-client/src/script.rs new file mode 100644 index 00000000000..c5d73133ecc --- /dev/null +++ b/litellm-rust/crates/redis-client/src/script.rs @@ -0,0 +1,88 @@ +use redis::{Cmd, ScriptInvocation}; + +#[derive(Clone)] +pub struct Script(redis::Script); + +impl Script { + pub fn new(source: &str) -> Self { + Self(redis::Script::new(source)) + } + + pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec]) -> Cmd { + redis::cmd("EVALSHA") + .arg(self.0.get_hash()) + .arg(keys.len()) + .arg(keys) + .arg(args) + .to_owned() + } + + pub(crate) fn invocation(&self, keys: &[String], args: &[Vec]) -> ScriptInvocation<'_> { + let mut invocation = self.0.prepare_invoke(); + invocation.key(keys).arg(args); + invocation + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::rstest; + + use super::Script; + use crate::{ + DEFAULT_TIMEOUT, Error, + test_support::{Server, Topology, flush_scripts}, + }; + + #[rstest] + #[tokio::test] + async fn script_keys_are_namespaced( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("a", DEFAULT_TIMEOUT).await; + let bare = server.client("", DEFAULT_TIMEOUT).await; + let read = Script::new("return redis.call('GET', KEYS[1])"); + bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap(); + + let value: Option = client + .eval(&read, &["x"], Vec::::new()) + .await + .unwrap(); + + assert_eq!(value.as_deref(), Some("1")); + } + + #[rstest] + #[tokio::test] + async fn a_failing_script_is_a_server_error_and_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.rpush("list", "item").await.unwrap(); + + let failed = client.eval::(&incr, &["list"], 1).await; + + assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); + assert_eq!(client.eval::(&incr, &["counter"], 3).await.unwrap(), 3); + } + + #[rstest] + #[tokio::test] + async fn a_flushed_script_is_loaded_again( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.eval::(&incr, &["counter"], 1).await.unwrap(); + + flush_scripts(&client).await; + + assert_eq!(client.eval::(&incr, &["counter"], 1).await.unwrap(), 2); + } +} diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs new file mode 100644 index 00000000000..34e78fc12d3 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/test_support.rs @@ -0,0 +1,101 @@ +use std::time::Duration; + +use litellm_redis_testing::{RedisCluster, RedisServer}; +use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; + +use crate::{ + ClusterEndpoint, ClusterNode, Endpoint, Namespace, RedisClient, RedisConfig, + connection::Connection, +}; + +#[derive(Clone, Copy, Debug)] +pub(crate) enum Topology { + Standalone, + Cluster, +} + +pub(crate) enum Server { + Standalone(RedisServer), + Cluster(RedisCluster), +} + +impl Server { + pub(crate) async fn start(topology: Topology) -> Self { + match topology { + Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()), + Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()), + } + } + + pub(crate) fn endpoint(&self) -> Endpoint { + match self { + Self::Standalone(server) => Endpoint::Url(server.url().to_owned()), + Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { + nodes: cluster + .nodes() + .map(|(host, port)| ClusterNode { + host: host.to_owned(), + port, + }) + .collect(), + username: None, + password: None, + tls: false, + }), + } + } + + pub(crate) async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient { + RedisClient::connect(RedisConfig { + endpoint: self.endpoint(), + namespace: Namespace::new(namespace), + timeout, + }) + .await + .unwrap() + } + + pub(crate) async fn restart(&self) { + match self { + Self::Standalone(server) => server.restart().await.unwrap(), + Self::Cluster(cluster) => cluster.restart().await.unwrap(), + } + } + + pub(crate) async fn freeze(&self) { + match self { + Self::Standalone(server) => server.freeze().await.unwrap(), + Self::Cluster(cluster) => cluster.freeze().await.unwrap(), + } + } + + pub(crate) async fn thaw(&self) { + match self { + Self::Standalone(server) => server.thaw().await.unwrap(), + Self::Cluster(cluster) => cluster.thaw().await.unwrap(), + } + } +} + +pub(crate) async fn client_id(client: &RedisClient, key: &str) -> i64 { + let command = redis::cmd("CLIENT").arg("ID").to_owned(); + match client.connection.clone() { + Connection::Standalone(mut connection) => { + command.query_async(&mut connection).await.unwrap() + } + Connection::Cluster(mut connection) => { + let node = Route::with_key(key, SlotAddr::Master); + let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node)); + redis::from_redis_value(connection.route_command(command, routing).await.unwrap()) + .unwrap() + } + } +} + +pub(crate) async fn flush_scripts(client: &RedisClient) { + redis::cmd("SCRIPT") + .arg("FLUSH") + .query_async::<()>(&mut client.connection.clone()) + .await + .unwrap(); +}