feat(redis-client): run Lua scripts and pipelines with per-op errors (#48273)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yujonglee 2026-10-26 18:41:46 -07:00 • committed by GitHub
parent a0b5c8b824
commit af0d29030e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 707 additions and 174 deletions

View file

@ -4299,6 +4299,7 @@ dependencies = [
name = "litellm-redis-client"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-redis-testing",
"redis",
"rstest",

View file

@ -21,6 +21,14 @@
- 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`)
## Scripts and pipelines
- A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production
- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async`
- A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip
- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order
- 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
## 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)

View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
futures-util.workspace = true
redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] }
thiserror.workspace = true

View file

@ -1,29 +1,18 @@
use std::time::Duration;
use redis::{
Cmd, ConnectionInfo, FromRedisValue, IntoConnectionInfo, RedisResult, ToRedisArgs,
ToSingleRedisArg,
aio::{ConnectionManager, ConnectionManagerConfig},
cluster::ClusterClientBuilder,
cluster_async::ClusterConnection,
use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value};
use crate::{
ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands,
connection::Connection,
};
use crate::{ClusterEndpoint, ConnectError, Endpoint, Error, Namespace, RedisConfig};
const RETRIES: u32 = 3;
#[derive(Clone)]
pub struct RedisClient {
connection: Connection,
pub(crate) connection: Connection,
namespace: Namespace,
}
#[derive(Clone)]
enum Connection {
Standalone(ConnectionManager),
Cluster(ClusterConnection),
}
#[derive(Debug, PartialEq, Eq)]
pub enum Ttl {
Missing,
@ -35,10 +24,12 @@ impl RedisClient {
pub async fn connect(config: RedisConfig) -> Result<Self, ConnectError> {
let connection = match &config.endpoint {
Endpoint::Url(url) => {
standalone(url.as_str().into_connection_info()?, config.timeout).await?
Connection::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?,
Endpoint::Host(host) => {
Connection::standalone(host.connection_info()?, config.timeout).await?
}
Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout).await?,
};
Ok(Self {
connection,
@ -47,14 +38,14 @@ impl RedisClient {
}
pub async fn get<T: FromRedisValue>(&self, key: &str) -> Result<Option<T>, Error> {
self.query(redis::cmd("GET").arg(self.key(key))).await
self.query(&commands::get(self.key(key))).await
}
pub async fn mget<T: FromRedisValue>(&self, keys: &[&str]) -> Result<Vec<Option<T>>, Error> {
if keys.is_empty() {
return Ok(Vec::new());
}
self.query(redis::cmd("MGET").arg(self.keys(keys))).await
self.query(&commands::mget(self.keys(keys))).await
}
pub async fn set(
@ -63,14 +54,7 @@ impl RedisClient {
value: impl ToSingleRedisArg,
ttl: Duration,
) -> Result<(), Error> {
self.query(
redis::cmd("SET")
.arg(self.key(key))
.arg(value)
.arg("PX")
.arg(millis(ttl)),
)
.await
self.query(&commands::set(self.key(key), value, ttl)).await
}
pub async fn set_nx(
@ -80,14 +64,7 @@ impl RedisClient {
ttl: Duration,
) -> Result<bool, Error> {
let reply: Option<String> = self
.query(
redis::cmd("SET")
.arg(self.key(key))
.arg(value)
.arg("NX")
.arg("PX")
.arg(millis(ttl)),
)
.query(&commands::set_nx(self.key(key), value, ttl))
.await?;
Ok(reply.is_some())
}
@ -96,16 +73,15 @@ impl RedisClient {
if keys.is_empty() {
return Ok(0);
}
self.query(redis::cmd("DEL").arg(self.keys(keys))).await
self.query(&commands::del(self.keys(keys))).await
}
pub async fn expire(&self, key: &str, ttl: Duration) -> Result<bool, Error> {
self.query(redis::cmd("PEXPIRE").arg(self.key(key)).arg(millis(ttl)))
.await
self.query(&commands::expire(self.key(key), ttl)).await
}
pub async fn ttl(&self, key: &str) -> Result<Ttl, Error> {
let reply: i64 = self.query(redis::cmd("PTTL").arg(self.key(key))).await?;
let reply: i64 = self.query(&commands::ttl(self.key(key))).await?;
Ok(match reply {
-2 => Ttl::Missing,
-1 => Ttl::Persistent,
@ -114,18 +90,38 @@ impl RedisClient {
}
pub async fn incr_by(&self, key: &str, delta: i64) -> Result<i64, Error> {
self.query(redis::cmd("INCRBY").arg(self.key(key)).arg(delta))
.await
self.query(&commands::incr_by(self.key(key), delta)).await
}
pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result<f64, Error> {
self.query(redis::cmd("INCRBYFLOAT").arg(self.key(key)).arg(delta))
self.query(&commands::incr_by_float(self.key(key), delta))
.await
}
pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result<usize, Error> {
self.query(redis::cmd("RPUSH").arg(self.key(key)).arg(values))
.await
self.query(&commands::rpush(self.key(key), values)).await
}
pub async fn eval<T: FromRedisValue>(
&self,
script: &Script,
keys: &[&str],
args: impl ToRedisArgs,
) -> Result<T, Error> {
let keys = self.keys(keys);
let args = args.to_redis_args();
Ok(script
.invocation(&keys, &args)
.invoke_async(&mut self.connection.clone())
.await?)
}
pub fn pipeline(&self) -> Pipeline {
Pipeline::new(self.namespace.clone())
}
pub async fn run(&self, pipeline: &Pipeline) -> Vec<Result<Value, Error>> {
pipeline.run(&self.connection).await
}
fn key(&self, key: &str) -> String {
@ -137,145 +133,23 @@ impl RedisClient {
}
async fn query<T: FromRedisValue>(&self, command: &Cmd) -> Result<T, Error> {
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?)
Ok(command.query_async(&mut self.connection.clone()).await?)
}
}
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)
}
#[cfg(test)]
mod tests {
use std::time::{Duration, Instant};
use litellm_redis_testing::{RedisCluster, RedisServer};
use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr};
use litellm_redis_testing::RedisServer;
use rstest::rstest;
use super::{Connection, RedisClient, Ttl};
use super::{RedisClient, Ttl};
use crate::{
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace,
RedisConfig,
DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig,
test_support::{Server, Topology, client_id},
};
#[derive(Clone, Copy, Debug)]
enum Topology {
Standalone,
Cluster,
}
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(
@ -500,6 +374,31 @@ mod tests {
assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75);
}
#[rstest]
#[tokio::test]
async fn concurrent_commands_on_one_cloned_client_all_apply(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
let tasks: Vec<_> = (0..100)
.map(|_| {
let client = client.clone();
tokio::spawn(async move {
for _ in 0..10 {
client.incr_by("shared", 1).await.unwrap();
}
})
})
.collect();
for task in tasks {
task.await.unwrap();
}
assert_eq!(client.get::<i64>("shared").await.unwrap(), Some(1000));
}
#[rstest]
#[tokio::test]
async fn a_host_endpoint_connects_to_its_database() {

View file

@ -0,0 +1,58 @@
use std::time::Duration;
use redis::{Cmd, ToRedisArgs, ToSingleRedisArg};
pub(crate) fn get(key: String) -> Cmd {
redis::cmd("GET").arg(key).to_owned()
}
pub(crate) fn mget(keys: Vec<String>) -> Cmd {
redis::cmd("MGET").arg(keys).to_owned()
}
pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("PX")
.arg(millis(ttl))
.to_owned()
}
pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("NX")
.arg("PX")
.arg(millis(ttl))
.to_owned()
}
pub(crate) fn del(keys: Vec<String>) -> Cmd {
redis::cmd("DEL").arg(keys).to_owned()
}
pub(crate) fn expire(key: String, ttl: Duration) -> Cmd {
redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned()
}
pub(crate) fn ttl(key: String) -> Cmd {
redis::cmd("PTTL").arg(key).to_owned()
}
pub(crate) fn incr_by(key: String, delta: i64) -> Cmd {
redis::cmd("INCRBY").arg(key).arg(delta).to_owned()
}
pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd {
redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned()
}
pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd {
redis::cmd("RPUSH").arg(key).arg(values).to_owned()
}
fn millis(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}

View file

@ -0,0 +1,80 @@
use std::time::Duration;
use redis::{
Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value,
aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig},
cluster::ClusterClientBuilder,
cluster_async::ClusterConnection,
};
use crate::ClusterEndpoint;
const RETRIES: u32 = 3;
#[derive(Clone)]
pub(crate) enum Connection {
Standalone(ConnectionManager),
Cluster(ClusterConnection),
}
impl Connection {
pub(crate) async fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult<Self> {
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(Self::Standalone(manager))
}
pub(crate) async fn cluster(
endpoint: &ClusterEndpoint,
timeout: Duration,
) -> RedisResult<Self> {
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(Self::Cluster(
builder.build()?.get_async_connection().await?,
))
}
}
impl ConnectionLike for Connection {
fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
match self {
Self::Standalone(connection) => connection.req_packed_command(cmd),
Self::Cluster(connection) => connection.req_packed_command(cmd),
}
}
fn req_packed_commands<'a>(
&'a mut self,
pipeline: &'a Pipeline,
offset: usize,
count: usize,
) -> RedisFuture<'a, Vec<Value>> {
match self {
Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count),
Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count),
}
}
fn get_db(&self) -> i64 {
match self {
Self::Standalone(connection) => connection.get_db(),
Self::Cluster(connection) => connection.get_db(),
}
}
}

