hrx/
plan_cache.rs

1//! Bounded, single-flight preparation with idle-only LRU eviction.
2
3use crate::{Error, Result};
4use std::{
5    collections::HashMap,
6    hash::Hash,
7    sync::{Arc, Condvar, Mutex},
8};
9
10/// Complete identity for a cache shared across artifacts and compiler profiles.
11/// Model-local caches may instead use shape keys when their owner fixes all
12/// other fields for the cache's entire lifetime.
13#[derive(Clone, Debug, PartialEq, Eq, Hash)]
14pub struct PlanKey {
15    /// Digest of the model artifact and its packing format.
16    pub artifact: String,
17    /// Actual device architecture and required features.
18    pub target: String,
19    /// Compiler version and compilation options identity.
20    pub compiler: String,
21    /// Model-specific specialization identity, including scalar arguments.
22    pub specialization: String,
23    /// Input shapes, element types, strides and layouts.
24    pub inputs: Vec<crate::tensor::TensorDesc>,
25    /// Output shapes, element types, strides and layouts.
26    pub outputs: Vec<crate::tensor::TensorDesc>,
27}
28
29type Outcome<V> = std::result::Result<Arc<V>, Arc<Error>>;
30struct Flight<V> {
31    result: Mutex<Option<Outcome<V>>>,
32    changed: Condvar,
33}
34enum Entry<V> {
35    Loading(Arc<Flight<V>>),
36    Ready { value: Arc<V>, used: u64 },
37}
38struct State<K, V> {
39    entries: HashMap<K, Entry<V>>,
40    clock: u64,
41    hits: u64,
42    misses: u64,
43    evictions: u64,
44}
45
46/// Cache accounting. A miss is one elected preparation, not every waiter.
47#[derive(Clone, Copy, Debug, Default, serde::Serialize)]
48pub struct CacheStatistics {
49    /// Current ready and in-progress entries.
50    pub entries: usize,
51    /// Requests reusing ready or in-progress preparation.
52    pub hits: u64,
53    /// Elected preparations, including failures.
54    pub misses: u64,
55    /// Idle entries removed to admit another shape.
56    pub evictions: u64,
57}
58
59/// Concurrent bounded preparation. Returned `Arc`s pin values. The additional
60/// idle predicate protects work whose output or completion outlives that Arc.
61/// Failed or panicking loaders release capacity and wake waiters for retry.
62pub struct PlanCache<K, V> {
63    capacity: usize,
64    idle: fn(&V) -> bool,
65    state: Mutex<State<K, V>>,
66}
67impl<K: Clone + Eq + Hash, V> PlanCache<K, V> {
68    /// Set the maximum number of resident/preparing plans, never zero.
69    pub fn new(capacity: usize, idle: fn(&V) -> bool) -> Result<Self> {
70        if capacity == 0 {
71            return Err(Error::Message("plan cache capacity must be nonzero".into()));
72        }
73        Ok(Self {
74            capacity,
75            idle,
76            state: Mutex::new(State {
77                entries: HashMap::new(),
78                clock: 0,
79                hits: 0,
80                misses: 0,
81                evictions: 0,
82            }),
83        })
84    }
85
86    /// Reuse or prepare a value. Concurrent requests for the same key share one
87    /// outcome; other keys may prepare concurrently. Returns `Busy` when all
88    /// capacity is pinned or preparing. No loader executes under the cache lock.
89    pub fn get_or_prepare(&self, key: K, load: impl FnOnce() -> Result<V>) -> Result<Arc<V>> {
90        let (flight, elected) = {
91            let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
92            state.clock = state.clock.saturating_add(1);
93            let used = state.clock;
94            if state.entries.contains_key(&key) {
95                state.hits += 1;
96                match state.entries.get_mut(&key).expect("entry checked") {
97                    Entry::Ready { value, used: last } => {
98                        *last = used;
99                        return Ok(value.clone());
100                    }
101                    Entry::Loading(flight) => (flight.clone(), false),
102                }
103            } else {
104                if state.entries.len() == self.capacity {
105                    let victim = state
106                        .entries
107                        .iter()
108                        .filter_map(|(key, entry)| match entry {
109                            Entry::Ready { value, used }
110                                if Arc::strong_count(value) == 1 && (self.idle)(value) =>
111                            {
112                                Some((key.clone(), *used))
113                            }
114                            _ => None,
115                        })
116                        .min_by_key(|(_, used)| *used)
117                        .map(|(key, _)| key);
118                    let Some(victim) = victim else {
119                        return Err(Error::Busy(
120                            "all shape plans are pinned or preparing".into(),
121                        ));
122                    };
123                    state.entries.remove(&victim);
124                    state.evictions += 1;
125                }
126                let flight = Arc::new(Flight {
127                    result: Mutex::new(None),
128                    changed: Condvar::new(),
129                });
130                state
131                    .entries
132                    .insert(key.clone(), Entry::Loading(flight.clone()));
133                state.misses += 1;
134                (flight, true)
135            }
136        };
137        if elected {
138            let result = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(load)) {
139                Ok(result) => result.map(Arc::new).map_err(Arc::new),
140                Err(_) => Err(Arc::new(Error::Message(
141                    "shape preparation panicked".into(),
142                ))),
143            };
144            {
145                let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
146                match &result {
147                    Ok(value) => {
148                        let used = state.clock;
149                        state.entries.insert(
150                            key,
151                            Entry::Ready {
152                                value: value.clone(),
153                                used,
154                            },
155                        );
156                    }
157                    Err(_) => {
158                        state.entries.remove(&key);
159                    }
160                }
161                *flight.result.lock().unwrap_or_else(|e| e.into_inner()) = Some(result.clone());
162            }
163            flight.changed.notify_all();
164            result.map_err(|source| Error::Execution { source })
165        } else {
166            let mut result = flight.result.lock().unwrap_or_else(|e| e.into_inner());
167            while result.is_none() {
168                result = flight
169                    .changed
170                    .wait(result)
171                    .unwrap_or_else(|e| e.into_inner());
172            }
173            result
174                .as_ref()
175                .expect("published above")
176                .clone()
177                .map_err(|source| Error::Execution { source })
178        }
179    }
180
181    /// Snapshot cache counters without waiting for preparation.
182    pub fn statistics(&self) -> CacheStatistics {
183        let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
184        CacheStatistics {
185            entries: state.entries.len(),
186            hits: state.hits,
187            misses: state.misses,
188            evictions: state.evictions,
189        }
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196    #[test]
197    fn pinned_values_apply_backpressure_and_idle_entries_are_lru() {
198        let cache = PlanCache::new(2, |_: &usize| true).unwrap();
199        let pinned = cache.get_or_prepare(1, || Ok(1)).unwrap();
200        drop(cache.get_or_prepare(2, || Ok(2)).unwrap());
201        drop(cache.get_or_prepare(3, || Ok(3)).unwrap());
202        assert_eq!(cache.statistics().evictions, 1);
203        let third = cache.get_or_prepare(3, || unreachable!()).unwrap();
204        assert!(matches!(
205            cache.get_or_prepare(4, || Ok(4)),
206            Err(Error::Busy(_))
207        ));
208        drop((pinned, third));
209        drop(cache.get_or_prepare(4, || Ok(4)).unwrap());
210        assert_eq!(cache.statistics().entries, 2);
211    }
212    #[test]
213    fn concurrent_preparation_is_single_flight_and_failures_retry() {
214        let cache = PlanCache::new(1, |_: &usize| true).unwrap();
215        let starts = std::sync::atomic::AtomicUsize::new(0);
216        let barrier = std::sync::Barrier::new(8);
217        std::thread::scope(|scope| {
218            for _ in 0..8 {
219                scope.spawn(|| {
220                    barrier.wait();
221                    assert_eq!(
222                        *cache
223                            .get_or_prepare(1, || {
224                                starts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
225                                Ok(7)
226                            })
227                            .unwrap(),
228                        7
229                    );
230                });
231            }
232        });
233        assert_eq!(starts.load(std::sync::atomic::Ordering::Relaxed), 1);
234        assert!(
235            cache
236                .get_or_prepare(2, || Err(Error::Message("retry".into())))
237                .is_err()
238        );
239        assert!(
240            cache
241                .get_or_prepare(2, || panic!("injected loader panic"))
242                .is_err()
243        );
244        assert_eq!(*cache.get_or_prepare(2, || Ok(9)).unwrap(), 9);
245    }
246}
hrx/plan_cache.rs

hrx/
plan_cache.rs
