ic_sqlite_vfs/stable/
memory.rs

1//! Byte-addressed SQLite memory wrapper.
2//!
3//! The crate stores SQLite inside a user-provided `VirtualMemory`, so it can
4//! coexist with other stable structures managed by the same MemoryManager.
5
6use crate::config::STABLE_PAGE_SIZE;
7use crate::stable::memory_manager::VirtualMemory;
8#[cfg(any(test, debug_assertions))]
9use crate::stable::memory_manager::{MemoryId, MemoryManager};
10use crate::stable::raw_memory::{DefaultMemoryImpl, Memory};
11use std::cell::{Cell, RefCell};
12#[cfg(any(test, feature = "canister-api-test-failpoints"))]
13use std::collections::BTreeMap;
14
15pub type DbMemory = VirtualMemory<DefaultMemoryImpl>;
16
17#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
18pub struct ContextId(u64);
19
20#[derive(Debug, thiserror::Error)]
21pub enum StableMemoryError {
22    #[error("stable memory backend is not initialized")]
23    NotInitialized,
24    #[error("stable memory backend is already initialized")]
25    AlreadyInitialized,
26    #[error(
27        "stable memory grow failed: current_pages={current_pages}, required_pages={required_pages}"
28    )]
29    GrowFailed {
30        current_pages: u64,
31        required_pages: u64,
32    },
33    #[error(
34        "stable memory read out of bounds: offset={offset}, len={len}, size_bytes={size_bytes}"
35    )]
36    ReadOutOfBounds {
37        offset: u64,
38        len: u64,
39        size_bytes: u64,
40    },
41    #[error("offset overflow")]
42    OffsetOverflow,
43    #[error("import session already started")]
44    ImportAlreadyStarted,
45    #[error("import session not started")]
46    ImportNotStarted,
47    #[error("database update already in progress")]
48    UpdateInProgress,
49    #[error("import chunk out of order: offset={offset}, expected={expected}")]
50    ImportOutOfOrder { offset: u64, expected: u64 },
51    #[error("import chunk out of bounds: offset={offset}, len={len}, db_size={db_size}")]
52    ImportOutOfBounds { offset: u64, len: u64, db_size: u64 },
53    #[error("import incomplete: written_until={written_until}, db_size={db_size}")]
54    ImportIncomplete { written_until: u64, db_size: u64 },
55    #[error("checksum mismatch: expected={expected}, actual={actual}")]
56    ChecksumMismatch { expected: u64, actual: u64 },
57    #[error("checksum refresh chunk size must be greater than zero")]
58    ChecksumRefreshChunkEmpty,
59    #[error("stable blob failpoint: {0}")]
60    Failpoint(&'static str),
61    #[error("superblock metadata checksum mismatch")]
62    MetaChecksumMismatch,
63    #[error("unsupported stable layout version: {0}")]
64    UnsupportedLayoutVersion(u64),
65}
66
67#[cfg(any(test, feature = "canister-api-test-failpoints"))]
68#[derive(Clone, Copy, Debug, Eq, PartialEq)]
69pub enum MemoryFailpoint {
70    GrowFailed { ordinal: u64 },
71    TrapAfterWrite { ordinal: u64 },
72}
73
74#[cfg(any(test, feature = "canister-api-test-failpoints"))]
75thread_local! {
76    static FAILPOINTS: RefCell<BTreeMap<ContextId, MemoryFailpointState>> = const { RefCell::new(BTreeMap::new()) };
77}
78
79thread_local! {
80    static NEXT_CONTEXT_ID: Cell<u64> = const { Cell::new(1) };
81    static DEFAULT_CONTEXT: Cell<Option<ContextId>> = const { Cell::new(None) };
82    static CURRENT_CONTEXT: Cell<Option<ContextId>> = const { Cell::new(None) };
83    static DB_MEMORY: RefCell<Vec<(ContextId, DbMemory)>> = const { RefCell::new(Vec::new()) };
84}
85
86pub fn init(memory: DbMemory) -> Result<ContextId, StableMemoryError> {
87    DEFAULT_CONTEXT.with(|default| {
88        if default.get().is_some() {
89            return Err(StableMemoryError::AlreadyInitialized);
90        }
91        let context = init_context(memory);
92        default.set(Some(context));
93        Ok(context)
94    })
95}
96
97pub fn init_context(memory: DbMemory) -> ContextId {
98    let context = NEXT_CONTEXT_ID.with(|next| {
99        let current = next.get();
100        let context = ContextId(current);
101        next.set(current.checked_add(1).expect("context id overflow"));
102        context
103    });
104    DB_MEMORY.with(|slot| {
105        slot.borrow_mut().push((context, memory));
106    });
107    context
108}
109
110pub fn is_initialized() -> bool {
111    DEFAULT_CONTEXT.with(|context| context.get().is_some())
112}
113
114pub fn default_context() -> Option<ContextId> {
115    DEFAULT_CONTEXT.with(Cell::get)
116}
117
118#[inline(always)]
119pub fn active_context_id() -> Result<ContextId, StableMemoryError> {
120    if let Some(context) = CURRENT_CONTEXT.with(Cell::get) {
121        return Ok(context);
122    }
123    default_context().ok_or(StableMemoryError::NotInitialized)
124}
125
126#[inline(always)]
127pub fn with_context<T>(context: ContextId, f: impl FnOnce() -> T) -> T {
128    let previous = CURRENT_CONTEXT.with(|current| {
129        let previous = current.get();
130        current.set(Some(context));
131        previous
132    });
133    let _guard = ContextGuard { previous };
134    f()
135}
136
137#[cfg(any(test, feature = "canister-api-test-failpoints"))]
138pub fn set_failpoint(failpoint: MemoryFailpoint) {
139    if let Ok(context) = active_context_id() {
140        FAILPOINTS.with(|slot| {
141            slot.borrow_mut().insert(
142                context,
143                MemoryFailpointState {
144                    failpoint: Some(failpoint),
145                    grow_count: 0,
146                    write_count: 0,
147                },
148            );
149        });
150    }
151}
152
153#[cfg(any(test, feature = "canister-api-test-failpoints"))]
154pub fn clear_failpoint() {
155    FAILPOINTS.with(|slot| slot.borrow_mut().clear());
156}
157
158pub fn size_pages() -> u64 {
159    with_memory(|memory| memory.size()).unwrap_or(0)
160}
161
162pub fn ensure_capacity(end_offset: u64) -> Result<(), StableMemoryError> {
163    with_memory(|memory| ensure_memory_capacity(memory, end_offset))?
164}
165
166pub fn read(offset: u64, dst: &mut [u8]) -> Result<(), StableMemoryError> {
167    if dst.is_empty() {
168        return Ok(());
169    }
170    let len = u64::try_from(dst.len()).map_err(|_| StableMemoryError::OffsetOverflow)?;
171    let end = offset
172        .checked_add(len)
173        .ok_or(StableMemoryError::OffsetOverflow)?;
174    with_memory(|memory| {
175        let size_bytes = memory
176            .size()
177            .checked_mul(STABLE_PAGE_SIZE)
178            .ok_or(StableMemoryError::OffsetOverflow)?;
179        if end > size_bytes {
180            return Err(StableMemoryError::ReadOutOfBounds {
181                offset,
182                len,
183                size_bytes,
184            });
185        }
186        memory.read(offset, dst);
187        Ok(())
188    })?
189}
190
191#[inline(always)]
192pub(crate) fn read_preallocated(offset: u64, dst: &mut [u8]) -> Result<(), StableMemoryError> {
193    checked_end(offset, dst.len())?;
194    with_memory(|memory| {
195        debug_assert_capacity(memory, offset, dst.len(), "read_preallocated");
196        memory.read(offset, dst);
197    })?;
198    Ok(())
199}
200
201pub fn write(offset: u64, bytes: &[u8]) -> Result<(), StableMemoryError> {
202    if bytes.is_empty() {
203        return Ok(());
204    }
205    let end = checked_end(offset, bytes.len())?;
206    with_memory(|memory| {
207        ensure_memory_capacity(memory, end)?;
208        memory.write(offset, bytes);
209        Ok(())
210    })??;
211
212    #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
213    crate::read_metrics::record_stable_data_write(bytes.len());
214
215    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
216    if hit_write_trap_failpoint() {
217        fail_after_stable_write();
218    }
219
220    Ok(())
221}
222
223pub(crate) fn write_preallocated(offset: u64, bytes: &[u8]) -> Result<(), StableMemoryError> {
224    if bytes.is_empty() {
225        return Ok(());
226    }
227    checked_end(offset, bytes.len())?;
228    with_memory(|memory| {
229        debug_assert_capacity(memory, offset, bytes.len(), "write_preallocated");
230        memory.write(offset, bytes);
231    })?;
232
233    #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
234    crate::read_metrics::record_stable_data_write(bytes.len());
235
236    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
237    if hit_write_trap_failpoint() {
238        fail_after_stable_write();
239    }
240
241    Ok(())
242}
243
244#[inline(always)]
245pub(crate) fn write_prechecked(offset: u64, bytes: &[u8]) -> Result<(), StableMemoryError> {
246    debug_assert!(checked_end(offset, bytes.len()).is_ok());
247    with_memory(|memory| {
248        debug_assert_capacity(memory, offset, bytes.len(), "write_prechecked");
249        memory.write(offset, bytes);
250    })?;
251
252    #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
253    crate::read_metrics::record_stable_data_write(bytes.len());
254
255    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
256    if hit_write_trap_failpoint() {
257        fail_after_stable_write();
258    }
259
260    Ok(())
261}
262
263#[inline(always)]
264pub(crate) fn write_prechecked_unmetered(
265    offset: u64,
266    bytes: &[u8],
267) -> Result<(), StableMemoryError> {
268    debug_assert!(checked_end(offset, bytes.len()).is_ok());
269    with_memory(|memory| {
270        debug_assert_capacity(memory, offset, bytes.len(), "write_prechecked_unmetered");
271        memory.write(offset, bytes);
272    })?;
273
274    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
275    if hit_write_trap_failpoint() {
276        fail_after_stable_write();
277    }
278
279    Ok(())
280}
281
282fn ensure_memory_capacity(memory: &DbMemory, end_offset: u64) -> Result<(), StableMemoryError> {
283    let current_pages = memory.size();
284    let current_bytes = current_pages
285        .checked_mul(STABLE_PAGE_SIZE)
286        .ok_or(StableMemoryError::OffsetOverflow)?;
287    if end_offset <= current_bytes {
288        return Ok(());
289    }
290
291    let missing = end_offset
292        .checked_sub(current_bytes)
293        .ok_or(StableMemoryError::OffsetOverflow)?;
294    let pages = missing.div_ceil(STABLE_PAGE_SIZE);
295
296    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
297    if hit_grow_failpoint() {
298        let required_pages = current_pages
299            .checked_add(pages)
300            .ok_or(StableMemoryError::OffsetOverflow)?;
301        return Err(StableMemoryError::GrowFailed {
302            current_pages,
303            required_pages,
304        });
305    }
306
307    let previous = memory.grow(pages);
308    if previous < 0 {
309        let required_pages = current_pages
310            .checked_add(pages)
311            .ok_or(StableMemoryError::OffsetOverflow)?;
312        return Err(StableMemoryError::GrowFailed {
313            current_pages,
314            required_pages,
315        });
316    }
317
318    #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
319    crate::read_metrics::record_stable_grow(pages);
320    Ok(())
321}
322
323#[inline(always)]
324fn debug_assert_capacity(memory: &DbMemory, offset: u64, len: usize, operation: &str) {
325    #[cfg(debug_assertions)]
326    {
327        let Ok(end) = checked_end(offset, len) else {
328            debug_assert!(false, "{operation} offset overflow");
329            return;
330        };
331        let Some(capacity) = memory.size().checked_mul(STABLE_PAGE_SIZE) else {
332            debug_assert!(false, "{operation} capacity overflow");
333            return;
334        };
335        debug_assert!(
336            end <= capacity,
337            "{operation} requires preallocated capacity: offset={offset}, len={len}, capacity={capacity}"
338        );
339    }
340    #[cfg(not(debug_assertions))]
341    {
342        let _ = (memory, offset, len, operation);
343    }
344}
345
346#[cfg(any(test, debug_assertions))]
347pub fn reset_for_tests() {
348    clear_initialization();
349    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
350    clear_failpoint();
351}
352
353#[cfg(any(test, debug_assertions))]
354pub fn set_next_context_id_for_tests(value: u64) {
355    NEXT_CONTEXT_ID.with(|next| next.set(value));
356}
357
358#[cfg(any(test, debug_assertions))]
359pub(crate) fn clear_initialization() {
360    DB_MEMORY.with(|memory| memory.borrow_mut().clear());
361    DEFAULT_CONTEXT.with(|context| context.set(None));
362    CURRENT_CONTEXT.with(|context| context.set(None));
363    NEXT_CONTEXT_ID.with(|next| next.set(1));
364    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
365    clear_failpoint();
366    crate::stable::meta::clear_superblock_cache();
367    crate::sqlite_vfs::stable_blob::invalidate_read_cache();
368}
369
370pub(crate) fn clear_failed_initialization(context: ContextId) {
371    DB_MEMORY.with(|memory| {
372        memory
373            .borrow_mut()
374            .retain(|(stored_context, _)| *stored_context != context);
375    });
376    DEFAULT_CONTEXT.with(|default| {
377        if default.get() == Some(context) {
378            default.set(None);
379        }
380    });
381    CURRENT_CONTEXT.with(|current| {
382        if current.get() == Some(context) {
383            current.set(None);
384        }
385    });
386    #[cfg(any(test, feature = "canister-api-test-failpoints"))]
387    FAILPOINTS.with(|slot| {
388        slot.borrow_mut().remove(&context);
389    });
390    crate::stable::meta::clear_superblock_cache();
391    crate::sqlite_vfs::stable_blob::invalidate_read_cache();
392}
393
394#[cfg(test)]
395pub fn snapshot_for_tests() -> Vec<u8> {
396    let len = usize::try_from(size_pages().saturating_mul(STABLE_PAGE_SIZE))
397        .expect("test memory size fits usize");
398    let mut out = vec![0_u8; len];
399    read(0, &mut out).expect("test memory snapshot succeeds");
400    out
401}
402
403#[cfg(test)]
404pub fn restore_for_tests(snapshot: Vec<u8>) -> DbMemory {
405    reset_for_tests();
406    let memory = memory_for_tests();
407    let pages = u64::try_from(snapshot.len())
408        .expect("snapshot len fits u64")
409        .div_ceil(STABLE_PAGE_SIZE);
410    if pages > 0 {
411        assert!(memory.grow(pages) >= 0, "snapshot memory grows");
412        memory.write(0, &snapshot);
413    }
414    crate::stable::meta::clear_superblock_cache();
415    memory
416}
417
418#[cfg(any(test, debug_assertions))]
419pub fn memory_for_tests() -> DbMemory {
420    MemoryManager::init(DefaultMemoryImpl::default()).get(MemoryId::new(42))
421}
422
423#[inline(always)]
424fn with_memory<T>(f: impl FnOnce(&DbMemory) -> T) -> Result<T, StableMemoryError> {
425    let context = active_context_id()?;
426    DB_MEMORY.with(|slot| {
427        let slot = slot.borrow();
428        if let Some((stored_context, memory)) = slot.first() {
429            if *stored_context == context {
430                return Ok(f(memory));
431            }
432        }
433        for (stored_context, memory) in slot.iter().skip(1) {
434            if *stored_context == context {
435                return Ok(f(memory));
436            }
437        }
438        Err(StableMemoryError::NotInitialized)
439    })
440}
441
442struct ContextGuard {
443    previous: Option<ContextId>,
444}
445
446impl Drop for ContextGuard {
447    fn drop(&mut self) {
448        CURRENT_CONTEXT.with(|current| current.set(self.previous));
449    }
450}
451
452#[cfg(any(test, feature = "canister-api-test-failpoints"))]
453#[derive(Clone, Copy, Debug)]
454struct MemoryFailpointState {
455    failpoint: Option<MemoryFailpoint>,
456    grow_count: u64,
457    write_count: u64,
458}
459
460fn checked_end(offset: u64, len: usize) -> Result<u64, StableMemoryError> {
461    let len = u64::try_from(len).map_err(|_| StableMemoryError::OffsetOverflow)?;
462    offset
463        .checked_add(len)
464        .ok_or(StableMemoryError::OffsetOverflow)
465}
466
467#[cfg(any(test, feature = "canister-api-test-failpoints"))]
468fn hit_grow_failpoint() -> bool {
469    let Ok(context) = active_context_id() else {
470        return false;
471    };
472    FAILPOINTS.with(|slot| {
473        let mut slot = slot.borrow_mut();
474        let Some(state) = slot.get_mut(&context) else {
475            return false;
476        };
477        state.grow_count += 1;
478        if state.failpoint
479            == Some(MemoryFailpoint::GrowFailed {
480                ordinal: state.grow_count,
481            })
482        {
483            state.failpoint = None;
484            true
485        } else {
486            false
487        }
488    })
489}
490
491#[cfg(any(test, feature = "canister-api-test-failpoints"))]
492fn hit_write_trap_failpoint() -> bool {
493    let Ok(context) = active_context_id() else {
494        return false;
495    };
496    FAILPOINTS.with(|slot| {
497        let mut slot = slot.borrow_mut();
498        let Some(state) = slot.get_mut(&context) else {
499            return false;
500        };
501        state.write_count += 1;
502        if state.failpoint
503            == Some(MemoryFailpoint::TrapAfterWrite {
504                ordinal: state.write_count,
505            })
506        {
507            state.failpoint = None;
508            true
509        } else {
510            false
511        }
512    })
513}
514
515#[cfg(all(target_arch = "wasm32", feature = "canister-api-test-failpoints"))]
516fn fail_after_stable_write() -> ! {
517    ic_cdk::trap("stable write failpoint");
518}
519
520#[cfg(all(
521    any(test, feature = "canister-api-test-failpoints"),
522    not(all(target_arch = "wasm32", feature = "canister-api-test-failpoints"))
523))]
524fn fail_after_stable_write() -> ! {
525    panic!("stable write failpoint");
526}
ic_sqlite_vfs/stable/memory.rs

ic_sqlite_vfs/stable/
memory.rs
