From f8b8225729d61b7bd73bf1353be6ffdc7f552f81 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 3 Nov 2026 14:23:10 +0900 Subject: [PATCH] feat(redis-client): publish and resubscribing subscribe for pub/sub Adds publish(channel, payload) through the guarded command path and subscribe(channel), a Stream of raw payloads on its own RESP3 connection that resubscribes after a lost connection with the Python subscriber's 5s to 60s backoff, on standalone, cluster and Sentinel. Channel names skip the namespace, matching the Python publisher, so the auth cache invalidation channel can move onto the Rust client Co-Authored-By: Claude Opus 5.5 --- litellm-rust/crates/redis-client/AGENTS.md | 8 + .../crates/redis-client/src/client.rs | 24 ++- .../crates/redis-client/src/commands.rs | 4 + .../crates/redis-client/src/connection.rs | 36 ++-- litellm-rust/crates/redis-client/src/lib.rs | 1 + .../crates/redis-client/src/pubsub.rs | 156 +++++++++++++++++ .../crates/redis-client/tests/pubsub.rs | 163 ++++++++++++++++++ 7 files changed, 379 insertions(+), 13 deletions(-) create mode 100644 litellm-rust/crates/redis-client/src/pubsub.rs create mode 100644 litellm-rust/crates/redis-client/tests/pubsub.rs diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index c42ec251872..f55a1fd1059 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -38,6 +38,14 @@ - 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 +## 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 RESP3 connection when first polled, 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 +- RESP3 needs Redis 6 or newer, while redis-py 5.3.1 subscribes over RESP2. redis-rs's RESP2 `PubSub` takes no credentials provider, timeouts or TCP_NODELAY + ## Batches - `RedisBatch` declares ops now and sends them later as one `Pipeline`, so it inherits every pipeline rule above. It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local - Each op returns a `Declared` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index cf11f269095..62301a46b10 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -4,6 +4,7 @@ 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}; @@ -12,11 +13,13 @@ use crate::{ RedisBatch, RedisConfig, RedisEvents, Script, commands::{self, ScriptCall}, connection::Connection, + pubsub::Subscriber, }; #[derive(Clone)] pub struct RedisClient { pub(crate) connection: Connection, + subscriber: Subscriber, namespace: Namespace, breaker: Option, events: Option>, @@ -33,8 +36,15 @@ pub enum Ttl { impl RedisClient { pub fn new(config: RedisConfig) -> Result { + let connection = Connection::new(&config.endpoint, &config.credentials, config.timeouts)?; Ok(Self { - connection: Connection::new(&config.endpoint, &config.credentials, config.timeouts)?, + subscriber: Subscriber::new( + &config.endpoint, + &config.credentials, + config.timeouts, + &connection, + )?, + connection, namespace: config.namespace, breaker: config.breaker, events: config.events, @@ -146,6 +156,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 095393dbe64..9f9929ff208 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -55,6 +55,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 5f343de4b58..21526a45e84 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -234,20 +234,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, }))) } @@ -268,11 +261,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 @@ -285,6 +276,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, .. } => { @@ -333,6 +335,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 f46720030e6..0dea8c1e8ac 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..5c0c67cabfc --- /dev/null +++ b/litellm-rust/crates/redis-client/src/pubsub.rs @@ -0,0 +1,156 @@ +use std::{sync::Arc, time::Duration}; + +use futures_util::{StreamExt, stream::BoxStream}; +use redis::{ + AsyncConnectionConfig, ConnectionInfo, Msg, ProtocolVersion, PushInfo, PushKind, RedisResult, + aio::MultiplexedConnection, +}; +use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + +use crate::{ + ConnectError, Credentials, Endpoint, HostEndpoint, Timeouts, config, + connection::{self, 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); + +#[derive(Clone)] +pub(crate) struct Subscriber { + source: Source, + config: AsyncConnectionConfig, +} + +#[derive(Clone)] +enum Source { + Nodes(Arc<[redis::Client]>), + Sentinel(SentinelConnection), +} + +struct Live { + _connection: MultiplexedConnection, + pushes: UnboundedReceiver, +} + +impl Subscriber { + pub(crate) fn new( + endpoint: &Endpoint, + credentials: &Credentials, + timeouts: Timeouts, + connection: &Connection, + ) -> Result { + let source = match (endpoint, connection) { + (_, Connection::Sentinel(sentinel)) => Source::Sentinel(sentinel.clone()), + (Endpoint::Cluster(cluster), _) => Source::Nodes( + cluster + .nodes + .iter() + .map(|node| { + let node = Endpoint::Host(HostEndpoint { + host: node.host.clone(), + port: node.port, + tls: cluster.tls.clone(), + db: 0, + unix_socket_path: None, + }); + Ok(resp3_client(config::connection_info(&node, credentials)?)?) + }) + .collect::>()?, + ), + _ => Source::Nodes(Arc::new([resp3_client(config::connection_info( + endpoint, + credentials, + )?)?])), + }; + Ok(Self { + source, + config: connection::async_config(credentials, timeouts), + }) + } + + 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.pushes.recv().await { + Some(push) if push.kind == PushKind::Disconnection => live = None, + None => live = None, + Some(push) => { + if let Some(message) = Msg::from_push_info(push) { + let payload = message.get_payload_bytes().to_vec(); + return Some((payload, (subscriber, channel, live, backoff))); + } + continue; + } + } + 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( + &resp3_client(master.get_connection_info().clone())?, + channel, + ) + .await + } + Source::Nodes(nodes) => { + let mut last = None; + for node in nodes.iter() { + match self.subscribe_on(node, 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, client: &redis::Client, channel: &str) -> RedisResult { + let (sender, pushes) = unbounded_channel(); + let config = self.config.clone().set_push_sender(sender); + let mut connection = client + .get_multiplexed_async_connection_with_config(&config) + .await?; + connection.subscribe(channel).await?; + Ok(Live { + _connection: connection, + pushes, + }) + } +} + +fn resp3_client(info: ConnectionInfo) -> RedisResult { + let settings = info + .redis_settings() + .clone() + .set_protocol(ProtocolVersion::RESP3); + redis::Client::open(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..2084c61388a --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -0,0 +1,163 @@ +mod support; + +use std::time::{Duration, Instant}; + +use futures_util::StreamExt; +use litellm_circuit_breaker::{Breaker, Config}; +use litellm_redis_client::{ + ClusterNode, Credentials, DEFAULT_TIMEOUT, Endpoint, Error, Namespace, RedisClient, + RedisConfig, SentinelEndpoint, 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 { + endpoint: Endpoint::Sentinel(SentinelEndpoint { + sentinels: vec![ClusterNode { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + sentinel_password: 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()); +}