Skip to main content

cranpose_services/
device_info.rs

1//! Device information (RAM, so far). Apps use it to size in-memory work (e.g.
2//! whether a large model fits).
3//!
4//! The default reads `/proc/meminfo` on Linux/Android and returns `None`
5//! elsewhere; the iOS backend installs a `ProcessInfo`-based implementation via
6//! [`set_platform_device_info`].
7
8use std::cell::RefCell;
9use std::rc::Rc;
10
11/// Provides device information.
12pub trait DeviceInfo {
13    /// Total physical memory in bytes, or `None` if unknown.
14    fn total_memory_bytes(&self) -> Option<u64>;
15}
16
17pub type DeviceInfoRef = Rc<dyn DeviceInfo>;
18
19struct DefaultDeviceInfo;
20
21impl DeviceInfo for DefaultDeviceInfo {
22    fn total_memory_bytes(&self) -> Option<u64> {
23        #[cfg(any(target_os = "linux", target_os = "android"))]
24        {
25            let text = std::fs::read_to_string("/proc/meminfo").ok()?;
26            for line in text.lines() {
27                if let Some(rest) = line.strip_prefix("MemTotal:") {
28                    let kb: u64 = rest.split_whitespace().next()?.parse().ok()?;
29                    return Some(kb * 1024);
30                }
31            }
32            None
33        }
34        #[cfg(not(any(target_os = "linux", target_os = "android")))]
35        {
36            None
37        }
38    }
39}
40
41thread_local! {
42    static PLATFORM_DEVICE_INFO: RefCell<Option<DeviceInfoRef>> = const { RefCell::new(None) };
43}
44
45/// Installs a platform device-info implementation, replacing any previous one.
46pub fn set_platform_device_info(info: DeviceInfoRef) {
47    PLATFORM_DEVICE_INFO.with(|cell| *cell.borrow_mut() = Some(info));
48}
49
50/// Removes any registered platform device info (tests and teardown).
51pub fn clear_platform_device_info() {
52    PLATFORM_DEVICE_INFO.with(|cell| *cell.borrow_mut() = None);
53}
54
55/// The active device info: the platform implementation if installed, otherwise
56/// the built-in default.
57pub fn device_info() -> DeviceInfoRef {
58    PLATFORM_DEVICE_INFO
59        .with(|cell| cell.borrow().clone())
60        .unwrap_or_else(|| Rc::new(DefaultDeviceInfo))
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn registered_device_info_takes_precedence() {
69        clear_platform_device_info();
70        struct Fake;
71        impl DeviceInfo for Fake {
72            fn total_memory_bytes(&self) -> Option<u64> {
73                Some(8 * 1024 * 1024 * 1024)
74            }
75        }
76        set_platform_device_info(Rc::new(Fake));
77        assert_eq!(device_info().total_memory_bytes(), Some(8 << 30));
78        clear_platform_device_info();
79    }
80}