scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 131

feat(cache-redis): expose the pooled connection handling for reuse

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-21 20:24:23 +00:00
parent 769917a7e8
commit 2b257bd9a3
3 changed files with 72 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

93
litellm-rust/crates/cache-redis/src/cache.rs
View file

@ -19,7 +19,7 @@ const DEFAULT_TTL: Duration = Duration::from_secs(600);
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
struct PooledConnection {
pub struct PooledConnection {
connection: redis::Connection,
failed: bool,
}
@ -27,16 +27,19 @@ struct PooledConnection {
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
struct ConnectionManager(redis::Client);
pub struct ConnectionManager {
client: redis::Client,
timeout: Duration,
}
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection;
type Error = redis::RedisError;
fn connect(&self) -> Result<PooledConnection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
let connection = self.client.get_connection()?;
connection.set_read_timeout(Some(self.timeout))?;
connection.set_write_timeout(Some(self.timeout))?;
Ok(PooledConnection {
connection,
failed: false,
@ -68,12 +71,12 @@ const CLAIM_SCRIPT: &str = concat!(
);
const CLAIM_ATTEMPTS: usize = 8;
enum Connections<C> {
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Fixed(Mutex<C>),
}
struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike);
pub struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike);
impl redis::ConnectionLike for ConnectionRef<'_> {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
@ -110,7 +113,23 @@ impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn execute<T>(
pub fn pooled(url: &str, timeout: Duration, pool_size: u32) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let pool = r2d2::Pool::builder()
.max_size(pool_size)
.min_idle(Some(0))
.connection_timeout(timeout)
.test_on_check_out(false)
.build(ConnectionManager { client, timeout })
.map_err(|_| Error::Unavailable)?;
Ok(Self::Pool(pool))
}
pub fn fixed(connection: C) -> Self {
Self::Fixed(Mutex::new(connection))
}
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
@ -127,6 +146,16 @@ where
}
}
}
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
}
pub struct RedisCache<S, C = redis::Connection> {
@ -138,16 +167,8 @@ pub struct RedisCache<S, C = redis::Connection> {
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let pool = r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(ConnectionManager(client))
.map_err(|_| Error::Unavailable)?;
Ok(Self {
connections: Arc::new(Connections::Pool(pool)),
connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
@ -162,7 +183,7 @@ where
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
Self {
connections: Arc::new(Connections::Fixed(Mutex::new(connection))),
connections: Arc::new(Connections::fixed(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
@ -241,20 +262,14 @@ where
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
ttl_seconds(ttl)
}
}
async fn run_blocking<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
pub fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
@ -313,7 +328,7 @@ where
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
@ -327,7 +342,7 @@ where
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
@ -350,7 +365,7 @@ where
})
.collect::<Result<Vec<_>, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, payload) in entries {
pipeline
@ -372,7 +387,7 @@ where
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Self::run_blocking(Arc::clone(&self.connections), |connection| {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match redis::cmd("PING").query::<String>(connection) {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
@ -433,7 +448,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -460,7 +475,7 @@ where
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
@ -480,7 +495,7 @@ where
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
@ -512,7 +527,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
.await
@ -623,7 +638,7 @@ where
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
let codec = self.codec.clone();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
})
.await

28
litellm-rust/crates/cache-redis/src/cache/operations.rs vendored
View file

@ -144,7 +144,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del(keys).map_err(|_| Error::Unavailable)
})
.await
@ -172,7 +172,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -192,7 +192,7 @@ where
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
Connections::run_blocking(Arc::clone(&self.connections), |connection| {
redis::cmd("PING")
.query::<String>(connection)
.map(|response| response == "PONG")
@ -203,7 +203,7 @@ where
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
let key = self.namespaced_key(key);
let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
@ -215,7 +215,7 @@ where
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut cursor = 0u64;
let mut matches = Vec::new();
loop {
@ -249,7 +249,7 @@ where
}
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
pipeline.cmd("SADD").arg(&key).arg(values);
pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore();
@ -266,7 +266,7 @@ where
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
@ -292,7 +292,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, values) in operations {
pipeline.cmd("RPUSH").arg(key).arg(values);
@ -309,7 +309,7 @@ where
) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let multiple = count.is_some();
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
@ -338,7 +338,7 @@ where
.iter()
.map(|(_, count)| count.is_some())
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, count) in operations {
let command = pipeline.cmd("LPOP").arg(key);
@ -368,7 +368,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
@ -440,7 +440,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, amount, ttl) in operations {
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount);
@ -461,7 +461,7 @@ where
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment_with_floor(connection, key, amount, ttl)
})
.await
@ -475,7 +475,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)

4
litellm-rust/crates/cache-redis/src/lib.rs
View file

@ -1,6 +1,10 @@
mod cache;
mod topology;
pub mod connection {
pub use crate::cache::{ConnectionRef, Connections, ttl_seconds};
}
pub use cache::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};