blue_lang_runtime/alloc.rs
1//! How many allocations a stretch of a blue program made, and how many bytes.
2//!
3//! `blue profile` reports time per definition; this is the other half of what
4//! a definition costs (`theory/BLUE-MEMORY.md` §VI M0: you cannot configure
5//! what you cannot account for). One counter, read at two points: a caller
6//! snapshots [`allocations`] when an activation enters and again when it
7//! leaves, and the difference is what that activation allocated, everything it
8//! called included. Subtracting what its callees allocated gives the
9//! exclusive figure, exactly as the profiler does with time.
10//!
11//! [`Counting`] is a `GlobalAlloc` over `System` that counts each allocation
12//! on the thread that made it. A binary opts in with one line:
13//!
14//! ```ignore
15//! #[global_allocator]
16//! static ALLOC: blue_lang_runtime::alloc::Counting = blue_lang_runtime::alloc::Counting;
17//! ```
18//!
19//! The count is per thread so a profile of one interpreter is not charged for
20//! another thread's work, and it costs two thread-local additions per
21//! allocation, nothing on a free. Until a binary installs [`Counting`],
22//! [`allocations`] answers `None`: a program that was never counted reads as
23//! unmeasured, never as having allocated nothing.
24
25use std::alloc::{GlobalAlloc, Layout, System};
26use std::cell::Cell;
27use std::sync::atomic::{AtomicBool, Ordering};
28
29/// Allocations made and bytes requested on one thread, from its start.
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct Allocations {
32 /// Allocation calls. A `realloc` that grows counts as one.
33 pub count: u64,
34 /// Bytes requested. A growing `realloc` counts the bytes it added.
35 pub bytes: u64,
36}
37
38impl Allocations {
39 /// What was allocated between `earlier` and `self`, on one thread.
40 #[must_use]
41 pub fn since(self, earlier: Self) -> Self {
42 Self {
43 count: self.count.saturating_sub(earlier.count),
44 bytes: self.bytes.saturating_sub(earlier.bytes),
45 }
46 }
47}
48
49impl std::ops::AddAssign for Allocations {
50 fn add_assign(&mut self, other: Self) {
51 self.count += other.count;
52 self.bytes += other.bytes;
53 }
54}
55
56/// Set by the first allocation [`Counting`] serves, so [`allocations`] can tell
57/// "not counted" from "counted zero".
58static INSTALLED: AtomicBool = AtomicBool::new(false);
59thread_local! {
60 // `const` initialisation: no lazy init, so no allocation from inside the
61 // allocator, and no destructor to run during thread teardown.
62 static COUNT: Cell<u64> = const { Cell::new(0) };
63 static BYTES: Cell<u64> = const { Cell::new(0) };
64}
65
66/// This thread's running totals, or `None` when [`Counting`] is not the global
67/// allocator of this binary.
68#[must_use]
69pub fn allocations() -> Option<Allocations> {
70 if !INSTALLED.load(Ordering::Relaxed) {
71 return None;
72 }
73 Some(Allocations {
74 count: COUNT.try_with(Cell::get).unwrap_or(0),
75 bytes: BYTES.try_with(Cell::get).unwrap_or(0),
76 })
77}
78
79fn charge(bytes: usize) {
80 // `try_with`: an allocation made while this thread's locals are being torn
81 // down is served and simply not counted.
82 let _ = COUNT.try_with(|c| c.set(c.get() + 1));
83 let _ = BYTES.try_with(|b| b.set(b.get() + bytes as u64));
84}
85
86/// The system allocator, counting what each thread allocates. See the module
87/// documentation for how a binary installs it.
88pub struct Counting;
89
90// SAFETY: every method forwards to `System` with the caller's arguments
91// unchanged; the counting touches only `Cell`s and an atomic flag, neither of
92// which allocates.
93unsafe impl GlobalAlloc for Counting {
94 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
95 INSTALLED.store(true, Ordering::Relaxed);
96 charge(layout.size());
97 unsafe { System.alloc(layout) }
98 }
99
100 unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
101 INSTALLED.store(true, Ordering::Relaxed);
102 charge(layout.size());
103 unsafe { System.alloc_zeroed(layout) }
104 }
105
106 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
107 unsafe { System.dealloc(ptr, layout) }
108 }
109
110 unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
111 if new_size > layout.size() {
112 charge(new_size - layout.size());
113 }
114 unsafe { System.realloc(ptr, layout, new_size) }
115 }
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121
122 /// This crate's own test binary does not install [`Counting`], so the
123 /// answer must be "unmeasured", not a zero that reads as "allocated
124 /// nothing". `tests/alloc_count.rs` installs it and measures.
125 #[test]
126 fn an_uncounted_binary_reads_as_unmeasured() {
127 let _ = vec![1u8; 64];
128 assert_eq!(allocations(), None);
129 }
130
131 #[test]
132 fn since_is_the_difference_and_never_underflows() {
133 let a = Allocations {
134 count: 3,
135 bytes: 100,
136 };
137 let b = Allocations {
138 count: 5,
139 bytes: 160,
140 };
141 assert_eq!(
142 b.since(a),
143 Allocations {
144 count: 2,
145 bytes: 60
146 }
147 );
148 assert_eq!(a.since(b), Allocations::default());
149 }
150}