From ed28888fd0bcf0e63bfe12a7f69eb7ead3d3b62e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 00:17:12 +0000 Subject: [PATCH] fix(redis-client): match redis-py command and failure behavior Every command is now the one RedisCache sends. Pipelined scripts go out as EVALSHA, and a NOSCRIPT op alone runs again after the pipeline like RedisBatch's _settle_alone. Only a ScriptCall whose reply was NOSCRIPT can be sent again, so nothing can resend a command that already ran. Increments are plain INCRBY and INCRBYFLOAT with TTL and EXPIRE in seconds, as increment_cache, async_increment and RedisBatch send them, instead of Lua. SET uses EX, and numbers are written the way redis-py encodes them. --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/gateway/src/redis.rs | 5 +- litellm-rust/crates/redis-client/AGENTS.md | 21 +- litellm-rust/crates/redis-client/Cargo.toml | 4 +- .../crates/redis-client/src/client.rs | 633 ++++---------- .../crates/redis-client/src/commands.rs | 55 +- .../crates/redis-client/src/config.rs | 29 +- .../crates/redis-client/src/connection.rs | 51 +- litellm-rust/crates/redis-client/src/error.rs | 136 ++- .../crates/redis-client/src/events.rs | 57 ++ litellm-rust/crates/redis-client/src/lib.rs | 9 +- .../crates/redis-client/src/pipeline.rs | 332 ++++---- .../crates/redis-client/src/script.rs | 76 +- .../crates/redis-client/src/test_support.rs | 126 --- .../crates/redis-client/tests/client.rs | 780 ++++++++++++++++++ .../crates/redis-client/tests/connection.rs | 63 ++ .../crates/redis-client/tests/pipeline.rs | 309 +++++++ .../crates/redis-client/tests/script.rs | 106 +++ .../crates/redis-client/tests/support/mod.rs | 92 +++ litellm-rust/crates/redis-testing/AGENTS.md | 4 + litellm-rust/crates/redis-testing/Cargo.toml | 2 +- .../crates/redis-testing/src/cluster.rs | 25 +- .../crates/redis-testing/src/deployment.rs | 113 +++ .../crates/redis-testing/src/error.rs | 2 + litellm-rust/crates/redis-testing/src/exec.rs | 31 + litellm-rust/crates/redis-testing/src/lib.rs | 3 + .../crates/redis-testing/src/standalone.rs | 14 +- 27 files changed, 2133 insertions(+), 947 deletions(-) create mode 100644 litellm-rust/crates/redis-client/src/events.rs delete mode 100644 litellm-rust/crates/redis-client/src/test_support.rs create mode 100644 litellm-rust/crates/redis-client/tests/client.rs create mode 100644 litellm-rust/crates/redis-client/tests/connection.rs create mode 100644 litellm-rust/crates/redis-client/tests/pipeline.rs create mode 100644 litellm-rust/crates/redis-client/tests/script.rs create mode 100644 litellm-rust/crates/redis-client/tests/support/mod.rs create mode 100644 litellm-rust/crates/redis-testing/src/deployment.rs create mode 100644 litellm-rust/crates/redis-testing/src/exec.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ba1131d7a92..a8e6b9c8548 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4361,10 +4361,12 @@ version = "0.1.0" dependencies = [ "futures-util", "litellm-circuit-breaker", + "litellm-python-compat", "litellm-redis-testing", "litellm-secret-value", "redis", "rstest", + "socket2 0.6.5", "thiserror 2.0.19", "tokio", ] diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs index c40b289029b..d0e9a838232 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -7,7 +7,7 @@ use litellm_config::{ }; use litellm_redis_client::{ ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient, - RedisConfig, + RedisConfig, Timeouts, }; use crate::Error; @@ -178,8 +178,9 @@ fn client(params: Params, breaker: Option) -> Result for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error +- Classify with `Error::classify(error, deployment)`, never a bare `From`, because redis-py (5.3.1) decides by topology. `Transport` is what Python's `_redis_health_error_types` counts: LOADING everywhere, CLUSTERDOWN and MASTERDOWN only on cluster (only `ClusterParser` maps them), plus NOAUTH, WRONGPASS and the `ERR` messages redis-py maps to `ConnectionError`, matched exactly like `parse_error`. TRYAGAIN is `Server`, since the cluster client already retried it. Each code gets its own rstest case built from a parsed RESP error - A `Server` error never drops the connection - `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome - Never swallow an error or turn it into a miss. The caller picks the fallback @@ -17,7 +17,6 @@ - Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it - `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success - An open breaker returns `Error::BreakerOpen` without touching the socket -- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count ## Namespace - Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace` @@ -26,18 +25,24 @@ ## Commands - One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET - 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`) +- Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same +- TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds +- Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat` +- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. In a pipeline `incr_by_float` is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when a TTL is given ## 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` +- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python +- `Entry` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice - 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 +- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed +- Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed - 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 +- 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 - 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) - Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables - Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index 3e9447a6bfa..17fc4d7ca44 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -8,11 +8,13 @@ repository.workspace = true [dependencies] litellm-circuit-breaker.workspace = true litellm-secret-value.workspace = true +litellm-python-compat.workspace = true futures-util.workspace = true redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } +socket2 = "0.6.5" thiserror.workspace = true +tokio = { workspace = true, features = ["sync", "time"] } [dev-dependencies] litellm-redis-testing.workspace = true rstest.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 48dbc2b2dfb..b9481512f8b 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,14 +1,17 @@ use std::{ future::Future, + sync::Arc, time::{Duration, Instant}, }; use litellm_circuit_breaker::{Admission, Breaker, Failure}; -use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value}; +use redis::{ + Cmd, FromRedisValue, IntoConnectionInfo, RedisError, ToRedisArgs, ToSingleRedisArg, Value, +}; use crate::{ - ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands, - connection::Connection, + Amount, CallEvent, ConnectError, Endpoint, Error, Expiry, FailureKind, Namespace, Operation, + Outcome, Pipeline, RedisConfig, RedisEvents, Script, commands, connection::Connection, }; #[derive(Clone)] @@ -16,6 +19,7 @@ pub struct RedisClient { pub(crate) connection: Connection, namespace: Namespace, breaker: Option, + events: Option>, } #[derive(Debug, PartialEq, Eq)] @@ -29,33 +33,34 @@ impl RedisClient { pub fn new(config: RedisConfig) -> Result { let connection = match &config.endpoint { Endpoint::Url(url) => { - Connection::standalone(url.expose().into_connection_info()?, config.timeout)? + Connection::standalone(url.expose().into_connection_info()?, config.timeouts)? } Endpoint::Host(host) => { - Connection::standalone(host.connection_info()?, config.timeout)? + Connection::standalone(host.connection_info()?, config.timeouts)? } - Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout)?, + Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeouts)?, }; Ok(Self { connection, namespace: config.namespace, breaker: config.breaker, + events: config.events, }) } pub async fn ping(&self) -> Result<(), Error> { - self.query(&commands::ping()).await + self.query("PING", &commands::ping()).await } pub async fn get(&self, key: &str) -> Result, Error> { - self.query(&commands::get(self.key(key))).await + self.query("GET", &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(&commands::mget(self.keys(keys))).await + self.query("MGET", &commands::mget(self.keys(keys))).await } pub async fn set( @@ -64,7 +69,8 @@ impl RedisClient { value: impl ToSingleRedisArg, ttl: Duration, ) -> Result<(), Error> { - self.query(&commands::set(self.key(key), value, ttl)).await + self.query("SET", &commands::set(self.key(key), value, ttl)) + .await } pub async fn set_nx( @@ -74,7 +80,7 @@ impl RedisClient { ttl: Duration, ) -> Result { let reply: Option = self - .query(&commands::set_nx(self.key(key), value, ttl)) + .query("SET", &commands::set_nx(self.key(key), value, ttl)) .await?; Ok(reply.is_some()) } @@ -83,33 +89,81 @@ impl RedisClient { if keys.is_empty() { return Ok(0); } - self.query(&commands::del(self.keys(keys))).await + self.query("DEL", &commands::del(self.keys(keys))).await } pub async fn expire(&self, key: &str, ttl: Duration) -> Result { - self.query(&commands::expire(self.key(key), ttl)).await - } - - pub async fn ttl(&self, key: &str) -> Result { - let reply: i64 = self.query(&commands::ttl(self.key(key))).await?; - Ok(match reply { - -2 => Ttl::Missing, - -1 => Ttl::Persistent, - millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())), - }) - } - - pub async fn incr_by(&self, key: &str, delta: i64) -> Result { - self.query(&commands::incr_by(self.key(key), delta)).await - } - - pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result { - self.query(&commands::incr_by_float(self.key(key), delta)) + self.query("EXPIRE", &commands::expire(self.key(key), ttl)) .await } - pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { - self.query(&commands::rpush(self.key(key), values)).await + pub async fn ttl(&self, key: &str) -> Result { + let reply: i64 = self.query("TTL", &commands::ttl(self.key(key))).await?; + Ok(match reply { + -2 => Ttl::Missing, + -1 => Ttl::Persistent, + seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())), + }) + } + + /// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no + /// TTL, each its own round trip. + pub async fn incr_by(&self, key: &str, delta: i64, ttl: Duration) -> Result { + let key = self.key(key); + self.guarded(Operation::Command("INCRBY"), async { + let mut connection = self.connection.clone(); + let value: i64 = commands::incr_by(key.clone(), delta) + .query_async(&mut connection) + .await?; + let current: i64 = commands::ttl(key.clone()) + .query_async(&mut connection) + .await?; + if current == -1 { + commands::expire(key, ttl) + .exec_async(&mut connection) + .await?; + } + Ok(value) + }) + .await + } + + /// `RedisCache.async_increment`: `INCRBYFLOAT` and its TTL command in one round trip, and for + /// [`Expiry::KeepExisting`] a third, `EXPIRE`, only when the key has no TTL. + pub async fn incr_by_float( + &self, + key: &str, + amount: Amount, + expiry: Expiry, + ) -> Result { + let key = self.key(key); + self.guarded(Operation::Command("INCRBYFLOAT"), async { + let mut connection = self.connection.clone(); + let mut pipeline = redis::pipe(); + pipeline.add_command(commands::incr_by_float(key.clone(), amount)); + match expiry { + Expiry::Refresh(ttl) => { + let (value, _): (f64, bool) = pipeline + .add_command(commands::expire(key, ttl)) + .query_async(&mut connection) + .await?; + Ok(value) + } + Expiry::KeepExisting(ttl) => { + let (value, current): (f64, i64) = pipeline + .add_command(commands::ttl(key.clone())) + .query_async(&mut connection) + .await?; + if current == -1 { + commands::expire(key, ttl) + .exec_async(&mut connection) + .await?; + } + Ok(value) + } + } + }) + .await } pub async fn eval( @@ -120,11 +174,11 @@ impl RedisClient { ) -> Result { let keys = self.keys(keys); let args = args.to_redis_args(); - self.guarded(async { - Ok(script + self.guarded(Operation::Script, async { + script .invocation(&keys, &args) .invoke_async(&mut self.connection.clone()) - .await?) + .await }) .await } @@ -134,9 +188,18 @@ impl RedisClient { } pub async fn run(&self, pipeline: &Pipeline) -> Vec> { + let operation = Operation::Pipeline { + ops: pipeline.len(), + }; + let started = Instant::now(); let admission = match self.admit() { Ok(admission) => admission, Err(_) => { + self.observe( + operation, + started, + Outcome::Failed(FailureKind::BreakerOpen), + ); return (0..pipeline.len()) .map(|_| Err(Error::BreakerOpen)) .collect(); @@ -150,6 +213,11 @@ impl RedisClient { .filter(|error| matches!(error, Error::Transport(_))) }); settle(admission, transport); + let failure = transport.or_else(|| results.iter().find_map(|result| result.as_ref().err())); + let outcome = failure.map_or(Outcome::Succeeded, |error| { + Outcome::Failed(FailureKind::of(error)) + }); + self.observe(operation, started, outcome); results } @@ -161,18 +229,47 @@ impl RedisClient { keys.iter().map(|key| self.key(key)).collect() } - async fn query(&self, command: &Cmd) -> Result { - self.guarded(async { Ok(command.query_async(&mut self.connection.clone()).await?) }) - .await + async fn query( + &self, + name: &'static str, + command: &Cmd, + ) -> Result { + self.guarded(Operation::Command(name), async { + command.query_async(&mut self.connection.clone()).await + }) + .await } - async fn guarded(&self, call: impl Future>) -> Result { - let admission = self.admit()?; - let result = call.await; - settle(admission, result.as_ref().err()); + async fn guarded( + &self, + operation: Operation, + call: impl Future>, + ) -> Result { + let started = Instant::now(); + let result = match self.admit() { + Ok(admission) => { + let result = call + .await + .map_err(|error| Error::classify(error, self.connection.deployment())); + settle(admission, result.as_ref().err()); + result + } + Err(error) => Err(error), + }; + self.observe(operation, started, Outcome::of(&result)); result } + fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) { + if let Some(events) = &self.events { + events.call(&CallEvent { + operation, + duration: started.elapsed(), + outcome, + }); + } + } + fn admit(&self) -> Result, Error> { self.breaker .as_ref() @@ -198,455 +295,3 @@ fn settle(admission: Option, error: Option<&Error>) { Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission), } } - -#[cfg(test)] -mod tests { - use std::time::{Duration, Instant}; - - use litellm_redis_testing::RedisServer; - use litellm_secret_value::SecretValue; - use rstest::rstest; - - use litellm_circuit_breaker::{Breaker, Config, State}; - - use super::{RedisClient, Ttl}; - use crate::{ - DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig, - test_support::{Server, Topology, client_id}, - }; - - #[rstest] - #[tokio::test] - async fn namespaces_isolate_keys_on_the_server( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let a = server.client("a", DEFAULT_TIMEOUT).await; - let b = server.client("b", DEFAULT_TIMEOUT).await; - let bare = server.client("", DEFAULT_TIMEOUT).await; - - a.set("x", "1", Duration::from_secs(60)).await.unwrap(); - - assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); - assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); - assert_eq!( - bare.get::("a:x").await.unwrap().as_deref(), - Some("1") - ); - assert_eq!(b.get::("x").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_rejected_command_keeps_the_connection( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.rpush("list", "item").await.unwrap(); - let before = client_id(&client, "list").await; - - let rejected = client.incr_by("list", 1).await; - - assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); - assert_eq!(client.rpush("list", "next").await.unwrap(), 2); - assert_eq!(client_id(&client, "list").await, before); - } - - #[rstest] - #[tokio::test] - async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client - .set("float", "15.0", Duration::from_secs(60)) - .await - .unwrap(); - - let decoded = client.get::("float").await; - - assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); - assert_eq!(client.get::("missing").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - client - .set("first", "1", Duration::from_secs(60)) - .await - .unwrap(); - client - .set("second", "2", Duration::from_secs(60)) - .await - .unwrap(); - - server.freeze().await; - let timed_out = client.get::("first").await; - server.thaw().await; - - assert!( - matches!(timed_out, Err(Error::Transport(_))), - "{timed_out:?}" - ); - assert_eq!( - client.get::("second").await.unwrap().as_deref(), - Some("2") - ); - } - - #[rstest] - #[tokio::test] - async fn the_same_client_recovers_after_redis_restarts( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - - server.restart().await; - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.incr_by("after_restart", 1).await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(1)); - } - - #[rstest] - #[tokio::test] - async fn a_client_created_while_redis_is_down_works_once_redis_is_up( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - server.stop().await; - - let client = server.client("", Duration::from_millis(500)).await; - let while_down = client.ping().await; - server.resume().await; - - assert!( - matches!(while_down, Err(Error::Transport(_))), - "{while_down:?}" - ); - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.ping().await { - Ok(()) => break true, - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break false, - } - }; - assert!(recovered); - } - - #[rstest] - #[tokio::test] - async fn empty_batches_succeed_without_asking_redis( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - assert_eq!( - client.mget::(&[]).await.unwrap(), - Vec::>::new() - ); - assert_eq!(client.del(&[]).await.unwrap(), 0); - } - - #[rstest] - #[tokio::test] - async fn mget_and_del_span_slots_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("a", "1", Duration::from_secs(60)).await.unwrap(); - client.set("c", "3", Duration::from_secs(60)).await.unwrap(); - - let values = client.mget::(&["c", "b", "a"]).await.unwrap(); - let deleted = client.del(&["a", "b", "c"]).await.unwrap(); - - assert_eq!( - values, - vec![Some("3".to_owned()), None, Some("1".to_owned())] - ); - assert_eq!(deleted, 2); - assert_eq!( - client.mget::(&["a", "c"]).await.unwrap(), - vec![None, None] - ); - } - - #[rstest] - #[tokio::test] - async fn set_nx_only_writes_a_missing_key( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - let first = client - .set_nx("lock", "owner-1", Duration::from_secs(60)) - .await - .unwrap(); - let second = client - .set_nx("lock", "owner-2", Duration::from_secs(60)) - .await - .unwrap(); - - assert!(first); - assert!(!second); - assert_eq!( - client.get::("lock").await.unwrap().as_deref(), - Some("owner-1") - ); - } - - #[rstest] - #[tokio::test] - async fn ttl_reports_missing_persistent_and_remaining_time( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.incr_by("persistent", 1).await.unwrap(); - client - .set("expiring", "1", Duration::from_secs(60)) - .await - .unwrap(); - client.incr_by_float("expired_later", 0.5).await.unwrap(); - assert!( - client - .expire("expired_later", Duration::from_secs(30)) - .await - .unwrap() - ); - - assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); - assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); - assert!(matches!( - client.ttl("expiring").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) - )); - assert!(matches!( - client.ttl("expired_later").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) - )); - assert!( - !client - .expire("missing", Duration::from_secs(30)) - .await - .unwrap() - ); - } - - #[rstest] - #[tokio::test] - async fn incr_by_float_accumulates_fractions( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - client.incr_by_float("spend", 0.25).await.unwrap(); - - 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)); - } - - fn breaker(failure_threshold: u32) -> Breaker { - Breaker::new( - Config { - failure_threshold, - recovery_timeout: Duration::from_millis(200), - timeout_min_duration: Duration::ZERO, - }, - Box::new(|_| {}), - ) - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker.clone())) - .await; - - server.freeze().await; - let first = client.get::("k").await; - let second = client.get::("k").await; - let started = Instant::now(); - let third = client.get::("k").await; - let rejected_in = started.elapsed(); - server.thaw().await; - - assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); - assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); - assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); - assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.get::("k").await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(50)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(None)); - assert!( - matches!(breaker.state(), State::Closed { .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn rejected_commands_never_open_the_breaker( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", DEFAULT_TIMEOUT, Some(breaker.clone())) - .await; - client.rpush("list", "item").await.unwrap(); - - let rejected = [ - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - ]; - - assert!( - matches!(rejected[2], Err(Error::Server(_))), - "{:?}", - rejected[2] - ); - assert!( - matches!(breaker.state(), State::Closed { hard: 0, .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_fails_every_pipeline_op_without_sending( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker(1))) - .await; - let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); - - server.freeze().await; - let failed = client.run(&pipeline).await; - let rejected = client.run(&pipeline).await; - server.thaw().await; - - assert!( - matches!(failed[0], Err(Error::Transport(_))), - "{:?}", - failed[0] - ); - assert!( - matches!(rejected[0], Err(Error::BreakerOpen)), - "{:?}", - rejected[0] - ); - assert!( - matches!(rejected[1], Err(Error::BreakerOpen)), - "{:?}", - rejected[1] - ); - } - - #[rstest] - #[tokio::test] - async fn a_host_endpoint_connects_to_its_database() { - let server = RedisServer::start().await.unwrap(); - let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); - let connect = |endpoint| { - RedisClient::new(RedisConfig { - endpoint, - namespace: Namespace::default(), - timeout: DEFAULT_TIMEOUT, - breaker: None, - }) - }; - let in_db_3 = connect(Endpoint::Host(HostEndpoint { - host: "127.0.0.1".to_owned(), - port, - username: None, - password: None, - tls: false, - db: 3, - })) - .unwrap(); - let in_db_0 = connect(Endpoint::Url(SecretValue::new(server.url().to_owned()))).unwrap(); - let via_url = connect(Endpoint::Url(SecretValue::new(format!( - "{}/3", - server.url() - )))) - .unwrap(); - - in_db_3 - .set("k", "3", Duration::from_secs(60)) - .await - .unwrap(); - - assert_eq!(in_db_0.get::("k").await.unwrap(), None); - assert_eq!( - via_url.get::("k").await.unwrap().as_deref(), - Some("3") - ); - } -} diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs index fb8a4a28c94..5cd7fd421a5 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -1,6 +1,34 @@ use std::time::Duration; -use redis::{Cmd, ToRedisArgs, ToSingleRedisArg}; +use litellm_python_compat::repr::float_repr; +use redis::{Cmd, RedisWrite, ToRedisArgs, ToSingleRedisArg}; + +/// How an increment treats the key's TTL, the two ways `RedisCache` increments with one. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Expiry { + /// `async_increment(refresh_ttl=True)` and every pipelined increment: `EXPIRE` after each write. + Refresh(Duration), + /// `async_increment(refresh_ttl=False)`: read the `TTL` and `EXPIRE` only a key that has none. + KeepExisting(Duration), +} + +/// A number written the way redis-py encodes it: an `int` with `str()`, a `float` with `repr()`. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Amount { + Int(i64), + Float(f64), +} + +impl ToRedisArgs for Amount { + fn write_redis_args(&self, out: &mut W) { + match self { + Self::Int(value) => out.write_arg(value.to_string().as_bytes()), + Self::Float(value) => out.write_arg(float_repr(*value).as_bytes()), + } + } +} + +impl ToSingleRedisArg for Amount {} pub(crate) fn ping() -> Cmd { redis::cmd("PING") @@ -14,12 +42,13 @@ pub(crate) fn mget(keys: Vec) -> Cmd { redis::cmd("MGET").arg(keys).to_owned() } +// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())` pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { redis::cmd("SET") .arg(key) .arg(value) - .arg("PX") - .arg(millis(ttl)) + .arg("EX") + .arg(ttl.as_secs()) .to_owned() } @@ -27,9 +56,9 @@ pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) - redis::cmd("SET") .arg(key) .arg(value) + .arg("EX") + .arg(ttl.as_secs()) .arg("NX") - .arg("PX") - .arg(millis(ttl)) .to_owned() } @@ -38,25 +67,17 @@ pub(crate) fn del(keys: Vec) -> Cmd { } pub(crate) fn expire(key: String, ttl: Duration) -> Cmd { - redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned() + redis::cmd("EXPIRE").arg(key).arg(ttl.as_secs()).to_owned() } pub(crate) fn ttl(key: String) -> Cmd { - redis::cmd("PTTL").arg(key).to_owned() + redis::cmd("TTL").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) +pub(crate) fn incr_by_float(key: String, amount: Amount) -> Cmd { + redis::cmd("INCRBYFLOAT").arg(key).arg(amount).to_owned() } diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index c40b9eb6b36..dd94e509e89 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,19 +1,42 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; use litellm_circuit_breaker::Breaker; use litellm_secret_value::SecretValue; -use crate::Namespace; +use crate::{Namespace, RedisEvents}; +// litellm/constants.py REDIS_SOCKET_TIMEOUT, which _coerce_redis_kwargs_types sets when socket_timeout is unset pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); pub struct RedisConfig { pub endpoint: Endpoint, pub namespace: Namespace, - pub timeout: Duration, + pub timeouts: Timeouts, pub breaker: Option, + pub events: Option>, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Timeouts { + pub response: Duration, + pub connect: Duration, +} + +impl Timeouts { + pub fn new(socket_timeout: Duration, connect_timeout: Option) -> Self { + Self { + response: socket_timeout, + connect: connect_timeout.unwrap_or(socket_timeout), + } + } +} + +impl Default for Timeouts { + fn default() -> Self { + Self::new(DEFAULT_TIMEOUT, None) + } } #[derive(Clone, Debug, PartialEq)] diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index 41729bfd107..68170757610 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -1,15 +1,20 @@ -use std::time::Duration; - use redis::{ Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value, aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig}, cluster::{ClusterClientBuilder, ClusterConfig}, cluster_async::ClusterConnection, + io::tcp::TcpSettings, }; -use crate::ClusterEndpoint; +use crate::{ClusterEndpoint, Timeouts, error::Deployment}; -const RETRIES: u32 = 3; +// redis-py 5.3.1 redis/asyncio/connection.py: without retry options a connection gets Retry(NoBackoff(), 0) +const RECONNECT_RETRIES: usize = 0; +// redis-py 5.3.1 redis/asyncio/cluster.py: cluster_error_retry_attempts=3 +const CLUSTER_RETRIES: u32 = 3; +// redis-py 5.3.1 redis/asyncio/cluster.py sleeps a flat 0.25s before retrying after CLUSTERDOWN. +// redis-rs draws the wait from [min, max), so a 1ms window pins it to 250ms +const CLUSTER_RETRY_WAIT_MS: u64 = 250; #[derive(Clone)] pub(crate) enum Connection { @@ -18,20 +23,36 @@ pub(crate) enum Connection { } impl Connection { - pub(crate) fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult { + pub(crate) fn deployment(&self) -> Deployment { + match self { + Self::Standalone(_) => Deployment::Standalone, + Self::Cluster(_) => Deployment::Cluster, + } + } + + pub(crate) fn standalone(info: ConnectionInfo, timeouts: Timeouts) -> RedisResult { let config = ConnectionManagerConfig::new() - .set_connection_timeout(Some(timeout)) - .set_response_timeout(Some(timeout)) - .set_number_of_retries(RETRIES as usize); + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)) + .set_number_of_retries(RECONNECT_RETRIES); + let info = info.set_tcp_settings(TcpSettings::default().set_nodelay(true)); let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?; Ok(Self::Standalone(manager)) } - pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult { + pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeouts: Timeouts) -> RedisResult { let builder = ClusterClientBuilder::new(endpoint.addrs()) - .connection_timeout(timeout) - .response_timeout(timeout) - .retries(RETRIES); + .tcp_settings( + TcpSettings::default() + .set_nodelay(true) + .set_keepalive(socket2::TcpKeepalive::new()), + ) + .connection_timeout(timeouts.connect) + .response_timeout(timeouts.response) + .retries(CLUSTER_RETRIES) + .retry_wait_formula(0, 2) + .min_retry_wait(CLUSTER_RETRY_WAIT_MS) + .max_retry_wait(CLUSTER_RETRY_WAIT_MS + 1); let builder = match &endpoint.username { Some(username) => builder.username(username.expose()), None => builder, @@ -40,9 +61,9 @@ impl Connection { Some(password) => builder.password(password.expose()), None => builder, }; - let config = ClusterConfig::new() - .set_connection_timeout(timeout) - .set_response_timeout(timeout); + // The builder's response timeout bounds the whole request, but with_config overwrites each node's + // response timeout with this one even when unset, so it has to be repeated here + let config = ClusterConfig::new().set_response_timeout(timeouts.response); Ok(Self::Cluster( builder .build()? diff --git a/litellm-rust/crates/redis-client/src/error.rs b/litellm-rust/crates/redis-client/src/error.rs index ed3b2c74004..683832bc078 100644 --- a/litellm-rust/crates/redis-client/src/error.rs +++ b/litellm-rust/crates/redis-client/src/error.rs @@ -16,29 +16,62 @@ pub enum Error { #[error("connecting to Redis")] pub struct ConnectError(#[from] RedisError); -impl From for Error { - fn from(error: RedisError) -> Self { - match error.kind() { - ErrorKind::Server( - ServerErrorKind::BusyLoading - | ServerErrorKind::ClusterDown - | ServerErrorKind::MasterDown - | ServerErrorKind::TryAgain - | ServerErrorKind::ReadOnly, - ) => Self::Transport(error), - ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error), - ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error), +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Deployment { + Standalone, + Cluster, +} + +// redis-py 5.3.1 redis/_parsers/base.py EXCEPTION_CLASSES["ERR"] entries that subclass ConnectionError +const CONNECTION_ERR_MESSAGES: &[&str] = &[ + "max number of clients reached", + "invalid password", + "Client sent AUTH, but no password is set", + "AUTH called without any password configured for the default user. Are you sure your configuration is correct?", +]; + +impl Error { + pub(crate) fn classify(error: RedisError, deployment: Deployment) -> Self { + match (error.kind(), deployment) { + (ErrorKind::Parse | ErrorKind::UnexpectedReturnType, _) => Self::Decode(error), + (ErrorKind::Server(kind), _) if is_unavailable(kind, deployment) => { + Self::Transport(error) + } + (ErrorKind::Server(ServerErrorKind::ResponseError), _) + if error + .detail() + .is_some_and(|detail| CONNECTION_ERR_MESSAGES.contains(&detail)) => + { + Self::Transport(error) + } + (ErrorKind::Extension, _) if matches!(error.code(), Some("NOAUTH" | "WRONGPASS")) => { + Self::Transport(error) + } + (ErrorKind::Server(_) | ErrorKind::Extension | ErrorKind::InvalidClientConfig, _) => { + Self::Server(error) + } _ => Self::Transport(error), } } } +fn is_unavailable(kind: ServerErrorKind, deployment: Deployment) -> bool { + matches!( + (kind, deployment), + (ServerErrorKind::BusyLoading, _) + | ( + ServerErrorKind::ClusterDown | ServerErrorKind::MasterDown, + Deployment::Cluster + ) + ) +} + #[cfg(test)] mod tests { use redis::{RedisError, Value}; use rstest::rstest; - use super::Error; + use super::{Deployment, Error}; fn reply(resp: &str) -> RedisError { match redis::parse_redis_value(resp.as_bytes()).unwrap() { @@ -47,23 +80,84 @@ mod tests { } } + fn is_transport(resp: &str, deployment: Deployment) -> bool { + match Error::classify(reply(resp), deployment) { + Error::Transport(_) => true, + Error::Server(_) => false, + other => panic!("{resp:?} classified as {other:?}"), + } + } + #[rstest] #[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")] #[case::no_script("-NOSCRIPT No matching script\r\n")] #[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")] #[case::generic("-ERR value is not an integer or out of range\r\n")] #[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")] - fn rejections_are_server_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Server(_))); + #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] + #[case::out_of_memory("-OOM command not allowed when used memory > 'maxmemory'.\r\n")] + #[case::auth_arity("-ERR wrong number of arguments for 'auth' command\r\n")] + #[case::max_clients_with_suffix("-ERR max number of clients reached, try later\r\n")] + fn rejections_are_server_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster)] deployment: Deployment, + ) { + assert!(!is_transport(resp, deployment)); } #[rstest] #[case::loading("-LOADING Redis is loading the dataset in memory\r\n")] - #[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")] - #[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")] - #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] - #[case::read_only("-READONLY You can't write against a read only replica.\r\n")] - fn unavailable_nodes_are_transport_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Transport(_))); + #[case::no_auth("-NOAUTH Authentication required.\r\n")] + #[case::wrong_pass("-WRONGPASS invalid username-password pair or user is disabled.\r\n")] + #[case::max_clients("-ERR max number of clients reached\r\n")] + #[case::invalid_password("-ERR invalid password\r\n")] + #[case::no_password_set("-ERR Client sent AUTH, but no password is set\r\n")] + #[case::no_default_password( + "-ERR AUTH called without any password configured for the default user. Are you sure your configuration is correct?\r\n" + )] + fn unavailable_nodes_are_transport_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster)] deployment: Deployment, + ) { + assert!(is_transport(resp, deployment)); + } + + #[rstest] + #[case::master_down_on_standalone( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Standalone, + false + )] + #[case::master_down_on_cluster( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Cluster, + true + )] + #[case::cluster_down_on_standalone( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Standalone, + false + )] + #[case::cluster_down_on_cluster( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Cluster, + true + )] + #[case::read_only_on_standalone( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Standalone, + false + )] + #[case::read_only_on_cluster( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Cluster, + false + )] + fn topology_decides_like_redis_py( + #[case] resp: &str, + #[case] deployment: Deployment, + #[case] transport: bool, + ) { + assert_eq!(is_transport(resp, deployment), transport); } } diff --git a/litellm-rust/crates/redis-client/src/events.rs b/litellm-rust/crates/redis-client/src/events.rs new file mode 100644 index 00000000000..e6a400cf87e --- /dev/null +++ b/litellm-rust/crates/redis-client/src/events.rs @@ -0,0 +1,57 @@ +use std::time::Duration; + +use crate::Error; + +pub trait RedisEvents: Send + Sync { + fn call(&self, event: &CallEvent); +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Operation { + Command(&'static str), + Script, + Pipeline { ops: usize }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Outcome { + Succeeded, + Failed(FailureKind), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FailureKind { + Timeout, + Transport, + Server, + Decode, + BreakerOpen, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallEvent { + pub operation: Operation, + pub duration: Duration, + pub outcome: Outcome, +} + +impl FailureKind { + pub fn of(error: &Error) -> Self { + match error { + Error::Transport(error) if error.is_timeout() => Self::Timeout, + Error::Transport(_) => Self::Transport, + Error::Server(_) => Self::Server, + Error::Decode(_) => Self::Decode, + Error::BreakerOpen => Self::BreakerOpen, + } + } +} + +impl Outcome { + pub(crate) fn of(result: &Result) -> Self { + match result { + Ok(_) => Self::Succeeded, + Err(error) => Self::Failed(FailureKind::of(error)), + } + } +} diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 47ff8677e06..1f6082d7b88 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -3,18 +3,19 @@ mod commands; mod config; mod connection; mod error; +mod events; mod namespace; mod pipeline; mod script; -#[cfg(test)] -mod test_support; pub use client::{RedisClient, Ttl}; +pub use commands::{Amount, Expiry}; pub use config::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig, + ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig, Timeouts, }; pub use error::{ConnectError, Error}; +pub use events::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; pub use namespace::Namespace; pub use pipeline::Pipeline; -pub use redis::Value; +pub use redis::{ErrorKind as RedisErrorKind, RedisError, 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 index e54c65e9605..c27ae56f4ce 100644 --- a/litellm-rust/crates/redis-client/src/pipeline.rs +++ b/litellm-rust/crates/redis-client/src/pipeline.rs @@ -2,11 +2,13 @@ use std::{collections::HashMap, time::Duration}; use futures_util::future::join_all; use redis::{ - Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, + Cmd, ErrorKind, RedisError, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, cluster_routing::Slot, }; -use crate::{Error, Namespace, Script, commands, connection::Connection}; +use crate::{ + Amount, Error, Namespace, Script, commands, connection::Connection, script::ScriptCall, +}; pub struct Pipeline { namespace: Namespace, @@ -14,15 +16,26 @@ pub struct Pipeline { } struct Op { - command: Cmd, + entry: Entry, key: String, - script: Option, + idempotent: bool, } -struct ScriptCall { - script: Script, - keys: Vec, - args: Vec>, +/// What a pipeline sends for one op. Only a [`ScriptCall`] can be sent again, so a retry can never +/// repeat a command that already ran. +enum Entry { + /// The op's reply is the first command's, like `RedisBatch`'s `INCRBYFLOAT` plus `EXPIRE`. + Commands(Vec), + Script(ScriptCall), +} + +impl Entry { + fn commands(&self) -> Vec { + match self { + Self::Commands(commands) => commands.clone(), + Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)], + } + } } type Settled = (usize, Result); @@ -37,59 +50,57 @@ impl Pipeline { pub fn get(self, key: &str) -> Self { let key = self.namespace.apply(key); - self.command(commands::get(key.clone()), key) + self.commands(vec![commands::get(key.clone())], key, true) } 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) + self.commands(vec![commands::set(key.clone(), value, ttl)], key, true) } pub fn del(self, key: &str) -> Self { let key = self.namespace.apply(key); - self.command(commands::del(vec![key.clone()]), key) + self.commands(vec![commands::del(vec![key.clone()])], key, true) } pub fn expire(self, key: &str, ttl: Duration) -> Self { let key = self.namespace.apply(key); - self.command(commands::expire(key.clone(), ttl), key) + self.commands(vec![commands::expire(key.clone(), ttl)], key, true) } - pub fn incr_by(self, key: &str, delta: i64) -> Self { + /// `INCRBYFLOAT`, then `EXPIRE` when a TTL is given, like `RedisBatch` and + /// `async_increment_pipeline` send it. + pub fn incr_by_float(self, key: &str, amount: Amount, ttl: Option) -> 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) + let increment = commands::incr_by_float(key.clone(), amount); + let commands = match ttl { + Some(ttl) => vec![increment, commands::expire(key.clone(), ttl)], + None => vec![increment], + }; + self.commands(commands, key, false) } 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, - }), + let key = keys.first().cloned().unwrap_or_default(); + let call = ScriptCall { + script: script.clone(), + keys, + args: args.to_redis_args(), }; - self.push(op) + self.push(Entry::Script(call), key, false) } - fn command(self, command: Cmd, key: String) -> Self { - self.push(Op { - command, + fn commands(self, commands: Vec, key: String, idempotent: bool) -> Self { + self.push(Entry::Commands(commands), key, idempotent) + } + + fn push(mut self, entry: Entry, key: String, idempotent: bool) -> Self { + self.ops.push(Op { + entry, key, - script: None, - }) - } - - fn push(mut self, op: Op) -> Self { - self.ops.push(op); + idempotent, + }); self } @@ -102,7 +113,7 @@ impl Pipeline { let settled = join_all( groups .iter() - .map(|group| self.run_group(connection.clone(), group)), + .map(|batches| self.run_in_order(connection, batches)), ) .await; let mut settled: Vec = settled.into_iter().flatten().collect(); @@ -110,29 +121,73 @@ impl Pipeline { settled.into_iter().map(|(_, result)| result).collect() } - fn groups(&self, connection: &Connection) -> Vec> { + async fn run_in_order(&self, connection: &Connection, batches: &[Vec]) -> Vec { + let mut settled = Vec::with_capacity(batches.iter().map(Vec::len).sum()); + for batch in batches { + settled.extend(self.run_group(connection.clone(), batch).await); + } + settled + } + + fn groups(&self, connection: &Connection) -> Vec>> { match connection { _ if self.ops.is_empty() => Vec::new(), - Connection::Standalone(_) => vec![(0..self.ops.len()).collect()], + Connection::Standalone(_) => vec![vec![(0..self.ops.len()).collect()]], Connection::Cluster(_) => { let mut by_slot: HashMap> = HashMap::new(); - for (index, op) in self.ops.iter().enumerate() { + for index in 0..self.ops.len() { by_slot - .entry(Slot::for_key(&op.key)) + .entry(Slot::for_key(&self.ops[index].key)) .or_default() .push(index); } - by_slot.into_values().collect() + by_slot + .into_values() + .map(|indices| self.retry_safe_batches(indices)) + .collect() } } } + fn retry_safe_batches(&self, indices: Vec) -> Vec> { + indices + .chunk_by(|&left, &right| self.ops[left].idempotent && self.ops[right].idempotent) + .map(<[usize]>::to_vec) + .collect() + } + + /// A script the server had not cached ran nothing, so it alone runs again, the way litellm's + /// `RedisBatch` reruns it through redis-py's `Script`: `EVALSHA`, then `SCRIPT LOAD` and + /// `EVALSHA` on `NOSCRIPT`. + async fn settle_script( + &self, + connection: Connection, + index: usize, + reply: RedisResult, + ) -> (usize, RedisResult) { + match (&self.ops[index].entry, reply) { + (Entry::Script(call), Err(error)) if is_missing_script(&error) => { + (index, rerun_missing_script(connection, call).await) + } + (_, reply) => (index, reply), + } + } + async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec { + let deployment = connection.deployment(); let mut pipeline = redis::pipe(); pipeline.ignore_errors(); - for &index in indices { - pipeline.add_command(self.ops[index].command.clone()); - } + let widths: Vec = indices + .iter() + .map(|&index| { + let commands = self.ops[index].entry.commands(); + let width = commands.len(); + for command in commands { + pipeline.add_command(command); + } + width + }) + .collect(); let replies: RedisResult>> = pipeline.query_async(&mut connection).await; let replies = match replies { @@ -140,153 +195,50 @@ impl Pipeline { Err(error) => { return indices .iter() - .map(|&index| (index, Err(Error::from(error.clone())))) + .map(|&index| (index, Err(Error::classify(error.clone(), deployment)))) .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?), - } + let mut replies = replies.into_iter(); + let per_op: Vec<(usize, RedisResult)> = indices + .iter() + .zip(widths) + .map(|(&index, width)| (index, first_reply(replies.by_ref().take(width)))) + .collect(); + join_all( + per_op + .into_iter() + .map(|(index, reply)| self.settle_script(connection.clone(), index, reply)), + ) + .await + .into_iter() + .map(|(index, reply)| { + ( + index, + reply.map_err(|error| Error::classify(error, deployment)), + ) + }) + .collect() } } -#[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:?}"); - } +async fn rerun_missing_script(mut connection: Connection, call: &ScriptCall) -> RedisResult { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await +} + +fn is_missing_script(error: &RedisError) -> bool { + error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) +} + +// RedisBatch settles an op from its first reply and fails it on any reply that failed +fn first_reply(replies: impl Iterator>) -> RedisResult { + replies + .reduce(|first, next| match (first, next) { + (Err(error), _) | (Ok(_), Err(error)) => Err(error), + (first, Ok(_)) => first, + }) + .unwrap_or(Ok(Value::Nil)) } diff --git a/litellm-rust/crates/redis-client/src/script.rs b/litellm-rust/crates/redis-client/src/script.rs index c5d73133ecc..6dbdbafe9c9 100644 --- a/litellm-rust/crates/redis-client/src/script.rs +++ b/litellm-rust/crates/redis-client/src/script.rs @@ -1,13 +1,18 @@ +use std::sync::Arc; + use redis::{Cmd, ScriptInvocation}; +/// A Lua script, cheap to clone into every call that runs it. #[derive(Clone)] -pub struct Script(redis::Script); +pub struct Script(Arc); impl Script { pub fn new(source: &str) -> Self { - Self(redis::Script::new(source)) + Self(Arc::new(redis::Script::new(source))) } + /// The `EVALSHA` a pipeline sends, like redis-py. A `NOSCRIPT` reply is settled by + /// [`Script::invocation`] alone, which loads the script and sends it again. pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec]) -> Cmd { redis::cmd("EVALSHA") .arg(self.0.get_hash()) @@ -24,65 +29,10 @@ impl Script { } } -#[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); - } +/// A script call in a pipeline, kept apart from plain commands because it is the one entry a +/// pipeline may send again. +pub(crate) struct ScriptCall { + pub(crate) script: Script, + pub(crate) keys: Vec, + pub(crate) args: Vec>, } diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs deleted file mode 100644 index e503739a942..00000000000 --- a/litellm-rust/crates/redis-client/src/test_support.rs +++ /dev/null @@ -1,126 +0,0 @@ -use std::time::Duration; - -use litellm_circuit_breaker::Breaker; -use litellm_redis_testing::{RedisCluster, RedisServer}; -use litellm_secret_value::SecretValue; -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(SecretValue::new(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 { - self.guarded_client(namespace, timeout, None).await - } - - pub(crate) async fn guarded_client( - &self, - namespace: &str, - timeout: Duration, - breaker: Option, - ) -> RedisClient { - RedisClient::new(RedisConfig { - endpoint: self.endpoint(), - namespace: Namespace::new(namespace), - timeout, - breaker, - }) - .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 stop(&self) { - match self { - Self::Standalone(server) => server.stop().await.unwrap(), - Self::Cluster(cluster) => cluster.stop().await.unwrap(), - } - } - - pub(crate) async fn resume(&self) { - match self { - Self::Standalone(server) => server.resume().await.unwrap(), - Self::Cluster(cluster) => cluster.resume().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(); -} diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs new file mode 100644 index 00000000000..1a441e6e787 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -0,0 +1,780 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use litellm_redis_testing::{RedisDeployment, RedisServer, Topology, eventually}; +use rstest::rstest; + +use litellm_circuit_breaker::{Breaker, Config, State}; + +use litellm_redis_client::{ + Amount, CallEvent, DEFAULT_TIMEOUT, Endpoint, Error, Expiry, FailureKind, HostEndpoint, + Namespace, Operation, Outcome, RedisClient, RedisConfig, RedisEvents, Script, Timeouts, Ttl, +}; +use litellm_secret_value::SecretValue; +use support::{KEEP, connect, connect_guarded, connections, endpoint}; + +async fn persist(client: &RedisClient, key: &str) { + let set = Script::new("return redis.call('SET', KEYS[1], '0')"); + client + .eval::<()>(&set, &[key], Vec::::new()) + .await + .unwrap(); +} + +#[derive(Clone, Copy, Debug)] +enum Counter { + IncrBy, + IncrByFloat, +} + +const MINUTE: Duration = Duration::from_secs(60); + +async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry) { + match (counter, expiry) { + (Counter::IncrBy, Expiry::KeepExisting(ttl)) => { + client.incr_by(key, 1, ttl).await.unwrap(); + } + (Counter::IncrBy, Expiry::Refresh(_)) => unreachable!("increment_cache never refreshes"), + (Counter::IncrByFloat, expiry) => { + client + .incr_by_float(key, Amount::Float(0.5), expiry) + .await + .unwrap(); + } + } +} + +async fn persist_as(client: &RedisClient, key: &str) { + let source = "return redis.call('SET', KEYS[1], '1')"; + client + .eval::(&Script::new(source), &[key], Vec::::new()) + .await + .unwrap(); +} + +fn remaining(ttl: Ttl) -> Duration { + match ttl { + Ttl::Expires(left) => left, + other => panic!("expected an expiring key, got {other:?}"), + } +} + +#[rstest] +#[case::increment_cache(Counter::IncrBy, Expiry::KeepExisting(MINUTE))] +#[case::async_increment(Counter::IncrByFloat, Expiry::KeepExisting(MINUTE))] +#[case::async_increment_refreshing(Counter::IncrByFloat, Expiry::Refresh(MINUTE))] +#[tokio::test] +async fn a_counter_write_expires_a_new_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[case] counter: Counter, + #[case] expiry: Expiry, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + write(&client, counter, "fresh", expiry).await; + + let left = remaining(client.ttl("fresh").await.unwrap()); + assert!(left > Duration::from_secs(50) && left <= MINUTE, "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_leaves_a_running_ttl_and_refresh_resets_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let short = Expiry::Refresh(Duration::from_secs(30)); + let long = Duration::from_secs(600); + write(&client, Counter::IncrByFloat, "kept", short).await; + write(&client, Counter::IncrByFloat, "refreshed", short).await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(Duration::from_secs(30)), + ) + .await; + + write( + &client, + Counter::IncrByFloat, + "kept", + Expiry::KeepExisting(long), + ) + .await; + write( + &client, + Counter::IncrByFloat, + "refreshed", + Expiry::Refresh(long), + ) + .await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(long), + ) + .await; + + assert!(remaining(client.ttl("kept").await.unwrap()) <= Duration::from_secs(30)); + assert!(remaining(client.ttl("kept_int").await.unwrap()) <= Duration::from_secs(30)); + assert!(remaining(client.ttl("refreshed").await.unwrap()) > Duration::from_secs(500)); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_expires_a_key_that_had_no_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(Counter::IncrBy, Counter::IncrByFloat)] counter: Counter, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + persist_as(&client, "persistent").await; + + write(&client, counter, "persistent", KEEP).await; + + let left = remaining(client.ttl("persistent").await.unwrap()); + assert!(left > Duration::from_secs(50), "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn counters_send_the_commands_redis_cache_sends( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + client.incr_by("int", 1, MINUTE).await.unwrap(); + client + .incr_by_float("kept", Amount::Int(1), KEEP) + .await + .unwrap(); + client + .incr_by_float("refreshed", Amount::Float(0.5), Expiry::Refresh(MINUTE)) + .await + .unwrap(); + client + .run( + &client + .pipeline() + .incr_by_float("piped", Amount::Float(0.25), Some(MINUTE)) + .incr_by_float("bare", Amount::Float(0.25), None), + ) + .await; + + let stats = support::command_stats(&server).await; + let calls = |command: &str| support::calls(&stats, command); + assert_eq!(calls("incrby"), 1); + assert_eq!(calls("incrbyfloat"), 4); + assert_eq!(calls("ttl"), 2); + assert_eq!(calls("expire"), 4); + assert_eq!( + calls("eval") + calls("evalsha") + calls("pttl") + calls("pexpire"), + 0 + ); +} + +#[derive(Default)] +struct Recorded(Mutex>); + +#[rstest] +#[case::beyond_double_precision(9_007_199_254_740_992_i64, 1_i64)] +#[case::negative_beyond_double_precision(-9_007_199_254_740_992_i64, -1_i64)] +#[case::maximum_integer(i64::MAX - 1, 1_i64)] +#[case::minimum_integer(i64::MIN + 1, -1_i64)] +#[tokio::test] +async fn integer_increments_preserve_every_bit( + #[case] initial: i64, + #[case] delta: i64, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .set("integer", initial, Duration::from_secs(60)) + .await + .unwrap(); + let expected = initial + delta; + let result = client.incr_by("integer", delta, MINUTE).await.unwrap(); + assert_eq!(result, expected); + assert_eq!(client.get::("integer").await.unwrap(), Some(expected)); +} + +#[rstest] +#[tokio::test] +async fn outage_events_report_the_failure_kind( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] timeout: bool, + #[values(false, true)] pipeline: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + endpoint: endpoint(&server), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(300), None), + breaker: None, + events: Some(recorded.clone()), + }) + .unwrap(); + client.ping().await.unwrap(); + if timeout { + server.freeze().await.unwrap(); + } else { + server.stop().await.unwrap(); + } + let result = if pipeline { + client + .run(&client.pipeline().get("missing")) + .await + .remove(0) + .map(drop) + } else { + client.get::("missing").await.map(drop) + }; + if timeout { + server.thaw().await.unwrap(); + } else { + server.resume().await.unwrap(); + } + assert!(matches!(result, Err(Error::Transport(_)))); + let expected = if timeout { + FailureKind::Timeout + } else { + FailureKind::Transport + }; + assert_eq!( + recorded.0.lock().unwrap().last().unwrap().outcome, + Outcome::Failed(expected) + ); +} + +impl RedisEvents for Recorded { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(*event); + } +} + +#[rstest] +#[tokio::test] +async fn every_call_reports_its_operation_and_outcome( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + endpoint: endpoint(&server), + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: Some(recorded.clone()), + }) + .unwrap(); + client + .set("list", "item", Duration::from_secs(60)) + .await + .unwrap(); + + let _ = client.incr_by("list", 1, Duration::from_secs(60)).await; + let _ = client.get::("missing").await; + let _ = client.run(&client.pipeline().get("a").get("b")).await; + let failures = client + .run(&client.pipeline().incr_by_float( + "list", + Amount::Int(1), + Some(Duration::from_secs(60)), + )) + .await; + assert!(matches!(failures[0], Err(Error::Server(_)))); + + let seen: Vec<(Operation, Outcome)> = recorded + .0 + .lock() + .unwrap() + .iter() + .map(|event| (event.operation, event.outcome)) + .collect(); + assert_eq!( + seen, + vec![ + (Operation::Command("SET"), Outcome::Succeeded), + ( + Operation::Command("INCRBY"), + Outcome::Failed(FailureKind::Server) + ), + (Operation::Command("GET"), Outcome::Succeeded), + (Operation::Pipeline { ops: 2 }, Outcome::Succeeded), + ( + Operation::Pipeline { ops: 1 }, + Outcome::Failed(FailureKind::Server) + ), + ] + ); +} + +#[rstest] +#[tokio::test] +async fn namespaces_isolate_keys_on_the_server( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let a = connect(&server, "a", DEFAULT_TIMEOUT); + let b = connect(&server, "b", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + + a.set("x", "1", Duration::from_secs(60)).await.unwrap(); + + assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); + assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); + assert_eq!( + bare.get::("a:x").await.unwrap().as_deref(), + Some("1") + ); + assert_eq!(b.get::("x").await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_command_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .set("list", "item", Duration::from_secs(60)) + .await + .unwrap(); + let before = connections(&server).await; + + let rejected = client.incr_by("list", 1, Duration::from_secs(60)).await; + + assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); + assert_eq!( + client.get::("list").await.unwrap().as_deref(), + Some("item") + ); + assert_eq!(connections(&server).await, before); +} + +#[rstest] +#[tokio::test] +async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .set("float", "15.0", Duration::from_secs(60)) + .await + .unwrap(); + + let decoded = client.get::("float").await; + + assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); + assert_eq!(client.get::("missing").await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + client + .set("first", "1", Duration::from_secs(60)) + .await + .unwrap(); + client + .set("second", "2", Duration::from_secs(60)) + .await + .unwrap(); + + server.freeze().await.unwrap(); + let timed_out = client.get::("first").await; + server.thaw().await.unwrap(); + + assert!( + matches!(timed_out, Err(Error::Transport(_))), + "{timed_out:?}" + ); + assert_eq!( + client.get::("second").await.unwrap().as_deref(), + Some("2") + ); +} + +#[rstest] +#[tokio::test] +async fn the_same_client_recovers_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + + server.restart().await.unwrap(); + + let recovered = eventually(async || { + client + .incr_by("after_restart", 1, Duration::from_secs(60)) + .await + .ok() + }) + .await; + assert_eq!(recovered, 1); +} + +#[rstest] +#[tokio::test] +async fn a_client_created_while_redis_is_down_works_once_redis_is_up( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server.stop().await.unwrap(); + + let client = connect(&server, "", Duration::from_millis(500)); + let started = Instant::now(); + let while_down = client.ping().await; + let failed_in = started.elapsed(); + server.resume().await.unwrap(); + + assert!(failed_in < Duration::from_secs(2), "{failed_in:?}"); + + assert!( + matches!(while_down, Err(Error::Transport(_))), + "{while_down:?}" + ); + let recovered = eventually(async || client.ping().await.ok()).await; + assert_eq!(recovered, ()); +} + +#[rstest] +#[tokio::test] +async fn empty_batches_succeed_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + assert_eq!( + client.mget::(&[]).await.unwrap(), + Vec::>::new() + ); + assert_eq!(client.del(&[]).await.unwrap(), 0); +} + +#[rstest] +#[tokio::test] +async fn mget_and_del_span_slots_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client.set("a", "1", Duration::from_secs(60)).await.unwrap(); + client.set("c", "3", Duration::from_secs(60)).await.unwrap(); + + let values = client.mget::(&["c", "b", "a"]).await.unwrap(); + let deleted = client.del(&["a", "b", "c"]).await.unwrap(); + + assert_eq!( + values, + vec![Some("3".to_owned()), None, Some("1".to_owned())] + ); + assert_eq!(deleted, 2); + assert_eq!( + client.mget::(&["a", "c"]).await.unwrap(), + vec![None, None] + ); +} + +#[rstest] +#[tokio::test] +async fn set_nx_only_writes_a_missing_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let first = client + .set_nx("lock", "owner-1", Duration::from_secs(60)) + .await + .unwrap(); + let second = client + .set_nx("lock", "owner-2", Duration::from_secs(60)) + .await + .unwrap(); + + assert!(first); + assert!(!second); + assert_eq!( + client.get::("lock").await.unwrap().as_deref(), + Some("owner-1") + ); +} + +#[rstest] +#[tokio::test] +async fn ttl_reports_missing_persistent_and_remaining_time( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + persist(&client, "persistent").await; + client + .set("expiring", "1", Duration::from_secs(60)) + .await + .unwrap(); + client + .incr_by_float("expired_later", Amount::Float(0.5), KEEP) + .await + .unwrap(); + assert!( + client + .expire("expired_later", Duration::from_secs(30)) + .await + .unwrap() + ); + + assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); + assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); + assert!(matches!( + client.ttl("expiring").await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); + assert!(matches!( + client.ttl("expired_later").await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) + )); + assert!( + !client + .expire("missing", Duration::from_secs(30)) + .await + .unwrap() + ); +} + +#[rstest] +#[tokio::test] +async fn incr_by_float_accumulates_fractions( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + client + .incr_by_float("spend", Amount::Float(0.25), KEEP) + .await + .unwrap(); + + assert_eq!( + client + .incr_by_float("spend", Amount::Float(0.5), KEEP) + .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 = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let tasks: Vec<_> = (0..100) + .map(|_| { + let client = client.clone(); + tokio::spawn(async move { + for _ in 0..10 { + client + .incr_by("shared", 1, Duration::from_secs(60)) + .await + .unwrap(); + } + }) + }) + .collect(); + for task in tasks { + task.await.unwrap(); + } + + assert_eq!(client.get::("shared").await.unwrap(), Some(1000)); +} + +fn breaker(failure_threshold: u32) -> Breaker { + Breaker::new( + Config { + failure_threshold, + recovery_timeout: Duration::from_millis(200), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ) +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + + server.freeze().await.unwrap(); + let first = client.get::("k").await; + let second = client.get::("k").await; + let started = Instant::now(); + let third = client.get::("k").await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); + assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); + assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + + let recovered = eventually(async || client.get::("k").await.ok()).await; + assert_eq!(recovered, None); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn rejected_commands_never_open_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded(&server, "", DEFAULT_TIMEOUT, Some(breaker.clone())); + client + .set("list", "item", Duration::from_secs(60)) + .await + .unwrap(); + + let rejected = [ + client.incr_by("list", 1, Duration::from_secs(60)).await, + client.incr_by("list", 1, Duration::from_secs(60)).await, + client.incr_by("list", 1, Duration::from_secs(60)).await, + ]; + + assert!( + matches!(rejected[2], Err(Error::Server(_))), + "{:?}", + rejected[2] + ); + assert!( + matches!(breaker.state(), State::Closed { hard: 0, .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_fails_every_pipeline_op_without_sending( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker(1))); + let pipeline = client + .pipeline() + .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) + .incr_by_float("c", Amount::Int(1), Some(Duration::from_secs(60))); + + server.freeze().await.unwrap(); + let failed = client.run(&pipeline).await; + let rejected = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + assert!( + matches!(failed[0], Err(Error::Transport(_))), + "{:?}", + failed[0] + ); + assert!( + matches!(rejected[0], Err(Error::BreakerOpen)), + "{:?}", + rejected[0] + ); + assert!( + matches!(rejected[1], Err(Error::BreakerOpen)), + "{:?}", + rejected[1] + ); +} + +#[rstest] +#[tokio::test] +async fn a_host_endpoint_connects_to_its_database() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let connect = |endpoint| { + RedisClient::new(RedisConfig { + endpoint, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + }) + }; + let url = |url: String| Endpoint::Url(SecretValue::new(url)); + let in_db_3 = connect(Endpoint::Host(HostEndpoint { + host: "127.0.0.1".to_owned(), + port, + username: None, + password: None, + tls: false, + db: 3, + })) + .unwrap(); + let in_db_0 = connect(url(server.url().to_owned())).unwrap(); + let via_url = connect(url(format!("{}/3", server.url()))).unwrap(); + + in_db_3 + .set("k", "3", Duration::from_secs(60)) + .await + .unwrap(); + + assert_eq!(in_db_0.get::("k").await.unwrap(), None); + assert_eq!( + via_url.get::("k").await.unwrap().as_deref(), + Some("3") + ); +} + +#[rstest] +#[case::int(Amount::Int(1), "1")] +#[case::float_keeps_its_point(Amount::Float(1.0), "1.0")] +#[case::small_float(Amount::Float(1e-5), "1e-05")] +#[case::large_float(Amount::Float(1e16), "1e+16")] +#[case::fraction(Amount::Float(0.1 + 0.2), "0.30000000000000004")] +fn amounts_are_written_like_redis_py_encodes_them(#[case] amount: Amount, #[case] expected: &str) { + assert_eq!( + redis::ToRedisArgs::to_redis_args(&amount), + vec![expected.as_bytes().to_vec()] + ); +} diff --git a/litellm-rust/crates/redis-client/tests/connection.rs b/litellm-rust/crates/redis-client/tests/connection.rs new file mode 100644 index 00000000000..44a7ded9335 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/connection.rs @@ -0,0 +1,63 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_client::{Endpoint, Error, HostEndpoint}; +use litellm_redis_testing::RedisServer; +use litellm_secret_value::SecretValue; +use rstest::rstest; + +use support::connect_to; + +#[rstest] +#[tokio::test] +async fn an_out_of_range_database_is_a_server_rejection() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let client = connect_to( + Endpoint::Host(HostEndpoint { + host: "127.0.0.1".into(), + port, + username: None, + password: None, + tls: false, + db: 999, + }), + "", + Duration::from_secs(2), + None, + ); + let result = client.ping().await; + assert!(matches!(result, Err(Error::Server(_))), "{result:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_password_is_a_transport_failure() { + let server = RedisServer::start().await.unwrap(); + server + .cli(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let with_password = |password: &str| { + connect_to( + Endpoint::Host(HostEndpoint { + host: "127.0.0.1".to_owned(), + port, + username: None, + password: Some(SecretValue::new(password)), + tls: false, + db: 0, + }), + "", + Duration::from_secs(2), + None, + ) + }; + with_password("right").ping().await.unwrap(); + + let rejected = with_password("wrong").ping().await; + + assert!(matches!(rejected, Err(Error::Transport(_))), "{rejected:?}"); +} diff --git a/litellm-rust/crates/redis-client/tests/pipeline.rs b/litellm-rust/crates/redis-client/tests/pipeline.rs new file mode 100644 index 00000000000..89d30d35d50 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pipeline.rs @@ -0,0 +1,309 @@ +mod support; + +use std::time::{Duration, Instant}; + +use litellm_redis_client::Value; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Script}; +use support::connect; + +async fn begin_migration(cluster: &litellm_redis_testing::RedisCluster, key: &str) { + let (_, first) = cluster.nodes().next().unwrap(); + let slot = cluster + .cli_on_node(first, &["CLUSTER", "KEYSLOT", key]) + .await + .unwrap(); + let slot = slot.trim(); + let number: u16 = slot.parse().unwrap(); + let nodes = cluster + .cli_on_node(first, &["CLUSTER", "NODES"]) + .await + .unwrap(); + let owner = nodes.lines().find(|line| line.split_whitespace().skip(8).any(|range| { + let (start, end) = range.split_once('-').unwrap_or((range, range)); + matches!((start.parse::(), end.parse::()), (Ok(start), Ok(end)) if (start..=end).contains(&number)) + })).unwrap(); + let target = nodes.lines().find(|line| *line != owner).unwrap(); + let address = |line: &str| { + let value = line.split_whitespace().nth(1).unwrap(); + value + .split_once('@') + .unwrap() + .0 + .rsplit_once(':') + .unwrap() + .1 + .parse::() + .unwrap() + }; + let source_port = address(owner); + let target_port = address(target); + let source_id = owner.split_whitespace().next().unwrap(); + let target_id = target.split_whitespace().next().unwrap(); + assert_eq!( + cluster + .cli_on_node(source_port, &["GET", key]) + .await + .unwrap() + .trim(), + "0" + ); + assert_eq!( + cluster + .cli_on_node( + target_port, + &["CLUSTER", "SETSLOT", slot, "IMPORTING", source_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + assert_eq!( + cluster + .cli_on_node( + source_port, + &["CLUSTER", "SETSLOT", slot, "MIGRATING", target_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + let state = cluster + .cli_on_node(source_port, &["CLUSTER", "NODES"]) + .await + .unwrap(); + assert!( + state.contains(&format!("[{slot}->-{target_id}]")), + "{state}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() { + let cluster = litellm_redis_testing::RedisCluster::start().await.unwrap(); + let server = RedisDeployment::Cluster(cluster); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let counter = "{migration}:counter"; + client + .set(counter, 0, Duration::from_secs(60)) + .await + .unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!() + }; + begin_migration(cluster, counter).await; + let both_keys = Script::new("return redis.call('MGET', KEYS[1], KEYS[2])"); + + let results = client + .run( + &client + .pipeline() + .incr_by_float(counter, Amount::Int(1), Some(Duration::from_secs(60))) + .eval( + &both_keys, + &[counter, "{migration}:missing"], + Vec::::new(), + ), + ) + .await; + + assert_eq!(client.get::(counter).await.unwrap(), Some(1)); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(&results[1], Err(Error::Server(error)) if error.code() == Some("TRYAGAIN")), + "{:?}", + results[1] + ); +} + +#[rstest] +#[tokio::test] +async fn each_op_settles_on_its_own_in_input_order( + #[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 pipeline = client + .pipeline() + .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) + .incr_by_float("text", Amount::Int(1), Some(Duration::from_secs(60))) + .incr_by_float("c", Amount::Int(5), Some(Duration::from_secs(60))) + .set("b", "2", Duration::from_secs(60)) + .get("b") + .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) + .get("missing"); + + let results = client.run(&pipeline).await; + + assert_eq!(results.len(), 7); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(results[1], Err(Error::Server(_))), + "{:?}", + results[1] + ); + assert_eq!( + results[2].as_ref().unwrap(), + &Value::BulkString(b"5".to_vec()) + ); + 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::BulkString(b"2".to_vec()) + ); + assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); +} + +#[rstest] +#[case::cached(true, Value::BulkString(b"2".to_vec()))] +#[case::lost_by_the_server(false, Value::Nil)] +#[tokio::test] +async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py( + #[case] cached: bool, + #[case] read_in_the_same_pipeline: Value, +) { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.eval::(&incr, &["warmup"], 1).await.unwrap(); + if !cached { + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + } + let pipeline = client + .pipeline() + .eval(&incr, &["counter"], 2) + .get("counter"); + + let results = client.run(&pipeline).await; + + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(2)); + assert_eq!(results[1].as_ref().unwrap(), &read_in_the_same_pipeline); + assert_eq!(client.get::("counter").await.unwrap(), Some(2)); +} + +#[rstest] +#[tokio::test] +async fn an_unreachable_server_fails_every_op_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let pipeline = client + .pipeline() + .incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60))) + .incr_by_float("c", Amount::Int(1), Some(Duration::from_secs(60))); + + server.freeze().await.unwrap(); + let results = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + 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 = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(2)); + + server.freeze().await.unwrap(); + let started = Instant::now(); + let results = client.run(&client.pipeline()).await; + let elapsed = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(results.is_empty()); + assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_lost( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&deployment, "scripts", DEFAULT_TIMEOUT); + let count = Script::new("return redis.call('INCR', KEYS[1])"); + deployment + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + deployment + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + let first = client + .run( + &client + .pipeline() + .incr_by_float( + "{a}spend", + Amount::Float(0.5), + Some(Duration::from_secs(60)), + ) + .eval(&count, &["{a}"], Vec::::new()) + .eval(&count, &["b"], Vec::::new()), + ) + .await; + let second = client + .run( + &client + .pipeline() + .eval(&count, &["{a}"], Vec::::new()), + ) + .await; + + let replies: Vec = first + .into_iter() + .chain(second) + .map(Result::unwrap) + .collect(); + assert_eq!( + replies, + [ + Value::BulkString(b"0.5".to_vec()), + Value::Int(1), + Value::Int(1), + Value::Int(2) + ] + ); + assert_eq!(client.get::("{a}spend").await.unwrap(), Some(0.5)); + let stats = support::command_stats(&deployment).await; + assert!(support::calls(&stats, "evalsha") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "eval"), 0, "{stats}"); +} diff --git a/litellm-rust/crates/redis-client/tests/script.rs b/litellm-rust/crates/redis-client/tests/script.rs new file mode 100644 index 00000000000..23c637989bf --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/script.rs @@ -0,0 +1,106 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::Script; +use litellm_redis_client::{DEFAULT_TIMEOUT, Error, Ttl}; +use support::connect; + +#[rstest] +#[tokio::test] +async fn script_keys_are_namespaced( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "a", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + 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 = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .set("list", "item", Duration::from_secs(60)) + .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 = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.eval::(&incr, &["counter"], 1).await.unwrap(); + + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + assert_eq!(client.eval::(&incr, &["counter"], 1).await.unwrap(), 2); +} + +const INCREMENT_WITH_FLOOR: &str = "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) \ +if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end \ +if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end \ +return count"; + +#[rstest] +#[tokio::test] +async fn pythons_increment_with_floor_script_floors_at_zero_and_expires_in_one_call( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] pipelined: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let floor = Script::new(INCREMENT_WITH_FLOOR); + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + let count = match pipelined { + false => client + .eval::(&floor, &["slots"], (-5, 60)) + .await + .unwrap(), + true => { + let results = client + .run(&client.pipeline().eval(&floor, &["slots"], (-5, 60))) + .await; + redis::from_redis_value(results[0].as_ref().unwrap().clone()).unwrap() + } + }; + + assert_eq!(count, 0); + assert_eq!(client.get::("slots").await.unwrap(), Some(0)); + assert!(matches!( + client.ttl("slots").await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); +} diff --git a/litellm-rust/crates/redis-client/tests/support/mod.rs b/litellm-rust/crates/redis-client/tests/support/mod.rs new file mode 100644 index 00000000000..42fec0c9df4 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/support/mod.rs @@ -0,0 +1,92 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::time::Duration; + +use litellm_circuit_breaker::Breaker; +use litellm_redis_client::{ + ClusterEndpoint, ClusterNode, Endpoint, Expiry, Namespace, RedisClient, RedisConfig, Timeouts, +}; +use litellm_redis_testing::{ClientConnection, RedisDeployment}; +use litellm_secret_value::SecretValue; + +pub const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60)); + +pub fn endpoint(deployment: &RedisDeployment) -> Endpoint { + match deployment { + RedisDeployment::Standalone(server) => Endpoint::Url(SecretValue::new(server.url())), + RedisDeployment::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 fn connect(deployment: &RedisDeployment, namespace: &str, timeout: Duration) -> RedisClient { + connect_guarded(deployment, namespace, timeout, None) +} + +pub fn connect_guarded( + deployment: &RedisDeployment, + namespace: &str, + timeout: Duration, + breaker: Option, +) -> RedisClient { + connect_to(endpoint(deployment), namespace, timeout, breaker) +} + +pub fn connect_to( + endpoint: Endpoint, + namespace: &str, + timeout: Duration, + breaker: Option, +) -> RedisClient { + RedisClient::new(RedisConfig { + endpoint, + namespace: Namespace::new(namespace), + timeouts: Timeouts::new(timeout, None), + breaker, + events: None, + }) + .unwrap() +} + +pub async fn connections(deployment: &RedisDeployment) -> Vec { + deployment.clients().await.unwrap() +} + +/// `INFO commandstats` from every node, so a test can count what was actually sent. +pub async fn command_stats(deployment: &RedisDeployment) -> String { + match deployment { + RedisDeployment::Standalone(server) => server.cli(&["INFO", "commandstats"]).await.unwrap(), + RedisDeployment::Cluster(cluster) => { + let mut stats = String::new(); + for (_, port) in cluster.nodes() { + stats.push_str( + &cluster + .cli_on_node(port, &["INFO", "commandstats"]) + .await + .unwrap(), + ); + } + stats + } + } +} + +/// How many times `command` reached a server, summed across the nodes in `stats`. +pub fn calls(stats: &str, command: &str) -> u64 { + let prefix = format!("cmdstat_{command}:calls="); + stats + .lines() + .filter_map(|line| line.trim().strip_prefix(&prefix)) + .filter_map(|rest| rest.split(',').next()?.parse::().ok()) + .sum() +} diff --git a/litellm-rust/crates/redis-testing/AGENTS.md b/litellm-rust/crates/redis-testing/AGENTS.md index fb7a4cd737a..225e4a71203 100644 --- a/litellm-rust/crates/redis-testing/AGENTS.md +++ b/litellm-rust/crates/redis-testing/AGENTS.md @@ -4,3 +4,7 @@ - `RedisCluster` runs 3 masters in one container. Each node listens on the same port inside and outside and announces 127.0.0.1, so MOVED redirects work from the test process - Outages come from the container (`restart`, `freeze` / `thaw`), never `DEBUG`, which the official image disables - No env gate that makes a test pass without Redis. A test that can't start its container fails +- `cli` / `cli_on_every_node` run `redis-cli -e` inside the container for setup such as ACL users, so tests don't need a second client. `-e` makes an error reply fail the call, and a test whose setup changes server state also asserts that state before relying on it +- Not `redis-test`'s server helpers, which need local `redis-server` and `openssl` binaries +- `RedisDeployment` runs either `Topology`, so a test can take `#[values(Topology::Standalone, Topology::Cluster)]`. `clients()` returns every client connection on every node from `CLIENT LIST`, leaving out the CLI's own +- Server fixtures only. This crate never depends on a Redis client crate, so the client crates can dev-depend on it without a cycle diff --git a/litellm-rust/crates/redis-testing/Cargo.toml b/litellm-rust/crates/redis-testing/Cargo.toml index c9bad227acc..5afef2f1aa9 100644 --- a/litellm-rust/crates/redis-testing/Cargo.toml +++ b/litellm-rust/crates/redis-testing/Cargo.toml @@ -8,4 +8,4 @@ repository.workspace = true [dependencies] testcontainers-modules = { workspace = true, features = ["redis"] } thiserror.workspace = true -tokio = { workspace = true, features = ["io-util"] } +tokio = { workspace = true, features = ["io-util", "time"] } diff --git a/litellm-rust/crates/redis-testing/src/cluster.rs b/litellm-rust/crates/redis-testing/src/cluster.rs index e61f7832b1c..b1373167050 100644 --- a/litellm-rust/crates/redis-testing/src/cluster.rs +++ b/litellm-rust/crates/redis-testing/src/cluster.rs @@ -5,7 +5,8 @@ use testcontainers_modules::testcontainers::{ }; use crate::{ - Error, + ClientConnection, Error, + deployment::parse_client_list, ports::{HOST, start_on_free_ports, wait_for_pong}, standalone::REDIS_TAG, }; @@ -70,6 +71,28 @@ impl RedisCluster { self.ports.iter().map(|&port| (HOST, port)) } + pub async fn cli_on_every_node(&self, args: &[&str]) -> Result<(), Error> { + for port in self.ports { + crate::exec::run(&self.container, crate::exec::cli(port, args)).await?; + } + Ok(()) + } + + pub async fn cli_on_node(&self, port: u16, args: &[&str]) -> Result { + crate::exec::run(&self.container, crate::exec::cli(port, args)).await + } + + pub async fn clients(&self) -> Result, Error> { + let mut clients = Vec::new(); + for port in self.ports { + let reply = + crate::exec::run(&self.container, crate::exec::cli(port, &["CLIENT", "LIST"])) + .await?; + clients.extend(parse_client_list(port, &reply)); + } + Ok(clients) + } + pub async fn restart(&self) -> Result<(), Error> { self.stop().await?; self.resume().await diff --git a/litellm-rust/crates/redis-testing/src/deployment.rs b/litellm-rust/crates/redis-testing/src/deployment.rs new file mode 100644 index 00000000000..5d090eae4ca --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/deployment.rs @@ -0,0 +1,113 @@ +use std::time::{Duration, Instant}; + +use crate::{Error, RedisCluster, RedisServer}; + +const EVENTUALLY_TIMEOUT: Duration = Duration::from_secs(30); + +#[derive(Clone, Copy, Debug)] +pub enum Topology { + Standalone, + Cluster, +} + +pub enum RedisDeployment { + Standalone(RedisServer), + Cluster(RedisCluster), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ClientConnection { + pub node: u16, + pub id: u64, + pub user: String, +} + +impl RedisDeployment { + pub async fn start(topology: Topology) -> Result { + Ok(match topology { + Topology::Standalone => Self::Standalone(RedisServer::start().await?), + Topology::Cluster => Self::Cluster(RedisCluster::start().await?), + }) + } + + pub async fn cli_on_every_node(&self, args: &[&str]) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.cli(args).await.map(drop), + Self::Cluster(cluster) => cluster.cli_on_every_node(args).await, + } + } + + pub async fn clients(&self) -> Result, Error> { + match self { + Self::Standalone(server) => server.clients().await, + Self::Cluster(cluster) => cluster.clients().await, + } + } + + pub async fn restart(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.restart().await, + Self::Cluster(cluster) => cluster.restart().await, + } + } + + pub async fn stop(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.stop().await, + Self::Cluster(cluster) => cluster.stop().await, + } + } + + pub async fn resume(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.resume().await, + Self::Cluster(cluster) => cluster.resume().await, + } + } + + pub async fn freeze(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.freeze().await, + Self::Cluster(cluster) => cluster.freeze().await, + } + } + + pub async fn thaw(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.thaw().await, + Self::Cluster(cluster) => cluster.thaw().await, + } + } +} + +pub(crate) fn parse_client_list(node: u16, reply: &str) -> Vec { + reply + .lines() + .filter(|line| !line.split(' ').any(|field| field == "cmd=client|list")) + .filter_map(|line| { + let field = |name: &str| { + line.split(' ') + .find_map(|field| field.strip_prefix(name)?.strip_prefix('=')) + }; + Some(ClientConnection { + node, + id: field("id")?.parse().ok()?, + user: field("user")?.to_owned(), + }) + }) + .collect() +} + +pub async fn eventually(mut attempt: impl AsyncFnMut() -> Option) -> T { + let deadline = Instant::now() + EVENTUALLY_TIMEOUT; + loop { + if let Some(value) = attempt().await { + return value; + } + assert!( + Instant::now() < deadline, + "condition not met within {EVENTUALLY_TIMEOUT:?}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } +} diff --git a/litellm-rust/crates/redis-testing/src/error.rs b/litellm-rust/crates/redis-testing/src/error.rs index f21cc959812..59a60f37a72 100644 --- a/litellm-rust/crates/redis-testing/src/error.rs +++ b/litellm-rust/crates/redis-testing/src/error.rs @@ -8,4 +8,6 @@ pub enum Error { Unreachable { port: u16 }, #[error("controlling the Redis container")] Container(#[from] testcontainers_modules::testcontainers::TestcontainersError), + #[error("{command} exited with {code}")] + Command { command: String, code: i64 }, } diff --git a/litellm-rust/crates/redis-testing/src/exec.rs b/litellm-rust/crates/redis-testing/src/exec.rs new file mode 100644 index 00000000000..74a94c085ec --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/exec.rs @@ -0,0 +1,31 @@ +use testcontainers_modules::testcontainers::{ + ContainerAsync, Image, + core::{CmdWaitFor, ExecCommand}, +}; + +use crate::Error; + +pub(crate) async fn run( + container: &ContainerAsync, + command: Vec, +) -> Result { + let mut result = container + .exec(ExecCommand::new(command.clone()).with_cmd_ready_condition(CmdWaitFor::exit())) + .await?; + let stdout = String::from_utf8_lossy(&result.stdout_to_vec().await?).into_owned(); + match result.exit_code().await? { + Some(0) | None => Ok(stdout), + Some(code) => Err(Error::Command { + command: command.join(" "), + code, + }), + } +} + +pub(crate) fn cli(port: u16, args: &[&str]) -> Vec { + ["redis-cli", "-e", "-p", &port.to_string()] + .into_iter() + .chain(args.iter().copied()) + .map(str::to_owned) + .collect() +} diff --git a/litellm-rust/crates/redis-testing/src/lib.rs b/litellm-rust/crates/redis-testing/src/lib.rs index 6c2afd7f275..4a1365c1a28 100644 --- a/litellm-rust/crates/redis-testing/src/lib.rs +++ b/litellm-rust/crates/redis-testing/src/lib.rs @@ -1,8 +1,11 @@ mod cluster; +mod deployment; mod error; +mod exec; mod ports; mod standalone; pub use cluster::RedisCluster; +pub use deployment::{ClientConnection, RedisDeployment, Topology, eventually}; pub use error::Error; pub use standalone::RedisServer; diff --git a/litellm-rust/crates/redis-testing/src/standalone.rs b/litellm-rust/crates/redis-testing/src/standalone.rs index 8060ccbf190..add0a791497 100644 --- a/litellm-rust/crates/redis-testing/src/standalone.rs +++ b/litellm-rust/crates/redis-testing/src/standalone.rs @@ -5,7 +5,8 @@ use testcontainers_modules::{ }; use crate::{ - Error, + ClientConnection, Error, + deployment::parse_client_list, ports::{HOST, start_on_free_ports, wait_for_pong}, }; @@ -39,6 +40,17 @@ impl RedisServer { &self.url } + pub async fn cli(&self, args: &[&str]) -> Result { + crate::exec::run(&self.container, crate::exec::cli(REDIS_PORT, args)).await + } + + pub async fn clients(&self) -> Result, Error> { + Ok(parse_client_list( + REDIS_PORT, + &self.cli(&["CLIENT", "LIST"]).await?, + )) + } + pub async fn restart(&self) -> Result<(), Error> { self.stop().await?; self.resume().await