1use crate::config::{SQLITE_PAGE_SIZE, STABLE_PAGE_SIZE, SUPERBLOCK_SIZE};
7use crate::sqlite_vfs::overlay::{self, Overlay};
8use crate::stable::memory::{self, ContextId, StableMemoryError};
9use crate::stable::meta::{
10 fnv1a64, Superblock, FLAG_CHECKSUM_REFRESHING, FLAG_CHECKSUM_STALE, FLAG_IMPORTING,
11 PAGE_MAP_LAYOUT_VERSION,
12};
13use std::cell::RefCell;
14use std::collections::BTreeMap;
15use std::mem::MaybeUninit;
16
17const CHECKSUM_CHUNK_LEN: u64 = 16 * 1024;
18const PAGE_TABLE_ENTRY_LEN: u64 = 8;
19const SEGMENT_PAGE_COUNT: u64 = 256;
20const SEGMENT_TABLE_BYTES: u64 = SEGMENT_PAGE_COUNT * PAGE_TABLE_ENTRY_LEN;
21const SINGLE_SEGMENT_PAGE_TABLE_BYTES: u64 = SEGMENT_TABLE_BYTES + PAGE_TABLE_ENTRY_LEN;
22const READ_SEGMENT_CACHE_CAPACITY: usize = 8;
23const FILE_PAGE_OFFSET_CACHE_CAPACITY: usize = 64;
24const FILE_PAGE_DATA_CACHE_CAPACITY: usize = 8;
25const COMPACT_MIN_ORPHAN_BYTES: u64 = 16 * 1024 * 1024;
26
27#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct ChecksumRefresh {
29 pub complete: bool,
30 pub checksum: u64,
31 pub scanned_bytes: u64,
32 pub db_size: u64,
33}
34
35#[derive(Clone, Debug, Eq, PartialEq)]
36pub struct StorageStats {
37 pub layout_version: u64,
38 pub page_count: u64,
39 pub page_table_bytes: u64,
40 pub active_bytes: u64,
41 pub allocated_bytes: u64,
42 pub orphan_bytes_estimate: u64,
43 pub orphan_ratio_basis_points: u64,
44 pub compact_recommended: bool,
45}
46
47#[derive(Clone, Copy, Debug, Eq, PartialEq)]
48pub(crate) enum StableBlobFailpoint {
49 OverlayWrite,
50 OverlayTruncate,
51 CommitCapacity,
52 CommitChunkWrite,
53 CommitPageTableWrite,
54 CommitSuperblockStore,
55}
56
57thread_local! {
58 #[cfg(test)]
59 static FAILPOINTS: RefCell<BTreeMap<ContextId, StableBlobFailpoint>> = const { RefCell::new(BTreeMap::new()) };
60 static READ_TABLE_CACHE: RefCell<Vec<(ContextId, ReadTableCache)>> = const { RefCell::new(Vec::new()) };
61 static COMMIT_SEGMENT_CACHE: RefCell<Vec<(ContextId, CommitSegmentCache)>> = const { RefCell::new(Vec::new()) };
62}
63
64#[derive(Clone, Copy, Debug, Eq, PartialEq)]
65struct ReadCacheKey {
66 page_table_offset: u64,
67 page_count: u64,
68 db_size: u64,
69 last_tx_id: u64,
70}
71
72#[derive(Debug)]
73struct ReadTableCache {
74 key: Option<ReadCacheKey>,
75 root: Vec<u64>,
76 segments: Vec<CachedSegment>,
77}
78
79#[derive(Debug)]
80struct CachedSegment {
81 segment_no: u64,
82 table: Vec<u64>,
83}
84
85#[derive(Debug)]
86struct CommitSegmentCache {
87 segment_no: u64,
88 segment_offset: u64,
89 table: Vec<u64>,
90}
91
92impl ReadTableCache {
93 fn new() -> Self {
94 Self {
95 key: None,
96 root: Vec::new(),
97 segments: Vec::new(),
98 }
99 }
100
101 fn clear(&mut self) {
102 self.key = None;
103 self.root.clear();
104 self.segments.clear();
105 }
106
107 fn ensure_key(&mut self, key: ReadCacheKey) {
108 if self.key == Some(key) {
109 return;
110 }
111 self.clear();
112 self.key = Some(key);
113 }
114
115 #[inline(always)]
116 fn segment_page_offset(&mut self, segment_no: u64, index: usize) -> Option<u64> {
117 if self.segments.is_empty() {
118 return None;
119 }
120 if self.segments.len() == 1 {
121 let segment = &self.segments[0];
122 if segment.segment_no == segment_no {
123 return Some(segment.table[index]);
124 }
125 return None;
126 }
127 let position = self
128 .segments
129 .iter()
130 .position(|segment| segment.segment_no == segment_no)?;
131 let offset = Some(self.segments[position].table[index]);
132 if position + 1 != self.segments.len() {
133 let segment = self.segments.remove(position);
134 self.segments.push(segment);
135 }
136 offset
137 }
138
139 fn insert_segment(&mut self, segment_no: u64, table: Vec<u64>) {
140 if let Some(position) = self
141 .segments
142 .iter()
143 .position(|segment| segment.segment_no == segment_no)
144 {
145 self.segments.remove(position);
146 }
147 self.segments.push(CachedSegment { segment_no, table });
148 while self.segments.len() > READ_SEGMENT_CACHE_CAPACITY {
149 self.segments.remove(0);
150 }
151 }
152}
153
154#[derive(Debug)]
155pub(crate) struct PageOffsetCache {
156 entries: Vec<(u64, u64)>,
157 pages: Vec<(u64, Vec<u8>)>,
158}
159
160impl PageOffsetCache {
161 pub(crate) fn new() -> Self {
162 Self {
163 entries: Vec::with_capacity(FILE_PAGE_OFFSET_CACHE_CAPACITY),
164 pages: Vec::new(),
165 }
166 }
167
168 fn get(&self, page_no: u64) -> Option<u64> {
169 match self.entries.as_slice() {
170 [] => None,
171 [(cached_page, physical)] => (*cached_page == page_no).then_some(*physical),
172 entries => {
173 for (cached_page, physical) in entries {
174 if *cached_page == page_no {
175 return Some(*physical);
176 }
177 }
178 None
179 }
180 }
181 }
182
183 fn insert(&mut self, page_no: u64, physical: u64) {
184 if self.entries.len() == FILE_PAGE_OFFSET_CACHE_CAPACITY {
185 self.entries.remove(0);
186 }
187 self.entries.push((page_no, physical));
188 }
189
190 #[inline(always)]
191 fn copy_page_slice(&self, page_no: u64, in_page: usize, dst: &mut [u8]) -> bool {
192 if self.pages.is_empty() {
193 return false;
194 }
195 if self.pages.len() == 1 {
196 let (cached_page, page) = &self.pages[0];
197 if *cached_page == page_no {
198 let end = in_page + dst.len();
199 dst.copy_from_slice(&page[in_page..end]);
200 return true;
201 }
202 return false;
203 }
204 for (cached_page, page) in &self.pages {
205 if *cached_page == page_no {
206 let end = in_page + dst.len();
207 dst.copy_from_slice(&page[in_page..end]);
208 return true;
209 }
210 }
211 false
212 }
213
214 fn insert_page(&mut self, page_no: u64, page: Vec<u8>) {
215 if self.pages.len() == FILE_PAGE_DATA_CACHE_CAPACITY {
216 self.pages.remove(0);
217 }
218 self.pages.push((page_no, page));
219 }
220}
221
222#[cfg(test)]
223pub(crate) fn set_failpoint(failpoint: StableBlobFailpoint) {
224 if let Ok(context) = memory::active_context_id() {
225 FAILPOINTS.with(|slot| {
226 slot.borrow_mut().insert(context, failpoint);
227 });
228 }
229}
230
231#[cfg(test)]
232pub(crate) fn clear_failpoint() {
233 FAILPOINTS.with(|slot| slot.borrow_mut().clear());
234}
235
236pub(crate) fn ensure_page_map_layout() -> Result<(), StableMemoryError> {
237 let block = Superblock::load()?;
238 if block.layout_version >= PAGE_MAP_LAYOUT_VERSION {
239 return Ok(());
240 }
241 Err(StableMemoryError::UnsupportedLayoutVersion(
242 block.layout_version,
243 ))
244}
245
246pub(crate) fn begin_update() -> Result<u64, StableMemoryError> {
247 let block = Superblock::load()?;
248 if block.layout_version < PAGE_MAP_LAYOUT_VERSION {
249 return Err(StableMemoryError::UnsupportedLayoutVersion(
250 block.layout_version,
251 ));
252 }
253 if block.is_importing() {
254 return Err(StableMemoryError::ImportAlreadyStarted);
255 }
256 overlay::begin(block.db_size)?;
257 Ok(block.db_size)
258}
259
260pub(crate) fn rollback_update() {
261 overlay::rollback();
262}
263
264#[doc(hidden)]
265pub fn invalidate_read_cache() {
266 READ_TABLE_CACHE.with(|cache| cache.borrow_mut().clear());
267 COMMIT_SEGMENT_CACHE.with(|cache| cache.borrow_mut().clear());
268}
269
270pub(crate) fn commit_update() -> Result<(), StableMemoryError> {
271 let Some(overlay) = overlay::take() else {
272 return Ok(());
273 };
274 if overlay.is_empty() {
275 return Ok(());
276 }
277 commit_overlay(overlay, true)
278}
279
280pub(crate) fn read_at(offset: u64, dst: &mut [u8]) -> Result<bool, StableMemoryError> {
281 if let Some(result) = overlay::read_at(offset, dst) {
282 return result;
283 }
284 read_base_at(offset, dst)
285}
286
287pub(crate) fn read_base_at(offset: u64, dst: &mut [u8]) -> Result<bool, StableMemoryError> {
288 if dst.is_empty() {
289 return Ok(true);
290 }
291 let block = Superblock::load()?;
292 read_base_at_with_block(&block, offset, dst)
293}
294
295pub(crate) fn read_base_at_with_block(
296 block: &Superblock,
297 offset: u64,
298 dst: &mut [u8],
299) -> Result<bool, StableMemoryError> {
300 if dst.is_empty() {
301 return Ok(true);
302 }
303 if offset >= block.db_size {
304 dst.fill(0);
305 return Ok(false);
306 }
307 let requested = u64::try_from(dst.len()).map_err(|_| StableMemoryError::OffsetOverflow)?;
308 if requested <= block.db_size - offset {
309 read_logical_range(block, offset, dst)?;
310 return Ok(true);
311 }
312 let copied = requested.min(block.db_size - offset);
313 let copied_len = usize::try_from(copied).map_err(|_| StableMemoryError::OffsetOverflow)?;
314 read_logical_range(block, offset, &mut dst[..copied_len])?;
315 dst[copied_len..].fill(0);
316 Ok(copied == requested)
317}
318
319#[inline(always)]
320pub(crate) fn read_base_at_with_page_cache(
321 block: &Superblock,
322 offset: u64,
323 dst: &mut [u8],
324 page_offsets: &mut PageOffsetCache,
325) -> Result<bool, StableMemoryError> {
326 if dst.is_empty() {
327 return Ok(true);
328 }
329 if offset >= block.db_size {
330 dst.fill(0);
331 return Ok(false);
332 }
333 let requested = u64::try_from(dst.len()).map_err(|_| StableMemoryError::OffsetOverflow)?;
334 if requested <= block.db_size - offset {
335 read_logical_range_with_page_cache(block, offset, dst, page_offsets)?;
336 return Ok(true);
337 }
338 let copied = requested.min(block.db_size - offset);
339 let copied_len = usize::try_from(copied).map_err(|_| StableMemoryError::OffsetOverflow)?;
340 read_logical_range_with_page_cache(block, offset, &mut dst[..copied_len], page_offsets)?;
341 dst[copied_len..].fill(0);
342 Ok(copied == requested)
343}
344
345pub(crate) fn read_base_page(page_no: u64) -> Result<Vec<u8>, StableMemoryError> {
346 let block = Superblock::load()?;
347 let mut page = zero_page();
348 if page_no >= active_page_count(&block)? {
349 return Ok(page);
350 }
351 let physical = page_offset_for(&block, page_no)?;
352 if physical != 0 {
353 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
354 crate::read_metrics::record_stable_data_read(page.len());
355 memory::read_preallocated(physical, &mut page)?;
356 }
357 Ok(page)
358}
359
360pub(crate) fn write_at(offset: u64, bytes: &[u8]) -> Result<(), StableMemoryError> {
361 if let Some(result) = overlay::write_at(offset, bytes) {
362 hit_failpoint(StableBlobFailpoint::OverlayWrite)?;
363 return result;
364 }
365 if bytes.is_empty() {
366 return Ok(());
367 }
368 ensure_page_map_layout()?;
369 let mut direct = Overlay::new(Superblock::load()?.db_size);
370 direct.write_at(offset, bytes)?;
371 commit_overlay(direct, false)
372}
373
374pub(crate) fn truncate(size: u64) -> Result<(), StableMemoryError> {
375 if let Some(result) = overlay::truncate(size) {
376 hit_failpoint(StableBlobFailpoint::OverlayTruncate)?;
377 return result;
378 }
379 ensure_page_map_layout()?;
380 let mut direct = Overlay::new(Superblock::load()?.db_size);
381 direct.truncate(size)?;
382 if direct.is_empty() {
383 return Ok(());
384 }
385 commit_overlay(direct, false)
386}
387
388pub(crate) fn file_size() -> Result<u64, StableMemoryError> {
389 if let Some(size) = overlay::file_size() {
390 return Ok(size);
391 }
392 Ok(Superblock::load()?.db_size)
393}
394
395pub fn export_chunk(offset: u64, len: u64) -> Result<Vec<u8>, StableMemoryError> {
396 reject_during_update()?;
397 let block = Superblock::load()?;
398 if offset >= block.db_size {
399 return Ok(Vec::new());
400 }
401 let copied = len.min(block.db_size - offset);
402 let copied_len = usize::try_from(copied).map_err(|_| StableMemoryError::OffsetOverflow)?;
403 let mut out = vec![0_u8; copied_len];
404 read_logical_range(&block, offset, &mut out)?;
405 Ok(out)
406}
407
408pub fn import_chunk(offset: u64, bytes: &[u8]) -> Result<(), StableMemoryError> {
409 reject_during_update()?;
410 let mut block = Superblock::load()?;
411 if !block.is_importing() {
412 return Err(StableMemoryError::ImportNotStarted);
413 }
414 let len = u64::try_from(bytes.len()).map_err(|_| StableMemoryError::OffsetOverflow)?;
415 if offset != block.import_written_until {
416 return Err(StableMemoryError::ImportOutOfOrder {
417 offset,
418 expected: block.import_written_until,
419 });
420 }
421 let end = checked_add(offset, len)?;
422 if end > block.import_total_size {
423 return Err(StableMemoryError::ImportOutOfBounds {
424 offset,
425 len,
426 db_size: block.import_total_size,
427 });
428 }
429 memory::write(import_offset(&block, offset)?, bytes)?;
430 block.import_written_until = end;
431 block.store()?;
432 invalidate_read_cache();
433 Ok(())
434}
435
436pub fn begin_import(total_size: u64, expected_checksum: u64) -> Result<(), StableMemoryError> {
437 reject_during_update()?;
438 let mut block = Superblock::load()?;
439 if block.is_importing() {
440 return Err(StableMemoryError::ImportAlreadyStarted);
441 }
442 let import_base_offset = append_base()?;
443 checked_add(import_base_offset, total_size)?;
444 block.flags |= FLAG_IMPORTING;
445 block.clear_checksum_refresh();
446 block.import_expected_checksum = expected_checksum;
447 block.import_written_until = 0;
448 block.import_total_size = total_size;
449 block.import_base_offset = import_base_offset;
450 block.store()?;
451 invalidate_read_cache();
452 Ok(())
453}
454
455pub fn finish_import() -> Result<(), StableMemoryError> {
456 reject_during_update()?;
457 let mut block = Superblock::load()?;
458 if !block.is_importing() {
459 return Err(StableMemoryError::ImportNotStarted);
460 }
461 if block.import_written_until != block.import_total_size {
462 return Err(StableMemoryError::ImportIncomplete {
463 written_until: block.import_written_until,
464 db_size: block.import_total_size,
465 });
466 }
467 let checksum = checksum_physical_range(block.import_base_offset, block.import_total_size)?;
468 if checksum != block.import_expected_checksum {
469 let expected = block.import_expected_checksum;
470 clear_import(&mut block)?;
471 return Err(StableMemoryError::ChecksumMismatch {
472 expected,
473 actual: checksum,
474 });
475 }
476 let entries = imported_page_table(&block)?;
477 let (root_offset, root_len) = write_segmented_tables(&entries)?;
478 block.db_size = block.import_total_size;
479 block.db_base_offset = block.import_base_offset;
480 block.page_table_offset = root_offset;
481 block.page_count = root_len;
482 block.layout_version = PAGE_MAP_LAYOUT_VERSION;
483 block.flags &= !FLAG_IMPORTING;
484 block.flags &= !FLAG_CHECKSUM_STALE;
485 block.clear_checksum_refresh();
486 block.checksum = checksum;
487 block.import_expected_checksum = 0;
488 block.import_written_until = 0;
489 block.import_total_size = 0;
490 block.import_base_offset = 0;
491 block.store()?;
492 invalidate_read_cache();
493 Ok(())
494}
495
496pub fn cancel_import() -> Result<(), StableMemoryError> {
497 reject_during_update()?;
498 let mut block = Superblock::load()?;
499 if !block.is_importing() {
500 return Err(StableMemoryError::ImportNotStarted);
501 }
502 clear_import(&mut block)
503}
504
505pub fn refresh_checksum() -> Result<u64, StableMemoryError> {
506 reject_during_update()?;
507 let checksum = checksum()?;
508 let mut block = Superblock::load()?;
509 block.checksum = checksum;
510 block.flags &= !FLAG_CHECKSUM_STALE;
511 block.clear_checksum_refresh();
512 block.store()?;
513 invalidate_read_cache();
514 Ok(checksum)
515}
516
517pub fn refresh_checksum_chunk(max_bytes: u64) -> Result<ChecksumRefresh, StableMemoryError> {
518 reject_during_update()?;
519 if max_bytes == 0 {
520 return Err(StableMemoryError::ChecksumRefreshChunkEmpty);
521 }
522
523 let mut block = Superblock::load()?;
524 if block.is_importing() {
525 return Err(StableMemoryError::ImportAlreadyStarted);
526 }
527 if !block.is_checksum_refreshing() {
528 block.flags |= FLAG_CHECKSUM_REFRESHING;
529 block.checksum_refresh_offset = 0;
530 block.checksum_refresh_hash = fnv1a64(&[]);
531 block.checksum_refresh_tx_id = block.last_tx_id;
532 }
533 if block.checksum_refresh_tx_id != block.last_tx_id {
534 block.clear_checksum_refresh();
535 block.store()?;
536 invalidate_read_cache();
537 return refresh_checksum_chunk(max_bytes);
538 }
539
540 let start = block.checksum_refresh_offset;
541 let end = block.db_size.min(start.saturating_add(max_bytes));
542 let mut offset = start;
543 let mut hash = block.checksum_refresh_hash;
544 while offset < end {
545 let len = (end - offset).min(CHECKSUM_CHUNK_LEN);
546 let copied_len = usize::try_from(len).map_err(|_| StableMemoryError::OffsetOverflow)?;
547 let mut bytes = vec![0_u8; copied_len];
548 read_logical_range(&block, offset, &mut bytes)?;
549 hash = fold_fnv1a64(hash, &bytes);
550 offset += len;
551 }
552
553 block.checksum_refresh_offset = offset;
554 block.checksum_refresh_hash = hash;
555 if offset == block.db_size {
556 block.checksum = hash;
557 block.flags &= !FLAG_CHECKSUM_STALE;
558 block.clear_checksum_refresh();
559 }
560 let out = ChecksumRefresh {
561 complete: offset == block.db_size,
562 checksum: hash,
563 scanned_bytes: offset,
564 db_size: block.db_size,
565 };
566 block.store()?;
567 invalidate_read_cache();
568 Ok(out)
569}
570
571pub fn checksum() -> Result<u64, StableMemoryError> {
572 reject_during_update()?;
573 let block = Superblock::load()?;
574 checksum_logical_range(&block, block.db_size)
575}
576
577pub fn compact() -> Result<(), StableMemoryError> {
578 reject_during_update()?;
579 ensure_page_map_layout()?;
580 let block = Superblock::load()?;
581 let table = read_page_table(&block)?;
582 let mut compacted = Vec::with_capacity(table.len());
583 let mut cursor = append_base()?;
584 let non_zero_pages = table.iter().filter(|offset| **offset != 0).count();
585 let data_bytes = u64::try_from(non_zero_pages)
586 .map_err(|_| StableMemoryError::OffsetOverflow)?
587 .checked_mul(page_size())
588 .ok_or(StableMemoryError::OffsetOverflow)?;
589 memory::ensure_capacity(checked_add(cursor, data_bytes)?)?;
590
591 for offset in table {
592 if offset == 0 {
593 compacted.push(0);
594 continue;
595 }
596 let mut page = zero_page();
597 memory::read_preallocated(offset, &mut page)?;
598 memory::write_preallocated(cursor, &page)?;
599 compacted.push(cursor);
600 cursor = checked_add(cursor, page_size())?;
601 }
602
603 let (root_offset, root_len) = write_segmented_tables(&compacted)?;
604 Superblock::store_page_map_without_tx(root_offset, root_len, block.db_size)?;
605 invalidate_read_cache();
606 Ok(())
607}
608
609pub fn storage_stats() -> Result<StorageStats, StableMemoryError> {
610 let block = Superblock::load()?;
611 let table = read_page_table(&block)?;
612 let non_zero_pages = u64::try_from(table.iter().filter(|offset| **offset != 0).count())
613 .map_err(|_| StableMemoryError::OffsetOverflow)?;
614 let segment_count = active_segment_count(&block)?;
615 let root_bytes = root_table_bytes(segment_count)?;
616 let segment_bytes = segment_count
617 .checked_mul(segment_table_bytes()?)
618 .ok_or(StableMemoryError::OffsetOverflow)?;
619 let page_table_bytes = checked_add(root_bytes, segment_bytes)?;
620 let active_bytes = SUPERBLOCK_SIZE
621 .checked_add(non_zero_pages.saturating_mul(page_size()))
622 .and_then(|value| value.checked_add(page_table_bytes))
623 .ok_or(StableMemoryError::OffsetOverflow)?;
624 let allocated_bytes = memory::size_pages()
625 .checked_mul(STABLE_PAGE_SIZE)
626 .ok_or(StableMemoryError::OffsetOverflow)?;
627 let orphan_bytes_estimate = allocated_bytes.saturating_sub(active_bytes);
628 let orphan_ratio_basis_points = orphan_bytes_estimate
629 .saturating_mul(10_000)
630 .checked_div(active_bytes)
631 .unwrap_or(0);
632 Ok(StorageStats {
633 layout_version: block.layout_version,
634 page_count: active_page_count(&block)?,
635 page_table_bytes,
636 active_bytes,
637 allocated_bytes,
638 orphan_bytes_estimate,
639 orphan_ratio_basis_points,
640 compact_recommended: orphan_bytes_estimate >= active_bytes
641 && orphan_bytes_estimate >= COMPACT_MIN_ORPHAN_BYTES,
642 })
643}
644
645pub(crate) fn page_count_for_size(size: u64) -> Result<u64, StableMemoryError> {
646 Ok(size.div_ceil(page_size()))
647}
648
649#[cfg(test)]
650pub(crate) fn debug_root_table_for_tests() -> Result<Vec<u64>, StableMemoryError> {
651 let block = Superblock::load()?;
652 read_root_table(&block)
653}
654
655fn commit_overlay(overlay: Overlay, advance_tx: bool) -> Result<(), StableMemoryError> {
656 hit_failpoint(StableBlobFailpoint::CommitCapacity)?;
657 let profile_enabled = commit_profile_enabled();
658 let block = Superblock::load()?;
659 let overlay_size = overlay.size();
660 let final_page_count = page_count_for_size(overlay_size)?;
661 let data_cursor = append_base()?;
662 debug_assert!(overlay
663 .dirty_pages()
664 .iter()
665 .all(|(page_no, _)| *page_no < final_page_count));
666 let dirty_pages = overlay.dirty_pages();
667 if let [(page_no, page)] = dirty_pages {
668 if overlay_size >= block.db_size
669 && *page_no < final_page_count
670 && final_page_count <= SEGMENT_PAGE_COUNT
671 {
672 let build_profile_start = commit_profile_start(profile_enabled);
673 let options = SinglePageCommitOptions {
674 advance_tx,
675 overlay_size,
676 data_cursor,
677 profile_enabled,
678 build_profile_start,
679 };
680 return commit_single_segment_page_overlay(&block, *page_no, page, options);
681 }
682 }
683
684 let final_segment_count = segment_count_for_pages(final_page_count)?;
685 let profile_start = commit_profile_start(profile_enabled);
686 let mut root = read_commit_root_table(&block)?;
687 commit_profile_record_load(profile_start);
688
689 let build_profile_start = commit_profile_start(profile_enabled);
690 let root_len =
691 usize::try_from(final_segment_count).map_err(|_| StableMemoryError::OffsetOverflow)?;
692 if root.len() != root_len {
693 root.resize(root_len, 0);
694 }
695
696 if let [(page_no, page)] = dirty_pages {
697 if overlay_size >= block.db_size && *page_no < final_page_count {
698 let options = SinglePageCommitOptions {
699 advance_tx,
700 overlay_size,
701 data_cursor,
702 profile_enabled,
703 build_profile_start,
704 };
705 return commit_single_page_overlay(
706 &block,
707 final_segment_count,
708 root,
709 *page_no,
710 page,
711 options,
712 );
713 }
714 }
715
716 let mut segment_updates = BTreeMap::<u64, Vec<u64>>::new();
717 let mut page_cursor = data_cursor;
718
719 for (page_no, _) in dirty_pages {
720 if *page_no >= final_page_count {
721 continue;
722 }
723 let segment_no = segment_no(*page_no);
724 let index = segment_index(*page_no)?;
725 let table = load_segment_for_update(&block, &root, &mut segment_updates, segment_no)?;
726 table[index] = page_cursor;
727 page_cursor = checked_add(page_cursor, page_size())?;
728 }
729
730 if overlay_size < block.db_size {
731 clear_truncated_tail(&block, &root, &mut segment_updates, final_page_count)?;
732 }
733 commit_profile_record_build_segments(build_profile_start);
734
735 let mut table_cursor = page_cursor;
736 let root_entries_len = final_segment_count;
737 let segment_table_writes = segment_updates.len();
738 let segment_table_bytes = u64::try_from(segment_table_writes)
739 .map_err(|_| StableMemoryError::OffsetOverflow)?
740 .checked_mul(segment_table_bytes()?)
741 .ok_or(StableMemoryError::OffsetOverflow)?;
742 let page_table_bytes = checked_add(segment_table_bytes, root_table_bytes(root_entries_len)?)?;
743 let profile_start = commit_profile_start(profile_enabled);
744 memory::ensure_capacity(checked_add(table_cursor, page_table_bytes)?)?;
745 commit_profile_record_capacity(profile_start);
746
747 let profile_start = commit_profile_start(profile_enabled);
748 let mut cursor = data_cursor;
749 for (_, page) in dirty_pages {
750 hit_failpoint(StableBlobFailpoint::CommitChunkWrite)?;
751 write_commit_page(cursor, page, profile_enabled)?;
752 cursor = checked_add(cursor, page_size())?;
753 }
754 commit_profile_record_page_write(profile_start);
755
756 hit_failpoint(StableBlobFailpoint::CommitPageTableWrite)?;
757 let profile_start = commit_profile_start(profile_enabled);
758 for (segment_no, table) in segment_updates {
759 let offset = write_commit_segment_table_at(&table, &mut table_cursor, profile_enabled)?;
760 let index = usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
761 root[index] = offset;
762 }
763 let root_offset = write_commit_root_table_at(&root, &mut table_cursor, profile_enabled)?;
764 commit_profile_record_table_write(profile_start);
765
766 hit_failpoint(StableBlobFailpoint::CommitSuperblockStore)?;
767 let profile_start = commit_profile_start(profile_enabled);
768 let result = store_commit_page_map(
769 advance_tx,
770 root_offset,
771 root_entries_len,
772 overlay_size,
773 profile_enabled,
774 );
775 commit_profile_record_superblock_store(profile_start);
776 result
777}
778
779#[derive(Clone, Copy)]
780struct SinglePageCommitOptions {
781 advance_tx: bool,
782 overlay_size: u64,
783 data_cursor: u64,
784 profile_enabled: bool,
785 build_profile_start: Option<u64>,
786}
787
788fn commit_single_page_overlay(
789 block: &Superblock,
790 final_segment_count: u64,
791 mut root: Vec<u64>,
792 page_no: u64,
793 page: &[u8],
794 options: SinglePageCommitOptions,
795) -> Result<(), StableMemoryError> {
796 let segment_no = segment_no(page_no);
797 let index = segment_index(page_no)?;
798 let mut table = read_commit_segment_table(block, &root, segment_no)?;
799 table[index] = options.data_cursor;
800 let page_cursor = checked_add(options.data_cursor, page_size())?;
801 commit_profile_record_build_segments(options.build_profile_start);
802
803 let root_entries_len = final_segment_count;
804 let page_table_bytes =
805 checked_add(segment_table_bytes()?, root_table_bytes(root_entries_len)?)?;
806 let profile_start = commit_profile_start(options.profile_enabled);
807 memory::ensure_capacity(checked_add(page_cursor, page_table_bytes)?)?;
808 commit_profile_record_capacity(profile_start);
809
810 hit_failpoint(StableBlobFailpoint::CommitChunkWrite)?;
811 let profile_start = commit_profile_start(options.profile_enabled);
812 write_commit_page(options.data_cursor, page, options.profile_enabled)?;
813 commit_profile_record_page_write(profile_start);
814
815 hit_failpoint(StableBlobFailpoint::CommitPageTableWrite)?;
816 let profile_start = commit_profile_start(options.profile_enabled);
817 let mut table_cursor = page_cursor;
818 let offset = write_commit_segment_table_at(&table, &mut table_cursor, options.profile_enabled)?;
819 let root_offset = if final_segment_count == 1 {
820 write_commit_root_table_at(&[offset], &mut table_cursor, options.profile_enabled)?
821 } else {
822 let root_index =
823 usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
824 root[root_index] = offset;
825 write_commit_root_table_at(&root, &mut table_cursor, options.profile_enabled)?
826 };
827 commit_profile_record_table_write(profile_start);
828
829 hit_failpoint(StableBlobFailpoint::CommitSuperblockStore)?;
830 let profile_start = commit_profile_start(options.profile_enabled);
831 let result = store_commit_page_map(
832 options.advance_tx,
833 root_offset,
834 root_entries_len,
835 options.overlay_size,
836 options.profile_enabled,
837 );
838 commit_profile_record_superblock_store(profile_start);
839 if result.is_ok() {
840 cache_commit_segment_table(segment_no, offset, table);
841 }
842 result
843}
844
845fn commit_single_segment_page_overlay(
846 block: &Superblock,
847 page_no: u64,
848 page: &[u8],
849 options: SinglePageCommitOptions,
850) -> Result<(), StableMemoryError> {
851 let index = segment_index(page_no)?;
852 let root = read_commit_root_table(block)?;
853 let mut table = read_commit_segment_table(block, &root, 0)?;
854 table[index] = options.data_cursor;
855 let page_cursor = checked_add(options.data_cursor, page_size())?;
856 commit_profile_record_build_segments(options.build_profile_start);
857
858 let profile_start = commit_profile_start(options.profile_enabled);
859 memory::ensure_capacity(checked_add(page_cursor, SINGLE_SEGMENT_PAGE_TABLE_BYTES)?)?;
860 commit_profile_record_capacity(profile_start);
861
862 hit_failpoint(StableBlobFailpoint::CommitChunkWrite)?;
863 let profile_start = commit_profile_start(options.profile_enabled);
864 memory::write_prechecked(options.data_cursor, page)?;
865 commit_profile_record_page_write(profile_start);
866
867 hit_failpoint(StableBlobFailpoint::CommitPageTableWrite)?;
868 let profile_start = commit_profile_start(options.profile_enabled);
869 let mut table_cursor = page_cursor;
870 let offset = write_commit_segment_table_at(&table, &mut table_cursor, options.profile_enabled)?;
871 let root_offset =
872 write_commit_root_table_at(&[offset], &mut table_cursor, options.profile_enabled)?;
873 commit_profile_record_table_write(profile_start);
874
875 hit_failpoint(StableBlobFailpoint::CommitSuperblockStore)?;
876 let profile_start = commit_profile_start(options.profile_enabled);
877 let result = store_commit_page_map(
878 options.advance_tx,
879 root_offset,
880 1,
881 options.overlay_size,
882 options.profile_enabled,
883 );
884 commit_profile_record_superblock_store(profile_start);
885 if result.is_ok() {
886 cache_commit_segment_table(0, offset, table);
887 }
888 result
889}
890
891#[cfg(any(test, debug_assertions, feature = "bench-profile"))]
892#[inline(always)]
893fn commit_profile_enabled() -> bool {
894 crate::read_metrics::metrics_enabled()
895}
896
897#[cfg(not(any(test, debug_assertions, feature = "bench-profile")))]
898#[inline(always)]
899fn commit_profile_enabled() -> bool {
900 false
901}
902
903#[cfg(any(test, debug_assertions, feature = "bench-profile"))]
904#[inline(always)]
905fn commit_profile_start(enabled: bool) -> Option<u64> {
906 if enabled {
907 Some(crate::read_metrics::instruction_counter())
908 } else {
909 None
910 }
911}
912
913#[cfg(not(any(test, debug_assertions, feature = "bench-profile")))]
914#[inline(always)]
915fn commit_profile_start(_enabled: bool) -> Option<u64> {
916 None
917}
918
919macro_rules! commit_profile_recorder {
920 ($name:ident, $record:ident) => {
921 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
922 #[inline(always)]
923 fn $name(start: Option<u64>) {
924 if let Some(start) = start {
925 crate::read_metrics::$record(
926 crate::read_metrics::instruction_counter().saturating_sub(start),
927 );
928 }
929 }
930
931 #[cfg(not(any(test, debug_assertions, feature = "bench-profile")))]
932 #[inline(always)]
933 fn $name(_start: Option<u64>) {}
934 };
935}
936
937commit_profile_recorder!(commit_profile_record_load, record_commit_load);
938commit_profile_recorder!(
939 commit_profile_record_build_segments,
940 record_commit_build_segments
941);
942commit_profile_recorder!(commit_profile_record_capacity, record_commit_capacity);
943commit_profile_recorder!(commit_profile_record_page_write, record_commit_page_write);
944commit_profile_recorder!(commit_profile_record_table_write, record_commit_table_write);
945commit_profile_recorder!(
946 commit_profile_record_superblock_store,
947 record_commit_superblock_store
948);
949
950#[inline(always)]
951fn write_commit_page(
952 offset: u64,
953 page: &[u8],
954 profile_enabled: bool,
955) -> Result<(), StableMemoryError> {
956 if profile_enabled {
957 memory::write_prechecked(offset, page)
958 } else {
959 memory::write_prechecked_unmetered(offset, page)
960 }
961}
962
963fn store_commit_page_map(
964 advance_tx: bool,
965 root_offset: u64,
966 root_entries_len: u64,
967 overlay_size: u64,
968 profile_enabled: bool,
969) -> Result<(), StableMemoryError> {
970 match (advance_tx, profile_enabled) {
971 (true, true) => Superblock::commit_page_map(root_offset, root_entries_len, overlay_size),
972 (true, false) => {
973 Superblock::commit_page_map_unmetered(root_offset, root_entries_len, overlay_size)
974 }
975 (false, true) => {
976 Superblock::store_page_map_without_tx(root_offset, root_entries_len, overlay_size)
977 }
978 (false, false) => Superblock::store_page_map_without_tx_unmetered(
979 root_offset,
980 root_entries_len,
981 overlay_size,
982 ),
983 }
984}
985
986fn load_segment_for_update<'a>(
987 block: &Superblock,
988 root: &[u64],
989 updates: &'a mut BTreeMap<u64, Vec<u64>>,
990 segment_no: u64,
991) -> Result<&'a mut Vec<u64>, StableMemoryError> {
992 match updates.entry(segment_no) {
993 std::collections::btree_map::Entry::Occupied(entry) => Ok(entry.into_mut()),
994 std::collections::btree_map::Entry::Vacant(entry) => {
995 let table = read_segment_table(block, root, segment_no)?;
996 Ok(entry.insert(table))
997 }
998 }
999}
1000
1001fn clear_truncated_tail(
1002 block: &Superblock,
1003 root: &[u64],
1004 updates: &mut BTreeMap<u64, Vec<u64>>,
1005 final_page_count: u64,
1006) -> Result<(), StableMemoryError> {
1007 let old_page_count = active_page_count(block)?;
1008 if final_page_count >= old_page_count || final_page_count == 0 {
1009 return Ok(());
1010 }
1011 let boundary_segment = segment_no(final_page_count);
1012 if boundary_segment >= segment_count_for_pages(final_page_count)? {
1013 return Ok(());
1014 }
1015 let start = segment_index(final_page_count)?;
1016 if start == 0 {
1017 return Ok(());
1018 }
1019 let table = load_segment_for_update(block, root, updates, boundary_segment)?;
1020 table[start..].fill(0);
1021 Ok(())
1022}
1023
1024fn reject_during_update() -> Result<(), StableMemoryError> {
1025 if overlay::is_active() {
1026 Err(StableMemoryError::UpdateInProgress)
1027 } else {
1028 Ok(())
1029 }
1030}
1031
1032fn read_logical_range(
1033 block: &Superblock,
1034 offset: u64,
1035 dst: &mut [u8],
1036) -> Result<(), StableMemoryError> {
1037 if dst.is_empty() {
1038 return Ok(());
1039 }
1040 let in_page =
1041 usize::try_from(offset % page_size()).map_err(|_| StableMemoryError::OffsetOverflow)?;
1042 if dst.len() <= page_len() - in_page {
1043 return read_logical_page_slice(block, offset / page_size(), in_page, dst);
1044 }
1045
1046 let mut copied_total = 0_usize;
1047 while copied_total < dst.len() {
1048 let absolute = checked_add(
1049 offset,
1050 u64::try_from(copied_total).map_err(|_| StableMemoryError::OffsetOverflow)?,
1051 )?;
1052 let page_no = absolute / page_size();
1053 let in_page = usize::try_from(absolute % page_size())
1054 .map_err(|_| StableMemoryError::OffsetOverflow)?;
1055 let copied = (page_len() - in_page).min(dst.len() - copied_total);
1056 read_logical_page_slice(
1057 block,
1058 page_no,
1059 in_page,
1060 &mut dst[copied_total..copied_total + copied],
1061 )?;
1062 copied_total += copied;
1063 }
1064 Ok(())
1065}
1066
1067fn read_logical_range_with_page_cache(
1068 block: &Superblock,
1069 offset: u64,
1070 dst: &mut [u8],
1071 page_offsets: &mut PageOffsetCache,
1072) -> Result<(), StableMemoryError> {
1073 let in_page =
1074 usize::try_from(offset % page_size()).map_err(|_| StableMemoryError::OffsetOverflow)?;
1075 if dst.len() <= page_len() - in_page {
1076 return read_logical_page_slice_with_page_cache(
1077 block,
1078 offset / page_size(),
1079 in_page,
1080 dst,
1081 page_offsets,
1082 );
1083 }
1084
1085 let mut copied_total = 0_usize;
1086 while copied_total < dst.len() {
1087 let absolute = checked_add(
1088 offset,
1089 u64::try_from(copied_total).map_err(|_| StableMemoryError::OffsetOverflow)?,
1090 )?;
1091 let page_no = absolute / page_size();
1092 let in_page = usize::try_from(absolute % page_size())
1093 .map_err(|_| StableMemoryError::OffsetOverflow)?;
1094 let copied = (page_len() - in_page).min(dst.len() - copied_total);
1095 read_logical_page_slice_with_page_cache(
1096 block,
1097 page_no,
1098 in_page,
1099 &mut dst[copied_total..copied_total + copied],
1100 page_offsets,
1101 )?;
1102 copied_total += copied;
1103 }
1104 Ok(())
1105}
1106
1107fn read_logical_page_slice(
1108 block: &Superblock,
1109 page_no: u64,
1110 in_page: usize,
1111 dst: &mut [u8],
1112) -> Result<(), StableMemoryError> {
1113 let physical = page_offset_for(block, page_no)?;
1114 if physical == 0 {
1115 dst.fill(0);
1116 return Ok(());
1117 }
1118 let stable_offset = checked_add(
1119 physical,
1120 u64::try_from(in_page).map_err(|_| StableMemoryError::OffsetOverflow)?,
1121 )?;
1122 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1123 crate::read_metrics::record_stable_data_read(dst.len());
1124 memory::read_preallocated(stable_offset, dst)
1125}
1126
1127#[inline(always)]
1128fn read_logical_page_slice_with_page_cache(
1129 block: &Superblock,
1130 page_no: u64,
1131 in_page: usize,
1132 dst: &mut [u8],
1133 page_offsets: &mut PageOffsetCache,
1134) -> Result<(), StableMemoryError> {
1135 if dst.len() < page_len() && page_offsets.copy_page_slice(page_no, in_page, dst) {
1136 return Ok(());
1137 }
1138 let physical = match page_offsets.get(page_no) {
1139 Some(physical) => physical,
1140 None => {
1141 let physical = if block.page_table_offset == 0 {
1142 0
1143 } else {
1144 cached_page_offset_for(block, page_no)?
1145 };
1146 page_offsets.insert(page_no, physical);
1147 physical
1148 }
1149 };
1150 if physical == 0 {
1151 dst.fill(0);
1152 return Ok(());
1153 }
1154 if in_page == 0 && dst.len() == page_len() {
1155 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1156 crate::read_metrics::record_stable_data_read(dst.len());
1157 return memory::read_preallocated(physical, dst);
1158 }
1159 if dst.len() < page_len() {
1160 let mut page = zero_page();
1161 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1162 crate::read_metrics::record_stable_data_read(page.len());
1163 memory::read_preallocated(physical, &mut page)?;
1164 let end = in_page + dst.len();
1165 dst.copy_from_slice(&page[in_page..end]);
1166 page_offsets.insert_page(page_no, page);
1167 return Ok(());
1168 }
1169 let stable_offset = checked_add(
1170 physical,
1171 u64::try_from(in_page).map_err(|_| StableMemoryError::OffsetOverflow)?,
1172 )?;
1173 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1174 crate::read_metrics::record_stable_data_read(dst.len());
1175 memory::read_preallocated(stable_offset, dst)
1176}
1177
1178fn page_offset_for(block: &Superblock, page_no: u64) -> Result<u64, StableMemoryError> {
1179 if page_no >= active_page_count(block)? || block.page_table_offset == 0 {
1180 return Ok(0);
1181 }
1182 cached_page_offset_for(block, page_no)
1183}
1184
1185fn read_page_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1186 let root = read_root_table(block)?;
1187 let count = active_page_count(block)?;
1188 let capacity = usize::try_from(count).map_err(|_| StableMemoryError::OffsetOverflow)?;
1189 let mut entries = Vec::with_capacity(capacity);
1190 for segment_no in 0..segment_count_for_pages(count)? {
1191 let table = read_segment_table(block, &root, segment_no)?;
1192 for entry in table {
1193 if entries.len() == capacity {
1194 break;
1195 }
1196 entries.push(entry);
1197 }
1198 }
1199 Ok(entries)
1200}
1201
1202fn cached_page_offset_for(block: &Superblock, page_no: u64) -> Result<u64, StableMemoryError> {
1203 let context = memory::active_context_id()?;
1204 let key = read_cache_key(block);
1205 let segment_no = segment_no(page_no);
1206 let index = segment_index(page_no)?;
1207 READ_TABLE_CACHE.with(|cache| {
1208 let mut caches = cache.borrow_mut();
1209 let cache = match read_table_cache_index(&caches, context) {
1210 Some(index) => &mut caches[index].1,
1211 None => {
1212 caches.push((context, ReadTableCache::new()));
1213 &mut caches
1214 .last_mut()
1215 .ok_or(StableMemoryError::OffsetOverflow)?
1216 .1
1217 }
1218 };
1219 cache.ensure_key(key);
1220 if cache.root.is_empty() {
1221 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1222 crate::read_metrics::record_page_table_root_miss();
1223 cache.root = read_root_table(block)?;
1224 } else {
1225 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1226 crate::read_metrics::record_page_table_root_hit();
1227 }
1228 let root_index =
1229 usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
1230 let segment_offset = cache.root[root_index];
1231 if segment_offset == 0 {
1232 return Ok(0);
1233 }
1234 if let Some(offset) = cache.segment_page_offset(segment_no, index) {
1235 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1236 crate::read_metrics::record_page_table_segment_hit();
1237 return Ok(offset);
1238 }
1239 #[cfg(any(test, debug_assertions, feature = "bench-profile"))]
1240 crate::read_metrics::record_page_table_segment_miss();
1241 let table = read_segment_table_at(segment_offset)?;
1242 let offset = table[index];
1243 cache.insert_segment(segment_no, table);
1244 Ok(offset)
1245 })
1246}
1247
1248fn read_table_cache_index(
1249 caches: &[(ContextId, ReadTableCache)],
1250 context: ContextId,
1251) -> Option<usize> {
1252 caches
1253 .iter()
1254 .position(|(stored_context, _)| *stored_context == context)
1255}
1256
1257fn read_root_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1258 if block.page_count == 0 {
1259 return Ok(Vec::new());
1260 }
1261 let entries_len =
1262 usize::try_from(block.page_count).map_err(|_| StableMemoryError::OffsetOverflow)?;
1263 read_u64_table_at(block.page_table_offset, entries_len)
1264}
1265
1266fn read_commit_root_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1267 read_root_table(block)
1268}
1269
1270fn read_segment_table(
1271 _block: &Superblock,
1272 root: &[u64],
1273 segment_no: u64,
1274) -> Result<Vec<u64>, StableMemoryError> {
1275 let index = usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
1276 let Some(offset) = root.get(index).copied() else {
1277 return Ok(vec![0_u64; segment_page_count_usize()]);
1278 };
1279 if offset == 0 {
1280 return Ok(vec![0_u64; segment_page_count_usize()]);
1281 }
1282 read_segment_table_at(offset)
1283}
1284
1285fn read_commit_segment_table(
1286 _block: &Superblock,
1287 root: &[u64],
1288 segment_no: u64,
1289) -> Result<Vec<u64>, StableMemoryError> {
1290 let index = usize::try_from(segment_no).map_err(|_| StableMemoryError::OffsetOverflow)?;
1291 let Some(offset) = root.get(index).copied() else {
1292 return Ok(vec![0_u64; segment_page_count_usize()]);
1293 };
1294 if offset == 0 {
1295 return Ok(vec![0_u64; segment_page_count_usize()]);
1296 }
1297 read_commit_segment_table_at(segment_no, offset)
1298}
1299
1300fn read_commit_segment_table_at(
1301 segment_no: u64,
1302 offset: u64,
1303) -> Result<Vec<u64>, StableMemoryError> {
1304 if offset == 0 {
1305 return Ok(vec![0_u64; segment_page_count_usize()]);
1306 }
1307 if let Some(table) = take_commit_segment_table(segment_no, offset) {
1308 return Ok(table);
1309 }
1310 read_segment_table_at(offset)
1311}
1312
1313fn take_commit_segment_table(segment_no: u64, segment_offset: u64) -> Option<Vec<u64>> {
1314 let Ok(context) = memory::active_context_id() else {
1315 return None;
1316 };
1317 COMMIT_SEGMENT_CACHE.with(|cache| {
1318 let mut cache = cache.borrow_mut();
1319 if cache.len() == 1 {
1320 let (stored_context, cached) = &cache[0];
1321 if *stored_context == context
1322 && cached.segment_no == segment_no
1323 && cached.segment_offset == segment_offset
1324 {
1325 return cache.pop().map(|(_, cached)| cached.table);
1326 }
1327 return None;
1328 }
1329 cache
1330 .iter()
1331 .position(|(stored_context, cached)| {
1332 *stored_context == context
1333 && cached.segment_no == segment_no
1334 && cached.segment_offset == segment_offset
1335 })
1336 .map(|position| cache.remove(position).1.table)
1337 })
1338}
1339
1340fn cache_commit_segment_table(segment_no: u64, segment_offset: u64, table: Vec<u64>) {
1341 let Ok(context) = memory::active_context_id() else {
1342 return;
1343 };
1344 COMMIT_SEGMENT_CACHE.with(|cache| {
1345 let mut cache = cache.borrow_mut();
1346 if cache.is_empty() {
1347 cache.push((
1348 context,
1349 CommitSegmentCache {
1350 segment_no,
1351 segment_offset,
1352 table,
1353 },
1354 ));
1355 return;
1356 }
1357 if cache.len() == 1 {
1358 let (stored_context, cached) = &mut cache[0];
1359 if *stored_context == context {
1360 cached.segment_no = segment_no;
1361 cached.segment_offset = segment_offset;
1362 cached.table = table;
1363 return;
1364 }
1365 } else if let Some((_, cached)) = cache
1366 .iter_mut()
1367 .find(|(stored_context, _)| *stored_context == context)
1368 {
1369 cached.segment_no = segment_no;
1370 cached.segment_offset = segment_offset;
1371 cached.table = table;
1372 return;
1373 }
1374 cache.push((
1375 context,
1376 CommitSegmentCache {
1377 segment_no,
1378 segment_offset,
1379 table,
1380 },
1381 ));
1382 });
1383}
1384
1385fn read_segment_table_at(offset: u64) -> Result<Vec<u64>, StableMemoryError> {
1386 read_u64_table_at(offset, segment_page_count_usize())
1387}
1388
1389fn write_segmented_tables(entries: &[u64]) -> Result<(u64, u64), StableMemoryError> {
1390 if entries.is_empty() {
1391 return Ok((0, 0));
1392 }
1393 let root_len = segment_count_for_pages(entries_len_u64(entries)?)?;
1394 let mut cursor = append_base()?;
1395 let segment_bytes = root_len
1396 .checked_mul(segment_table_bytes()?)
1397 .ok_or(StableMemoryError::OffsetOverflow)?;
1398 let page_table_bytes = checked_add(segment_bytes, root_table_bytes(root_len)?)?;
1399 memory::ensure_capacity(checked_add(cursor, page_table_bytes)?)?;
1400 let mut root = Vec::with_capacity(
1401 usize::try_from(root_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1402 );
1403 for segment_no in 0..root_len {
1404 let start = usize::try_from(
1405 segment_no
1406 .checked_mul(SEGMENT_PAGE_COUNT)
1407 .ok_or(StableMemoryError::OffsetOverflow)?,
1408 )
1409 .map_err(|_| StableMemoryError::OffsetOverflow)?;
1410 let mut table = vec![0_u64; segment_page_count_usize()];
1411 for (offset, entry) in entries[start..]
1412 .iter()
1413 .take(segment_page_count_usize())
1414 .enumerate()
1415 {
1416 table[offset] = *entry;
1417 }
1418 root.push(write_segment_table_at(&table, &mut cursor)?);
1419 }
1420 let root_offset = write_root_table_at(&root, &mut cursor)?;
1421 Ok((root_offset, entries_len_u64(&root)?))
1422}
1423
1424#[inline(always)]
1425fn write_segment_table_at(entries: &[u64], cursor: &mut u64) -> Result<u64, StableMemoryError> {
1426 if entries.len() == segment_page_count_usize() {
1427 return write_u64_table_at(entries, cursor);
1428 }
1429
1430 let mut table = vec![0_u64; segment_page_count_usize()];
1431 for (index, entry) in entries.iter().take(segment_page_count_usize()).enumerate() {
1432 table[index] = *entry;
1433 }
1434 write_u64_table_at(&table, cursor)
1435}
1436
1437fn write_root_table_at(entries: &[u64], cursor: &mut u64) -> Result<u64, StableMemoryError> {
1438 write_u64_table_at(entries, cursor)
1439}
1440
1441#[inline(always)]
1442fn write_commit_segment_table_at(
1443 entries: &[u64],
1444 cursor: &mut u64,
1445 profile_enabled: bool,
1446) -> Result<u64, StableMemoryError> {
1447 if profile_enabled {
1448 write_segment_table_at(entries, cursor)
1449 } else {
1450 write_segment_table_at_unmetered(entries, cursor)
1451 }
1452}
1453
1454#[inline(always)]
1455fn write_commit_root_table_at(
1456 entries: &[u64],
1457 cursor: &mut u64,
1458 profile_enabled: bool,
1459) -> Result<u64, StableMemoryError> {
1460 if profile_enabled {
1461 write_root_table_at(entries, cursor)
1462 } else {
1463 write_u64_table_at_unmetered(entries, cursor)
1464 }
1465}
1466
1467fn write_segment_table_at_unmetered(
1468 entries: &[u64],
1469 cursor: &mut u64,
1470) -> Result<u64, StableMemoryError> {
1471 if entries.len() == segment_page_count_usize() {
1472 return write_u64_table_at_unmetered(entries, cursor);
1473 }
1474
1475 let mut table = vec![0_u64; segment_page_count_usize()];
1476 for (index, entry) in entries.iter().take(segment_page_count_usize()).enumerate() {
1477 table[index] = *entry;
1478 }
1479 write_u64_table_at_unmetered(&table, cursor)
1480}
1481
1482fn write_u64_table_at(entries: &[u64], cursor: &mut u64) -> Result<u64, StableMemoryError> {
1483 if entries.is_empty() {
1484 return Ok(0);
1485 }
1486 let offset = *cursor;
1487 let byte_len = entries
1488 .len()
1489 .checked_mul(8)
1490 .ok_or(StableMemoryError::OffsetOverflow)?;
1491 #[cfg(target_endian = "little")]
1492 {
1493 let bytes = unsafe { std::slice::from_raw_parts(entries.as_ptr().cast::<u8>(), byte_len) };
1495 memory::write_prechecked(offset, bytes)?;
1496 *cursor = checked_add(
1497 offset,
1498 u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1499 )?;
1500 Ok(offset)
1501 }
1502
1503 #[cfg(not(target_endian = "little"))]
1504 {
1505 let mut bytes = vec![0_u8; byte_len];
1506 for (chunk, entry) in bytes.chunks_exact_mut(8).zip(entries) {
1507 chunk.copy_from_slice(&entry.to_le_bytes());
1508 }
1509 memory::write_prechecked(offset, &bytes)?;
1510 *cursor = checked_add(
1511 offset,
1512 u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1513 )?;
1514 Ok(offset)
1515 }
1516}
1517
1518fn read_u64_table_at(offset: u64, entries_len: usize) -> Result<Vec<u64>, StableMemoryError> {
1519 if entries_len == 0 {
1520 return Ok(Vec::new());
1521 }
1522 let byte_len = entries_len
1523 .checked_mul(8)
1524 .ok_or(StableMemoryError::OffsetOverflow)?;
1525 #[cfg(target_endian = "little")]
1526 {
1527 let mut entries = Vec::<MaybeUninit<u64>>::with_capacity(entries_len);
1528 unsafe {
1529 entries.set_len(entries_len);
1530 }
1531 let bytes =
1534 unsafe { std::slice::from_raw_parts_mut(entries.as_mut_ptr().cast::<u8>(), byte_len) };
1535 memory::read_preallocated(offset, bytes)?;
1536 let ptr = entries.as_mut_ptr().cast::<u64>();
1537 let len = entries.len();
1538 let capacity = entries.capacity();
1539 std::mem::forget(entries);
1540 unsafe { Ok(Vec::from_raw_parts(ptr, len, capacity)) }
1543 }
1544
1545 #[cfg(not(target_endian = "little"))]
1546 {
1547 let mut bytes = vec![0_u8; byte_len];
1548 memory::read_preallocated(offset, &mut bytes)?;
1549 decode_u64_table(&bytes)
1550 }
1551}
1552
1553fn write_u64_table_at_unmetered(
1554 entries: &[u64],
1555 cursor: &mut u64,
1556) -> Result<u64, StableMemoryError> {
1557 if entries.is_empty() {
1558 return Ok(0);
1559 }
1560 let offset = *cursor;
1561 let byte_len = entries
1562 .len()
1563 .checked_mul(8)
1564 .ok_or(StableMemoryError::OffsetOverflow)?;
1565 #[cfg(target_endian = "little")]
1566 {
1567 let bytes = unsafe { std::slice::from_raw_parts(entries.as_ptr().cast::<u8>(), byte_len) };
1569 memory::write_prechecked_unmetered(offset, bytes)?;
1570 *cursor = checked_add(
1571 offset,
1572 u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1573 )?;
1574 Ok(offset)
1575 }
1576
1577 #[cfg(not(target_endian = "little"))]
1578 {
1579 let mut bytes = vec![0_u8; byte_len];
1580 for (chunk, entry) in bytes.chunks_exact_mut(8).zip(entries) {
1581 chunk.copy_from_slice(&entry.to_le_bytes());
1582 }
1583 memory::write_prechecked_unmetered(offset, &bytes)?;
1584 *cursor = checked_add(
1585 offset,
1586 u64::try_from(byte_len).map_err(|_| StableMemoryError::OffsetOverflow)?,
1587 )?;
1588 Ok(offset)
1589 }
1590}
1591
1592#[cfg(not(target_endian = "little"))]
1593fn decode_u64_table(bytes: &[u8]) -> Result<Vec<u64>, StableMemoryError> {
1594 if !bytes.len().is_multiple_of(8) {
1595 return Err(StableMemoryError::OffsetOverflow);
1596 }
1597 let mut entries = Vec::with_capacity(bytes.len() / 8);
1598 for chunk in bytes.chunks_exact(8) {
1599 entries.push(u64::from_le_bytes([
1600 chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7],
1601 ]));
1602 }
1603 Ok(entries)
1604}
1605
1606fn imported_page_table(block: &Superblock) -> Result<Vec<u64>, StableMemoryError> {
1607 let count = page_count_for_size(block.import_total_size)?;
1608 let capacity = usize::try_from(count).map_err(|_| StableMemoryError::OffsetOverflow)?;
1609 let mut entries = Vec::with_capacity(capacity);
1610 for page_no in 0..count {
1611 entries.push(checked_add(
1612 block.import_base_offset,
1613 page_no
1614 .checked_mul(page_size())
1615 .ok_or(StableMemoryError::OffsetOverflow)?,
1616 )?);
1617 }
1618 Ok(entries)
1619}
1620
1621fn checksum_logical_range(block: &Superblock, len: u64) -> Result<u64, StableMemoryError> {
1622 let mut offset = 0_u64;
1623 let mut hash = fnv1a64(&[]);
1624 while offset < len {
1625 let chunk_len = (len - offset).min(CHECKSUM_CHUNK_LEN);
1626 let copied_len =
1627 usize::try_from(chunk_len).map_err(|_| StableMemoryError::OffsetOverflow)?;
1628 let mut bytes = vec![0_u8; copied_len];
1629 read_logical_range(block, offset, &mut bytes)?;
1630 hash = fold_fnv1a64(hash, &bytes);
1631 offset += chunk_len;
1632 }
1633 Ok(hash)
1634}
1635
1636fn checksum_physical_range(base_offset: u64, len: u64) -> Result<u64, StableMemoryError> {
1637 let mut offset = 0_u64;
1638 let mut hash = fnv1a64(&[]);
1639 while offset < len {
1640 let chunk_len = (len - offset).min(CHECKSUM_CHUNK_LEN);
1641 let copied_len =
1642 usize::try_from(chunk_len).map_err(|_| StableMemoryError::OffsetOverflow)?;
1643 let mut bytes = vec![0_u8; copied_len];
1644 memory::read_preallocated(checked_add(base_offset, offset)?, &mut bytes)?;
1645 hash = fold_fnv1a64(hash, &bytes);
1646 offset += chunk_len;
1647 }
1648 Ok(hash)
1649}
1650
1651fn clear_import(block: &mut Superblock) -> Result<(), StableMemoryError> {
1652 block.flags &= !FLAG_IMPORTING;
1653 block.import_expected_checksum = 0;
1654 block.import_written_until = 0;
1655 block.import_total_size = 0;
1656 block.import_base_offset = 0;
1657 block.store()?;
1658 invalidate_read_cache();
1659 Ok(())
1660}
1661
1662fn import_offset(block: &Superblock, offset: u64) -> Result<u64, StableMemoryError> {
1663 checked_add(block.import_base_offset, offset)
1664}
1665
1666fn active_page_count(block: &Superblock) -> Result<u64, StableMemoryError> {
1667 page_count_for_size(block.db_size)
1668}
1669
1670fn active_segment_count(block: &Superblock) -> Result<u64, StableMemoryError> {
1671 Ok(block.page_count)
1672}
1673
1674fn read_cache_key(block: &Superblock) -> ReadCacheKey {
1675 ReadCacheKey {
1676 page_table_offset: block.page_table_offset,
1677 page_count: block.page_count,
1678 db_size: block.db_size,
1679 last_tx_id: block.last_tx_id,
1680 }
1681}
1682
1683fn segment_count_for_pages(page_count: u64) -> Result<u64, StableMemoryError> {
1684 Ok(page_count.div_ceil(SEGMENT_PAGE_COUNT))
1685}
1686
1687fn segment_no(page_no: u64) -> u64 {
1688 page_no / SEGMENT_PAGE_COUNT
1689}
1690
1691fn segment_index(page_no: u64) -> Result<usize, StableMemoryError> {
1692 usize::try_from(page_no % SEGMENT_PAGE_COUNT).map_err(|_| StableMemoryError::OffsetOverflow)
1693}
1694
1695fn segment_page_count_usize() -> usize {
1696 usize::try_from(SEGMENT_PAGE_COUNT).expect("segment page count fits usize")
1697}
1698
1699fn segment_table_len() -> usize {
1700 segment_page_count_usize() * 8
1701}
1702
1703fn segment_table_bytes() -> Result<u64, StableMemoryError> {
1704 u64::try_from(segment_table_len()).map_err(|_| StableMemoryError::OffsetOverflow)
1705}
1706
1707fn root_table_bytes(entry_count: u64) -> Result<u64, StableMemoryError> {
1708 entry_count
1709 .checked_mul(PAGE_TABLE_ENTRY_LEN)
1710 .ok_or(StableMemoryError::OffsetOverflow)
1711}
1712
1713fn entries_len_u64<T>(entries: &[T]) -> Result<u64, StableMemoryError> {
1714 u64::try_from(entries.len()).map_err(|_| StableMemoryError::OffsetOverflow)
1715}
1716
1717fn append_base() -> Result<u64, StableMemoryError> {
1718 memory::size_pages()
1719 .checked_mul(STABLE_PAGE_SIZE)
1720 .ok_or(StableMemoryError::OffsetOverflow)
1721}
1722
1723fn page_size() -> u64 {
1724 u64::from(SQLITE_PAGE_SIZE)
1725}
1726
1727fn page_len() -> usize {
1728 usize::try_from(SQLITE_PAGE_SIZE).expect("SQLite page size fits usize")
1729}
1730
1731fn zero_page() -> Vec<u8> {
1732 vec![0_u8; page_len()]
1733}
1734
1735fn checked_add(left: u64, right: u64) -> Result<u64, StableMemoryError> {
1736 left.checked_add(right)
1737 .ok_or(StableMemoryError::OffsetOverflow)
1738}
1739
1740fn fold_fnv1a64(mut hash: u64, bytes: &[u8]) -> u64 {
1741 for byte in bytes {
1742 hash ^= u64::from(*byte);
1743 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
1744 }
1745 hash
1746}
1747
1748#[cfg(test)]
1749fn hit_failpoint(failpoint: StableBlobFailpoint) -> Result<(), StableMemoryError> {
1750 let Ok(context) = memory::active_context_id() else {
1751 return Ok(());
1752 };
1753 FAILPOINTS.with(|slot| {
1754 let mut slot = slot.borrow_mut();
1755 if slot.get(&context).copied() == Some(failpoint) {
1756 slot.remove(&context);
1757 Err(StableMemoryError::Failpoint(failpoint.name()))
1758 } else {
1759 Ok(())
1760 }
1761 })
1762}
1763
1764#[cfg(not(test))]
1765fn hit_failpoint(_failpoint: StableBlobFailpoint) -> Result<(), StableMemoryError> {
1766 Ok(())
1767}
1768
1769#[cfg(test)]
1770impl StableBlobFailpoint {
1771 fn name(self) -> &'static str {
1772 match self {
1773 Self::OverlayWrite => "before overlay write",
1774 Self::OverlayTruncate => "before overlay truncate",
1775 Self::CommitCapacity => "before commit capacity",
1776 Self::CommitChunkWrite => "before commit page write",
1777 Self::CommitPageTableWrite => "before commit page table write",
1778 Self::CommitSuperblockStore => "before commit superblock store",
1779 }
1780 }
1781}
1782
1783#[cfg(test)]
1784mod tests {
1785 use super::*;
1786 use proptest::prelude::*;
1787 use proptest::test_runner::{Config, TestRunner};
1788 use std::collections::BTreeSet;
1789
1790 #[test]
1791 fn layout_math_matches_expected_boundaries() {
1792 assert_eq!(page_count_for_size(0).unwrap(), 0);
1793 assert_eq!(page_count_for_size(1).unwrap(), 1);
1794 assert_eq!(page_count_for_size(page_size()).unwrap(), 1);
1795 assert_eq!(page_count_for_size(page_size() + 1).unwrap(), 2);
1796
1797 assert_eq!(segment_count_for_pages(0).unwrap(), 0);
1798 assert_eq!(segment_count_for_pages(1).unwrap(), 1);
1799 assert_eq!(segment_count_for_pages(SEGMENT_PAGE_COUNT).unwrap(), 1);
1800 assert_eq!(segment_count_for_pages(SEGMENT_PAGE_COUNT + 1).unwrap(), 2);
1801
1802 assert_eq!(segment_no(SEGMENT_PAGE_COUNT), 1);
1803 assert_eq!(segment_index(SEGMENT_PAGE_COUNT - 1).unwrap(), 255);
1804 assert_eq!(segment_index(SEGMENT_PAGE_COUNT).unwrap(), 0);
1805 assert_eq!(root_table_bytes(2).unwrap(), 16);
1806 }
1807
1808 #[test]
1809 fn layout_math_rejects_u64_max_overflow_boundaries() {
1810 assert!(matches!(
1811 root_table_bytes(u64::MAX),
1812 Err(StableMemoryError::OffsetOverflow)
1813 ));
1814 assert!(matches!(
1815 checked_add(u64::MAX, 1),
1816 Err(StableMemoryError::OffsetOverflow)
1817 ));
1818
1819 let mut block = Superblock::fresh();
1820 block.import_base_offset = u64::MAX;
1821 assert!(matches!(
1822 import_offset(&block, 1),
1823 Err(StableMemoryError::OffsetOverflow)
1824 ));
1825
1826 block.import_base_offset = u64::MAX - page_size() + 1;
1827 block.import_total_size = page_size() + 1;
1828 assert!(matches!(
1829 imported_page_table(&block),
1830 Err(StableMemoryError::OffsetOverflow)
1831 ));
1832 }
1833
1834 #[test]
1835 fn pbt_layout_math_matches_verus_model() {
1836 let mut runner = TestRunner::new(Config {
1837 cases: 512,
1838 ..Config::default()
1839 });
1840
1841 runner
1842 .run(
1843 &(
1844 boundary_size_strategy(),
1845 boundary_page_strategy(),
1846 boundary_entry_strategy(),
1847 ),
1848 |(size, page_no, entries)| {
1849 let page_count = page_count_for_size(size).unwrap();
1850 let page_size = u128::from(page_size());
1851 if size == 0 {
1852 prop_assert_eq!(page_count, 0);
1853 } else {
1854 prop_assert!(u128::from(page_count - 1) * page_size < u128::from(size));
1855 prop_assert!(u128::from(size) <= u128::from(page_count) * page_size);
1856 }
1857
1858 let segment_count = segment_count_for_pages(page_count).unwrap();
1859 if page_count == 0 {
1860 prop_assert_eq!(segment_count, 0);
1861 } else {
1862 prop_assert!(
1863 u128::from(segment_count - 1) * u128::from(SEGMENT_PAGE_COUNT)
1864 < u128::from(page_count)
1865 );
1866 prop_assert!(
1867 u128::from(page_count)
1868 <= u128::from(segment_count) * u128::from(SEGMENT_PAGE_COUNT)
1869 );
1870 }
1871
1872 let index = segment_index(page_no).unwrap();
1873 prop_assert!(index < segment_page_count_usize());
1874 prop_assert_eq!(
1875 u128::from(segment_no(page_no)) * u128::from(SEGMENT_PAGE_COUNT)
1876 + index as u128,
1877 u128::from(page_no)
1878 );
1879
1880 match root_table_bytes(entries) {
1881 Ok(bytes) => prop_assert_eq!(bytes, entries * PAGE_TABLE_ENTRY_LEN),
1882 Err(StableMemoryError::OffsetOverflow) => {
1883 prop_assert!(entries.checked_mul(PAGE_TABLE_ENTRY_LEN).is_none());
1884 }
1885 Err(error) => return Err(TestCaseError::fail(error.to_string())),
1886 }
1887 Ok(())
1888 },
1889 )
1890 .unwrap();
1891 }
1892
1893 fn boundary_size_strategy() -> impl Strategy<Value = u64> {
1894 let page = page_size();
1895 let segment_bytes = SEGMENT_PAGE_COUNT * page;
1896 prop_oneof![
1897 any::<u64>(),
1898 prop::sample::select(boundary_values(&[
1899 0,
1900 1,
1901 page - 1,
1902 page,
1903 page + 1,
1904 segment_bytes - 1,
1905 segment_bytes,
1906 segment_bytes + 1,
1907 u64::MAX,
1908 ])),
1909 ]
1910 }
1911
1912 fn boundary_page_strategy() -> impl Strategy<Value = u64> {
1913 prop_oneof![
1914 any::<u64>(),
1915 prop::sample::select(boundary_values(&[
1916 0,
1917 1,
1918 SEGMENT_PAGE_COUNT - 1,
1919 SEGMENT_PAGE_COUNT,
1920 SEGMENT_PAGE_COUNT + 1,
1921 u64::MAX,
1922 ])),
1923 ]
1924 }
1925
1926 fn boundary_entry_strategy() -> impl Strategy<Value = u64> {
1927 let max_without_overflow = u64::MAX / PAGE_TABLE_ENTRY_LEN;
1928 prop_oneof![
1929 any::<u64>(),
1930 prop::sample::select(boundary_values(&[
1931 0,
1932 1,
1933 SEGMENT_PAGE_COUNT - 1,
1934 SEGMENT_PAGE_COUNT,
1935 SEGMENT_PAGE_COUNT + 1,
1936 max_without_overflow - 1,
1937 max_without_overflow,
1938 max_without_overflow + 1,
1939 u64::MAX - 1,
1940 u64::MAX,
1941 ])),
1942 ]
1943 }
1944
1945 fn boundary_values(values: &[u64]) -> Vec<u64> {
1946 values
1947 .iter()
1948 .flat_map(|value| [value.saturating_sub(1), *value, value.saturating_add(1)])
1949 .collect()
1950 }
1951
1952 #[test]
1953 fn fnv_fold_matches_one_pass_for_multiple_partitions() {
1954 let bytes: Vec<u8> = (0..97)
1955 .map(|index| (index as u8).wrapping_mul(37).wrapping_add(11))
1956 .collect();
1957 let expected = fnv1a64(&bytes);
1958
1959 for split in [0_usize, 1, 2, 7, 31, 64, bytes.len()] {
1960 let split = split.min(bytes.len());
1961 let mut hash = fnv1a64(&[]);
1962 hash = fold_fnv1a64(hash, &bytes[..split]);
1963 hash = fold_fnv1a64(hash, &bytes[split..]);
1964 assert_eq!(hash, expected);
1965 }
1966
1967 let mut hash = fnv1a64(&[]);
1968 for chunk in bytes.chunks(13) {
1969 hash = fold_fnv1a64(hash, chunk);
1970 }
1971 assert_eq!(hash, expected);
1972 }
1973
1974 #[test]
1975 #[serial_test::serial]
1976 fn page_map_commit_tracks_dirty_page_offsets() {
1977 crate::stable::memory::reset_for_tests();
1978 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
1979 invalidate_read_cache();
1980
1981 let page_zero = vec![1_u8; page_len()];
1982 let page_later = vec![2_u8; page_len()];
1983 let later_page_no = SEGMENT_PAGE_COUNT + 1;
1984 write_at(0, &page_zero).unwrap();
1985 write_at(later_page_no * page_size(), &page_later).unwrap();
1986
1987 let block = Superblock::load().unwrap();
1988 let root = read_root_table(&block).unwrap();
1989 let table = read_page_table(&block).unwrap();
1990 let expected_pages = active_page_count(&block).unwrap();
1991 let expected_segments = segment_count_for_pages(expected_pages).unwrap();
1992
1993 assert_eq!(root.len() as u64, expected_segments);
1994 assert_eq!(table.len() as u64, expected_pages);
1995 assert_ne!(table[0], 0);
1996 assert_ne!(table[later_page_no as usize], 0);
1997
1998 let old_page_zero_offset = table[0];
1999 let updated_page_zero = vec![3_u8; page_len()];
2000 write_at(0, &updated_page_zero).unwrap();
2001 let updated_table = read_page_table(&Superblock::load().unwrap()).unwrap();
2002 let mut out = vec![0_u8; page_len()];
2003 read_base_at(0, &mut out).unwrap();
2004
2005 assert_ne!(updated_table[0], old_page_zero_offset);
2006 assert_eq!(out, updated_page_zero);
2007 }
2008
2009 #[test]
2010 #[serial_test::serial]
2011 fn page_map_commit_tracks_multi_segment_dirty_and_clean_pages() {
2012 crate::stable::memory::reset_for_tests();
2013 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2014 invalidate_read_cache();
2015
2016 let clean_page_no = 1;
2017 let later_page_no = SEGMENT_PAGE_COUNT + 1;
2018 write_at(0, &vec![1_u8; page_len()]).unwrap();
2019 write_at(clean_page_no * page_size(), &vec![2_u8; page_len()]).unwrap();
2020 write_at(later_page_no * page_size(), &vec![3_u8; page_len()]).unwrap();
2021
2022 let before = Superblock::load().unwrap();
2023 let before_root = read_root_table(&before).unwrap();
2024 let before_table = read_page_table(&before).unwrap();
2025
2026 begin_update().unwrap();
2027 write_at(0, &vec![4_u8; page_len()]).unwrap();
2028 write_at(later_page_no * page_size(), &vec![5_u8; page_len()]).unwrap();
2029 commit_update().unwrap();
2030
2031 let after = Superblock::load().unwrap();
2032 let after_root = read_root_table(&after).unwrap();
2033 let after_table = read_page_table(&after).unwrap();
2034
2035 assert_eq!(after_root.len(), after.page_count as usize);
2036 assert_eq!(after_root.len(), before_root.len());
2037 assert_ne!(after_table[0], before_table[0]);
2038 assert_eq!(
2039 after_table[clean_page_no as usize],
2040 before_table[clean_page_no as usize]
2041 );
2042 assert_ne!(
2043 after_table[later_page_no as usize],
2044 before_table[later_page_no as usize]
2045 );
2046 }
2047
2048 #[test]
2049 #[serial_test::serial]
2050 fn page_map_commit_zeroes_truncated_tail_slots() {
2051 crate::stable::memory::reset_for_tests();
2052 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2053 invalidate_read_cache();
2054
2055 write_at(0, &vec![1_u8; page_len()]).unwrap();
2056 write_at(page_size(), &vec![2_u8; page_len()]).unwrap();
2057 write_at(2 * page_size(), &vec![3_u8; page_len()]).unwrap();
2058 truncate(page_size()).unwrap();
2059
2060 let block = Superblock::load().unwrap();
2061 let root = read_root_table(&block).unwrap();
2062 let segment = read_segment_table(&block, &root, 0).unwrap();
2063
2064 assert_eq!(block.db_size, page_size());
2065 assert!(segment[0] != 0);
2066 assert_eq!(segment[1], 0);
2067 assert_eq!(segment[2], 0);
2068 }
2069
2070 #[test]
2071 #[serial_test::serial]
2072 fn compact_keeps_zero_pages_and_densifies_offsets_across_segments() {
2073 crate::stable::memory::reset_for_tests();
2074 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2075 invalidate_read_cache();
2076
2077 let later_page_no = SEGMENT_PAGE_COUNT + 2;
2078 let first_page = vec![7_u8; page_len()];
2079 let later_page = vec![9_u8; page_len()];
2080 write_at(0, &first_page).unwrap();
2081 write_at(later_page_no * page_size(), &later_page).unwrap();
2082
2083 compact().unwrap();
2084
2085 let block = Superblock::load().unwrap();
2086 let root = read_root_table(&block).unwrap();
2087 let table = read_page_table(&block).unwrap();
2088 let mut first_out = vec![0_u8; page_len()];
2089 let mut later_out = vec![0_u8; page_len()];
2090
2091 read_base_at(0, &mut first_out).unwrap();
2092 read_base_at(later_page_no * page_size(), &mut later_out).unwrap();
2093
2094 assert_eq!(root.len() as u64, block.page_count);
2095 assert_eq!(table.len() as u64, active_page_count(&block).unwrap());
2096 assert_ne!(table[0], 0);
2097 assert_eq!(table[1], 0);
2098 assert_eq!(table[later_page_no as usize], table[0] + page_size());
2099 assert_eq!(first_out, first_page);
2100 assert_eq!(later_out, later_page);
2101 }
2102
2103 #[test]
2104 #[serial_test::serial]
2105 fn single_segment_fast_path_preserves_table_after_expand_only_commit() {
2106 crate::stable::memory::reset_for_tests();
2107 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2108 invalidate_read_cache();
2109
2110 write_at(0, &[0]).unwrap();
2111 truncate(page_size() * 4).unwrap();
2112 truncate(page_size() * 4 + 1).unwrap();
2113
2114 let block = Superblock::load().unwrap();
2115 let table = read_page_table(&block).unwrap();
2116 let mut first = [1_u8; 1];
2117 let mut expanded_tail = [1_u8; 1];
2118
2119 read_base_at(0, &mut first).unwrap();
2120 read_base_at(page_size() * 4, &mut expanded_tail).unwrap();
2121
2122 assert_eq!(block.db_size, page_size() * 4 + 1);
2123 assert_ne!(table[0], 0);
2124 assert_eq!(table[1], 0);
2125 assert_ne!(table[4], 0);
2126 assert_eq!(first, [0]);
2127 assert_eq!(expanded_tail, [0]);
2128 }
2129
2130 #[test]
2131 #[serial_test::serial]
2132 fn page_table_u64_encoding_is_little_endian_and_round_trips() {
2133 crate::stable::memory::reset_for_tests();
2134 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2135 invalidate_read_cache();
2136
2137 let entries = [
2138 0_u64,
2139 1,
2140 0x0102_0304_0506_0708,
2141 0xf1f2_f3f4_f5f6_f7f8,
2142 u64::MAX,
2143 ];
2144 let mut cursor = 128_u64;
2145 let expected_len = u64::try_from(entries.len() * 8).unwrap();
2146 crate::stable::memory::ensure_capacity(cursor + expected_len).unwrap();
2147
2148 let offset = write_u64_table_at(&entries, &mut cursor).unwrap();
2149 let decoded = read_u64_table_at(offset, entries.len()).unwrap();
2150 let mut encoded = vec![0_u8; entries.len() * 8];
2151 crate::stable::memory::read_preallocated(offset, &mut encoded).unwrap();
2152 let expected = entries
2153 .iter()
2154 .flat_map(|entry| entry.to_le_bytes())
2155 .collect::<Vec<_>>();
2156
2157 assert_eq!(offset, 128);
2158 assert_eq!(cursor, 128 + expected_len);
2159 assert_eq!(decoded, entries);
2160 assert_eq!(encoded, expected);
2161
2162 let mut empty_cursor = cursor;
2163 assert_eq!(write_u64_table_at(&[], &mut empty_cursor).unwrap(), 0);
2164 assert_eq!(empty_cursor, cursor);
2165 assert!(read_u64_table_at(cursor, 0).unwrap().is_empty());
2166 }
2167
2168 #[test]
2169 #[serial_test::serial]
2170 fn pbt_page_table_u64_encoding_round_trips() {
2171 let mut runner = TestRunner::new(Config {
2172 cases: 128,
2173 ..Config::default()
2174 });
2175
2176 runner
2177 .run(
2178 &proptest::collection::vec(any::<u64>(), 0..=512),
2179 |entries| {
2180 crate::stable::memory::reset_for_tests();
2181 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2182 invalidate_read_cache();
2183
2184 let mut cursor = 128_u64;
2185 let byte_len = entries.len().checked_mul(8).unwrap();
2186 let end = cursor + u64::try_from(byte_len).unwrap();
2187 crate::stable::memory::ensure_capacity(end).unwrap();
2188
2189 let offset = write_u64_table_at(&entries, &mut cursor).unwrap();
2190 let decoded = read_u64_table_at(offset, entries.len()).unwrap();
2191 prop_assert_eq!(decoded, entries.clone());
2192 prop_assert_eq!(cursor, end);
2193
2194 let mut encoded = vec![0_u8; byte_len];
2195 crate::stable::memory::read_preallocated(offset, &mut encoded).unwrap();
2196 let expected = entries
2197 .iter()
2198 .flat_map(|entry| entry.to_le_bytes())
2199 .collect::<Vec<_>>();
2200 prop_assert_eq!(encoded, expected);
2201 Ok(())
2202 },
2203 )
2204 .unwrap();
2205 }
2206
2207 #[test]
2208 #[serial_test::serial]
2209 fn pbt_compact_preserves_sparse_page_model() {
2210 let mut runner = TestRunner::new(Config {
2211 cases: 32,
2212 ..Config::default()
2213 });
2214
2215 runner
2216 .run(
2217 &proptest::collection::vec(prop::option::of(any::<u8>()), 0..=300),
2218 |pages| {
2219 crate::stable::memory::reset_for_tests();
2220 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2221 invalidate_read_cache();
2222
2223 let active_len = pages
2224 .iter()
2225 .rposition(Option::is_some)
2226 .map(|index| index + 1)
2227 .unwrap_or(0);
2228 for (page_no, byte) in pages.iter().take(active_len).enumerate() {
2229 if let Some(byte) = byte {
2230 write_at(
2231 u64::try_from(page_no).unwrap() * page_size(),
2232 &vec![*byte; page_len()],
2233 )
2234 .unwrap();
2235 }
2236 }
2237
2238 compact().unwrap();
2239 let block = Superblock::load().unwrap();
2240 prop_assert_eq!(
2241 block.db_size,
2242 u64::try_from(active_len).unwrap() * page_size()
2243 );
2244 let table = read_page_table(&block).unwrap();
2245 prop_assert_eq!(table.len(), active_len);
2246
2247 let mut first_compacted_offset = None;
2248 let mut non_zero_seen = 0_u64;
2249 for (page_no, byte) in pages.iter().take(active_len).enumerate() {
2250 let entry = table[page_no];
2251 let mut page = vec![0_u8; page_len()];
2252 read_base_at(u64::try_from(page_no).unwrap() * page_size(), &mut page)
2253 .unwrap();
2254
2255 if let Some(byte) = byte {
2256 let base = *first_compacted_offset.get_or_insert(entry);
2257 prop_assert_ne!(entry, 0);
2258 prop_assert_eq!(entry, base + non_zero_seen * page_size());
2259 prop_assert_eq!(page, vec![*byte; page_len()]);
2260 non_zero_seen += 1;
2261 } else {
2262 prop_assert_eq!(entry, 0);
2263 prop_assert_eq!(page, vec![0_u8; page_len()]);
2264 }
2265 }
2266 Ok(())
2267 },
2268 )
2269 .unwrap();
2270 }
2271
2272 #[derive(Clone, Debug)]
2273 enum BlobOp {
2274 Write { offset: u64, len: usize, byte: u8 },
2275 Truncate { size: u64 },
2276 Compact,
2277 }
2278
2279 #[test]
2280 #[serial_test::serial]
2281 fn pbt_blob_operations_match_logical_model_across_compact() {
2282 let mut runner = TestRunner::new(Config {
2283 cases: 48,
2284 ..Config::default()
2285 });
2286
2287 runner
2288 .run(&blob_operation_sequence(), |operations| {
2289 crate::stable::memory::reset_for_tests();
2290 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2291 invalidate_read_cache();
2292
2293 let mut model = Vec::new();
2294 let mut materialized = BTreeSet::new();
2295 assert_blob_model(&model, &materialized, false)?;
2296
2297 for operation in operations {
2298 let compacted = apply_blob_op(operation, &mut model, &mut materialized)?;
2299 assert_blob_model(&model, &materialized, compacted)?;
2300 }
2301 Ok(())
2302 })
2303 .unwrap();
2304 }
2305
2306 fn blob_operation_sequence() -> impl Strategy<Value = Vec<BlobOp>> {
2307 let write = (blob_offset_strategy(), blob_len_strategy(), any::<u8>())
2308 .prop_map(|(offset, len, byte)| BlobOp::Write { offset, len, byte });
2309 let truncate = blob_offset_strategy().prop_map(|size| BlobOp::Truncate { size });
2310 proptest::collection::vec(prop_oneof![write, truncate, Just(BlobOp::Compact)], 0..=48)
2311 }
2312
2313 fn blob_offset_strategy() -> impl Strategy<Value = u64> {
2314 let limit = blob_model_limit();
2315 let page = page_size();
2316 let segment = SEGMENT_PAGE_COUNT * page;
2317 prop_oneof![
2318 0_u64..=limit,
2319 prop::sample::select(boundary_values(&[
2320 0,
2321 1,
2322 page - 1,
2323 page,
2324 page + 1,
2325 segment - 1,
2326 segment,
2327 segment + 1,
2328 limit - 1,
2329 limit,
2330 ]))
2331 .prop_map(move |value| value.min(limit)),
2332 ]
2333 }
2334
2335 fn blob_len_strategy() -> impl Strategy<Value = usize> {
2336 prop_oneof![
2337 0_usize..=(page_len() * 2 + 17),
2338 prop::sample::select(vec![
2339 0,
2340 1,
2341 page_len() - 1,
2342 page_len(),
2343 page_len() + 1,
2344 page_len() * 2 + 1,
2345 ]),
2346 ]
2347 }
2348
2349 fn blob_model_limit() -> u64 {
2350 (SEGMENT_PAGE_COUNT + 3) * page_size()
2351 }
2352
2353 fn apply_blob_op(
2354 operation: BlobOp,
2355 model: &mut Vec<u8>,
2356 materialized: &mut BTreeSet<u64>,
2357 ) -> Result<bool, TestCaseError> {
2358 match operation {
2359 BlobOp::Write { offset, len, byte } => {
2360 let len = len.min(usize::try_from(blob_model_limit() - offset).unwrap());
2361 let bytes = vec![byte; len];
2362 write_at(offset, &bytes).map_err(|error| TestCaseError::fail(error.to_string()))?;
2363 if len == 0 {
2364 return Ok(false);
2365 }
2366
2367 let start = usize::try_from(offset).unwrap();
2368 let end = start + len;
2369 if model.len() < start {
2370 model.resize(start, 0);
2371 }
2372 if model.len() < end {
2373 model.resize(end, 0);
2374 }
2375 model[start..end].copy_from_slice(&bytes);
2376 mark_materialized_range(offset, len, materialized);
2377 Ok(false)
2378 }
2379 BlobOp::Truncate { size } => {
2380 truncate(size).map_err(|error| TestCaseError::fail(error.to_string()))?;
2381 let new_len = usize::try_from(size).unwrap();
2382 model.resize(new_len, 0);
2383 let active_pages = page_count_for_size(size)
2384 .map_err(|error| TestCaseError::fail(error.to_string()))?;
2385 materialized.retain(|page_no| *page_no < active_pages);
2386 if size > 0 && !size.is_multiple_of(page_size()) {
2387 materialized.insert(size / page_size());
2388 }
2389 Ok(false)
2390 }
2391 BlobOp::Compact => {
2392 compact().map_err(|error| TestCaseError::fail(error.to_string()))?;
2393 Ok(true)
2394 }
2395 }
2396 }
2397
2398 fn mark_materialized_range(offset: u64, len: usize, materialized: &mut BTreeSet<u64>) {
2399 let end = offset + u64::try_from(len).unwrap();
2400 let first_page = offset / page_size();
2401 let last_page = (end - 1) / page_size();
2402 for page_no in first_page..=last_page {
2403 materialized.insert(page_no);
2404 }
2405 }
2406
2407 fn assert_blob_model(
2408 model: &[u8],
2409 materialized: &BTreeSet<u64>,
2410 expect_compacted: bool,
2411 ) -> Result<(), TestCaseError> {
2412 let block = Superblock::load().map_err(|error| TestCaseError::fail(error.to_string()))?;
2413 prop_assert_eq!(block.db_size, u64::try_from(model.len()).unwrap());
2414
2415 if !model.is_empty() {
2416 let mut out = vec![0_u8; model.len()];
2417 read_base_at(0, &mut out).map_err(|error| TestCaseError::fail(error.to_string()))?;
2418 prop_assert_eq!(out, model);
2419 }
2420
2421 let mut tail = vec![1_u8; 32];
2422 read_base_at(u64::try_from(model.len()).unwrap(), &mut tail)
2423 .map_err(|error| TestCaseError::fail(error.to_string()))?;
2424 prop_assert_eq!(tail, vec![0_u8; 32]);
2425
2426 let table =
2427 read_page_table(&block).map_err(|error| TestCaseError::fail(error.to_string()))?;
2428 let active_pages = page_count_for_size(u64::try_from(model.len()).unwrap())
2429 .map_err(|error| TestCaseError::fail(error.to_string()))?;
2430 prop_assert_eq!(table.len(), usize::try_from(active_pages).unwrap());
2431
2432 let mut first_compacted_offset = None;
2433 let mut non_zero_seen = 0_u64;
2434 for (index, entry) in table.iter().enumerate() {
2435 let page_no = u64::try_from(index).unwrap();
2436 if materialized.contains(&page_no) {
2437 prop_assert_ne!(*entry, 0);
2438 if expect_compacted {
2439 let base = *first_compacted_offset.get_or_insert(*entry);
2440 prop_assert_eq!(*entry, base + non_zero_seen * page_size());
2441 }
2442 non_zero_seen += 1;
2443 } else {
2444 prop_assert_eq!(*entry, 0);
2445 }
2446 }
2447 Ok(())
2448 }
2449
2450 #[test]
2451 #[serial_test::serial]
2452 fn read_metrics_separate_table_cache_from_data_reads() {
2453 crate::stable::memory::reset_for_tests();
2454 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2455 invalidate_read_cache();
2456
2457 let page = vec![7_u8; page_len()];
2458 write_at(0, &page).unwrap();
2459 invalidate_read_cache();
2460 crate::read_metrics::reset_read_metrics();
2461
2462 let first = read_base_page(0).unwrap();
2463 let second = read_base_page(0).unwrap();
2464 let metrics = crate::read_metrics::read_metrics_snapshot();
2465
2466 assert_eq!(first, page);
2467 assert_eq!(second, page);
2468 assert!(metrics.stable_data_read_calls >= 2);
2469 assert!(metrics.stable_data_read_bytes >= page_size() * 2);
2470 assert!(metrics.page_table_root_misses >= 1);
2471 assert!(metrics.page_table_root_hits >= 1);
2472 assert!(metrics.page_table_segment_misses >= 1);
2473 assert!(metrics.page_table_segment_hits >= 1);
2474 #[cfg(feature = "bench-profile")]
2475 assert!(metrics.superblock_loads <= 1);
2476 #[cfg(not(feature = "bench-profile"))]
2477 assert_eq!(metrics.superblock_loads, 0);
2478 }
2479
2480 #[test]
2481 #[serial_test::serial]
2482 fn page_offset_cache_reuses_page_data_for_small_reads() {
2483 crate::stable::memory::reset_for_tests();
2484 crate::stable::memory::init(crate::stable::memory::memory_for_tests()).unwrap();
2485 invalidate_read_cache();
2486
2487 let page = vec![9_u8; page_len()];
2488 write_at(0, &page).unwrap();
2489 let block = Superblock::load().unwrap();
2490 let mut cache = PageOffsetCache::new();
2491 let mut first = [0_u8; 16];
2492 let mut second = [0_u8; 16];
2493
2494 crate::read_metrics::reset_read_metrics();
2495 read_base_at_with_page_cache(&block, 0, &mut first, &mut cache).unwrap();
2496 read_base_at_with_page_cache(&block, 8, &mut second, &mut cache).unwrap();
2497 let metrics = crate::read_metrics::read_metrics_snapshot();
2498
2499 assert_eq!(first, [9_u8; 16]);
2500 assert_eq!(second, [9_u8; 16]);
2501 assert_eq!(metrics.stable_data_read_calls, 1);
2502 assert_eq!(metrics.stable_data_read_bytes, page_size());
2503 }
2504}