diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9a2b2e79364..bfbe543e235 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4418,6 +4418,7 @@ dependencies = [ "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` @@ -27,17 +26,22 @@ - 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`) +- `incr_by`, `incr_by_float` and `rpush` take an `Expiry` and apply it in the same script as the write, so a key never outlives a lost follow-up `EXPIRE`. `KeepExisting` sets the TTL only on a key without one, like Python's `_incrbyfloat_with_ttl`, and `Refresh` always resets it +- Atomic writes require scripting permission (`EVAL`, `EVALSHA`, `SCRIPT LOAD`) and permission for the underlying write, `PTTL`, `PEXPIRE` and `GET` for integer increments. There is no non-atomic fallback for restricted ACLs. Ordinary reads and SET do not require scripting +- Integer increments read their exact decimal result inside the script before expiry, without converting through Lua's double. The public method and pipeline return an exact integer ## 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 `EVAL` with its source, so it can never miss the script cache and run out of order with the ops around it - 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 +- 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 - 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..4a5de772cfe 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -10,9 +10,10 @@ litellm-circuit-breaker.workspace = true litellm-secret-value.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..a3382e6d5ef 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,13 +1,18 @@ 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, + CallEvent, ConnectError, Endpoint, Error, Expiry, FailureKind, Namespace, Operation, Outcome, + Pipeline, RedisConfig, RedisEvents, Script, + commands::{self, ScriptCall}, connection::Connection, }; @@ -16,6 +21,7 @@ pub struct RedisClient { pub(crate) connection: Connection, namespace: Namespace, breaker: Option, + events: Option>, } #[derive(Debug, PartialEq, Eq)] @@ -29,33 +35,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 +71,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 +82,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,15 +91,16 @@ 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 + self.query("PEXPIRE", &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?; + let reply: i64 = self.query("PTTL", &commands::ttl(self.key(key))).await?; Ok(match reply { -2 => Ttl::Missing, -1 => Ttl::Persistent, @@ -99,17 +108,27 @@ impl RedisClient { }) } - 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)) + pub async fn incr_by(&self, key: &str, delta: i64, expiry: Expiry) -> Result { + self.call("INCRBY", commands::incr_by(self.key(key), delta, expiry)) .await } - pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { - self.query(&commands::rpush(self.key(key), values)).await + pub async fn incr_by_float(&self, key: &str, delta: f64, expiry: Expiry) -> Result { + self.call( + "INCRBYFLOAT", + commands::incr_by_float(self.key(key), delta, expiry), + ) + .await + } + + pub async fn rpush( + &self, + key: &str, + values: impl ToRedisArgs, + expiry: Expiry, + ) -> Result { + self.call("RPUSH", commands::rpush(self.key(key), values, expiry)) + .await } pub async fn eval( @@ -120,11 +139,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 +153,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 +178,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 +194,61 @@ 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 call( + &self, + name: &'static str, + call: ScriptCall, + ) -> Result { + self.guarded(Operation::Command(name), async { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut self.connection.clone()) + .await + }) + .await + } + + 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 +274,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..095393dbe64 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -1,7 +1,56 @@ -use std::time::Duration; +use std::{sync::LazyLock, time::Duration}; use redis::{Cmd, ToRedisArgs, ToSingleRedisArg}; +use crate::Script; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Expiry { + Refresh(Duration), + KeepExisting(Duration), +} + +pub(crate) struct ScriptCall { + pub(crate) script: &'static Script, + pub(crate) keys: Vec, + pub(crate) args: Vec>, +} + +macro_rules! with_expiry { + ($write:literal) => { + with_expiry!($write, "value") + }; + ($write:literal, $reply:literal) => { + concat!( + "local value = ", + $write, + "\nlocal reply = ", + $reply, + "\nif ARGV[1] == 'refresh' or redis.call('PTTL', KEYS[1]) == -1 then\n", + " redis.call('PEXPIRE', KEYS[1], ARGV[2])\n", + "end\n", + "return reply" + ) + }; +} + +static INCR_BY: LazyLock