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/// Default logical retained-byte bound for extension cache entries.
17pub const DEFAULT_EXTENSION_CACHE_RETAINED_BYTES: usize = 64 * 1024 * 1024;
18
19const DEFAULT_EXTENSION_CACHE_RETAINED_BYTES_NONZERO: NonZeroUsize =
20    match NonZeroUsize::new(DEFAULT_EXTENSION_CACHE_RETAINED_BYTES) {
21        Some(value) => value,
22        None => NonZeroUsize::MIN,
23    };
24
25/// A stable key for one extension-owned runtime cache entry.
26///
27/// `family_id` names the extension family, `cache_name` names the specific
28/// cache within that family, and `discriminator` is chosen by the extension
29/// executor from shape, dtype, device, or other runtime planning inputs.
30#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
31pub struct ExtensionCacheKey {
32    /// Extension family that owns this cache entry.
33    pub family_id: &'static str,
34    /// Cache namespace within the extension family.
35    pub cache_name: &'static str,
36    /// Extension-defined stable discriminator for this runtime entry.
37    pub discriminator: u64,
38}
39
40impl ExtensionCacheKey {
41    /// Build an extension cache key.
42    ///
43    /// # Examples
44    ///
45    /// ```
46    /// use tenferro_runtime::ExtensionCacheKey;
47    ///
48    /// let key = ExtensionCacheKey::new("example.identity.v1", "plans", 7);
49    /// assert_eq!(key.cache_name, "plans");
50    /// ```
51    pub const fn new(
52        family_id: &'static str,
53        cache_name: &'static str,
54        discriminator: u64,
55    ) -> Self {
56        Self {
57            family_id,
58            cache_name,
59            discriminator,
60        }
61    }
62}
63
64/// Selector used to inspect or clear extension cache entries.
65#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
66pub enum ExtensionCacheSelector {
67    /// Select every extension cache entry.
68    All,
69    /// Select every entry owned by an extension family.
70    Family { family_id: &'static str },
71    /// Select entries in one named cache for one extension family.
72    Cache {
73        family_id: &'static str,
74        cache_name: &'static str,
75    },
76}
77
78impl ExtensionCacheSelector {
79    /// Return `true` when this selector includes `key`.
80    ///
81    /// # Examples
82    ///
83    /// ```
84    /// use tenferro_runtime::{ExtensionCacheKey, ExtensionCacheSelector};
85    ///
86    /// let key = ExtensionCacheKey::new("example.identity.v1", "plans", 0);
87    /// assert!(ExtensionCacheSelector::Family {
88    ///     family_id: "example.identity.v1",
89    /// }.matches(&key));
90    /// ```
91    pub fn matches(&self, key: &ExtensionCacheKey) -> bool {
92        match *self {
93            Self::All => true,
94            Self::Family { family_id } => key.family_id == family_id,
95            Self::Cache {
96                family_id,
97                cache_name,
98            } => key.family_id == family_id && key.cache_name == cache_name,
99        }
100    }
101}
102
103/// Bounded retention limits for extension runtime caches.
104#[derive(Clone, Copy, Debug, PartialEq, Eq)]
105pub struct ExtensionCacheLimits {
106    max_entries: NonZeroUsize,
107    max_retained_bytes: Option<NonZeroUsize>,
108}
109
110impl ExtensionCacheLimits {
111    /// Build limits from a maximum entry count and the default byte bound.
112    ///
113    /// # Examples
114    ///
115    /// ```
116    /// use std::num::NonZeroUsize;
117    /// use tenferro_runtime::ExtensionCacheLimits;
118    ///
119    /// let limits = ExtensionCacheLimits::new(NonZeroUsize::new(4).unwrap());
120    /// assert_eq!(limits.max_entries().get(), 4);
121    /// assert!(limits.max_retained_bytes().is_some());
122    /// ```
123    pub const fn new(max_entries: NonZeroUsize) -> Self {
124        Self {
125            max_entries,
126            max_retained_bytes: Some(DEFAULT_EXTENSION_CACHE_RETAINED_BYTES_NONZERO),
127        }
128    }
129
130    /// Maximum entries retained by the store.
131    pub const fn max_entries(self) -> NonZeroUsize {
132        self.max_entries
133    }
134
135    /// Maximum logical retained bytes, when configured.
136    pub const fn max_retained_bytes(self) -> Option<NonZeroUsize> {
137        self.max_retained_bytes
138    }
139
140    /// Return limits with a new logical retained-byte bound.
141    ///
142    /// # Examples
143    ///
144    /// ```
145    /// use std::num::NonZeroUsize;
146    /// use tenferro_runtime::ExtensionCacheLimits;
147    ///
148    /// let limits = ExtensionCacheLimits::default()
149    ///     .with_max_retained_bytes(NonZeroUsize::new(1024).unwrap());
150    /// assert_eq!(limits.max_retained_bytes().unwrap().get(), 1024);
151    /// ```
152    pub const fn with_max_retained_bytes(mut self, max_retained_bytes: NonZeroUsize) -> Self {
153        self.max_retained_bytes = Some(max_retained_bytes);
154        self
155    }
156}
157
158impl Default for ExtensionCacheLimits {
159    fn default() -> Self {
160        Self {
161            max_entries: NonZeroUsize::new(DEFAULT_EXTENSION_CACHE_CAPACITY)
162                .unwrap_or(NonZeroUsize::MIN),
163            max_retained_bytes: Some(DEFAULT_EXTENSION_CACHE_RETAINED_BYTES_NONZERO),
164        }
165    }
166}
167
168trait ExtensionCacheValue: Send + Sync {
169    fn as_any(&self) -> &(dyn Any + Send + Sync);
170    fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync);
171    fn retained_bytes(&self) -> usize;
172}
173
174struct FixedRetainedBytes<T> {
175    value: T,
176    retained_bytes: usize,
177}
178
179impl<T> ExtensionCacheValue for FixedRetainedBytes<T>
180where
181    T: Any + Send + Sync + 'static,
182{
183    fn as_any(&self) -> &(dyn Any + Send + Sync) {
184        &self.value
185    }
186
187    fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) {
188        &mut self.value
189    }
190
191    fn retained_bytes(&self) -> usize {
192        self.retained_bytes
193    }
194}
195
196struct DynamicRetainedBytes<T, F> {
197    value: T,
198    retained_bytes: F,
199}
200
201impl<T, F> ExtensionCacheValue for DynamicRetainedBytes<T, F>
202where
203    T: Any + Send + Sync + 'static,
204    F: Fn(&T) -> usize + Send + Sync + 'static,
205{
206    fn as_any(&self) -> &(dyn Any + Send + Sync) {
207        &self.value
208    }
209
210    fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync) {
211        &mut self.value
212    }
213
214    fn retained_bytes(&self) -> usize {
215        (self.retained_bytes)(&self.value)
216    }
217}
218
219struct ExtensionCacheEntry {
220    value: Box<dyn ExtensionCacheValue>,
221}
222
223#[derive(Clone, Copy, Debug, Default)]
224struct ExtensionCacheEventStats {
225    hits: u64,
226    misses: u64,
227    evictions: u64,
228    clears: u64,
229}
230
231impl ExtensionCacheEventStats {
232    fn to_cache_stats(self, entries: usize, retained_bytes: usize) -> CacheStats {
233        CacheStats {
234            entries,
235            retained_bytes,
236            hits: self.hits,
237            misses: self.misses,
238            evictions: self.evictions,
239            clears: self.clears,
240        }
241    }
242}
243
244#[derive(Clone, Copy, Debug)]
245struct FamilyEventStats {
246    family_id: &'static str,
247    stats: ExtensionCacheEventStats,
248}
249
250#[derive(Clone, Copy, Debug)]
251struct CacheEventStats {
252    family_id: &'static str,
253    cache_name: &'static str,
254    stats: ExtensionCacheEventStats,
255}
256
257/// Bounded type-erased cache storage owned by an extension executor.
258pub struct ExtensionCacheStore {
259    limits: ExtensionCacheLimits,
260    entries: LruCache<ExtensionCacheKey, ExtensionCacheEntry>,
261    events: ExtensionCacheEventStats,
262    family_events: Vec<FamilyEventStats>,
263    cache_events: Vec<CacheEventStats>,
264}
265
266impl fmt::Debug for ExtensionCacheStore {
267    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
268        f.debug_struct("ExtensionCacheStore")
269            .field("limits", &self.limits)
270            .field("stats", &self.stats(ExtensionCacheSelector::All))
271            .finish_non_exhaustive()
272    }
273}
274
275impl ExtensionCacheStore {
276    /// Create an empty cache store with default limits.
277    ///
278    /// # Examples
279    ///
280    /// ```
281    /// use tenferro_runtime::ExtensionCacheStore;
282    ///
283    /// let store = ExtensionCacheStore::new();
284    /// assert_eq!(store.len(), 0);
285    /// ```
286    pub fn new() -> Self {
287        Self::with_limits(ExtensionCacheLimits::default())
288    }
289
290    /// Create an empty cache store with explicit limits.
291    pub fn with_limits(limits: ExtensionCacheLimits) -> Self {
292        Self {
293            entries: LruCache::new(limits.max_entries()),
294            limits,
295            events: ExtensionCacheEventStats::default(),
296            family_events: Vec::new(),
297            cache_events: Vec::new(),
298        }
299    }
300
301    /// Return the active cache limits.
302    pub const fn limits(&self) -> ExtensionCacheLimits {
303        self.limits
304    }
305
306    /// Resize the store and evict least-recently-used entries if needed.
307    pub fn set_limits(&mut self, limits: ExtensionCacheLimits) {
308        self.limits = limits;
309        self.evict_to_limits();
310        self.entries.resize(limits.max_entries());
311    }
312
313    /// Current entry count.
314    pub fn len(&self) -> usize {
315        self.entries.len()
316    }
317
318    /// Return whether the store contains no entries.
319    pub fn is_empty(&self) -> bool {
320        self.entries.is_empty()
321    }
322
323    /// Insert or replace a typed cache entry.
324    pub fn put<T>(&mut self, key: ExtensionCacheKey, value: T, retained_bytes: usize)
325    where
326        T: Any + Send + Sync + 'static,
327    {
328        self.put_entry(
329            key,
330            ExtensionCacheEntry {
331                value: Box::new(FixedRetainedBytes {
332                    value,
333                    retained_bytes,
334                }),
335            },
336        );
337    }
338
339    /// Insert or replace a typed cache entry whose retained bytes are computed
340    /// from the current value whenever stats are requested.
341    ///
342    /// Use this for entries that mutate after insertion, such as compiled
343    /// execution plans with backend-owned nested caches.
344    ///
345    /// # Examples
346    ///
347    /// ```
348    /// use tenferro_runtime::{ExtensionCacheKey, ExtensionCacheStore, ExtensionCacheSelector};
349    ///
350    /// let mut store = ExtensionCacheStore::new();
351    /// let key = ExtensionCacheKey::new("example.cache.v1", "plans", 0);
352    /// store.put_with_retained_bytes(key, Vec::<usize>::with_capacity(2), |values| {
353    ///     values.capacity().saturating_mul(std::mem::size_of::<usize>())
354    /// });
355    /// let values = store.get_mut::<Vec<usize>>(&key).unwrap();
356    /// values.reserve_exact(4);
357    /// let retained_capacity = values.capacity();
358    ///
359    /// assert_eq!(
360    ///     store.stats(ExtensionCacheSelector::All).retained_bytes,
361    ///     retained_capacity * std::mem::size_of::<usize>()
362    /// );
363    /// ```
364    pub fn put_with_retained_bytes<T, F>(
365        &mut self,
366        key: ExtensionCacheKey,
367        value: T,
368        retained_bytes: F,
369    ) where
370        T: Any + Send + Sync + 'static,
371        F: Fn(&T) -> usize + Send + Sync + 'static,
372    {
373        self.put_entry(
374            key,
375            ExtensionCacheEntry {
376                value: Box::new(DynamicRetainedBytes {
377                    value,
378                    retained_bytes,
379                }),
380            },
381        );
382    }
383
384    /// Get a typed cache entry, updating its LRU position.
385    pub fn get<T>(&mut self, key: &ExtensionCacheKey) -> Option<&T>
386    where
387        T: Any + Send + Sync + 'static,
388    {
389        let Self {
390            entries,
391            events,
392            family_events,
393            cache_events,
394            ..
395        } = self;
396        let Some(entry) = entries.get(key) else {
397            record_miss(events, family_events, cache_events, key);
398            return None;
399        };
400        if entry.value.as_any().is::<T>() {
401            record_hit(events, family_events, cache_events, key);
402            entry.value.as_any().downcast_ref::<T>()
403        } else {
404            record_miss(events, family_events, cache_events, key);
405            None
406        }
407    }
408
409    /// Get a mutable typed cache entry, updating its LRU position.
410    pub fn get_mut<T>(&mut self, key: &ExtensionCacheKey) -> Option<&mut T>
411    where
412        T: Any + Send + Sync + 'static,
413    {
414        let Self {
415            entries,
416            events,
417            family_events,
418            cache_events,
419            ..
420        } = self;
421        let Some(entry) = entries.get_mut(key) else {
422            record_miss(events, family_events, cache_events, key);
423            return None;
424        };
425        if entry.value.as_any().is::<T>() {
426            record_hit(events, family_events, cache_events, key);
427            entry.value.as_any_mut().downcast_mut::<T>()
428        } else {
429            record_miss(events, family_events, cache_events, key);
430            None
431        }
432    }
433
434    /// Clear entries selected by `selector`.
435    pub fn clear_selected(&mut self, selector: ExtensionCacheSelector) {
436        self.record_clear(selector);
437        if selector == ExtensionCacheSelector::All {
438            self.entries.clear();
439            return;
440        }
441
442        let keys: Vec<_> = self
443            .entries
444            .iter()
445            .map(|(key, _)| *key)
446            .filter(|key| selector.matches(key))
447            .collect();
448        for key in keys {
449            self.entries.pop(&key);
450        }
451    }
452
453    /// Clear every extension cache entry.
454    pub fn clear(&mut self) {
455        self.record_clear(ExtensionCacheSelector::All);
456        self.entries.clear();
457    }
458
459    /// Return cache-style stats for entries selected by `selector`.
460    pub fn stats(&self, selector: ExtensionCacheSelector) -> CacheStats {
461        let entries = self
462            .entries
463            .iter()
464            .filter(|(key, _)| selector.matches(key))
465            .count();
466        let retained_bytes = self
467            .entries
468            .iter()
469            .filter(|(key, _)| selector.matches(key))
470            .map(|(_, entry)| entry.value.retained_bytes())
471            .fold(0usize, usize::saturating_add);
472        self.event_stats(selector)
473            .to_cache_stats(entries, retained_bytes)
474    }
475
476    fn put_entry(&mut self, key: ExtensionCacheKey, entry: ExtensionCacheEntry) {
477        if let Some((removed_key, _)) = self.entries.push(key, entry) {
478            if removed_key != key {
479                self.record_eviction(&removed_key);
480            }
481        }
482        self.evict_to_limits();
483    }
484
485    fn evict_to_limits(&mut self) {
486        while self.entries.len() > self.limits.max_entries().get()
487            || self.retained_bytes_exceeds_limit()
488        {
489            let Some((key, _)) = self.entries.pop_lru() else {
490                break;
491            };
492            self.record_eviction(&key);
493        }
494    }
495
496    fn retained_bytes_exceeds_limit(&self) -> bool {
497        self.limits
498            .max_retained_bytes
499            .is_some_and(|limit| self.retained_bytes() > limit.get())
500    }
501
502    fn retained_bytes(&self) -> usize {
503        self.entries
504            .iter()
505            .map(|(_, entry)| entry.value.retained_bytes())
506            .fold(0usize, usize::saturating_add)
507    }
508
509    fn event_stats(&self, selector: ExtensionCacheSelector) -> ExtensionCacheEventStats {
510        match selector {
511            ExtensionCacheSelector::All => self.events,
512            ExtensionCacheSelector::Family { family_id } => {
513                family_event_stats(&self.family_events, family_id)
514                    .copied()
515                    .unwrap_or_default()
516            }
517            ExtensionCacheSelector::Cache {
518                family_id,
519                cache_name,
520            } => cache_event_stats(&self.cache_events, family_id, cache_name)
521                .copied()
522                .unwrap_or_default(),
523        }
524    }
525
526    fn record_eviction(&mut self, key: &ExtensionCacheKey) {
527        record_eviction(
528            &mut self.events,
529            &mut self.family_events,
530            &mut self.cache_events,
531            key,
532        );
533    }
534
535    fn record_clear(&mut self, selector: ExtensionCacheSelector) {
536        self.ensure_event_scopes_for_current_entries(selector);
537        self.events.clears = self.events.clears.saturating_add(1);
538        match selector {
539            ExtensionCacheSelector::All => {
540                for scoped in &mut self.family_events {
541                    scoped.stats.clears = scoped.stats.clears.saturating_add(1);
542                }
543                for scoped in &mut self.cache_events {
544                    scoped.stats.clears = scoped.stats.clears.saturating_add(1);
545                }
546            }
547            ExtensionCacheSelector::Family { family_id } => {
548                self.record_family_clear(family_id);
549                for scoped in &mut self.cache_events {
550                    if scoped.family_id == family_id {
551                        scoped.stats.clears = scoped.stats.clears.saturating_add(1);
552                    }
553                }
554            }
555            ExtensionCacheSelector::Cache {
556                family_id,
557                cache_name,
558            } => {
559                self.record_family_clear(family_id);
560                self.record_cache_clear(family_id, cache_name);
561            }
562        }
563    }
564
565    fn ensure_event_scopes_for_current_entries(&mut self, selector: ExtensionCacheSelector) {
566        let keys: Vec<_> = self
567            .entries
568            .iter()
569            .map(|(key, _)| *key)
570            .filter(|key| selector.matches(key))
571            .collect();
572        for key in keys {
573            ensure_family_event_stats(&mut self.family_events, key.family_id);
574            ensure_cache_event_stats(&mut self.cache_events, key.family_id, key.cache_name);
575        }
576    }
577
578    fn record_family_clear(&mut self, family_id: &'static str) {
579        let stats = ensure_family_event_stats(&mut self.family_events, family_id);
580        stats.clears = stats.clears.saturating_add(1);
581    }
582
583    fn record_cache_clear(&mut self, family_id: &'static str, cache_name: &'static str) {
584        let stats = ensure_cache_event_stats(&mut self.cache_events, family_id, cache_name);
585        stats.clears = stats.clears.saturating_add(1);
586    }
587}
588
589fn family_event_stats<'a>(
590    events: &'a [FamilyEventStats],
591    family_id: &'static str,
592) -> Option<&'a ExtensionCacheEventStats> {
593    events
594        .iter()
595        .find(|event| event.family_id == family_id)
596        .map(|event| &event.stats)
597}
598
599fn cache_event_stats<'a>(
600    events: &'a [CacheEventStats],
601    family_id: &'static str,
602    cache_name: &'static str,
603) -> Option<&'a ExtensionCacheEventStats> {
604    events
605        .iter()
606        .find(|event| event.family_id == family_id && event.cache_name == cache_name)
607        .map(|event| &event.stats)
608}
609
610fn ensure_family_event_stats<'a>(
611    events: &'a mut Vec<FamilyEventStats>,
612    family_id: &'static str,
613) -> &'a mut ExtensionCacheEventStats {
614    if let Some(index) = events.iter().position(|event| event.family_id == family_id) {
615        return &mut events[index].stats;
616    }
617    events.push(FamilyEventStats {
618        family_id,
619        stats: ExtensionCacheEventStats::default(),
620    });
621    &mut events
622        .last_mut()
623        .expect("just-pushed family event stats")
624        .stats
625}
626
627fn ensure_cache_event_stats<'a>(
628    events: &'a mut Vec<CacheEventStats>,
629    family_id: &'static str,
630    cache_name: &'static str,
631) -> &'a mut ExtensionCacheEventStats {
632    if let Some(index) = events
633        .iter()
634        .position(|event| event.family_id == family_id && event.cache_name == cache_name)
635    {
636        return &mut events[index].stats;
637    }
638    events.push(CacheEventStats {
639        family_id,
640        cache_name,
641        stats: ExtensionCacheEventStats::default(),
642    });
643    &mut events
644        .last_mut()
645        .expect("just-pushed cache event stats")
646        .stats
647}
648
649fn record_hit(
650    events: &mut ExtensionCacheEventStats,
651    family_events: &mut Vec<FamilyEventStats>,
652    cache_events: &mut Vec<CacheEventStats>,
653    key: &ExtensionCacheKey,
654) {
655    events.hits = events.hits.saturating_add(1);
656    let family = ensure_family_event_stats(family_events, key.family_id);
657    family.hits = family.hits.saturating_add(1);
658    let cache = ensure_cache_event_stats(cache_events, key.family_id, key.cache_name);
659    cache.hits = cache.hits.saturating_add(1);
660}
661
662fn record_miss(
663    events: &mut ExtensionCacheEventStats,
664    family_events: &mut Vec<FamilyEventStats>,
665    cache_events: &mut Vec<CacheEventStats>,
666    key: &ExtensionCacheKey,
667) {
668    events.misses = events.misses.saturating_add(1);
669    let family = ensure_family_event_stats(family_events, key.family_id);
670    family.misses = family.misses.saturating_add(1);
671    let cache = ensure_cache_event_stats(cache_events, key.family_id, key.cache_name);
672    cache.misses = cache.misses.saturating_add(1);
673}
674
675fn record_eviction(
676    events: &mut ExtensionCacheEventStats,
677    family_events: &mut Vec<FamilyEventStats>,
678    cache_events: &mut Vec<CacheEventStats>,
679    key: &ExtensionCacheKey,
680) {
681    events.evictions = events.evictions.saturating_add(1);
682    let family = ensure_family_event_stats(family_events, key.family_id);
683    family.evictions = family.evictions.saturating_add(1);
684    let cache = ensure_cache_event_stats(cache_events, key.family_id, key.cache_name);
685    cache.evictions = cache.evictions.saturating_add(1);
686}
687
688impl Default for ExtensionCacheStore {
689    fn default() -> Self {
690        Self::new()
691    }
692}
693
694#[cfg(test)]
695mod tests;
tenferro_runtime/extension_cache.rs

tenferro_runtime/
extension_cache.rs