View file

@ -1,7 +1,13 @@
mod client;
mod commands;
mod config;
mod connection;
mod error;
mod namespace;
mod pipeline;
mod script;
#[cfg(test)]
mod test_support;
pub use client::{RedisClient, Ttl};
pub use config::{
@ -9,3 +15,6 @@ pub use config::{
};
pub use error::{ConnectError, Error};
pub use namespace::Namespace;
pub use pipeline::Pipeline;
pub use redis::Value;
pub use script::Script;

View file

@ -0,0 +1,288 @@
use std::{collections::HashMap, time::Duration};
use futures_util::future::join_all;
use redis::{
Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value,
cluster_routing::Slot,
};
use crate::{Error, Namespace, Script, commands, connection::Connection};
pub struct Pipeline {
namespace: Namespace,
ops: Vec<Op>,
}
struct Op {
command: Cmd,
key: String,
script: Option<ScriptCall>,
}
struct ScriptCall {
script: Script,
keys: Vec<String>,
args: Vec<Vec<u8>>,
}
type Settled = (usize, Result<Value, Error>);
impl Pipeline {
pub(crate) fn new(namespace: Namespace) -> Self {
Self {
namespace,
ops: Vec::new(),
}
}
pub fn get(self, key: &str) -> Self {
let key = self.namespace.apply(key);
self.command(commands::get(key.clone()), key)
}
pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self {
let key = self.namespace.apply(key);
self.command(commands::set(key.clone(), value, ttl), key)
}
pub fn del(self, key: &str) -> Self {
let key = self.namespace.apply(key);
self.command(commands::del(vec![key.clone()]), key)
}
pub fn expire(self, key: &str, ttl: Duration) -> Self {
let key = self.namespace.apply(key);
self.command(commands::expire(key.clone(), ttl), key)
}
pub fn incr_by(self, key: &str, delta: i64) -> Self {
let key = self.namespace.apply(key);
self.command(commands::incr_by(key.clone(), delta), key)
}
pub fn incr_by_float(self, key: &str, delta: f64) -> Self {
let key = self.namespace.apply(key);
self.command(commands::incr_by_float(key.clone(), delta), key)
}
pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self {
let keys: Vec<String> = keys.iter().map(|key| self.namespace.apply(key)).collect();
let args = args.to_redis_args();
let op = Op {
command: script.evalsha(&keys, &args),
key: keys.first().cloned().unwrap_or_default(),
script: Some(ScriptCall {
script: script.clone(),
keys,
args,
}),
};
self.push(op)
}
fn command(self, command: Cmd, key: String) -> Self {
self.push(Op {
command,
key,
script: None,
})
}
fn push(mut self, op: Op) -> Self {
self.ops.push(op);
self
}
pub(crate) async fn run(&self, connection: &Connection) -> Vec<Result<Value, Error>> {
let groups = self.groups(connection);
let settled = join_all(
groups
.iter()
.map(|group| self.run_group(connection.clone(), group)),
)
.await;
let mut settled: Vec<Settled> = settled.into_iter().flatten().collect();
settled.sort_by_key(|(index, _)| *index);
settled.into_iter().map(|(_, result)| result).collect()
}
fn groups(&self, connection: &Connection) -> Vec<Vec<usize>> {
match connection {
_ if self.ops.is_empty() => Vec::new(),
Connection::Standalone(_) => vec![(0..self.ops.len()).collect()],
Connection::Cluster(_) => {
let mut by_slot: HashMap<Slot, Vec<usize>> = HashMap::new();
for (index, op) in self.ops.iter().enumerate() {
by_slot
.entry(Slot::for_key(&op.key))
.or_default()
.push(index);
}
by_slot.into_values().collect()
}
}
}
async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec<Settled> {
let mut pipeline = redis::pipe();
pipeline.ignore_errors();
for &index in indices {
pipeline.add_command(self.ops[index].command.clone());
}
let replies: RedisResult<Vec<RedisResult<Value>>> =
pipeline.query_async(&mut connection).await;
let replies = match replies {
Ok(replies) => replies,
Err(error) => {
return indices
.iter()
.map(|&index| (index, Err(Error::from(error.clone()))))
.collect();
}
};
let mut settled = Vec::with_capacity(indices.len());
for (&index, reply) in indices.iter().zip(replies) {
settled.push((index, self.settle(&mut connection, index, reply).await));
}
settled
}
async fn settle(
&self,
connection: &mut Connection,
index: usize,
reply: RedisResult<Value>,
) -> Result<Value, Error> {
match (reply, &self.ops[index].script) {
(Err(error), Some(call))
if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) =>
{
Ok(call
.script
.invocation(&call.keys, &call.args)
.invoke_async(connection)
.await?)
}
(reply, _) => Ok(reply?),
}
}
}
#[cfg(test)]
mod tests {
use std::time::{Duration, Instant};
use redis::Value;
use rstest::rstest;
use crate::{
DEFAULT_TIMEOUT, Error, Script,
test_support::{Server, Topology, flush_scripts},
};
#[rstest]
#[tokio::test]
async fn each_op_settles_on_its_own_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("text", "not a number", Duration::from_secs(60))
.await
.unwrap();
let pipeline = client
.pipeline()
.incr_by("a", 1)
.incr_by("text", 1)
.incr_by("c", 5)
.set("b", "2", Duration::from_secs(60))
.get("b")
.incr_by("a", 1)
.get("missing");
let results = client.run(&pipeline).await;
assert_eq!(results.len(), 7);
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1));
assert!(
matches!(results[1], Err(Error::Server(_))),
"{:?}",
results[1]
);
assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5));
assert_eq!(results[3].as_ref().unwrap(), &Value::Okay);
assert_eq!(
results[4].as_ref().unwrap(),
&Value::BulkString(b"2".to_vec())
);
assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2));
assert_eq!(results[6].as_ref().unwrap(), &Value::Nil);
}
#[rstest]
#[tokio::test]
async fn a_flushed_script_in_a_pipeline_reruns_alone(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("ns", DEFAULT_TIMEOUT).await;
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.eval::<i64>(&incr, &["counter"], 1).await.unwrap();
flush_scripts(&client).await;
let pipeline = client
.pipeline()
.eval(&incr, &["counter"], 2)
.incr_by("other", 7);
let results = client.run(&pipeline).await;
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3));
assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7));
assert_eq!(client.get::<i64>("counter").await.unwrap(), Some(3));
}
#[rstest]
#[tokio::test]
async fn an_unreachable_server_fails_every_op_as_transport(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", Duration::from_millis(500)).await;
let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1);
server.freeze().await;
let results = client.run(&pipeline).await;
server.thaw().await;
assert_eq!(results.len(), 2);
assert!(
matches!(results[0], Err(Error::Transport(_))),
"{:?}",
results[0]
);
assert!(
matches!(results[1], Err(Error::Transport(_))),
"{:?}",
results[1]
);
}
#[rstest]
#[tokio::test]
async fn an_empty_pipeline_returns_without_asking_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", Duration::from_secs(2)).await;
server.freeze().await;
let started = Instant::now();
let results = client.run(&client.pipeline()).await;
let elapsed = started.elapsed();
server.thaw().await;
assert!(results.is_empty());
assert!(elapsed < Duration::from_secs(1), "{elapsed:?}");
}
}

