mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
feat(auth): add bounded Moka cache and native invalidation subscriber
This commit is contained in:
Devin AI
parent
77979758cd
commit
13fa09b04a
22 changed files with 874 additions and 218 deletions
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -4535,6 +4535,7 @@ dependencies = [
|
|||
"litellm-redis-testing",
|
||||
"litellm-secret-value",
|
||||
"litellm-tracing",
|
||||
"moka",
|
||||
"rstest",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
|
|
|
|||
|
|
@ -43,7 +43,7 @@ mod _native {
|
|||
#[pymodule_export]
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes};
|
||||
use crate::redis::store::auth_cache::{NativeAuthCache, NativeAuthSubscription, PyLocalSizes};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::rate_limits::{
|
||||
NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired,
|
||||
|
|
@ -273,6 +273,7 @@ mod tests {
|
|||
"NativeRequestRedis",
|
||||
"NativeSend",
|
||||
"NativeAuthCache",
|
||||
"NativeAuthSubscription",
|
||||
"NativeRateLimits",
|
||||
"WindowCount",
|
||||
"AtomicApplied",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,9 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use litellm_redis_client::SubscriptionEvent;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
use litellm_redis_store::auth_cache::{AuthCache, Settings};
|
||||
use pyo3::prelude::*;
|
||||
|
|
@ -44,6 +48,35 @@ impl NativeAuthCache {
|
|||
self.0.has_redis()
|
||||
}
|
||||
|
||||
fn attach(&self, client: &NativeRedisClient) -> bool {
|
||||
self.0.attach(remote(client))
|
||||
}
|
||||
|
||||
fn current_generation(&self) -> u64 {
|
||||
self.0.current_generation()
|
||||
}
|
||||
|
||||
#[pyo3(signature = (key, value, generation, ttl_seconds = None))]
|
||||
fn set_if_generation(
|
||||
&self,
|
||||
key: &str,
|
||||
value: PyValue,
|
||||
generation: u64,
|
||||
ttl_seconds: Option<f64>,
|
||||
) -> PyResult<bool> {
|
||||
Ok(self
|
||||
.0
|
||||
.set_if_generation(key, value.0, generation, ttl(ttl_seconds)?))
|
||||
}
|
||||
|
||||
fn subscribe(&self, client: &NativeRedisClient, channel: &str) -> NativeAuthSubscription {
|
||||
self.0.subscriber_connected(false);
|
||||
NativeAuthSubscription {
|
||||
cache: self.0.clone(),
|
||||
stream: Arc::new(Mutex::new(Some(client.client.subscribe_events(channel)))),
|
||||
}
|
||||
}
|
||||
|
||||
/// `request` answers a key its pre-call reads already found absent without another `GET`.
|
||||
#[pyo3(signature = (key, request = None))]
|
||||
fn get<'py>(
|
||||
|
|
@ -245,6 +278,65 @@ impl NativeAuthCache {
|
|||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen)]
|
||||
pub(crate) struct NativeAuthSubscription {
|
||||
cache: Arc<AuthCache>,
|
||||
stream: Arc<Mutex<Option<BoxStream<'static, SubscriptionEvent>>>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeAuthSubscription {
|
||||
fn is_connected(&self) -> bool {
|
||||
self.cache.local_ready()
|
||||
}
|
||||
|
||||
fn next<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let cache = self.cache.clone();
|
||||
let stream = self.stream.clone();
|
||||
litellm_host_python::run_async_value(py, async move {
|
||||
let mut guard = stream.lock().await;
|
||||
let Some(stream) = guard.as_mut() else {
|
||||
return Ok(None::<Py<pyo3::types::PyBytes>>);
|
||||
};
|
||||
match stream.next().await {
|
||||
Some(SubscriptionEvent::Connected) => cache.subscriber_connected(true),
|
||||
Some(SubscriptionEvent::Message(payload)) => {
|
||||
if let Ok(message) = serde_json::from_slice::<serde_json::Value>(&payload)
|
||||
&& let Some(key) =
|
||||
message.get("cache_key").and_then(serde_json::Value::as_str)
|
||||
&& message
|
||||
.get("new_value")
|
||||
.and_then(serde_json::Value::as_f64)
|
||||
.is_none()
|
||||
{
|
||||
cache.delete_local(key).map_err(kind_error)?;
|
||||
}
|
||||
return Ok(Some(Python::attach(|py| {
|
||||
pyo3::types::PyBytes::new(py, &payload).unbind()
|
||||
})));
|
||||
}
|
||||
_ => cache.subscriber_connected(false),
|
||||
}
|
||||
Ok(None)
|
||||
})
|
||||
}
|
||||
|
||||
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
self.cache.subscriber_connected(false);
|
||||
let stream = self.stream.clone();
|
||||
litellm_host_python::run_async_value(py, async move {
|
||||
stream.lock().await.take();
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for NativeAuthSubscription {
|
||||
fn drop(&mut self) {
|
||||
self.cache.subscriber_connected(false);
|
||||
}
|
||||
}
|
||||
|
||||
/// Entries the management-object and the user-key partitions hold in memory.
|
||||
#[pyclass(frozen, eq, get_all, name = "LocalSizes")]
|
||||
#[derive(PartialEq)]
|
||||
|
|
|
|||
|
|
@ -155,6 +155,10 @@ impl RedisClient {
|
|||
self.0.subscriber.clone().messages(channel.to_owned())
|
||||
}
|
||||
|
||||
pub fn subscribe_events(&self, channel: &str) -> BoxStream<'static, crate::SubscriptionEvent> {
|
||||
self.0.subscriber.clone().events(channel.to_owned())
|
||||
}
|
||||
|
||||
pub fn pipeline(&self) -> Pipeline {
|
||||
Pipeline::new(self.0.namespace.clone())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ pub use litellm_redis_types::{
|
|||
};
|
||||
pub use namespace::Namespace;
|
||||
pub use pipeline::Pipeline;
|
||||
pub use pubsub::SubscriptionEvent;
|
||||
pub use redis::{
|
||||
BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value,
|
||||
cluster_routing::Slot,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,10 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use redis::{ConnectionInfo, RedisResult, aio::PubSubStream};
|
||||
use redis::{
|
||||
ConnectionInfo, RedisResult,
|
||||
aio::{PubSubSink, PubSubStream},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
ConnectError, Credentials, Standalone, Target, Timeouts, config,
|
||||
|
|
@ -11,6 +14,18 @@ use crate::{
|
|||
// 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);
|
||||
const IDLE_PROBE: Duration = Duration::from_secs(15);
|
||||
|
||||
pub enum SubscriptionEvent {
|
||||
Connected,
|
||||
Disconnected,
|
||||
Message(Vec<u8>),
|
||||
}
|
||||
|
||||
struct Live {
|
||||
sink: PubSubSink,
|
||||
stream: PubSubStream,
|
||||
}
|
||||
|
||||
/// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends
|
||||
/// `SUBSCRIBE`, authenticated once when it connects.
|
||||
|
|
@ -66,43 +81,76 @@ impl Subscriber {
|
|||
}
|
||||
|
||||
pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec<u8>> {
|
||||
self.events(channel)
|
||||
.filter_map(|event| async move {
|
||||
match event {
|
||||
SubscriptionEvent::Message(payload) => Some(payload),
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
.boxed()
|
||||
}
|
||||
|
||||
pub(crate) fn events(self, channel: String) -> BoxStream<'static, SubscriptionEvent> {
|
||||
futures_util::stream::unfold(
|
||||
(self, channel, None::<PubSubStream>, BACKOFF_INITIAL),
|
||||
|(subscriber, channel, live, backoff)| async move {
|
||||
(self, channel, None::<Live>, BACKOFF_INITIAL, false),
|
||||
|(subscriber, channel, live, backoff, retry)| async move {
|
||||
let mut live = live;
|
||||
let mut backoff = backoff;
|
||||
if retry {
|
||||
tokio::time::sleep(backoff).await;
|
||||
}
|
||||
loop {
|
||||
let Some(current) = live.as_mut() else {
|
||||
match subscriber.subscribe(&channel).await {
|
||||
let event = match subscriber.subscribe(&channel).await {
|
||||
Ok(subscribed) => {
|
||||
live = Some(subscribed);
|
||||
backoff = BACKOFF_INITIAL;
|
||||
SubscriptionEvent::Connected
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying");
|
||||
tokio::time::sleep(backoff).await;
|
||||
backoff = (backoff * 2).min(BACKOFF_MAX);
|
||||
Err(_) => {
|
||||
if retry {
|
||||
backoff = (backoff * 2).min(BACKOFF_MAX);
|
||||
}
|
||||
SubscriptionEvent::Disconnected
|
||||
}
|
||||
}
|
||||
continue;
|
||||
};
|
||||
let retry = live.is_none();
|
||||
return Some((event, (subscriber, channel, live, backoff, retry)));
|
||||
};
|
||||
match current.next().await {
|
||||
Some(message) => {
|
||||
match tokio::time::timeout(IDLE_PROBE, current.stream.next()).await {
|
||||
Ok(Some(message)) => {
|
||||
let payload = message.get_payload_bytes().to_vec();
|
||||
return Some((payload, (subscriber, channel, live, backoff)));
|
||||
return Some((
|
||||
SubscriptionEvent::Message(payload),
|
||||
(subscriber, channel, live, backoff, false),
|
||||
));
|
||||
}
|
||||
None => live = None,
|
||||
Err(_) => {
|
||||
if matches!(
|
||||
tokio::time::timeout(
|
||||
subscriber.connect_timeout,
|
||||
current.sink.ping::<String>()
|
||||
)
|
||||
.await,
|
||||
Ok(Ok(_))
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
Ok(None) => {}
|
||||
}
|
||||
tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing");
|
||||
tokio::time::sleep(backoff).await;
|
||||
backoff = (backoff * 2).min(BACKOFF_MAX);
|
||||
return Some((
|
||||
SubscriptionEvent::Disconnected,
|
||||
(subscriber, channel, None, backoff, true),
|
||||
));
|
||||
}
|
||||
},
|
||||
)
|
||||
.boxed()
|
||||
}
|
||||
|
||||
async fn subscribe(&self, channel: &str) -> RedisResult<PubSubStream> {
|
||||
async fn subscribe(&self, channel: &str) -> RedisResult<Live> {
|
||||
match &self.source {
|
||||
Source::Sentinel(sentinel) => {
|
||||
let master = sentinel.master().await?;
|
||||
|
|
@ -122,7 +170,7 @@ impl Subscriber {
|
|||
}
|
||||
}
|
||||
|
||||
async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult<PubSubStream> {
|
||||
async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult<Live> {
|
||||
let settings =
|
||||
config::authenticated(info.redis_settings().clone(), &self.credentials).await?;
|
||||
let client = redis::Client::open(info.set_redis_settings(settings))?;
|
||||
|
|
@ -135,7 +183,8 @@ impl Subscriber {
|
|||
))
|
||||
})??;
|
||||
subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?;
|
||||
Ok(pubsub.into_on_message())
|
||||
let (sink, stream) = pubsub.split();
|
||||
Ok(Live { sink, stream })
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ use litellm_circuit_breaker::{Breaker, Config};
|
|||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{
|
||||
Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel,
|
||||
Target, Timeouts,
|
||||
SubscriptionEvent, Target, Timeouts,
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually};
|
||||
use rstest::rstest;
|
||||
|
|
@ -16,6 +16,85 @@ use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel};
|
|||
|
||||
const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation";
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_silent_dead_subscription_is_detected_and_reconnected(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let proxy = support::proxied(&server).await;
|
||||
let client = RedisClient::new(RedisConfig {
|
||||
target: proxy.target.clone(),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::new(Duration::from_millis(500), None),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap();
|
||||
let mut events = client.subscribe_events(CHANNEL);
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(Duration::from_secs(5), events.next())
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(SubscriptionEvent::Connected)
|
||||
));
|
||||
proxy.blackhole_open_connections();
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(Duration::from_secs(20), events.next())
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(SubscriptionEvent::Disconnected)
|
||||
));
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(Duration::from_secs(10), events.next())
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(SubscriptionEvent::Connected)
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn lifecycle_reports_disconnect_before_reconnect(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", Duration::from_millis(500));
|
||||
let mut events = client.subscribe_events(CHANNEL);
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(Duration::from_secs(5), events.next())
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(SubscriptionEvent::Connected)
|
||||
));
|
||||
server.restart().await.unwrap();
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(Duration::from_secs(5), events.next())
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(SubscriptionEvent::Disconnected)
|
||||
));
|
||||
assert!(matches!(
|
||||
tokio::time::timeout(Duration::from_secs(15), events.next())
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(SubscriptionEvent::Connected)
|
||||
));
|
||||
eventually(async || {
|
||||
client
|
||||
.send(commands::publish(CHANNEL, b"after".as_slice()))
|
||||
.await
|
||||
.ok()
|
||||
})
|
||||
.await;
|
||||
assert!(
|
||||
matches!(tokio::time::timeout(Duration::from_secs(5), events.next()).await.unwrap(), Some(SubscriptionEvent::Message(bytes)) if bytes == b"after")
|
||||
);
|
||||
}
|
||||
|
||||
fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver<Vec<u8>> {
|
||||
let (sender, receiver) = unbounded_channel();
|
||||
let mut messages = client.subscribe(channel);
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ litellm-cache.workspace = true
|
|||
litellm-cache-memory.workspace = true
|
||||
litellm-python-compat.workspace = true
|
||||
litellm-redis-client.workspace = true
|
||||
moka = { workspace = true, features = ["sync"] }
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
|
||||
|
|
|
|||
108
litellm-rust/crates/redis-store/src/auth_cache/local.rs
Normal file
108
litellm-rust/crates/redis-store/src/auth_cache/local.rs
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
use moka::sync::Cache;
|
||||
|
||||
use super::{Settings, is_key_object};
|
||||
|
||||
const MAX_AUTH_TTL: Duration = Duration::from_secs(60);
|
||||
|
||||
#[derive(Clone)]
|
||||
struct Entry {
|
||||
value: Value,
|
||||
expires: Duration,
|
||||
}
|
||||
|
||||
pub(super) struct Local {
|
||||
objects: Cache<String, Entry>,
|
||||
keys: Cache<String, Entry>,
|
||||
pub generation: u64,
|
||||
pub healthy: bool,
|
||||
}
|
||||
|
||||
impl Local {
|
||||
pub fn new(settings: Settings) -> Self {
|
||||
let partition = || {
|
||||
Cache::builder()
|
||||
.max_capacity(settings.memory_entries as u64)
|
||||
.build()
|
||||
};
|
||||
Self {
|
||||
objects: partition(),
|
||||
keys: partition(),
|
||||
generation: 0,
|
||||
healthy: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn partition(&self, key: &str) -> &Cache<String, Entry> {
|
||||
if is_key_object(key) {
|
||||
&self.keys
|
||||
} else {
|
||||
&self.objects
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(&self, key: &str, now: Duration) -> Option<Value> {
|
||||
if !self.healthy {
|
||||
return None;
|
||||
}
|
||||
let entry = self.partition(key).get(key)?;
|
||||
if now >= entry.expires {
|
||||
self.partition(key).invalidate(key);
|
||||
return None;
|
||||
}
|
||||
Some(entry.value)
|
||||
}
|
||||
|
||||
pub fn set(&self, key: &str, value: Value, ttl: Option<Duration>, now: Duration) {
|
||||
if !self.healthy || key.starts_with("sk-") {
|
||||
return;
|
||||
}
|
||||
let ttl = ttl.unwrap_or(Duration::from_secs(600));
|
||||
let ttl = if is_key_object(key) {
|
||||
ttl.min(MAX_AUTH_TTL)
|
||||
} else {
|
||||
ttl
|
||||
};
|
||||
self.partition(key).insert(
|
||||
key.to_owned(),
|
||||
Entry {
|
||||
value,
|
||||
expires: now.saturating_add(ttl),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn delete(&mut self, key: &str) {
|
||||
self.generation = self.generation.wrapping_add(1);
|
||||
self.partition(key).invalidate(key);
|
||||
}
|
||||
|
||||
pub fn flush(&mut self) {
|
||||
self.generation = self.generation.wrapping_add(1);
|
||||
self.keys.invalidate_all();
|
||||
self.objects.invalidate_all();
|
||||
}
|
||||
|
||||
pub fn reconfigure(&mut self, settings: Settings) {
|
||||
let replacement = Self::new(settings);
|
||||
for (key, value) in &self.keys {
|
||||
replacement.keys.insert((*key).clone(), value);
|
||||
}
|
||||
for (key, value) in &self.objects {
|
||||
replacement.objects.insert((*key).clone(), value);
|
||||
}
|
||||
self.keys = replacement.keys;
|
||||
self.objects = replacement.objects;
|
||||
}
|
||||
|
||||
pub fn len(&self) -> (usize, usize) {
|
||||
self.objects.run_pending_tasks();
|
||||
self.keys.run_pending_tasks();
|
||||
(
|
||||
self.objects.entry_count() as usize,
|
||||
self.keys.entry_count() as usize,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
mod keys;
|
||||
mod local;
|
||||
mod types;
|
||||
|
||||
use std::{
|
||||
|
|
@ -12,23 +13,19 @@ use litellm_redis_client::{Declared, RedisBatch, commands};
|
|||
pub use self::keys::is_key_object;
|
||||
pub use self::types::Settings;
|
||||
use crate::{
|
||||
Error, Number, Pending, RequestRedis, Send,
|
||||
dual_cache::DualCache,
|
||||
Error, Number, Pending, RequestRedis, Send, codec,
|
||||
failure_log::redis_failure,
|
||||
memory::{Clock, lock, system_clock},
|
||||
remote::{Remote, RemoteSlot},
|
||||
request::ready,
|
||||
request::{pending, ready},
|
||||
};
|
||||
|
||||
/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it.
|
||||
const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])";
|
||||
|
||||
/// The proxy's auth object cache with `enable_redis_auth_cache` on: `UserApiKeyCache`'s two
|
||||
/// `DualCache` partitions, so churn in management objects cannot evict user-key objects, in front of
|
||||
/// one Redis shared by every pod. Keys are the logical cache keys `UserApiKeyCache` callers already
|
||||
/// use; values are what `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`.
|
||||
pub struct AuthCache {
|
||||
key_objects: DualCache<Value>,
|
||||
objects: DualCache<Value>,
|
||||
local: Mutex<local::Local>,
|
||||
now: Clock,
|
||||
redis: Arc<RemoteSlot>,
|
||||
settings: Mutex<Settings>,
|
||||
}
|
||||
|
|
@ -40,10 +37,9 @@ impl AuthCache {
|
|||
|
||||
pub fn with_clock(remote: Option<Remote>, settings: Settings, now: Clock) -> Self {
|
||||
let redis = Arc::new(RemoteSlot::new(remote));
|
||||
let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone());
|
||||
Self {
|
||||
key_objects: partition(),
|
||||
objects: partition(),
|
||||
local: Mutex::new(local::Local::new(settings)),
|
||||
now,
|
||||
redis,
|
||||
settings: Mutex::new(settings),
|
||||
}
|
||||
|
|
@ -53,6 +49,40 @@ impl AuthCache {
|
|||
self.redis.get().is_some()
|
||||
}
|
||||
|
||||
pub fn attach(&self, remote: Remote) -> bool {
|
||||
self.redis.attach(remote)
|
||||
}
|
||||
|
||||
pub fn current_generation(&self) -> u64 {
|
||||
lock(&self.local).generation
|
||||
}
|
||||
|
||||
pub fn local_ready(&self) -> bool {
|
||||
lock(&self.local).healthy
|
||||
}
|
||||
|
||||
pub fn subscriber_connected(&self, connected: bool) {
|
||||
let mut local = lock(&self.local);
|
||||
local.flush();
|
||||
local.healthy = connected;
|
||||
}
|
||||
|
||||
pub fn set_if_generation(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Value,
|
||||
generation: u64,
|
||||
ttl: Option<Duration>,
|
||||
) -> bool {
|
||||
let ttl = ttl.or(self.settings().memory_ttl);
|
||||
let local = lock(&self.local);
|
||||
if local.generation != generation || !local.healthy || key.starts_with("sk-") {
|
||||
return false;
|
||||
}
|
||||
local.set(key, value, ttl, (self.now)());
|
||||
true
|
||||
}
|
||||
|
||||
pub fn settings(&self) -> Settings {
|
||||
*lock(&self.settings)
|
||||
}
|
||||
|
|
@ -61,37 +91,33 @@ impl AuthCache {
|
|||
/// keep what they hold.
|
||||
pub fn reconfigure(&self, settings: Settings) {
|
||||
*lock(&self.settings) = settings;
|
||||
self.key_objects.reconfigure(settings.memory());
|
||||
self.objects.reconfigure(settings.memory());
|
||||
lock(&self.local).reconfigure(settings);
|
||||
}
|
||||
|
||||
pub fn get_local(&self, key: &str) -> Result<Option<Value>, Error> {
|
||||
Ok(self.partition(key).memory().get(key)?)
|
||||
Ok(lock(&self.local).get(key, (self.now)()))
|
||||
}
|
||||
|
||||
/// A memory write with `ttl`, or `default_in_memory_ttl` when it names none.
|
||||
pub fn set_local(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
|
||||
let partition = self.partition(key);
|
||||
Ok(partition
|
||||
.memory()
|
||||
.set(key, value, ttl.or(partition.default_in_memory_ttl()))?)
|
||||
let ttl = ttl.or(self.settings().memory_ttl);
|
||||
lock(&self.local).set(key, value, ttl, (self.now)());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_local(&self, key: &str) -> Result<(), Error> {
|
||||
Ok(self.partition(key).memory().delete(key)?)
|
||||
lock(&self.local).delete(key);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn flush_local(&self) -> Result<(), Error> {
|
||||
self.key_objects.memory().flush()?;
|
||||
Ok(self.objects.memory().flush()?)
|
||||
lock(&self.local).flush();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Entries held by the management-object and the user-key partitions.
|
||||
pub fn len_local(&self) -> Result<(usize, usize), Error> {
|
||||
Ok((
|
||||
self.objects.memory().len()?,
|
||||
self.key_objects.memory().len()?,
|
||||
))
|
||||
Ok(lock(&self.local).len())
|
||||
}
|
||||
|
||||
/// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read
|
||||
|
|
@ -101,12 +127,42 @@ impl AuthCache {
|
|||
request: Option<&RequestRedis>,
|
||||
key: &str,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
self.partition(key).get(request, key, None).await
|
||||
let generation = self.current_generation();
|
||||
if let Some(value) = self.get_local(key)? {
|
||||
return Ok(Some(value));
|
||||
}
|
||||
if !lock(&self.local).healthy || key.starts_with("sk-") {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(client) = self.redis.client() else {
|
||||
return Ok(None);
|
||||
};
|
||||
if request.is_some_and(|request| request.misses(client).contains(key)) {
|
||||
return Ok(None);
|
||||
}
|
||||
let value = client
|
||||
.send(commands::get::<Vec<u8>>(key))
|
||||
.await
|
||||
.map_err(Error::from)
|
||||
.and_then(|raw| Ok(types::decode(raw)?))
|
||||
.unwrap_or_else(|error| {
|
||||
redis_failure!(&error, "auth cache read failed");
|
||||
None
|
||||
});
|
||||
if let Some(value) = &value
|
||||
&& !self.set_if_generation(key, value.clone(), generation, None)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
/// `DualCache.async_set_cache`, writing `json.dumps(value)`.
|
||||
pub async fn set(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
|
||||
self.partition(key).set(key, value, ttl).await
|
||||
if let Err(error) = self.declare_set(Send::Now, key, value, ttl)?.await {
|
||||
redis_failure!(&error, "auth cache write failed");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`.
|
||||
|
|
@ -116,7 +172,14 @@ impl AuthCache {
|
|||
value: Value,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
self.partition(key).set_str(key, value, ttl).await
|
||||
let encoded = codec::str(value.clone());
|
||||
if let Err(error) = self
|
||||
.declare_encoded(Send::Now, key, value, ttl, Ok(encoded))?
|
||||
.await
|
||||
{
|
||||
redis_failure!(&error, "auth cache write failed");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks.
|
||||
|
|
@ -127,12 +190,56 @@ impl AuthCache {
|
|||
value: Value,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Pending<()>, Error> {
|
||||
self.partition(key).declare_set(send, key, value, ttl)
|
||||
let encoded = codec::dumps(value.clone());
|
||||
self.declare_encoded(send, key, value, ttl, encoded)
|
||||
}
|
||||
|
||||
fn declare_encoded(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
key: &str,
|
||||
value: Value,
|
||||
ttl: Option<Duration>,
|
||||
encoded: Result<String, Error>,
|
||||
) -> Result<Pending<()>, Error> {
|
||||
let ttl = ttl.or(self.settings().memory_ttl);
|
||||
self.set_local(key, value, ttl)?;
|
||||
if key.starts_with("sk-") {
|
||||
return Ok(ready(()));
|
||||
}
|
||||
let Some(remote) = self.redis.get() else {
|
||||
return Ok(ready(()));
|
||||
};
|
||||
let encoded = match encoded {
|
||||
Ok(encoded) => encoded,
|
||||
Err(error) => {
|
||||
redis_failure!(&error, "auth cache value did not encode");
|
||||
return Ok(ready(()));
|
||||
}
|
||||
};
|
||||
if let Some(misses) = send.misses(&remote.client) {
|
||||
misses.forget(key);
|
||||
}
|
||||
Ok(pending(send.batch(&remote.client).add(commands::set(
|
||||
key,
|
||||
encoded,
|
||||
ttl.unwrap_or(remote.default_ttl),
|
||||
))))
|
||||
}
|
||||
|
||||
/// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks.
|
||||
pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result<Pending<()>, Error> {
|
||||
self.partition(key).declare_delete(send, key)
|
||||
self.delete_local(key)?;
|
||||
let Some(client) = self.redis.client() else {
|
||||
return Ok(ready(()));
|
||||
};
|
||||
if let Some(misses) = send.misses(client) {
|
||||
misses.note([key.to_owned()]);
|
||||
}
|
||||
let deleted = send.batch(client).add(commands::del(&[key]));
|
||||
Ok(Box::pin(
|
||||
async move { Ok(deleted.resolve().await.map(drop)?) },
|
||||
))
|
||||
}
|
||||
|
||||
/// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then
|
||||
|
|
@ -142,10 +249,30 @@ impl AuthCache {
|
|||
entries: Vec<(String, Value)>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let (key_objects, others): (Vec<_>, Vec<_>) =
|
||||
let (keys, objects): (Vec<_>, Vec<_>) =
|
||||
entries.into_iter().partition(|(key, _)| is_key_object(key));
|
||||
self.key_objects.set_pipeline(key_objects, ttl).await?;
|
||||
self.objects.set_pipeline(others, ttl).await
|
||||
self.set_partition(keys, ttl).await?;
|
||||
self.set_partition(objects, ttl).await
|
||||
}
|
||||
|
||||
async fn set_partition(
|
||||
&self,
|
||||
entries: Vec<(String, Value)>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let ttl = ttl.or(self.settings().memory_ttl);
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| (key, value, ttl))
|
||||
.collect();
|
||||
let result = match self.write_back(Send::Now, entries) {
|
||||
Ok(written) => written.await,
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
if let Err(error) = result {
|
||||
redis_failure!(&error, "auth cache pipeline write failed");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The auth prefetch write-back: memory first, an entry without a TTL taking
|
||||
|
|
@ -158,6 +285,10 @@ impl AuthCache {
|
|||
send: Send<'_>,
|
||||
entries: Vec<(String, Value, Option<Duration>)>,
|
||||
) -> Result<Pending<()>, Error> {
|
||||
let entries: Vec<_> = entries
|
||||
.into_iter()
|
||||
.filter(|(key, _, _)| !key.starts_with("sk-"))
|
||||
.collect();
|
||||
let encoded = self
|
||||
.redis
|
||||
.get()
|
||||
|
|
@ -199,13 +330,30 @@ impl AuthCache {
|
|||
/// `UserApiKeyCache.async_delete_cache_keys`: the key-object partition and the rest as two
|
||||
/// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails.
|
||||
pub async fn delete(&self, keys: &[String]) -> Result<(), Error> {
|
||||
let (key_objects, others): (Vec<String>, Vec<String>) =
|
||||
for key in keys {
|
||||
self.delete_local(key)?;
|
||||
}
|
||||
let (keys, objects): (Vec<_>, Vec<_>) =
|
||||
keys.iter().cloned().partition(|key| is_key_object(key));
|
||||
let (key_objects, others) = tokio::join!(
|
||||
self.key_objects.delete_keys(&key_objects),
|
||||
self.objects.delete_keys(&others)
|
||||
let (keys, objects) = tokio::join!(
|
||||
self.delete_partition(&keys),
|
||||
self.delete_partition(&objects)
|
||||
);
|
||||
key_objects.and(others)
|
||||
keys.and(objects)
|
||||
}
|
||||
|
||||
async fn delete_partition(&self, keys: &[String]) -> Result<(), Error> {
|
||||
let Some(client) = self.redis.client() else {
|
||||
return Ok(());
|
||||
};
|
||||
for chunk in keys.chunks(1000) {
|
||||
client
|
||||
.send(commands::del(
|
||||
&chunk.iter().map(String::as_str).collect::<Vec<_>>(),
|
||||
))
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on
|
||||
|
|
@ -247,12 +395,67 @@ impl AuthCache {
|
|||
.send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount))
|
||||
.await?)
|
||||
}
|
||||
}
|
||||
|
||||
fn partition(&self, key: &str) -> &DualCache<Value> {
|
||||
if is_key_object(key) {
|
||||
&self.key_objects
|
||||
} else {
|
||||
&self.objects
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
};
|
||||
|
||||
const KEY: &str = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b";
|
||||
|
||||
#[rstest]
|
||||
fn invalidation_fences_an_in_flight_fill_and_allows_a_new_one() {
|
||||
let cache = AuthCache::new(None, Settings::default());
|
||||
let before = cache.current_generation();
|
||||
cache.delete_local(KEY).unwrap();
|
||||
assert!(!cache.set_if_generation(KEY, Value::from("revoked"), before, None));
|
||||
assert_eq!(cache.get_local(KEY).unwrap(), None);
|
||||
let after = cache.current_generation();
|
||||
assert!(cache.set_if_generation(KEY, Value::from("fresh"), after, None));
|
||||
assert_eq!(cache.get_local(KEY).unwrap(), Some(Value::from("fresh")));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn disconnect_bypasses_writes_and_reconnect_fences_old_fills() {
|
||||
let cache = AuthCache::new(None, Settings::default());
|
||||
cache.set_local(KEY, Value::Bool(true), None).unwrap();
|
||||
cache.subscriber_connected(false);
|
||||
assert_eq!(cache.get_local(KEY).unwrap(), None);
|
||||
cache.set_local(KEY, Value::Bool(true), None).unwrap();
|
||||
let disconnected = cache.current_generation();
|
||||
assert!(!cache.set_if_generation(KEY, Value::Bool(true), disconnected, None));
|
||||
cache.subscriber_connected(true);
|
||||
assert_eq!(cache.get_local(KEY).unwrap(), None);
|
||||
assert!(!cache.set_if_generation(KEY, Value::Bool(true), disconnected, None));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn auth_entries_expire_even_without_a_configured_ttl() {
|
||||
let seconds = Arc::new(AtomicU64::new(100));
|
||||
let clock = seconds.clone();
|
||||
let cache = AuthCache::with_clock(
|
||||
None,
|
||||
Settings::default(),
|
||||
Arc::new(move || Duration::from_secs(clock.load(Ordering::SeqCst))),
|
||||
);
|
||||
cache.set_local(KEY, Value::Bool(true), None).unwrap();
|
||||
assert_eq!(cache.get_local(KEY).unwrap(), Some(Value::Bool(true)));
|
||||
seconds.store(160, Ordering::SeqCst);
|
||||
assert_eq!(cache.get_local(KEY).unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn raw_credentials_are_not_cached() {
|
||||
let cache = AuthCache::new(None, Settings::default());
|
||||
cache
|
||||
.set_local("sk-do-not-cache", Value::Bool(true), None)
|
||||
.unwrap();
|
||||
assert_eq!(cache.get_local("sk-do-not-cache").unwrap(), None);
|
||||
assert_eq!(cache.len_local().unwrap(), (0, 0));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use std::time::Duration;
|
|||
|
||||
use litellm_python_compat::Value;
|
||||
|
||||
use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings};
|
||||
use crate::{DecodeError, codec::python_value};
|
||||
|
||||
/// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
|
|
@ -21,15 +21,6 @@ impl Default for Settings {
|
|||
}
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
pub(super) fn memory(self) -> MemorySettings {
|
||||
MemorySettings {
|
||||
entries: Some(self.memory_entries),
|
||||
ttl: self.memory_ttl,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Value>, DecodeError> {
|
||||
raw.as_deref().map(python_value).transpose()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -88,6 +88,7 @@ impl<V: Clone, R> DualCache<V, R> {
|
|||
|
||||
/// `DualCache.update_cache_ttl` and `update_in_memory_max_size` together: both apply in place,
|
||||
/// and the entries memory already holds stay.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn reconfigure(&self, memory: MemorySettings) {
|
||||
*lock(&self.default_in_memory_ttl) = memory.ttl;
|
||||
self.memory.update_max_size(memory.entries);
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_client::{Declared, RedisBatch, commands};
|
||||
use litellm_redis_client::commands;
|
||||
|
||||
use super::DualCache;
|
||||
use crate::{
|
||||
|
|
@ -35,69 +35,6 @@ impl<V: Clone + IntoValue, R> DualCache<V, R> {
|
|||
.await
|
||||
}
|
||||
|
||||
/// `DualCache._set_on_batch`: memory now, the `SET` on the batch `send` picks. Memory and Redis
|
||||
/// both take `ttl`, else `default_in_memory_ttl`, and Redis falls back to its own default.
|
||||
pub(crate) fn declare_set(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
key: &str,
|
||||
value: V,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Pending<()>, Error> {
|
||||
self.declare_encoded(send, key, value, ttl, codec::dumps)
|
||||
}
|
||||
|
||||
/// `DualCache.async_set_cache_pipeline`: memory, then one pipeline of `SET`s. A value
|
||||
/// `json.dumps` refuses aborts the whole pipeline, and any failure only logs, as `RedisCache`
|
||||
/// does.
|
||||
pub(crate) async fn set_pipeline(
|
||||
&self,
|
||||
entries: Vec<(String, V)>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let ttl = ttl.or(self.default_in_memory_ttl());
|
||||
let encoded = self.redis.get().map(|_| {
|
||||
entries
|
||||
.iter()
|
||||
.map(|(_, value)| codec::dumps(value.clone()))
|
||||
.collect::<Result<Vec<_>, Error>>()
|
||||
});
|
||||
let keys = entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
self.memory.set(&key, value, ttl)?;
|
||||
Ok(key)
|
||||
})
|
||||
.collect::<Result<Vec<_>, Error>>()?;
|
||||
let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else {
|
||||
return Ok(());
|
||||
};
|
||||
let failure = match encoded {
|
||||
Err(error) => Some(error),
|
||||
Ok(values) => {
|
||||
let pipeline = RedisBatch::pipelined(remote.client.clone());
|
||||
let ttl = ttl.unwrap_or(remote.default_ttl);
|
||||
let written = keys
|
||||
.iter()
|
||||
.zip(values)
|
||||
.map(|(key, value)| pipeline.add(commands::set(key, value, ttl)))
|
||||
.collect::<Vec<_>>();
|
||||
Declared::all(written)
|
||||
.await
|
||||
.into_iter()
|
||||
.find_map(Result::err)
|
||||
.map(Error::from)
|
||||
}
|
||||
};
|
||||
if let Some(error) = failure {
|
||||
redis_failure!(
|
||||
&error,
|
||||
"LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS",
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn set_now(
|
||||
&self,
|
||||
key: &str,
|
||||
|
|
@ -147,22 +84,6 @@ impl<V: Clone + IntoValue, R> DualCache<V, R> {
|
|||
}
|
||||
|
||||
impl<V: Clone, R> DualCache<V, R> {
|
||||
/// `DualCache.async_delete_cache_pre_call`: memory now, the `DEL` on the batch `send` picks,
|
||||
/// after which the request reads the key as missing.
|
||||
pub(crate) fn declare_delete(&self, send: Send<'_>, key: &str) -> Result<Pending<()>, Error> {
|
||||
self.memory.delete(key)?;
|
||||
let Some(client) = self.redis.client() else {
|
||||
return Ok(ready(()));
|
||||
};
|
||||
if let Some(misses) = send.misses(client) {
|
||||
misses.note([key.to_owned()]);
|
||||
}
|
||||
let deleted = send.batch(client).add(commands::del(&[key]));
|
||||
Ok(Box::pin(
|
||||
async move { Ok(deleted.resolve().await.map(drop)?) },
|
||||
))
|
||||
}
|
||||
|
||||
/// `DualCache.async_delete_cache_keys`: memory, then one multi-key `DEL` per chunk, in order.
|
||||
pub(crate) async fn delete_keys(&self, keys: &[String]) -> Result<(), Error> {
|
||||
keys.iter().try_for_each(|key| self.memory.delete(key))?;
|
||||
|
|
|
|||
|
|
@ -50,10 +50,6 @@ impl<V: Clone> Memory<V> {
|
|||
self.cache.flush_cache()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> Result<usize, Error> {
|
||||
self.cache.len()
|
||||
}
|
||||
|
||||
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
|
||||
self.cache.expires_at(key)
|
||||
}
|
||||
|
|
@ -63,6 +59,7 @@ impl<V: Clone> Memory<V> {
|
|||
}
|
||||
|
||||
/// `DualCache.update_in_memory_max_size`: `None` restores the 200-entry default.
|
||||
#[cfg(test)]
|
||||
pub fn update_max_size(&self, entries: Option<usize>) {
|
||||
self.cache.update_max_size_in_memory(entries);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import asdict, dataclass
|
||||
from types import MappingProxyType
|
||||
|
|
@ -18,6 +19,8 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # leg
|
|||
pubsub_capable_client,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
|
||||
from litellm.redis.stores.connection import connect
|
||||
from litellm.rust_bridge._native import NativeAuthSubscription
|
||||
|
||||
|
||||
class SpendCounterCopies(Protocol):
|
||||
|
|
@ -65,7 +68,7 @@ def _cache_invalidation_message_json(cache_key: str, new_value: float | None = N
|
|||
|
||||
|
||||
def _finite_number_or_none(value: object) -> float | None:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
|
||||
return None
|
||||
return float(value)
|
||||
|
||||
|
|
@ -239,7 +242,14 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
|
|||
|
||||
|
||||
class AuthCacheInvalidationSubscriber:
|
||||
__slots__ = ("_additional_in_memory_caches", "_redis_cache", "_spend_counters", "_task", "_user_api_key_cache")
|
||||
__slots__ = (
|
||||
"_additional_in_memory_caches",
|
||||
"_native",
|
||||
"_redis_cache",
|
||||
"_spend_counters",
|
||||
"_task",
|
||||
"_user_api_key_cache",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -253,12 +263,31 @@ class AuthCacheInvalidationSubscriber:
|
|||
self._additional_in_memory_caches = tuple(additional_in_memory_caches)
|
||||
self._spend_counters: Final = spend_counters
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
self._native: NativeAuthSubscription | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
if self._task is not None:
|
||||
return
|
||||
shared: Final = self._user_api_key_cache.shared
|
||||
client: Final = connect(self._redis_cache).client
|
||||
if shared is not None and client is not None:
|
||||
self._native = shared.subscribe(client, auth_cache_invalidation_channel(self._redis_cache))
|
||||
self._task = asyncio.create_task(self._consume_native(self._native))
|
||||
return
|
||||
self._task = asyncio.create_task(self._run())
|
||||
|
||||
async def _consume_native(self, subscription: NativeAuthSubscription) -> None:
|
||||
try:
|
||||
while True:
|
||||
if (data := await subscription.next()) is not None:
|
||||
self._apply_message({"data": data})
|
||||
finally:
|
||||
await subscription.close()
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
return self._native is not None and self._native.is_connected()
|
||||
|
||||
async def stop(self) -> None:
|
||||
task: Final = self._task
|
||||
if task is None:
|
||||
|
|
@ -269,6 +298,7 @@ class AuthCacheInvalidationSubscriber:
|
|||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._native = None
|
||||
|
||||
async def _run(self) -> None:
|
||||
backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects
|
||||
|
|
|
|||
|
|
@ -68,9 +68,8 @@ class UserApiKeyCache(DualCache):
|
|||
User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition,
|
||||
``key_object_cache``, so churn in the other management objects cannot evict them.
|
||||
|
||||
Without Redis both partitions are this class's Python memory. ``attach_redis_cache`` moves the whole
|
||||
keyspace, both memory partitions included, onto the native ``AuthCacheStore`` (``shared``); redis-py never
|
||||
reads or writes it, and ``redis_cache`` stays None.
|
||||
Both partitions live in native Moka, including without Redis. ``attach_redis_cache`` adds the shared
|
||||
Redis tier without creating another local authority; ``redis_cache`` stays None.
|
||||
|
||||
``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous
|
||||
(no other methods in between) so mypy resolves ``@overload`` + implementation correctly.
|
||||
|
|
@ -103,7 +102,7 @@ class UserApiKeyCache(DualCache):
|
|||
)
|
||||
self._auth_cache: Final = auth_cache
|
||||
self._redis_connection: RedisCache | None = None
|
||||
self._shared: AuthCacheStore | None = None
|
||||
self._shared: AuthCacheStore | None = AuthCacheStore(connect(None).auth_cache(self._settings()))
|
||||
self.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl)
|
||||
|
||||
@property
|
||||
|
|
@ -113,7 +112,7 @@ class UserApiKeyCache(DualCache):
|
|||
|
||||
@property
|
||||
def has_redis(self) -> bool:
|
||||
return self._shared is not None
|
||||
return self._redis_connection is not None
|
||||
|
||||
@property
|
||||
def redis_connection(self) -> RedisCache | None:
|
||||
|
|
@ -150,12 +149,29 @@ class UserApiKeyCache(DualCache):
|
|||
def attach_redis_cache(
|
||||
self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None
|
||||
) -> None:
|
||||
if redis_cache is None or self._shared is not None:
|
||||
if redis_cache is None or self._redis_connection is not None:
|
||||
return
|
||||
if default_redis_ttl is not None:
|
||||
self.update_cache_ttl(default_in_memory_ttl=None, default_redis_ttl=default_redis_ttl)
|
||||
self._redis_connection = redis_cache
|
||||
self._shared = self._auth_cache(redis_cache, self._settings())
|
||||
if self._auth_cache is _connected_auth_cache and self._shared is not None:
|
||||
client: Final = connect(redis_cache).client
|
||||
if client is not None:
|
||||
self._shared.attach(client)
|
||||
else:
|
||||
self._shared = self._auth_cache(redis_cache, self._settings())
|
||||
|
||||
def begin_fill(self, key: str) -> int:
|
||||
assert self._shared is not None
|
||||
return self._shared.begin_fill(key)
|
||||
|
||||
def current_generation(self) -> int:
|
||||
assert self._shared is not None
|
||||
return self._shared.current_generation()
|
||||
|
||||
def set_if_generation(self, key: str, value: CacheValue, generation: int, ttl: float | None = None) -> bool:
|
||||
assert self._shared is not None
|
||||
return self._shared.set_if_generation(key, CacheCodec.serialize(value, model_type=None), generation, ttl)
|
||||
|
||||
def get_local(self, key: str) -> object:
|
||||
"""This worker's memory copy alone."""
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from litellm.caching.redis_batch import Batch, active_request_redis_batches
|
|||
from litellm.redis.stores._route import route
|
||||
from litellm.redis.stores.auth_cache.handle import AuthCacheHandle
|
||||
from litellm.redis.stores.auth_cache.types import AUTH_CACHE_TARGET, AuthCacheSettings, WriteBack
|
||||
from litellm.rust_bridge._native import LocalSizes
|
||||
from litellm.rust_bridge._native import LocalSizes, NativeAuthSubscription, NativeRedisClient
|
||||
|
||||
__all__ = [
|
||||
"AUTH_CACHE_TARGET",
|
||||
|
|
@ -36,6 +36,22 @@ class AuthCacheStore:
|
|||
def __init__(self, handle: AuthCacheHandle) -> None:
|
||||
self._handle: Final = handle
|
||||
|
||||
def attach(self, client: NativeRedisClient) -> bool:
|
||||
return self._handle.attach(client)
|
||||
|
||||
def begin_fill(self, key: str) -> int:
|
||||
return self._handle.current_generation()
|
||||
|
||||
def current_generation(self) -> int:
|
||||
return self._handle.current_generation()
|
||||
|
||||
def set_if_generation(self, key: str, value: object, generation: int, ttl: float | None = None) -> bool:
|
||||
self._forget_miss(key)
|
||||
return self._handle.set_if_generation(key, value, generation, _redis_ttl(ttl))
|
||||
|
||||
def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription:
|
||||
return self._handle.subscribe(client, channel)
|
||||
|
||||
def reconfigure(self, settings: AuthCacheSettings) -> None:
|
||||
"""Both memory partitions keep what they hold, as ``update_cache_ttl`` and ``update_in_memory_max_size`` do."""
|
||||
self._handle.reconfigure(memory_entries=settings.memory_entries, memory_ttl_seconds=settings.memory_ttl)
|
||||
|
|
|
|||
|
|
@ -5,10 +5,14 @@ from typing import Protocol
|
|||
|
||||
from litellm.caching.redis_batch import NativePending
|
||||
from litellm.redis.stores.auth_cache.types import WriteBack
|
||||
from litellm.rust_bridge._native import LocalSizes, NativeSend
|
||||
from litellm.rust_bridge._native import LocalSizes, NativeAuthSubscription, NativeRedisClient, NativeSend
|
||||
|
||||
|
||||
class AuthCacheHandle(Protocol):
|
||||
def attach(self, client: NativeRedisClient) -> bool: ...
|
||||
def current_generation(self) -> int: ...
|
||||
def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: ...
|
||||
def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: ...
|
||||
def get(self, key: str) -> Awaitable[object]: ...
|
||||
def get_sync(self, key: str) -> object: ...
|
||||
def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Awaitable[None]: ...
|
||||
|
|
|
|||
|
|
@ -395,6 +395,10 @@ class NativeAuthCache:
|
|||
) -> NativeAuthCache: ...
|
||||
@property
|
||||
def has_redis(self) -> bool: ...
|
||||
def attach(self, client: NativeRedisClient) -> bool: ...
|
||||
def current_generation(self) -> int: ...
|
||||
def set_if_generation(self, key: str, value: object, generation: int, ttl_seconds: float | None = None) -> bool: ...
|
||||
def subscribe(self, client: NativeRedisClient, channel: str) -> NativeAuthSubscription: ...
|
||||
def get(self, key: str, request: NativeRequestRedis | None = None) -> Future[object]: ...
|
||||
def reconfigure(self, *, memory_entries: int, memory_ttl_seconds: float | None = None) -> None: ...
|
||||
def get_sync(self, key: str) -> object: ...
|
||||
|
|
@ -416,6 +420,12 @@ class NativeAuthCache:
|
|||
def flush_local(self) -> None: ...
|
||||
def len_local(self) -> LocalSizes: ...
|
||||
|
||||
@final
|
||||
class NativeAuthSubscription:
|
||||
def is_connected(self) -> bool: ...
|
||||
def next(self) -> Future[bytes | None]: ...
|
||||
def close(self) -> Future[None]: ...
|
||||
|
||||
@final
|
||||
class SpendRead:
|
||||
def __new__(cls, value: float | None, from_redis: bool) -> SpendRead: ...
|
||||
|
|
|
|||
125
tests/integration/caching/test_auth_cache_native.py
Normal file
125
tests/integration/caching/test_auth_cache_native.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
AuthCacheInvalidationSubscriber,
|
||||
auth_cache_invalidation_channel,
|
||||
evict_local,
|
||||
evict_shared,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.redis.stores.connection import connect
|
||||
from litellm.redis.stores.spend_counters import SpendCounterStore
|
||||
from tests.integration._support.redis_process import owned_redis
|
||||
|
||||
KEY: Final = hashlib.sha256(b"native-cache-test-key").hexdigest()
|
||||
|
||||
|
||||
async def until(predicate: Callable[[], bool], seconds: float = 12) -> None:
|
||||
async with asyncio.timeout(seconds):
|
||||
while not predicate():
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("namespace", (None, "phase3"))
|
||||
@pytest.mark.parametrize("shared", (False, True))
|
||||
async def test_native_subscriber_invalidates_and_recovers_without_serving_outage_fills(
|
||||
tmp_path: Path, namespace: str | None, shared: bool
|
||||
) -> None:
|
||||
with owned_redis(tmp_path) as server:
|
||||
redis_cache: Final = RedisCache(
|
||||
host=server.host,
|
||||
port=server.port,
|
||||
namespace=namespace,
|
||||
socket_timeout=0.3,
|
||||
socket_connect_timeout=0.3,
|
||||
)
|
||||
cache: Final = UserApiKeyCache(redis_cache=redis_cache if shared else None)
|
||||
additional: Final = InMemoryCache()
|
||||
spend: Final = SpendCounterStore(connect(None).spend_counters())
|
||||
subscriber: Final = AuthCacheInvalidationSubscriber(redis_cache, cache, (additional,), spend)
|
||||
publisher: Final = Redis(host=server.host, port=server.port)
|
||||
channel: Final = auth_cache_invalidation_channel(redis_cache)
|
||||
subscriber.start()
|
||||
started: Final = cache.current_generation()
|
||||
try:
|
||||
assert not subscriber.is_ready
|
||||
await until(lambda: subscriber.is_ready)
|
||||
assert cache.current_generation() > started
|
||||
fill: Final = cache.begin_fill(KEY)
|
||||
assert cache.set_if_generation(KEY, {"token": KEY, "models": ["allowed"]}, fill)
|
||||
assert cache.get_local(KEY) == {"token": KEY, "models": ["allowed"]}
|
||||
spend.set_local("spend:key:k-1", 99.0)
|
||||
assert (
|
||||
await publisher.publish(
|
||||
channel, json.dumps({"cache_key": "spend:key:k-1", "new_value": 0.0, "ttl": 60})
|
||||
)
|
||||
== 1
|
||||
)
|
||||
await until(lambda: spend.get_local("spend:key:k-1") == 0.0)
|
||||
assert cache.get_local(KEY) == {"token": KEY, "models": ["allowed"]}
|
||||
assert (
|
||||
await publisher.publish(
|
||||
channel, json.dumps({"cache_key": "spend_db_floor:spend:key:k-1", "new_value": 3.0})
|
||||
)
|
||||
== 1
|
||||
)
|
||||
await until(lambda: additional.get_cache("spend_db_floor:spend:key:k-1") == 3.0)
|
||||
assert await publisher.publish(channel, json.dumps({"cache_key": "spend:key:k-1"})) == 1
|
||||
await until(lambda: spend.get_local("spend:key:k-1") is None)
|
||||
additional.set_cache("model-cache", "stale")
|
||||
assert await publisher.publish(channel, json.dumps({"cache_key": KEY})) == 1
|
||||
await until(lambda: cache.get_local(KEY) is None)
|
||||
assert not cache.set_if_generation(KEY, {"token": KEY}, fill)
|
||||
assert await publisher.publish(channel, json.dumps({"cache_key": "model-cache"})) == 1
|
||||
await until(lambda: additional.get_cache("model-cache") is None)
|
||||
|
||||
cache.set_local(KEY, {"token": KEY}, None)
|
||||
before_stop: Final = cache.current_generation()
|
||||
server.stop()
|
||||
await until(lambda: cache.current_generation() > before_stop)
|
||||
assert not subscriber.is_ready
|
||||
assert cache.get_local(KEY) is None
|
||||
outage: Final = cache.begin_fill(KEY)
|
||||
assert not cache.set_if_generation(KEY, {"token": KEY}, outage)
|
||||
cache.set_local(KEY, {"token": KEY}, None)
|
||||
assert cache.get_local(KEY) is None
|
||||
server.start()
|
||||
await until(lambda: subscriber.is_ready)
|
||||
assert cache.current_generation() > outage
|
||||
assert not cache.set_if_generation(KEY, {"token": KEY}, outage)
|
||||
assert cache.set_if_generation(KEY, {"token": KEY}, cache.begin_fill(KEY))
|
||||
assert await publisher.publish(channel, json.dumps({"cache_key": KEY})) == 1
|
||||
await until(lambda: cache.get_local(KEY) is None)
|
||||
finally:
|
||||
await subscriber.stop()
|
||||
await publisher.aclose()
|
||||
assert cache.get_local(KEY) is None
|
||||
assert not cache.set_if_generation(KEY, {"token": KEY}, cache.begin_fill(KEY))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_management_eviction_reaches_native_memory_and_shared_redis(tmp_path: Path) -> None:
|
||||
with owned_redis(tmp_path) as server:
|
||||
redis_cache: Final = RedisCache(host=server.host, port=server.port, namespace="management")
|
||||
first: Final = UserApiKeyCache(redis_cache=redis_cache)
|
||||
second: Final = UserApiKeyCache(redis_cache=redis_cache)
|
||||
await first.async_set_cache(KEY, {"token": KEY})
|
||||
assert await second.async_get_cache(KEY) == {"token": KEY}
|
||||
evict_local((KEY,), first)
|
||||
assert first.get_local(KEY) is None
|
||||
await evict_shared((KEY,), first)
|
||||
second.delete_local(KEY)
|
||||
assert await second.async_get_cache(KEY) is None
|
||||
|
|
@ -145,6 +145,8 @@ class _ScriptedPubSubRedisClient(Redis):
|
|||
|
||||
|
||||
class _FakeRedisCache:
|
||||
store_settings = None
|
||||
|
||||
def __init__(self, client: object, namespace: str | None = None) -> None:
|
||||
self._client = client
|
||||
self.namespace = namespace
|
||||
|
|
@ -207,8 +209,8 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
|
|||
instead of serving the stale object until the TTL expires.
|
||||
"""
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("project_id:p-1", {"models": []})
|
||||
assert cache.in_memory_cache.get_cache("project_id:p-1") is not None
|
||||
cache.set_local("project_id:p-1", {"models": []}, None)
|
||||
assert cache.get_local("project_id:p-1") is not None
|
||||
|
||||
pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")])
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
|
|
@ -221,7 +223,7 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
|
|||
finally:
|
||||
await subscriber.stop()
|
||||
|
||||
assert cache.in_memory_cache.get_cache("project_id:p-1") is None
|
||||
assert cache.get_local("project_id:p-1") is None
|
||||
assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL]
|
||||
|
||||
|
||||
|
|
@ -357,7 +359,7 @@ async def test_subscriber_writes_and_drops_the_native_spend_counter_copies_a_bro
|
|||
@pytest.mark.asyncio
|
||||
async def test_subscriber_ignores_malformed_messages() -> None:
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("project_id:p-1", {"models": []})
|
||||
cache.set_local("project_id:p-1", {"models": []}, None)
|
||||
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])),
|
||||
|
|
@ -368,7 +370,7 @@ async def test_subscriber_ignores_malformed_messages() -> None:
|
|||
subscriber._apply_message("raw string")
|
||||
subscriber._apply_message(None)
|
||||
|
||||
assert cache.in_memory_cache.get_cache("project_id:p-1") is not None
|
||||
assert cache.get_local("project_id:p-1") is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import datetime
|
||||
import hashlib
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -17,7 +17,6 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
get_management_object_ttl,
|
||||
is_user_key_cache_key,
|
||||
)
|
||||
from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum
|
||||
from litellm.redis.stores.auth_cache import AuthCacheStore
|
||||
from litellm.rust_bridge._native import NativeAuthCache
|
||||
from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache
|
||||
|
|
@ -25,18 +24,6 @@ from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cach
|
|||
HASHED_TOKEN = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest()
|
||||
|
||||
|
||||
class CapturingInMemoryCache(InMemoryCache):
|
||||
"""Records ``ttl`` passed into ``set_cache`` (what DualCache injects)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.last_ttl: Any = None
|
||||
|
||||
def set_cache(self, key, value, **kwargs): # type: ignore[override]
|
||||
self.last_ttl = kwargs.get("ttl")
|
||||
super().set_cache(key, value, **kwargs)
|
||||
|
||||
|
||||
class _ManualClock:
|
||||
def __init__(self) -> None:
|
||||
self.seconds = 1000.0
|
||||
|
|
@ -54,39 +41,55 @@ def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth:
|
|||
|
||||
|
||||
class TestUserApiKeyCache:
|
||||
def test_native_fill_fence_and_request_local_projection_without_redis(self):
|
||||
cache: Final = UserApiKeyCache()
|
||||
stale: Final = cache.begin_fill(HASHED_TOKEN)
|
||||
cache.delete_local(HASHED_TOKEN)
|
||||
assert not cache.set_if_generation(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), stale)
|
||||
assert cache.set_if_generation(
|
||||
HASHED_TOKEN, UserAPIKeyAuth(token=HASHED_TOKEN, models=["original"]), cache.begin_fill(HASHED_TOKEN)
|
||||
)
|
||||
first: Final = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth)
|
||||
assert first is not None
|
||||
first.models.append("request-only")
|
||||
second: Final = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth)
|
||||
assert second is not None and second is not first
|
||||
assert second.models == ["original"]
|
||||
cache.flush_cache()
|
||||
assert cache.get_local(HASHED_TOKEN) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted(
|
||||
async def test_async_set_with_zero_default_ttl_expires_without_a_python_copy(
|
||||
self,
|
||||
):
|
||||
"""
|
||||
If ``general_settings.user_api_key_cache_ttl`` is absent, the proxy never
|
||||
calls ``update_cache_ttl``; ``user_api_key_cache`` keeps
|
||||
``default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl``.
|
||||
DualCache must forward that as the in-memory ``ttl`` kwarg on each set.
|
||||
"""
|
||||
mem = CapturingInMemoryCache()
|
||||
mem = InMemoryCache()
|
||||
cache = UserApiKeyCache(
|
||||
in_memory_cache=mem,
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
|
||||
default_in_memory_ttl=0,
|
||||
)
|
||||
cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=None)
|
||||
await cache.async_set_cache(
|
||||
"k",
|
||||
_make_key_obj("t"),
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
expected = UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
assert mem.last_ttl == expected
|
||||
assert cache.get_local("k") is None
|
||||
assert mem.get_cache("k") is None
|
||||
cache.set_cache("k", {"token": "t"}, ttl=60)
|
||||
assert cache.get_local("k") == {"token": "t"}
|
||||
|
||||
def test_sync_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted(
|
||||
def test_sync_set_with_zero_default_ttl_expires_without_a_python_copy(
|
||||
self,
|
||||
):
|
||||
mem = CapturingInMemoryCache()
|
||||
mem = InMemoryCache()
|
||||
cache = UserApiKeyCache(
|
||||
in_memory_cache=mem,
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
|
||||
default_in_memory_ttl=0,
|
||||
)
|
||||
cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=None)
|
||||
cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth)
|
||||
assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
assert cache.get_local("sk") is None
|
||||
assert mem.get_cache("sk") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_shared_store_gets_default_in_memory_ttl_and_writes_without_a_ttl(self):
|
||||
|
|
@ -173,7 +176,9 @@ class TestUserApiKeyCache:
|
|||
async def test_a_model_holding_a_datetime_is_cached_as_its_json_dump(self):
|
||||
cache = UserApiKeyCache(
|
||||
redis_cache=MagicMock(spec=RedisCache),
|
||||
auth_cache=lambda _redis, settings: AuthCacheStore(NativeAuthCache(None, memory_entries=settings.memory_entries)),
|
||||
auth_cache=lambda _redis, settings: AuthCacheStore(
|
||||
NativeAuthCache(None, memory_entries=settings.memory_entries)
|
||||
),
|
||||
)
|
||||
key_obj = UserAPIKeyAuth(token="abc", expires=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc))
|
||||
|
||||
|
|
@ -273,13 +278,14 @@ class TestUserKeyObjectPartition:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_ttl_applies_to_key_object_partition(self):
|
||||
key_partition = CapturingInMemoryCache()
|
||||
key_partition = InMemoryCache()
|
||||
cache = UserApiKeyCache(default_in_memory_ttl=60, key_object_in_memory_cache=key_partition)
|
||||
cache.update_cache_ttl(default_in_memory_ttl=7, default_redis_ttl=7)
|
||||
cache.update_cache_ttl(default_in_memory_ttl=0, default_redis_ttl=7)
|
||||
|
||||
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
|
||||
|
||||
assert key_partition.last_ttl == 7
|
||||
assert cache.get_local(HASHED_TOKEN) is None
|
||||
assert key_partition.get_cache(HASHED_TOKEN) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attach_redis_cache_moves_both_partitions_onto_the_shared_store(self):
|
||||
|
|
@ -372,8 +378,10 @@ class TestUserKeyObjectPartition:
|
|||
cache.set_local(HASHED_TOKEN, {"token": HASHED_TOKEN}, None)
|
||||
cache.set_local(end_user_cache_key("u1"), {"user_id": "u1"}, 5)
|
||||
|
||||
assert cache.key_object_cache.in_memory_cache.get_cache(HASHED_TOKEN) == {"token": HASHED_TOKEN}
|
||||
assert cache.in_memory_cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"}
|
||||
assert cache.get_local(HASHED_TOKEN) == {"token": HASHED_TOKEN}
|
||||
assert cache.get_local(end_user_cache_key("u1")) == {"user_id": "u1"}
|
||||
assert cache.key_object_cache.in_memory_cache.get_cache(HASHED_TOKEN) is None
|
||||
assert cache.in_memory_cache.get_cache(end_user_cache_key("u1")) is None
|
||||
cache.delete_local(HASHED_TOKEN)
|
||||
assert cache.get_local(HASHED_TOKEN) is None
|
||||
|
||||
|
|
@ -461,11 +469,7 @@ class TestManagementObjectTTL:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_management_write_uses_configured_ttl_over_constant(self):
|
||||
mem = CapturingInMemoryCache()
|
||||
cache = UserApiKeyCache(
|
||||
in_memory_cache=mem,
|
||||
default_in_memory_ttl=300,
|
||||
)
|
||||
cache, fake = shared_cache(default_in_memory_ttl=300)
|
||||
assert get_management_object_ttl(cache) != (DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL)
|
||||
|
||||
await cache.async_set_cache(
|
||||
|
|
@ -475,4 +479,4 @@ class TestManagementObjectTTL:
|
|||
ttl=get_management_object_ttl(cache),
|
||||
)
|
||||
|
||||
assert mem.last_ttl == 300
|
||||
assert fake.commands == [("SET", "team_id:abc", 300)]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue