Skip to main content

canic_memory/runtime/
mod.rs

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