Skip to main content

canic_memory/runtime/
mod.rs

1pub mod registry;
2
3use std::cell::{Cell, RefCell};
4
5// -----------------------------------------------------------------------------
6// CANIC_EAGER_TLS
7// -----------------------------------------------------------------------------
8// Internal registry of "TLS touch" functions.
9//
10// Each function must be a plain `fn()` pointer (not a closure). When invoked,
11// the function must perform a `.with(|_| {})` on a thread_local! static.
12// This guarantees that the TLS slot is *initialized eagerly*, not lazily, so
13// stable memory pages or other backing buffers are allocated in a deterministic
14// order before any canister entry points are executed.
15//
16// These functions are registered by the `eager_static!` macro via
17// `defer_tls_initializer()`, and run once during process startup by
18// `init_eager_tls()`.
19// -----------------------------------------------------------------------------
20
21thread_local! {
22    static CANIC_EAGER_TLS: RefCell<Vec<fn()>> = const {
23        RefCell::new(Vec::new())
24    };
25    static CANIC_EAGER_TLS_RUNNING: Cell<bool> = const { Cell::new(false) };
26}
27
28/// Run all deferred TLS initializers and clear the registry.
29///
30/// This drains the internal queue of initializer functions and invokes
31/// each *exactly once*. The use of `std::mem::take` ensures:
32///
33/// - the vector is fully emptied before we run any initializers
34/// - we drop the borrow before calling user code (prevents borrow panics)
35/// - functions cannot be re-run accidentally
36/// - reentrant modifications of the queue become visible *after* this call
37///
38/// This should be invoked before any IC canister lifecycle hooks (init, update,
39/// heartbeat, etc.) so that thread-local caches are in a fully-initialized state
40/// before the canister performs memory-dependent work.
41pub fn init_eager_tls() {
42    ///
43    /// RunningGuard
44    ///
45
46    struct RunningGuard;
47
48    impl Drop for RunningGuard {
49        fn drop(&mut self) {
50            CANIC_EAGER_TLS_RUNNING.with(|running| running.set(false));
51        }
52    }
53
54    CANIC_EAGER_TLS_RUNNING.with(|running| running.set(true));
55    let _running_guard = RunningGuard;
56
57    let funcs = CANIC_EAGER_TLS.with(|v| {
58        let mut v = v.borrow_mut();
59        std::mem::take(&mut *v)
60    });
61
62    debug_assert!(
63        CANIC_EAGER_TLS.with(|v| v.borrow().is_empty()),
64        "CANIC_EAGER_TLS was modified during init_eager_tls() execution"
65    );
66
67    for f in funcs {
68        f();
69    }
70}
71
72#[must_use]
73pub fn is_eager_tls_initializing() -> bool {
74    CANIC_EAGER_TLS_RUNNING.with(Cell::get)
75}
76
77#[must_use]
78pub fn is_memory_bootstrap_ready() -> bool {
79    is_eager_tls_initializing() || registry::MemoryRegistryRuntime::is_initialized()
80}
81
82pub fn assert_memory_bootstrap_ready(label: &str, id: u8) {
83    if is_memory_bootstrap_ready() {
84        return;
85    }
86
87    panic!(
88        "stable memory slot '{label}' (id {id}) accessed before memory bootstrap; call init_eager_tls() and MemoryRegistryRuntime::init(...) first"
89    );
90}
91
92/// Register a TLS initializer function for eager execution.
93///
94/// This is called by the `eager_static!` macro. The function pointer `f`
95/// must be a zero-argument function (`fn()`) that performs a `.with(|_| {})`
96/// on the thread-local static it is meant to initialize.
97pub fn defer_tls_initializer(f: fn()) {
98    CANIC_EAGER_TLS.with_borrow_mut(|v| v.push(f));
99}
100
101///
102/// TESTS
103///
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use std::cell::Cell;
109
110    thread_local! {
111        static COUNT: Cell<u32> = const { Cell::new(0) };
112    }
113
114    fn bump() {
115        COUNT.with(|c| c.set(c.get() + 1));
116    }
117
118    #[test]
119    fn init_eager_tls_runs_and_clears_queue() {
120        COUNT.with(|c| c.set(0));
121        CANIC_EAGER_TLS.with(|v| v.borrow_mut().push(bump));
122        init_eager_tls();
123        let first = COUNT.with(Cell::get);
124        assert_eq!(first, 1);
125
126        // second call sees empty queue
127        init_eager_tls();
128        let second = COUNT.with(Cell::get);
129        assert_eq!(second, 1);
130    }
131}