liter_llm/tenant/
in_memory.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use dashmap::DashMap;
6
7use super::resolver::{KeyResolver, KeyResolverError, ResolvedKey};
8
9/// Thread-safe in-memory [`KeyResolver`] backed by a [`DashMap`].
10///
11/// Suitable for tests and single-process deployments where the full key set
12/// fits in memory.  For database- or cache-backed production deployments,
13/// implement [`KeyResolver`] directly.
14///
15/// The map is stored behind an [`Arc`] so that the futures returned by
16/// [`KeyResolver::resolve`] are `'static` (they capture the `Arc` by clone,
17/// not a borrow of `self`).
18pub struct InMemoryKeyResolver {
19    keys: Arc<DashMap<String, ResolvedKey>>,
20}
21
22impl InMemoryKeyResolver {
23    /// Create an empty resolver.
24    #[must_use]
25    pub fn new() -> Self {
26        Self {
27            keys: Arc::new(DashMap::new()),
28        }
29    }
30
31    /// Create a resolver pre-populated with the given entries.
32    #[must_use]
33    pub fn with_entries(entries: impl IntoIterator<Item = (String, ResolvedKey)>) -> Self {
34        let keys = DashMap::new();
35        for (k, v) in entries {
36            keys.insert(k, v);
37        }
38        Self { keys: Arc::new(keys) }
39    }
40
41    /// Insert or replace a key record.
42    pub fn insert(&self, api_key: impl Into<String>, resolved: ResolvedKey) {
43        self.keys.insert(api_key.into(), resolved);
44    }
45
46    /// Remove a key record.  Returns the removed record if it existed.
47    pub fn remove(&self, api_key: &str) -> Option<ResolvedKey> {
48        self.keys.remove(api_key).map(|(_, v)| v)
49    }
50}
51
52impl Default for InMemoryKeyResolver {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58impl KeyResolver for InMemoryKeyResolver {
59    fn resolve(
60        &self,
61        api_key: String,
62    ) -> Pin<Box<dyn Future<Output = Result<ResolvedKey, KeyResolverError>> + Send + 'static>> {
63        let keys = self.keys.clone();
64        Box::pin(async move {
65            match keys.get(&api_key) {
66                None => Err(KeyResolverError::NotFound),
67                Some(entry) => {
68                    if !entry.active {
69                        Err(KeyResolverError::Inactive)
70                    } else {
71                        Ok(entry.clone())
72                    }
73                }
74            }
75        })
76    }
77}
liter_llm/tenant/in_memory.rs

liter_llm/tenant/
in_memory.rs
