diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index f10a507b561..2af2dcf20b8 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -39,6 +39,14 @@ - 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 +## Pub/sub +- Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel +- Channel names never get the namespace. Python builds `{namespace}:{channel}` itself, so the store does too +- `publish` is a command like any other, guarded by the breaker and reported to events +- `subscribe` returns a `Stream` that opens its own RESP2 `PubSub` connection when first polled, like redis-py's `PubSub`, to the endpoint, the first reachable cluster node, or the master Sentinel names at that moment. It never touches the breaker or events, like Python's subscriber, which reads off its own connection +- After a lost connection or a failed subscribe it waits and resubscribes, 5s doubling to 60s and reset by a successful subscribe, like `AuthCacheInvalidationSubscriber._run`. Messages published while it is down are lost. The connection has no idle ping, so a socket that dies silently is only noticed by TCP +- redis-rs's `PubSub` takes no credentials provider, so a provider is asked for its current credentials once when the subscription connects, which is when redis-py's pubsub connection asks it. The connect is bounded by the connect timeout + ## Batches - `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment` - `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`. Use it instead of one `get` per key whenever Python reads several keys in one call It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 79ff37609b3..c63d2121081 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -4,17 +4,20 @@ use std::{ time::{Duration, Instant}, }; +use futures_util::stream::BoxStream; use litellm_circuit_breaker::{Admission, Breaker, Failure, State}; use redis::{Cmd, FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value}; use crate::{ Amount, CallEvent, ConnectError, Error, Expiry, FailureKind, Namespace, Operation, Outcome, Pipeline, RedisBatch, RedisConfig, RedisEvents, Script, commands, connection::Connection, + pubsub::Subscriber, }; #[derive(Clone)] pub struct RedisClient { pub(crate) connection: Connection, + subscriber: Subscriber, namespace: Namespace, breaker: Option, events: Option>, @@ -31,8 +34,15 @@ pub enum Ttl { impl RedisClient { pub fn new(config: RedisConfig) -> Result { + let connection = Connection::new(&config.target, &config.credentials, config.timeouts)?; Ok(Self { - connection: Connection::new(&config.target, &config.credentials, config.timeouts)?, + subscriber: Subscriber::new( + &config.target, + &config.credentials, + config.timeouts, + &connection, + )?, + connection, namespace: config.namespace, breaker: config.breaker, events: config.events, @@ -194,6 +204,18 @@ impl RedisClient { .await } + /// The channel is sent as given, without the namespace, like redis-py's `publish`. + pub async fn publish(&self, channel: &str, payload: &[u8]) -> Result<(), Error> { + self.query("PUBLISH", &commands::publish(channel, payload)) + .await + } + + /// Messages published on `channel`, on a connection of its own. After the connection is lost + /// it reconnects and resubscribes, and messages published meanwhile are not delivered. + pub fn subscribe(&self, channel: &str) -> BoxStream<'static, Vec> { + self.subscriber.clone().messages(channel.to_owned()) + } + pub fn pipeline(&self) -> Pipeline { Pipeline::new(self.namespace.clone()) } diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs index 5cd7fd421a5..6d7485c7bdb 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -34,6 +34,10 @@ pub(crate) fn ping() -> Cmd { redis::cmd("PING") } +pub(crate) fn publish(channel: &str, payload: &[u8]) -> Cmd { + redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned() +} + pub(crate) fn get(key: String) -> Cmd { redis::cmd("GET").arg(key).to_owned() } diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index 2734be82b25..9440c3e432b 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -233,20 +233,13 @@ impl SentinelConnection { .set_client_to_sentinel_certificates(config::tls_certificates(tls)?), None => builder, }; - let config = AsyncConnectionConfig::new() - .set_connection_timeout(Some(timeouts.connect)) - .set_response_timeout(Some(timeouts.response)); - let config = match credentials { - Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), - Credentials::Static { .. } => config, - }; Ok(Self(Arc::new(SentinelState { discovery: tokio::sync::Mutex::new(builder.build()?), current: Mutex::new(None), generations: AtomicU64::new(0), endpoint: endpoint.clone(), credentials: credentials.clone(), - config, + config: async_config(credentials, timeouts), connect_timeout: timeouts.connect, }))) } @@ -267,11 +260,9 @@ impl SentinelConnection { if let Some(current) = self.cached() { return Ok(current); } - let master = tokio::time::timeout(self.0.connect_timeout, discovery.async_get_client()) - .await - .map_err(|_| timed_out("asking the Sentinels for the master timed out"))??; let connection = self - .master_client(master.get_connection_info().addr())? + .discover(&mut discovery) + .await? .get_multiplexed_async_connection_with_config(&self.0.config) .await?; let mut current = self @@ -284,6 +275,17 @@ impl SentinelConnection { Ok((generation, connection)) } + pub(crate) async fn master(&self) -> RedisResult { + self.discover(&mut *self.0.discovery.lock().await).await + } + + async fn discover(&self, discovery: &mut SentinelClient) -> RedisResult { + let master = tokio::time::timeout(self.0.connect_timeout, discovery.async_get_client()) + .await + .map_err(|_| timed_out("asking the Sentinels for the master timed out"))??; + self.master_client(master.get_connection_info().addr()) + } + fn master_client(&self, master: &ConnectionAddr) -> Result { let (host, port) = match master { ConnectionAddr::Tcp(host, port) | ConnectionAddr::TcpTls { host, port, .. } => { @@ -336,6 +338,16 @@ impl SentinelConnection { } } +pub(crate) fn async_config(credentials: &Credentials, timeouts: Timeouts) -> AsyncConnectionConfig { + let config = AsyncConnectionConfig::new() + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)); + match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + } +} + fn timed_out(message: &'static str) -> RedisError { RedisError::from(std::io::Error::new(std::io::ErrorKind::TimedOut, message)) } diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 6a06ad4b2fd..a5ed803dca8 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -7,6 +7,7 @@ mod error; mod events; mod namespace; mod pipeline; +mod pubsub; mod script; pub use batch::{Declared, RedisBatch}; diff --git a/litellm-rust/crates/redis-client/src/pubsub.rs b/litellm-rust/crates/redis-client/src/pubsub.rs new file mode 100644 index 00000000000..9de617843a7 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/pubsub.rs @@ -0,0 +1,153 @@ +use std::{sync::Arc, time::Duration}; + +use futures_util::{StreamExt, stream::BoxStream}; +use redis::{ConnectionInfo, RedisResult, StreamingCredentialsProvider, aio::PubSubStream}; + +use crate::{ + ConnectError, Credentials, Standalone, Target, Timeouts, config, + connection::{Connection, SentinelConnection}, +}; + +// litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS +const BACKOFF_INITIAL: Duration = Duration::from_secs(5); +const BACKOFF_MAX: Duration = Duration::from_secs(60); + +/// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends +/// `SUBSCRIBE`, authenticated once when it connects. +#[derive(Clone)] +pub(crate) struct Subscriber { + source: Source, + credentials: Credentials, + connect_timeout: Duration, +} + +#[derive(Clone)] +enum Source { + Nodes(Arc<[ConnectionInfo]>), + Sentinel(SentinelConnection), +} + +impl Subscriber { + pub(crate) fn new( + endpoint: &Target, + credentials: &Credentials, + timeouts: Timeouts, + connection: &Connection, + ) -> Result { + let source = match (endpoint, connection) { + (_, Connection::Sentinel(sentinel)) => Source::Sentinel(sentinel.clone()), + (Target::Cluster(cluster), _) => Source::Nodes( + cluster + .nodes + .iter() + .map(|node| { + let node = Target::Standalone(Standalone { + host: node.host.clone(), + port: node.port, + tls: cluster.tls.clone(), + db: 0, + }); + config::connection_info(&node, credentials) + }) + .collect::>()?, + ), + _ => Source::Nodes(Arc::new([config::connection_info(endpoint, credentials)?])), + }; + Ok(Self { + source, + credentials: credentials.clone(), + connect_timeout: timeouts.connect, + }) + } + + pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec> { + futures_util::stream::unfold( + (self, channel, None::, BACKOFF_INITIAL), + |(subscriber, channel, live, backoff)| async move { + let mut live = live; + let mut backoff = backoff; + loop { + let Some(current) = live.as_mut() else { + match subscriber.subscribe(&channel).await { + Ok(subscribed) => { + live = Some(subscribed); + backoff = BACKOFF_INITIAL; + } + Err(error) => { + tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + } + continue; + }; + match current.next().await { + Some(message) => { + let payload = message.get_payload_bytes().to_vec(); + return Some((payload, (subscriber, channel, live, backoff))); + } + None => live = None, + } + tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + }, + ) + .boxed() + } + + async fn subscribe(&self, channel: &str) -> RedisResult { + match &self.source { + Source::Sentinel(sentinel) => { + let master = sentinel.master().await?; + self.subscribe_on(master.get_connection_info().clone(), channel) + .await + } + Source::Nodes(nodes) => { + let mut last = None; + for node in nodes.iter() { + match self.subscribe_on(node.clone(), channel).await { + Ok(live) => return Ok(live), + Err(error) => last = Some(error), + } + } + Err(last.expect("an endpoint always has at least one node")) + } + } + } + + async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { + let info = self.authenticated(info).await?; + let client = redis::Client::open(info)?; + let mut pubsub = tokio::time::timeout(self.connect_timeout, client.get_async_pubsub()) + .await + .map_err(|_| { + redis::RedisError::from(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "connecting the Redis subscription timed out", + )) + })??; + pubsub.subscribe(channel).await?; + Ok(pubsub.into_on_message()) + } + + // redis-py's pubsub connection asks the credential provider once, on connect + async fn authenticated(&self, info: ConnectionInfo) -> RedisResult { + let Credentials::Provider(provider) = &self.credentials else { + return Ok(info); + }; + let current = provider.subscribe().next().await.ok_or_else(|| { + redis::RedisError::from(( + redis::ErrorKind::AuthenticationFailed, + "the Redis credential provider ended before giving credentials", + )) + })??; + let settings = info + .redis_settings() + .clone() + .set_username(current.username()) + .set_password(current.password()); + Ok(info.set_redis_settings(settings)) + } +} diff --git a/litellm-rust/crates/redis-client/tests/pubsub.rs b/litellm-rust/crates/redis-client/tests/pubsub.rs new file mode 100644 index 00000000000..c899bfbc4bf --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -0,0 +1,184 @@ +mod support; + +use std::time::{Duration, Instant}; + +use futures_util::StreamExt; +use litellm_circuit_breaker::{Breaker, Config}; +use litellm_redis_client::{ + Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, + Target, Timeouts, +}; +use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually}; +use rstest::rstest; +use support::{connect, connect_guarded}; +use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + +const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation"; + +fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver> { + let (sender, receiver) = unbounded_channel(); + let mut messages = client.subscribe(channel); + tokio::spawn(async move { + while let Some(message) = messages.next().await { + if sender.send(message).is_err() { + return; + } + } + }); + receiver +} + +async fn delivered( + publisher: &RedisClient, + channel: &str, + received: &mut UnboundedReceiver>, + payload: &[u8], +) { + eventually(async || { + publisher.publish(channel, payload).await.ok()?; + let deadline = tokio::time::Instant::now() + Duration::from_millis(200); + while let Ok(Some(message)) = tokio::time::timeout_at(deadline, received.recv()).await { + if message == payload { + return Some(()); + } + } + None + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_receives_raw_bytes_on_the_channel_name_without_the_namespace( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let publisher = connect(&server, "ns", DEFAULT_TIMEOUT); + let plain = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&plain, CHANNEL); + + delivered( + &publisher, + CHANNEL, + &mut received, + b"{\"cache_key\":\"k\"}\xff", + ) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_resubscribes_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)); + let mut received = listen(&client, CHANNEL); + delivered(&client, CHANNEL, &mut received, b"before").await; + + server.restart().await.unwrap(); + + delivered(&client, CHANNEL, &mut received, b"after").await; +} + +#[rstest] +#[tokio::test] +async fn a_sentinel_subscriber_receives_from_the_master() { + let server = RedisSentinel::start().await.unwrap(); + let (host, port) = server.sentinel(); + let client = RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + }) + .unwrap(); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"payload").await; +} + +async fn publish_calls(server: &RedisDeployment) -> Vec { + let replies = match server { + RedisDeployment::Standalone(server) => { + vec![server.cli(&["INFO", "commandstats"]).await.unwrap()] + } + RedisDeployment::Cluster(cluster) => cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap(), + }; + replies + .iter() + .flat_map(|reply| reply.lines()) + .filter(|line| line.starts_with("cmdstat_publish:")) + .map(str::to_owned) + .collect() +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_a_publish_without_sending_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = Breaker::new( + Config { + failure_threshold: 2, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker)); + + server.freeze().await.unwrap(); + let opened = [client.ping().await, client.ping().await]; + let started = Instant::now(); + let rejected = client.publish(CHANNEL, b"payload").await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!( + opened + .iter() + .all(|ping| matches!(ping, Err(Error::Transport(_)))), + "{opened:?}" + ); + assert!(matches!(rejected, Err(Error::BreakerOpen)), "{rejected:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + assert_eq!(publish_calls(&server).await, Vec::::new()); +} + +// redis-py's PubSub subscribes on a RESP2 connection, so it never sends HELLO +#[rstest] +#[tokio::test] +async fn a_subscription_speaks_resp2_like_redis_py( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"evict").await; + + let stats = support::command_stats(&server).await; + assert!(support::calls(&stats, "subscribe") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "hello"), 0, "{stats}"); +}