Skip to main content

canic_core/memory/runtime/
mod.rs

1use std::sync::Mutex;
2
3// -----------------------------------------------------------------------------
4// CANIC_EAGER_TLS
5// -----------------------------------------------------------------------------
6// Internal registry of "TLS touch" functions.
7//
8// Each function must be a plain `fn()` pointer (not a closure). When invoked,
9// the function must perform a `.with(|_| {})` on a thread_local! static.
10// This guarantees that the TLS slot is *initialized eagerly*, not lazily, so
11// stable memory pages or other backing buffers are allocated in a deterministic
12// order before any canister entry points are executed.
13//
14// These functions are registered by the `eager_static!` macro via
15// `defer_tls_initializer()`, and run once during process startup by
16// `init_eager_tls()`.
17// -----------------------------------------------------------------------------
18
19static CANIC_EAGER_TLS: Mutex<Vec<fn()>> = Mutex::new(Vec::new());
20#[cfg(any(test, debug_assertions))]
21static TEST_BOOTSTRAP_HOOK: Mutex<Option<fn()>> = Mutex::new(None);
22
23/// Run all deferred TLS initializers and clear the registry.
24///
25/// This drains the internal queue of initializer functions and invokes
26/// each *exactly once*. The use of `std::mem::take` ensures:
27///
28/// - the vector is fully emptied before we run any initializers
29/// - we drop the borrow before calling user code (prevents borrow panics)
30/// - functions cannot be re-run accidentally
31/// - reentrant modifications of the queue become visible *after* this call
32///
33/// This should be invoked before any IC canister lifecycle hooks (init, update,
34/// heartbeat, etc.) so that thread-local caches are in a fully-initialized state
35/// before the canister performs memory-dependent work.
36pub fn init_eager_tls() {
37    let funcs = {
38        let mut funcs = CANIC_EAGER_TLS.lock().expect("eager tls queue poisoned");
39        std::mem::take(&mut *funcs)
40    };
41
42    debug_assert!(
43        CANIC_EAGER_TLS
44            .lock()
45            .expect("eager tls queue poisoned")
46            .is_empty(),
47        "CANIC_EAGER_TLS was modified during init_eager_tls() execution"
48    );
49
50    for f in funcs {
51        f();
52    }
53}
54
55/// Return whether memory access is currently allowed during bootstrap.
56#[must_use]
57pub fn is_memory_bootstrap_ready() -> bool {
58    ic_memory::runtime::is_default_memory_manager_bootstrapped()
59}
60
61/// Panic if a stable-memory slot is touched before memory bootstrap is ready.
62pub fn assert_memory_bootstrap_ready(label: &str, id: u8) {
63    if is_memory_bootstrap_ready() {
64        return;
65    }
66
67    #[cfg(any(test, debug_assertions))]
68    {
69        run_test_bootstrap_hook();
70        if is_memory_bootstrap_ready() {
71            return;
72        }
73    }
74
75    panic!(
76        "stable memory slot '{label}' (id {id}) accessed before memory bootstrap; call ic_memory::bootstrap_default_memory_manager_with_policy(...) first"
77    );
78}
79
80/// Register a TLS initializer function for eager execution.
81///
82/// This is called by the `eager_static!` macro. The function pointer `f`
83/// must be a zero-argument function (`fn()`) that performs a `.with(|_| {})`
84/// on the thread-local static it is meant to initialize.
85pub fn defer_tls_initializer(f: fn()) {
86    CANIC_EAGER_TLS
87        .lock()
88        .expect("eager tls queue poisoned")
89        .push(f);
90}
91
92/// Install a test-only hook that can run the crate's normal memory bootstrap
93/// before host unit tests first touch macro-backed stable memory.
94#[cfg(any(test, debug_assertions))]
95pub fn install_test_bootstrap_hook(hook: fn()) {
96    *TEST_BOOTSTRAP_HOOK
97        .lock()
98        .expect("test bootstrap hook poisoned") = Some(hook);
99}
100
101/// Return whether a test bootstrap hook has been installed.
102#[cfg(any(test, debug_assertions))]
103#[must_use]
104pub fn has_test_bootstrap_hook() -> bool {
105    TEST_BOOTSTRAP_HOOK
106        .lock()
107        .expect("test bootstrap hook poisoned")
108        .is_some()
109}
110
111#[cfg(any(test, debug_assertions))]
112fn run_test_bootstrap_hook() {
113    let hook = *TEST_BOOTSTRAP_HOOK
114        .lock()
115        .expect("test bootstrap hook poisoned");
116    if let Some(hook) = hook {
117        hook();
118    }
119}
120
121///
122/// TESTS
123///
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use std::sync::{
129        Mutex,
130        atomic::{AtomicU32, Ordering},
131    };
132
133    static COUNT: AtomicU32 = AtomicU32::new(0);
134    static TEST_LOCK: Mutex<()> = Mutex::new(());
135
136    fn clear_test_queues() {
137        CANIC_EAGER_TLS
138            .lock()
139            .expect("eager tls queue poisoned")
140            .clear();
141    }
142
143    fn bump() {
144        COUNT.fetch_add(1, Ordering::SeqCst);
145    }
146
147    #[test]
148    fn init_eager_tls_runs_and_clears_queue() {
149        let _guard = TEST_LOCK.lock().expect("test lock poisoned");
150        clear_test_queues();
151        COUNT.store(0, Ordering::SeqCst);
152        CANIC_EAGER_TLS
153            .lock()
154            .expect("eager tls queue poisoned")
155            .push(bump);
156        init_eager_tls();
157        let first = COUNT.load(Ordering::SeqCst);
158        assert_eq!(first, 1);
159
160        // second call sees empty queue
161        init_eager_tls();
162        let second = COUNT.load(Ordering::SeqCst);
163        assert_eq!(second, 1);
164    }
165}