tenferro_runtime/
extension_cache.rs

1//! Generic runtime caches for extension executors.
2//!
3//! Extension payloads describe operation semantics. Runtime plans, vendor
4//! handles, and other mutable execution state belong here instead, behind
5//! explicit bounded cache ownership.
6
7use std::any::Any;
8use std::fmt;
9use std::num::NonZeroUsize;
10
11use lru::LruCache;
12use tenferro_tensor::CacheStats;
13
14/// Default number of type-erased extension cache entries retained per owner.
15pub const DEFAULT_EXTENSION_CACHE_CAPACITY: usize = 256;
16
17/// A stable key for one extension-owned runtime cache entry.
18///
19/// `family_id` names the extension family, `cache_name` names the specific
20/// cache within that family, and `discriminator` is chosen by the extension
21/// executor from shape, dtype, device, or other runtime planning inputs.
22#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
23pub struct ExtensionCacheKey {
24    /// Extension family that owns this cache entry.
25    pub family_id: &'static str,
26    /// Cache namespace within the extension family.
27    pub cache_name: &'static str,
28    /// Extension-defined stable discriminator for this runtime entry.
29    pub discriminator: u64,
30}
31
32impl ExtensionCacheKey {
33    /// Build an extension cache key.
34    ///
35    /// # Examples
36    ///
37    /// ```
38    /// use tenferro_runtime::ExtensionCacheKey;
39    ///
40    /// let key = ExtensionCacheKey::new("example.identity.v1", "plans", 7);
41    /// assert_eq!(key.cache_name, "plans");
42    /// ```
43    pub const fn new(
44        family_id: &'static str,
45        cache_name: &'static str,
46        discriminator: u64,
47    ) -> Self {
48        Self {
49            family_id,
50            cache_name,
51            discriminator,
52        }
53    }
54}
55
56/// Selector used to inspect or clear extension cache entries.
57#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
58pub enum ExtensionCacheSelector {
59    /// Select every extension cache entry.
60    All,
61    /// Select every entry owned by an extension family.
62    Family { family_id: &'static str },
63    /// Select entries in one named cache for one extension family.
64    Cache {
65        family_id: &'static str,
66        cache_name: &'static str,
67    },
68}
69
70impl ExtensionCacheSelector {
71    /// Return `true` when this selector includes `key`.
72    ///
73    /// # Examples
74    ///
75    /// ```
76    /// use tenferro_runtime::{ExtensionCacheKey, ExtensionCacheSelector};
77    ///
78    /// let key = ExtensionCacheKey::new("example.identity.v1", "plans", 0);
79    /// assert!(ExtensionCacheSelector::Family {
80    ///     family_id: "example.identity.v1",
81    /// }.matches(&key));
82    /// ```
83    pub fn matches(&self, key: &ExtensionCacheKey) -> bool {
84        match *self {
85            Self::All => true,
86            Self::Family { family_id } => key.family_id == family_id,
87            Self::Cache {
88                family_id,
89                cache_name,
90            } => key.family_id == family_id && key.cache_name == cache_name,
91        }
92    }
93}
94
95/// Bounded retention limits for extension runtime caches.
96#[derive(Clone, Copy, Debug, PartialEq, Eq)]
97pub struct ExtensionCacheLimits {
98    max_entries: NonZeroUsize,
99}
100
101impl ExtensionCacheLimits {
102    /// Build limits from a maximum entry count.
103    ///
104    /// # Examples
105    ///
106    /// ```
107    /// use std::num::NonZeroUsize;
108    /// use tenferro_runtime::ExtensionCacheLimits;
109    ///
110    /// let limits = ExtensionCacheLimits::new(NonZeroUsize::new(4).unwrap());
111    /// assert_eq!(limits.max_entries().get(), 4);
112    /// ```
113    pub const fn new(max_entries: NonZeroUsize) -> Self {
114        Self { max_entries }
115    }
116
117    /// Maximum entries retained by the store.
118    pub const fn max_entries(self) -> NonZeroUsize {
119        self.max_entries
120    }
121}
122
123impl Default for ExtensionCacheLimits {
124    fn default() -> Self {
125        Self {
126            max_entries: NonZeroUsize::new(DEFAULT_EXTENSION_CACHE_CAPACITY)
127                .unwrap_or(NonZeroUsize::MIN),
128        }
129    }
130}
131
132trait ExtensionCacheValue: Send + Sync {
133    fn as_any(&self) -> &(dyn Any + Send + Sync);
134    fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync);
135    fn retained_bytes(&self) -> usize;
136}
137
138struct FixedRetainedBytes<T> {
139    value: T,
140    retained_bytes: usize,
141}
142
143impl<T> ExtensionCacheValue for FixedRetainedBytes<T>
144where
145    T: Any + Send + Sync + 'static,
146{
147    fn as_any(&self) -> &(dyn Any + Send + Sync) {
148        &self.value
149    }
150
151    fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) {
152        &mut self.value
153    }
154
155    fn retained_bytes(&self) -> usize {
156        self.retained_bytes
157    }
158}
159
160struct DynamicRetainedBytes<T, F> {
161    value: T,
162    retained_bytes: F,
163}
164
165impl<T, F> ExtensionCacheValue for DynamicRetainedBytes<T, F>
166where
167    T: Any + Send + Sync + 'static,
168    F: Fn(&T) -> usize + Send + Sync + 'static,
169{
170    fn as_any(&self) -> &(dyn Any + Send + Sync) {
171        &self.value
172    }
173
174    fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) {
175        &mut self.value
176    }
177
178    fn retained_bytes(&self) -> usize {
179        (self.retained_bytes)(&self.value)
180    }
181}
182
183struct ExtensionCacheEntry {
184    value: Box<dyn ExtensionCacheValue>,
185}
186
187/// Bounded type-erased cache storage owned by an extension executor.
188pub struct ExtensionCacheStore {
189    limits: ExtensionCacheLimits,
190    entries: LruCache<ExtensionCacheKey, ExtensionCacheEntry>,
191}
192
193impl fmt::Debug for ExtensionCacheStore {
194    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
195        f.debug_struct("ExtensionCacheStore")
196            .field("limits", &self.limits)
197            .field("stats", &self.stats(ExtensionCacheSelector::All))
198            .finish_non_exhaustive()
199    }
200}
201
202impl ExtensionCacheStore {
203    /// Create an empty cache store with default limits.
204    ///
205    /// # Examples
206    ///
207    /// ```
208    /// use tenferro_runtime::ExtensionCacheStore;
209    ///
210    /// let store = ExtensionCacheStore::new();
211    /// assert_eq!(store.len(), 0);
212    /// ```
213    pub fn new() -> Self {
214        Self::with_limits(ExtensionCacheLimits::default())
215    }
216
217    /// Create an empty cache store with explicit limits.
218    pub fn with_limits(limits: ExtensionCacheLimits) -> Self {
219        Self {
220            entries: LruCache::new(limits.max_entries()),
221            limits,
222        }
223    }
224
225    /// Return the active cache limits.
226    pub const fn limits(&self) -> ExtensionCacheLimits {
227        self.limits
228    }
229
230    /// Resize the store and evict least-recently-used entries if needed.
231    pub fn set_limits(&mut self, limits: ExtensionCacheLimits) {
232        self.entries.resize(limits.max_entries());
233        self.limits = limits;
234    }
235
236    /// Current entry count.
237    pub fn len(&self) -> usize {
238        self.entries.len()
239    }
240
241    /// Return whether the store contains no entries.
242    pub fn is_empty(&self) -> bool {
243        self.entries.is_empty()
244    }
245
246    /// Insert or replace a typed cache entry.
247    pub fn put<T>(&mut self, key: ExtensionCacheKey, value: T, retained_bytes: usize)
248    where
249        T: Any + Send + Sync + 'static,
250    {
251        self.entries.put(
252            key,
253            ExtensionCacheEntry {
254                value: Box::new(FixedRetainedBytes {
255                    value,
256                    retained_bytes,
257                }),
258            },
259        );
260    }
261
262    /// Insert or replace a typed cache entry whose retained bytes are computed
263    /// from the current value whenever stats are requested.
264    ///
265    /// Use this for entries that mutate after insertion, such as compiled
266    /// execution plans with backend-owned nested caches.
267    ///
268    /// # Examples
269    ///
270    /// ```
271    /// use tenferro_runtime::{ExtensionCacheKey, ExtensionCacheStore, ExtensionCacheSelector};
272    ///
273    /// let mut store = ExtensionCacheStore::new();
274    /// let key = ExtensionCacheKey::new("example.cache.v1", "plans", 0);
275    /// store.put_with_retained_bytes(key, Vec::<usize>::with_capacity(2), |values| {
276    ///     values.capacity().saturating_mul(std::mem::size_of::<usize>())
277    /// });
278    /// let values = store.get_mut::<Vec<usize>>(&key).unwrap();
279    /// values.reserve_exact(4);
280    /// let retained_capacity = values.capacity();
281    ///
282    /// assert_eq!(
283    ///     store.stats(ExtensionCacheSelector::All).retained_bytes,
284    ///     retained_capacity * std::mem::size_of::<usize>()
285    /// );
286    /// ```
287    pub fn put_with_retained_bytes<T, F>(
288        &mut self,
289        key: ExtensionCacheKey,
290        value: T,
291        retained_bytes: F,
292    ) where
293        T: Any + Send + Sync + 'static,
294        F: Fn(&T) -> usize + Send + Sync + 'static,
295    {
296        self.entries.put(
297            key,
298            ExtensionCacheEntry {
299                value: Box::new(DynamicRetainedBytes {
300                    value,
301                    retained_bytes,
302                }),
303            },
304        );
305    }
306
307    /// Get a typed cache entry, updating its LRU position.
308    pub fn get<T>(&mut self, key: &ExtensionCacheKey) -> Option<&T>
309    where
310        T: Any + Send + Sync + 'static,
311    {
312        self.entries
313            .get(key)
314            .and_then(|entry| entry.value.as_any().downcast_ref::<T>())
315    }
316
317    /// Get a mutable typed cache entry, updating its LRU position.
318    pub fn get_mut<T>(&mut self, key: &ExtensionCacheKey) -> Option<&mut T>
319    where
320        T: Any + Send + Sync + 'static,
321    {
322        self.entries
323            .get_mut(key)
324            .and_then(|entry| entry.value.as_any_mut().downcast_mut::<T>())
325    }
326
327    /// Clear entries selected by `selector`.
328    pub fn clear_selected(&mut self, selector: ExtensionCacheSelector) {
329        if selector == ExtensionCacheSelector::All {
330            self.entries.clear();
331            return;
332        }
333
334        let keys: Vec<_> = self
335            .entries
336            .iter()
337            .map(|(key, _)| *key)
338            .filter(|key| selector.matches(key))
339            .collect();
340        for key in keys {
341            self.entries.pop(&key);
342        }
343    }
344
345    /// Clear every extension cache entry.
346    pub fn clear(&mut self) {
347        self.entries.clear();
348    }
349
350    /// Return cache-style stats for entries selected by `selector`.
351    pub fn stats(&self, selector: ExtensionCacheSelector) -> CacheStats {
352        CacheStats {
353            entries: self
354                .entries
355                .iter()
356                .filter(|(key, _)| selector.matches(key))
357                .count(),
358            retained_bytes: self
359                .entries
360                .iter()
361                .filter(|(key, _)| selector.matches(key))
362                .map(|(_, entry)| entry.value.retained_bytes())
363                .fold(0usize, usize::saturating_add),
364        }
365    }
366}
367
368impl Default for ExtensionCacheStore {
369    fn default() -> Self {
370        Self::new()
371    }
372}
373
374#[cfg(test)]
375mod tests;
tenferro_runtime/extension_cache.rs

tenferro_runtime/
extension_cache.rs