View file

@ -0,0 +1,88 @@
use redis::{Cmd, ScriptInvocation};
#[derive(Clone)]
pub struct Script(redis::Script);
impl Script {
pub fn new(source: &str) -> Self {
Self(redis::Script::new(source))
}
pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec<u8>]) -> Cmd {
redis::cmd("EVALSHA")
.arg(self.0.get_hash())
.arg(keys.len())
.arg(keys)
.arg(args)
.to_owned()
}
pub(crate) fn invocation(&self, keys: &[String], args: &[Vec<u8>]) -> ScriptInvocation<'_> {
let mut invocation = self.0.prepare_invoke();
invocation.key(keys).arg(args);
invocation
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use rstest::rstest;
use super::Script;
use crate::{
DEFAULT_TIMEOUT, Error,
test_support::{Server, Topology, flush_scripts},
};
#[rstest]
#[tokio::test]
async fn script_keys_are_namespaced(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("a", DEFAULT_TIMEOUT).await;
let bare = server.client("", DEFAULT_TIMEOUT).await;
let read = Script::new("return redis.call('GET', KEYS[1])");
bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap();
let value: Option<String> = client
.eval(&read, &["x"], Vec::<String>::new())
.await
.unwrap();
assert_eq!(value.as_deref(), Some("1"));
}
#[rstest]
#[tokio::test]
async fn a_failing_script_is_a_server_error_and_keeps_the_connection(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.rpush("list", "item").await.unwrap();
let failed = client.eval::<i64>(&incr, &["list"], 1).await;
assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}");
assert_eq!(client.eval::<i64>(&incr, &["counter"], 3).await.unwrap(), 3);
}
#[rstest]
#[tokio::test]
async fn a_flushed_script_is_loaded_again(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.eval::<i64>(&incr, &["counter"], 1).await.unwrap();
flush_scripts(&client).await;
assert_eq!(client.eval::<i64>(&incr, &["counter"], 1).await.unwrap(), 2);
}
}

