1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
//! What the device has, and what this process is using of it.
//!
//! Applications ask two different questions here. *How much memory does this
//! device have* sizes a decision made once — whether a large model fits at all.
//! *How much is this process holding, and how much may it still have* is asked
//! while work is running, because the answer moves and because the platform
//! kills a process that gets it wrong.
//!
//! Every reading is optional. A platform that will not say reports `None`
//! rather than a zero an application would divide by, and
//! [`release_free_memory`] reports whether the platform has such a call at all
//! instead of quietly doing nothing.
//!
//! The default reads what can be read without leaving safe Rust:
//! `/proc/meminfo` and `/proc/self/statm` on Linux and Android. The platform
//! backends install richer implementations through
//! [`set_platform_device_info`].
use std::{rc::Rc, time::Duration};
/// Provides device and process information.
pub trait DeviceInfo {
/// Total physical memory in bytes, or `None` if unknown.
fn total_memory_bytes(&self) -> Option<u64>;
/// Memory this process currently holds resident, or `None` where the
/// platform will not say.
fn resident_memory_bytes(&self) -> Option<u64> {
None
}
/// Memory this process may still allocate before the platform stops it.
///
/// Not the same as free system memory: a device with gigabytes free may
/// still refuse this process another hundred megabytes, and it is the
/// second number that decides whether the next allocation is the one that
/// gets the application killed.
fn available_memory_bytes(&self) -> Option<u64> {
None
}
/// Processor time this process has used, user and system together.
///
/// Wall-clock time says how long something took; this says how much of a
/// core it took, which is what a background lane that must not heat the
/// device up is actually rationing.
fn process_cpu_time(&self) -> Option<Duration> {
None
}
/// Asks the allocator to return free pages to the system.
///
/// Returns whether the platform has such a call. Freeing a large buffer
/// does not necessarily shrink the process — an allocator keeps the pages
/// for the next allocation — and on a platform that kills by resident size
/// that is the difference between finishing and being killed.
fn release_free_memory(&self) -> bool {
false
}
}
pub type DeviceInfoRef = Rc<dyn DeviceInfo>;
struct DefaultDeviceInfo;
impl DeviceInfo for DefaultDeviceInfo {
fn total_memory_bytes(&self) -> Option<u64> {
#[cfg(any(target_os = "linux", target_os = "android"))]
{
let text = std::fs::read_to_string("/proc/meminfo").ok()?;
for line in text.lines() {
if let Some(rest) = line.strip_prefix("MemTotal:") {
let kb: u64 = rest.split_whitespace().next()?.parse().ok()?;
return Some(kb * 1024);
}
}
None
}
#[cfg(not(any(target_os = "linux", target_os = "android")))]
{
None
}
}
fn resident_memory_bytes(&self) -> Option<u64> {
#[cfg(any(target_os = "linux", target_os = "android"))]
{
let text = std::fs::read_to_string("/proc/self/statm").ok()?;
resident_bytes_from_statm(&text, page_size_bytes())
}
#[cfg(not(any(target_os = "linux", target_os = "android")))]
{
None
}
}
}
#[cfg(any(target_os = "linux", target_os = "android", test))]
fn resident_bytes_from_statm(text: &str, page_size: u64) -> Option<u64> {
text.split_whitespace()
.nth(1)?
.parse::<u64>()
.ok()?
.checked_mul(page_size)
}
#[cfg(any(target_os = "linux", target_os = "android", test))]
const fn page_size_bytes() -> u64 {
4096
}
/// Installs a platform device-info implementation, replacing any previous one.
pub fn set_platform_device_info(info: DeviceInfoRef) {
crate::thread_state::with(|state| *state.device_info.borrow_mut() = Some(info));
}
/// Removes any registered platform device info (tests and teardown).
pub fn clear_platform_device_info() {
crate::thread_state::with(|state| *state.device_info.borrow_mut() = None);
}
/// The active device info: the platform implementation if installed, otherwise
/// the built-in default.
pub fn device_info() -> DeviceInfoRef {
crate::thread_state::with(|state| state.device_info.borrow().clone())
.unwrap_or_else(|| Rc::new(DefaultDeviceInfo))
}
/// Asks the allocator to return free pages to the system.
///
/// Returns whether the platform has such a call — see
/// [`DeviceInfo::release_free_memory`]. Worth doing after finishing with a
/// large buffer on a platform that kills by resident size, and worth doing
/// nowhere else: it walks the allocator's arenas.
pub fn release_free_memory() -> bool {
device_info().release_free_memory()
}
#[cfg(test)]
#[path = "tests/device_info_tests.rs"]
mod tests;