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 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-03 14:23:10 +09:00
parent e96e1009cb
commit dfce567d66
7 changed files with 379 additions and 13 deletions

View file

@ -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<T>` 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

View file

@ -4,6 +4,7 @@ use std::{
time::{Duration, Instant},
};
use futures_util::stream::BoxStream;
use litellm_circuit_breaker::{Admission, Breaker, Failure};
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<Breaker>,
events: Option<Arc<dyn RedisEvents>>,
@ -33,8 +36,15 @@ pub enum Ttl {
impl RedisClient {
pub fn new(config: RedisConfig) -> Result<Self, ConnectError> {
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,
@ -141,6 +151,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<u8>> {
self.subscriber.clone().messages(channel.to_owned())
}
pub fn pipeline(&self) -> Pipeline {
Pipeline::new(self.namespace.clone())
}

View file

@ -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()
}

View file

@ -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<redis::Client> {
self.discover(&mut *self.0.discovery.lock().await).await
}
async fn discover(&self, discovery: &mut SentinelClient) -> RedisResult<redis::Client> {
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<redis::Client, RedisError> {
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))
}

View file

@ -7,6 +7,7 @@ mod error;
mod events;
mod namespace;
mod pipeline;
mod pubsub;
mod script;
pub use batch::{Declared, RedisBatch};

View file

@ -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<PushInfo>,
}
impl Subscriber {
pub(crate) fn new(
endpoint: &Endpoint,
credentials: &Credentials,
timeouts: Timeouts,
connection: &Connection,
) -> Result<Self, ConnectError> {
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::<Result<_, ConnectError>>()?,
),
_ => 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<u8>> {
futures_util::stream::unfold(
(self, channel, None::<Live>, 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<Live> {
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<Live> {
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<redis::Client> {
let settings = info
.redis_settings()
.clone()
.set_protocol(ProtocolVersion::RESP3);
redis::Client::open(info.set_redis_settings(settings))
}

View file

@ -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<Vec<u8>> {
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<Vec<u8>>,
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<String> {
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::<String>::new());
}