View file

@ -0,0 +1,101 @@
use std::time::Duration;
use litellm_redis_testing::{RedisCluster, RedisServer};
use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr};
use crate::{
ClusterEndpoint, ClusterNode, Endpoint, Namespace, RedisClient, RedisConfig,
connection::Connection,
};
#[derive(Clone, Copy, Debug)]
pub(crate) enum Topology {
Standalone,
Cluster,
}
pub(crate) enum Server {
Standalone(RedisServer),
Cluster(RedisCluster),
}
impl Server {
pub(crate) async fn start(topology: Topology) -> Self {
match topology {
Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()),
Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()),
}
}
pub(crate) 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,
}),
}
}
pub(crate) async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient {
RedisClient::connect(RedisConfig {
endpoint: self.endpoint(),
namespace: Namespace::new(namespace),
timeout,
})
.await
.unwrap()
}
pub(crate) async fn restart(&self) {
match self {
Self::Standalone(server) => server.restart().await.unwrap(),
Self::Cluster(cluster) => cluster.restart().await.unwrap(),
}
}
pub(crate) async fn freeze(&self) {
match self {
Self::Standalone(server) => server.freeze().await.unwrap(),
Self::Cluster(cluster) => cluster.freeze().await.unwrap(),
}
}
pub(crate) async fn thaw(&self) {
match self {
Self::Standalone(server) => server.thaw().await.unwrap(),
Self::Cluster(cluster) => cluster.thaw().await.unwrap(),
}
}
}
pub(crate) 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()
}
}
}
pub(crate) async fn flush_scripts(client: &RedisClient) {
redis::cmd("SCRIPT")
.arg("FLUSH")
.query_async::<()>(&mut client.connection.clone())
.await
.unwrap();
}