laddu_memory/
state.rs

1use std::{
2    collections::BTreeMap,
3    sync::{
4        Arc, Mutex, OnceLock,
5        atomic::{AtomicU64, Ordering},
6    },
7};
8
9use crate::{
10    budget::MemoryBudget,
11    discovery::{MemoryProbe, SystemMemoryProbe, discover_host, discover_process_memory},
12    error::{MemoryError, MemoryResult},
13    pool::{MemoryPool, MemoryPoolInner, ReservationAccount, ResourceLedger},
14    report::{MemoryReport, MemoryResourceReport, ProcessMemoryReport},
15    resource::{DeviceIdentity, MemoryResource, MemoryResourceKind},
16};
17
18/// Live resource discovery and process-wide laddu reservation state.
19#[derive(Clone, Debug)]
20pub struct MemoryState {
21    inner: Arc<MemoryStateInner>,
22}
23
24#[derive(Debug)]
25pub(crate) struct MemoryStateInner {
26    pub(crate) resources: Mutex<BTreeMap<String, ResourceLedger>>,
27    process_high_water: AtomicU64,
28}
29
30impl MemoryState {
31    /// Discovers host memory and creates an independent reservation state.
32    pub fn discover() -> Self {
33        let host = discover_host();
34        let mut resources = BTreeMap::new();
35        resources.insert(host.id.clone(), ResourceLedger::new(host));
36        Self {
37            inner: Arc::new(MemoryStateInner {
38                resources: Mutex::new(resources),
39                process_high_water: AtomicU64::new(0),
40            }),
41        }
42    }
43
44    /// Returns the process-wide default state.
45    pub fn current() -> Self {
46        static CURRENT: OnceLock<MemoryState> = OnceLock::new();
47        CURRENT.get_or_init(Self::discover).clone()
48    }
49
50    /// Refreshes host total and available memory.
51    pub fn refresh(&self) {
52        self.refresh_inner();
53    }
54    fn refresh_inner(&self) -> Option<ProcessMemoryReport> {
55        self.refresh_with_probe(&SystemMemoryProbe)
56    }
57
58    fn refresh_with_probe(&self, probe: &dyn MemoryProbe) -> Option<ProcessMemoryReport> {
59        let host = discover_host();
60        let process = self.sample_process_memory();
61        let targets = {
62            let mut resources = self
63                .inner
64                .resources
65                .lock()
66                .unwrap_or_else(|e| e.into_inner());
67            let ledger = resources
68                .entry(host.id.clone())
69                .or_insert_with(|| ResourceLedger::new(host.clone()));
70            ledger.update_snapshot(host);
71            resources
72                .values()
73                .filter(|ledger| ledger.snapshot.is_refreshable())
74                .map(|ledger| ledger.snapshot.clone())
75                .collect::<Vec<_>>()
76        };
77        let outcomes = targets
78            .into_iter()
79            .map(|target| {
80                let outcome = probe.probe_device(&target);
81                (target, outcome)
82            })
83            .collect::<Vec<_>>();
84        let mut resources = self
85            .inner
86            .resources
87            .lock()
88            .unwrap_or_else(|e| e.into_inner());
89        for (target, outcome) in outcomes {
90            if let Ok(snapshot) = outcome
91                && let Some(ledger) = resources.get_mut(&target.id)
92            {
93                ledger.apply_telemetry(&target, snapshot);
94            }
95        }
96        process
97    }
98
99    /// Returns the current host snapshot.
100    pub fn host(&self) -> MemoryResource {
101        self.resource("host").unwrap_or_else(discover_host)
102    }
103
104    /// Registers capacity telemetry for a runtime-selected accelerator.
105    ///
106    /// The runtime owns the stable identifier and adapter identity. Platform
107    /// telemetry is used when available; otherwise `fallback_bytes` becomes an
108    /// adaptive capacity estimate.
109    pub fn register_discovered_device(
110        &self,
111        id: impl Into<String>,
112        name: impl Into<String>,
113        identity: DeviceIdentity,
114        fallback_bytes: u64,
115    ) {
116        self.insert_device_snapshot(MemoryResource::discover_device(
117            id,
118            name,
119            identity,
120            fallback_bytes,
121        ));
122    }
123
124    /// Overrides capacity telemetry for a device until another user override.
125    pub fn override_device_capacity(
126        &self,
127        id: impl Into<String>,
128        name: impl Into<String>,
129        total_bytes: u64,
130        available_bytes: Option<u64>,
131    ) {
132        self.insert_device_snapshot(MemoryResource::user_device(
133            id,
134            name,
135            total_bytes,
136            available_bytes,
137        ));
138    }
139
140    pub(crate) fn insert_device_snapshot(&self, resource: MemoryResource) {
141        let mut resources = self
142            .inner
143            .resources
144            .lock()
145            .unwrap_or_else(|e| e.into_inner());
146        if let Some(ledger) = resources.get_mut(&resource.id) {
147            ledger.update_snapshot(resource);
148        } else {
149            resources.insert(resource.id.clone(), ResourceLedger::new(resource));
150        }
151    }
152
153    /// Returns one resource by stable identifier.
154    pub fn resource(&self, id: &str) -> Option<MemoryResource> {
155        self.inner
156            .resources
157            .lock()
158            .unwrap_or_else(|e| e.into_inner())
159            .get(id)
160            .map(|ledger| ledger.snapshot.clone())
161    }
162
163    /// Returns all registered accelerator resources.
164    pub fn devices(&self) -> Vec<MemoryResource> {
165        self.inner
166            .resources
167            .lock()
168            .unwrap_or_else(|e| e.into_inner())
169            .values()
170            .filter(|ledger| ledger.snapshot.kind == MemoryResourceKind::Device)
171            .map(|ledger| ledger.snapshot.clone())
172            .collect()
173    }
174
175    /// Resolves a budget and creates a local pool for a resource.
176    ///
177    /// # Errors
178    ///
179    /// Returns an error when the resource is unknown or the budget cannot be resolved.
180    pub fn pool(&self, resource_id: &str, budget: MemoryBudget) -> MemoryResult<MemoryPool> {
181        let resource = self.resource(resource_id).ok_or_else(|| {
182            MemoryError::InvalidBudget(format!("unknown memory resource {resource_id:?}"))
183        })?;
184        resource.validate().map_err(|error| {
185            MemoryError::InvalidBudget(format!(
186                "invalid memory resource {resource_id:?}: {error:?}"
187            ))
188        })?;
189        let capacity = budget.resolve(&resource)?;
190        Ok(MemoryPool {
191            inner: Arc::new(MemoryPoolInner {
192                requested: budget,
193                accounting: Arc::new(ReservationAccount::new(
194                    Arc::downgrade(&self.inner),
195                    resource_id.to_owned(),
196                    capacity,
197                )),
198            }),
199        })
200    }
201
202    /// Returns a structured report for all resources.
203    pub fn report(&self) -> MemoryReport {
204        let process = self.refresh_inner();
205        let resources = self
206            .inner
207            .resources
208            .lock()
209            .unwrap_or_else(|e| e.into_inner());
210        MemoryReport {
211            process,
212            resources: resources
213                .values()
214                .map(|ledger| MemoryResourceReport {
215                    resource: ledger.snapshot.clone(),
216                    laddu_reserved_bytes: ledger.reserved,
217                    laddu_high_water_bytes: ledger.high_water,
218                })
219                .collect(),
220        }
221    }
222
223    fn sample_process_memory(&self) -> Option<ProcessMemoryReport> {
224        let snapshot = discover_process_memory().ok()?;
225        self.inner
226            .process_high_water
227            .fetch_max(snapshot.resident_bytes, Ordering::AcqRel);
228        Some(ProcessMemoryReport {
229            resident_bytes: snapshot.resident_bytes,
230            virtual_bytes: snapshot.virtual_bytes,
231            sampled_high_water_bytes: self.inner.process_high_water.load(Ordering::Acquire),
232        })
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239    use crate::{
240        CapacitySource, DeviceIdentity, MemoryResourceKind,
241        discovery::{CapacitySnapshot, ProbeFailure, ProbeOutcome},
242    };
243    use std::{sync::mpsc, thread, time::Duration};
244
245    fn resource() -> MemoryResource {
246        MemoryResource {
247            id: "test".into(),
248            name: "Test".into(),
249            kind: MemoryResourceKind::Device,
250            total_bytes: Some(1_000),
251            available_bytes: Some(500),
252            capacity_source: CapacitySource::User,
253            device_identity: None,
254        }
255    }
256    fn refreshable(identity: usize) -> MemoryResource {
257        MemoryResource {
258            capacity_source: CapacitySource::Adaptive,
259            device_identity: Some(DeviceIdentity {
260                adapter_index: identity,
261                vendor_id: 1,
262                device_id: 2,
263                pci_bus_id: format!("bus-{identity}"),
264            }),
265            ..resource()
266        }
267    }
268    fn snapshot(total_bytes: u64, available_bytes: u64) -> CapacitySnapshot {
269        CapacitySnapshot {
270            total_bytes,
271            available_bytes,
272            source: CapacitySource::Drm,
273        }
274    }
275
276    #[derive(Clone, Copy)]
277    struct FixedProbe(ProbeOutcome);
278    impl MemoryProbe for FixedProbe {
279        fn probe_device(&self, _: &MemoryResource) -> ProbeOutcome {
280            self.0
281        }
282    }
283    struct BlockingProbe {
284        entered: mpsc::SyncSender<()>,
285        release: Mutex<mpsc::Receiver<()>>,
286        outcome: ProbeOutcome,
287    }
288    impl MemoryProbe for BlockingProbe {
289        fn probe_device(&self, _: &MemoryResource) -> ProbeOutcome {
290            self.entered.send(()).unwrap();
291            self.release.lock().unwrap().recv().unwrap();
292            self.outcome
293        }
294    }
295    struct ReplacingProbe {
296        state: MemoryState,
297        replacement: MemoryResource,
298        outcome: ProbeOutcome,
299    }
300    impl MemoryProbe for ReplacingProbe {
301        fn probe_device(&self, _: &MemoryResource) -> ProbeOutcome {
302            self.state.insert_device_snapshot(self.replacement.clone());
303            self.outcome
304        }
305    }
306
307    #[test]
308    fn ledgers_clamp_capacity_and_keep_user_overrides_sticky() {
309        let state = MemoryState::discover();
310        let discovered = MemoryResource {
311            total_bytes: Some(1_000),
312            available_bytes: Some(1_100),
313            ..refreshable(0)
314        };
315        state.insert_device_snapshot(discovered);
316        assert_eq!(state.resource("test").unwrap().available_bytes, Some(1_000));
317        let user = MemoryResource::user_device("test", "Test", 700, Some(650));
318        state.override_device_capacity("test", "Test", 700, Some(650));
319        state.insert_device_snapshot(MemoryResource {
320            total_bytes: Some(900),
321            available_bytes: Some(800),
322            ..refreshable(0)
323        });
324        assert_eq!(state.resource("test"), Some(user));
325
326        let replacement = MemoryResource::user_device("test", "Test", 600, Some(550));
327        state.override_device_capacity("test", "Test", 600, Some(550));
328        assert_eq!(state.resource("test"), Some(replacement));
329        assert_eq!(
330            state
331                .devices()
332                .into_iter()
333                .filter(|resource| resource.id == "test")
334                .count(),
335            1
336        );
337    }
338
339    #[test]
340    fn refresh_probes_without_holding_the_resource_lock() {
341        let state = MemoryState::discover();
342        state.insert_device_snapshot(refreshable(0));
343        let (entered_tx, entered_rx) = mpsc::sync_channel(0);
344        let (release_tx, release_rx) = mpsc::sync_channel(0);
345        let probe = BlockingProbe {
346            entered: entered_tx,
347            release: Mutex::new(release_rx),
348            outcome: Ok(snapshot(900, 700)),
349        };
350        let refresh_state = state.clone();
351        let refresh = thread::spawn(move || refresh_state.refresh_with_probe(&probe));
352        entered_rx.recv_timeout(Duration::from_secs(2)).unwrap();
353        let access_state = state.clone();
354        let (done_tx, done_rx) = mpsc::sync_channel(0);
355        let access = thread::spawn(move || {
356            assert!(access_state.resource("test").is_some());
357            done_tx.send(()).unwrap();
358        });
359        let result = done_rx.recv_timeout(Duration::from_secs(2));
360        release_tx.send(()).unwrap();
361        refresh.join().unwrap();
362        access.join().unwrap();
363        assert!(result.is_ok(), "resource access blocked during telemetry");
364    }
365
366    #[test]
367    fn refresh_retains_stale_snapshot_on_failure() {
368        let state = MemoryState::discover();
369        state.insert_device_snapshot(refreshable(0));
370        state.refresh_with_probe(&FixedProbe(Ok(snapshot(800, 600))));
371        let refreshed = state.resource("test").unwrap();
372        state.refresh_with_probe(&FixedProbe(Err(ProbeFailure::Unavailable)));
373        assert_eq!(state.resource("test"), Some(refreshed));
374    }
375
376    #[test]
377    fn refresh_does_not_overwrite_user_override_or_replaced_device() {
378        let state = MemoryState::discover();
379        state.insert_device_snapshot(refreshable(0));
380        let user = MemoryResource::user_device("test", "Test", 700, Some(650));
381        state.refresh_with_probe(&ReplacingProbe {
382            state: state.clone(),
383            replacement: user.clone(),
384            outcome: Ok(snapshot(900, 800)),
385        });
386        assert_eq!(state.resource("test"), Some(user));
387
388        let state = MemoryState::discover();
389        state.insert_device_snapshot(refreshable(0));
390        let replacement = MemoryResource {
391            total_bytes: Some(400),
392            available_bytes: Some(300),
393            ..refreshable(1)
394        };
395        state.refresh_with_probe(&ReplacingProbe {
396            state: state.clone(),
397            replacement: replacement.clone(),
398            outcome: Ok(snapshot(900, 800)),
399        });
400        assert_eq!(state.resource("test"), Some(replacement));
401    }
402
403    #[test]
404    #[cfg(target_os = "linux")]
405    fn reports_current_process_memory() {
406        let state = MemoryState::discover();
407        let first = state.report().process.unwrap();
408        let second = state.report().process.unwrap();
409        assert!(first.resident_bytes > 0);
410        assert!(first.virtual_bytes >= first.resident_bytes);
411        assert!(second.sampled_high_water_bytes >= second.resident_bytes);
412    }
413}
laddu_memory/state.rs

laddu_memory/
state.rs
