mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
feat(redis-client): support Redis Cluster (#48271)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
6edbd27f74
commit
a0b5c8b824
10 changed files with 356 additions and 106 deletions
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -5805,6 +5805,7 @@ dependencies = [
|
|||
"futures-channel",
|
||||
"futures-util",
|
||||
"itoa",
|
||||
"log",
|
||||
"num-bigint 0.5.1",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
|
|
|
|||
|
|
@ -18,9 +18,11 @@
|
|||
|
||||
## Commands
|
||||
- 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`)
|
||||
|
||||
## Tests
|
||||
- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes
|
||||
- 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
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
redis = { workspace = true, features = ["connection-manager", "tokio-rustls-comp"] }
|
||||
redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] }
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
|
|
|||
|
|
@ -1,20 +1,29 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use redis::{
|
||||
Cmd, FromRedisValue, ToRedisArgs, ToSingleRedisArg,
|
||||
Cmd, ConnectionInfo, FromRedisValue, IntoConnectionInfo, RedisResult, ToRedisArgs,
|
||||
ToSingleRedisArg,
|
||||
aio::{ConnectionManager, ConnectionManagerConfig},
|
||||
cluster::ClusterClientBuilder,
|
||||
cluster_async::ClusterConnection,
|
||||
};
|
||||
|
||||
use crate::{ConnectError, Error, Namespace, RedisConfig};
|
||||
use crate::{ClusterEndpoint, ConnectError, Endpoint, Error, Namespace, RedisConfig};
|
||||
|
||||
const RECONNECT_RETRIES: usize = 3;
|
||||
const RETRIES: u32 = 3;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RedisClient {
|
||||
connection: ConnectionManager,
|
||||
connection: Connection,
|
||||
namespace: Namespace,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum Connection {
|
||||
Standalone(ConnectionManager),
|
||||
Cluster(ClusterConnection),
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum Ttl {
|
||||
Missing,
|
||||
|
|
@ -24,15 +33,15 @@ pub enum Ttl {
|
|||
|
||||
impl RedisClient {
|
||||
pub async fn connect(config: RedisConfig) -> Result<Self, ConnectError> {
|
||||
let client = redis::Client::open(config.endpoint.connection_info()?)?;
|
||||
let manager_config = ConnectionManagerConfig::new()
|
||||
.set_connection_timeout(Some(config.timeout))
|
||||
.set_response_timeout(Some(config.timeout))
|
||||
.set_number_of_retries(RECONNECT_RETRIES);
|
||||
let connection = match &config.endpoint {
|
||||
Endpoint::Url(url) => {
|
||||
standalone(url.as_str().into_connection_info()?, config.timeout).await?
|
||||
}
|
||||
Endpoint::Host(host) => standalone(host.connection_info()?, config.timeout).await?,
|
||||
Endpoint::Cluster(cluster) => self::cluster(cluster, config.timeout).await?,
|
||||
};
|
||||
Ok(Self {
|
||||
connection: client
|
||||
.get_connection_manager_with_config(manager_config)
|
||||
.await?,
|
||||
connection,
|
||||
namespace: config.namespace,
|
||||
})
|
||||
}
|
||||
|
|
@ -128,10 +137,43 @@ impl RedisClient {
|
|||
}
|
||||
|
||||
async fn query<T: FromRedisValue>(&self, command: &Cmd) -> Result<T, Error> {
|
||||
Ok(command.query_async(&mut self.connection.clone()).await?)
|
||||
let reply = match self.connection.clone() {
|
||||
Connection::Standalone(mut connection) => command.query_async(&mut connection).await,
|
||||
Connection::Cluster(mut connection) => command.query_async(&mut connection).await,
|
||||
};
|
||||
Ok(reply?)
|
||||
}
|
||||
}
|
||||
|
||||
async fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult<Connection> {
|
||||
let config = ConnectionManagerConfig::new()
|
||||
.set_connection_timeout(Some(timeout))
|
||||
.set_response_timeout(Some(timeout))
|
||||
.set_number_of_retries(RETRIES as usize);
|
||||
let manager = redis::Client::open(info)?
|
||||
.get_connection_manager_with_config(config)
|
||||
.await?;
|
||||
Ok(Connection::Standalone(manager))
|
||||
}
|
||||
|
||||
async fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult<Connection> {
|
||||
let builder = ClusterClientBuilder::new(endpoint.addrs())
|
||||
.connection_timeout(timeout)
|
||||
.response_timeout(timeout)
|
||||
.retries(RETRIES);
|
||||
let builder = match &endpoint.username {
|
||||
Some(username) => builder.username(username),
|
||||
None => builder,
|
||||
};
|
||||
let builder = match &endpoint.password {
|
||||
Some(password) => builder.password(password),
|
||||
None => builder,
|
||||
};
|
||||
Ok(Connection::Cluster(
|
||||
builder.build()?.get_async_connection().await?,
|
||||
))
|
||||
}
|
||||
|
||||
fn millis(duration: Duration) -> u64 {
|
||||
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
|
||||
}
|
||||
|
|
@ -140,37 +182,109 @@ fn millis(duration: Duration) -> u64 {
|
|||
mod tests {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use litellm_redis_testing::RedisServer;
|
||||
use litellm_redis_testing::{RedisCluster, RedisServer};
|
||||
use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{RedisClient, Ttl};
|
||||
use crate::{DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig};
|
||||
use super::{Connection, RedisClient, Ttl};
|
||||
use crate::{
|
||||
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace,
|
||||
RedisConfig,
|
||||
};
|
||||
|
||||
async fn client(server: &RedisServer, namespace: &str, timeout: Duration) -> RedisClient {
|
||||
RedisClient::connect(RedisConfig {
|
||||
endpoint: Endpoint::Url(server.url().to_owned()),
|
||||
namespace: Namespace::new(namespace),
|
||||
timeout,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Topology {
|
||||
Standalone,
|
||||
Cluster,
|
||||
}
|
||||
|
||||
async fn client_id(client: &RedisClient) -> i64 {
|
||||
redis::cmd("CLIENT")
|
||||
.arg("ID")
|
||||
.query_async(&mut client.connection.clone())
|
||||
enum Server {
|
||||
Standalone(RedisServer),
|
||||
Cluster(RedisCluster),
|
||||
}
|
||||
|
||||
impl Server {
|
||||
async fn start(topology: Topology) -> Self {
|
||||
match topology {
|
||||
Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()),
|
||||
Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()),
|
||||
}
|
||||
}
|
||||
|
||||
fn endpoint(&self) -> Endpoint {
|
||||
match self {
|
||||
Self::Standalone(server) => Endpoint::Url(server.url().to_owned()),
|
||||
Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
|
||||
nodes: cluster
|
||||
.nodes()
|
||||
.map(|(host, port)| ClusterNode {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
})
|
||||
.collect(),
|
||||
username: None,
|
||||
password: None,
|
||||
tls: false,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient {
|
||||
RedisClient::connect(RedisConfig {
|
||||
endpoint: self.endpoint(),
|
||||
namespace: Namespace::new(namespace),
|
||||
timeout,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn restart(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.restart().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.restart().await.unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn freeze(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.freeze().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.freeze().await.unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn thaw(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.thaw().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.thaw().await.unwrap(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn client_id(client: &RedisClient, key: &str) -> i64 {
|
||||
let command = redis::cmd("CLIENT").arg("ID").to_owned();
|
||||
match client.connection.clone() {
|
||||
Connection::Standalone(mut connection) => {
|
||||
command.query_async(&mut connection).await.unwrap()
|
||||
}
|
||||
Connection::Cluster(mut connection) => {
|
||||
let node = Route::with_key(key, SlotAddr::Master);
|
||||
let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node));
|
||||
redis::from_redis_value(connection.route_command(command, routing).await.unwrap())
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn namespaces_isolate_keys_on_the_server() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let a = client(&server, "a", DEFAULT_TIMEOUT).await;
|
||||
let b = client(&server, "b", DEFAULT_TIMEOUT).await;
|
||||
let bare = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
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();
|
||||
|
||||
|
|
@ -185,24 +299,28 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_rejected_command_keeps_the_connection() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
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).await;
|
||||
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.incr_by("counter", 2).await.unwrap(), 2);
|
||||
assert_eq!(client_id(&client).await, before);
|
||||
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() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
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
|
||||
|
|
@ -216,9 +334,11 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", Duration::from_millis(500)).await;
|
||||
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
|
||||
|
|
@ -228,9 +348,9 @@ mod tests {
|
|||
.await
|
||||
.unwrap();
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
server.freeze().await;
|
||||
let timed_out = client.get::<String>("first").await;
|
||||
server.thaw().await.unwrap();
|
||||
server.thaw().await;
|
||||
|
||||
assert!(
|
||||
matches!(timed_out, Err(Error::Transport(_))),
|
||||
|
|
@ -244,11 +364,13 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_same_client_recovers_after_redis_restarts() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", Duration::from_millis(500)).await;
|
||||
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.unwrap();
|
||||
server.restart().await;
|
||||
|
||||
let deadline = Instant::now() + Duration::from_secs(30);
|
||||
let recovered = loop {
|
||||
|
|
@ -265,9 +387,11 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn empty_batches_succeed_without_asking_redis() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
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::<String>(&[]).await.unwrap(),
|
||||
|
|
@ -278,25 +402,35 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn mget_returns_values_in_input_order_with_gaps() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "ns", DEFAULT_TIMEOUT).await;
|
||||
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::<String>(&["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::<String>(&["a", "c"]).await.unwrap(),
|
||||
vec![None, None]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn set_nx_only_writes_a_missing_key() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
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))
|
||||
|
|
@ -317,9 +451,11 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ttl_reports_missing_persistent_and_remaining_time() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
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))
|
||||
|
|
@ -353,9 +489,11 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn incr_by_float_accumulates_fractions() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let client = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
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();
|
||||
|
||||
|
|
@ -367,21 +505,29 @@ mod tests {
|
|||
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 in_db_3 = RedisClient::connect(RedisConfig {
|
||||
endpoint: Endpoint::Host(HostEndpoint {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port,
|
||||
username: None,
|
||||
password: None,
|
||||
tls: false,
|
||||
db: 3,
|
||||
}),
|
||||
namespace: Namespace::default(),
|
||||
timeout: DEFAULT_TIMEOUT,
|
||||
})
|
||||
let connect = |endpoint| {
|
||||
RedisClient::connect(RedisConfig {
|
||||
endpoint,
|
||||
namespace: Namespace::default(),
|
||||
timeout: DEFAULT_TIMEOUT,
|
||||
})
|
||||
};
|
||||
let in_db_3 = connect(Endpoint::Host(HostEndpoint {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port,
|
||||
username: None,
|
||||
password: None,
|
||||
tls: false,
|
||||
db: 3,
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
let in_db_0 = client(&server, "", DEFAULT_TIMEOUT).await;
|
||||
let in_db_0 = connect(Endpoint::Url(server.url().to_owned()))
|
||||
.await
|
||||
.unwrap();
|
||||
let via_url = connect(Endpoint::Url(format!("{}/3", server.url())))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
in_db_3
|
||||
.set("k", "3", Duration::from_secs(60))
|
||||
|
|
@ -389,13 +535,6 @@ mod tests {
|
|||
.unwrap();
|
||||
|
||||
assert_eq!(in_db_0.get::<String>("k").await.unwrap(), None);
|
||||
let via_url = RedisClient::connect(RedisConfig {
|
||||
endpoint: Endpoint::Url(format!("{}/3", server.url())),
|
||||
namespace: Namespace::default(),
|
||||
timeout: DEFAULT_TIMEOUT,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
via_url.get::<String>("k").await.unwrap().as_deref(),
|
||||
Some("3")
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ pub struct RedisConfig {
|
|||
pub enum Endpoint {
|
||||
Url(String),
|
||||
Host(HostEndpoint),
|
||||
Cluster(ClusterEndpoint),
|
||||
}
|
||||
|
||||
pub struct HostEndpoint {
|
||||
|
|
@ -26,26 +27,28 @@ pub struct HostEndpoint {
|
|||
pub db: i64,
|
||||
}
|
||||
|
||||
impl Endpoint {
|
||||
pub(crate) fn connection_info(&self) -> RedisResult<ConnectionInfo> {
|
||||
match self {
|
||||
Self::Url(url) => url.as_str().into_connection_info(),
|
||||
Self::Host(host) => host.connection_info(),
|
||||
}
|
||||
pub struct ClusterEndpoint {
|
||||
pub nodes: Vec<ClusterNode>,
|
||||
pub username: Option<String>,
|
||||
pub password: Option<String>,
|
||||
pub tls: bool,
|
||||
}
|
||||
|
||||
pub struct ClusterNode {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
impl ClusterEndpoint {
|
||||
pub(crate) fn addrs(&self) -> impl Iterator<Item = ConnectionAddr> + '_ {
|
||||
self.nodes
|
||||
.iter()
|
||||
.map(|node| addr(&node.host, node.port, self.tls))
|
||||
}
|
||||
}
|
||||
|
||||
impl HostEndpoint {
|
||||
fn connection_info(&self) -> RedisResult<ConnectionInfo> {
|
||||
let addr = match self.tls {
|
||||
true => ConnectionAddr::TcpTls {
|
||||
host: self.host.clone(),
|
||||
port: self.port,
|
||||
insecure: false,
|
||||
tls_params: None,
|
||||
},
|
||||
false => ConnectionAddr::Tcp(self.host.clone(), self.port),
|
||||
};
|
||||
pub(crate) fn connection_info(&self) -> RedisResult<ConnectionInfo> {
|
||||
let settings = RedisConnectionInfo::default().set_db(self.db);
|
||||
let settings = match &self.username {
|
||||
Some(username) => settings.set_username(username),
|
||||
|
|
@ -55,6 +58,20 @@ impl HostEndpoint {
|
|||
Some(password) => settings.set_password(password),
|
||||
None => settings,
|
||||
};
|
||||
Ok(addr.into_connection_info()?.set_redis_settings(settings))
|
||||
Ok(addr(&self.host, self.port, self.tls)
|
||||
.into_connection_info()?
|
||||
.set_redis_settings(settings))
|
||||
}
|
||||
}
|
||||
|
||||
fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr {
|
||||
match tls {
|
||||
true => ConnectionAddr::TcpTls {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
insecure: false,
|
||||
tls_params: None,
|
||||
},
|
||||
false => ConnectionAddr::Tcp(host.to_owned(), port),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ mod error;
|
|||
mod namespace;
|
||||
|
||||
pub use client::{RedisClient, Ttl};
|
||||
pub use config::{DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig};
|
||||
pub use config::{
|
||||
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig,
|
||||
};
|
||||
pub use error::{ConnectError, Error};
|
||||
pub use namespace::Namespace;
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
- Dev-only crate that starts real Redis for tests. Never a dependency outside `[dev-dependencies]`
|
||||
- Same shape as `db-testing`: one container per test, a handle that removes it on drop
|
||||
- The host port is fixed at start so `restart` keeps the same URL and a client can reconnect to it
|
||||
- `RedisCluster` runs 3 masters in one container. Each node listens on the same port inside and outside and announces 127.0.0.1, so MOVED redirects work from the test process
|
||||
- Outages come from the container (`restart`, `freeze` / `thaw`), never `DEBUG`, which the official image disables
|
||||
- No env gate that makes a test pass without Redis. A test that can't start its container fails
|
||||
|
|
|
|||
86
litellm-rust/crates/redis-testing/src/cluster.rs
Normal file
86
litellm-rust/crates/redis-testing/src/cluster.rs
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
use testcontainers_modules::testcontainers::{
|
||||
ContainerAsync, GenericImage, ImageExt,
|
||||
core::{IntoContainerPort, WaitFor},
|
||||
runners::AsyncRunner,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
ports::{HOST, start_on_free_ports, wait_for_pong},
|
||||
standalone::REDIS_TAG,
|
||||
};
|
||||
|
||||
const READY: &str = "redis-testing cluster ready";
|
||||
|
||||
const START_CLUSTER: &str = r#"
|
||||
set -e
|
||||
bus=17000
|
||||
for port in "$@"; do
|
||||
mkdir -p "/data/$port"
|
||||
redis-server --port "$port" --cluster-enabled yes --cluster-port "$bus" \
|
||||
--cluster-config-file "/data/$port/nodes.conf" --cluster-announce-ip 127.0.0.1 \
|
||||
--dir "/data/$port" --save '' --appendonly no --daemonize yes
|
||||
bus=$((bus + 1))
|
||||
done
|
||||
for port in "$@"; do
|
||||
until redis-cli -p "$port" ping >/dev/null 2>&1; do sleep 0.1; done
|
||||
done
|
||||
if [ ! -f /data/created ]; then
|
||||
redis-cli --cluster create $(for port in "$@"; do printf '127.0.0.1:%s ' "$port"; done) --cluster-yes >/dev/null
|
||||
touch /data/created
|
||||
fi
|
||||
for port in "$@"; do
|
||||
until redis-cli -p "$port" cluster info | grep -q cluster_state:ok; do sleep 0.1; done
|
||||
done
|
||||
echo "redis-testing cluster ready"
|
||||
exec sleep infinity
|
||||
"#;
|
||||
|
||||
pub struct RedisCluster {
|
||||
container: ContainerAsync<GenericImage>,
|
||||
ports: [u16; 3],
|
||||
}
|
||||
|
||||
impl RedisCluster {
|
||||
pub async fn start() -> Result<Self, Error> {
|
||||
start_on_free_ports(Self::start_on).await
|
||||
}
|
||||
|
||||
async fn start_on(ports: [u16; 3]) -> Result<Self, Error> {
|
||||
let image = GenericImage::new("redis", REDIS_TAG)
|
||||
.with_wait_for(WaitFor::message_on_stdout(READY))
|
||||
.with_entrypoint("sh");
|
||||
let request = ports.iter().fold(
|
||||
image.with_cmd(
|
||||
["-c", START_CLUSTER, "sh"]
|
||||
.into_iter()
|
||||
.map(str::to_owned)
|
||||
.chain(ports.iter().map(u16::to_string)),
|
||||
),
|
||||
|request, &port| request.with_mapped_port(port, port.tcp()),
|
||||
);
|
||||
let container = request.start().await?;
|
||||
for port in ports {
|
||||
wait_for_pong(port).await?;
|
||||
}
|
||||
Ok(Self { container, ports })
|
||||
}
|
||||
|
||||
pub fn nodes(&self) -> impl Iterator<Item = (&'static str, u16)> + '_ {
|
||||
self.ports.iter().map(|&port| (HOST, port))
|
||||
}
|
||||
|
||||
pub async fn restart(&self) -> Result<(), Error> {
|
||||
self.container.stop().await?;
|
||||
self.container.start().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn freeze(&self) -> Result<(), Error> {
|
||||
Ok(self.container.pause().await?)
|
||||
}
|
||||
|
||||
pub async fn thaw(&self) -> Result<(), Error> {
|
||||
Ok(self.container.unpause().await?)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
mod cluster;
|
||||
mod error;
|
||||
mod ports;
|
||||
mod standalone;
|
||||
|
||||
pub use cluster::RedisCluster;
|
||||
pub use error::Error;
|
||||
pub use standalone::RedisServer;
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ use crate::{
|
|||
ports::{HOST, start_on_free_ports, wait_for_pong},
|
||||
};
|
||||
|
||||
const REDIS_TAG: &str =
|
||||
pub(crate) const REDIS_TAG: &str =
|
||||
"7.4@sha256:44960b3100c48159a66d28c18b0588c193492e25f7551d60e51dd464fc3eca5b";
|
||||
|
||||
pub struct RedisServer {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue