mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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:
parent
d9556120aa
commit
f8b8225729
7 changed files with 379 additions and 13 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<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,
|
||||
|
|
@ -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<u8>> {
|
||||
self.subscriber.clone().messages(channel.to_owned())
|
||||
}
|
||||
|
||||
pub fn pipeline(&self) -> Pipeline {
|
||||
Pipeline::new(self.namespace.clone())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ mod error;
|
|||
mod events;
|
||||
mod namespace;
|
||||
mod pipeline;
|
||||
mod pubsub;
|
||||
mod script;
|
||||
|
||||
pub use batch::{Declared, RedisBatch};
|
||||
|
|
|
|||
156
litellm-rust/crates/redis-client/src/pubsub.rs
Normal file
156
litellm-rust/crates/redis-client/src/pubsub.rs
Normal 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))
|
||||
}
|
||||
163
litellm-rust/crates/redis-client/tests/pubsub.rs
Normal file
163
litellm-rust/crates/redis-client/tests/pubsub.rs
Normal 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());
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue