use ic_sqlite_vfs::config::STABLE_PAGE_SIZE;
use ic_sqlite_vfs::stable::memory_manager::{MemoryId, MemoryManager};
use ic_sqlite_vfs::stable::raw_memory::{DefaultMemoryImpl, Memory};
use proptest::prelude::*;
use proptest::test_runner::{Config, TestRunner};
use std::cell::{Cell, RefCell};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::rc::Rc;
use upstream_ic_stable_structures::memory_manager::{
MemoryId as UpstreamMemoryId, MemoryManager as UpstreamMemoryManager,
};
use upstream_ic_stable_structures::Memory as UpstreamMemory;
#[test]
fn memory_manager_reloads_interleaved_bucket_layout() {
let backing = DefaultMemoryImpl::default();
let manager = MemoryManager::init_with_bucket_size(backing.clone(), 1);
let first = manager.get(MemoryId::new(1));
let second = manager.get(MemoryId::new(2));
assert_eq!(first.grow(2), 0);
assert_eq!(second.grow(1), 0);
first.write(STABLE_PAGE_SIZE - 1, &[1, 2, 3]);
second.write(0, &[4, 5, 6]);
let mut magic = [0_u8; 3];
Memory::read(&backing, 0, &mut magic);
assert_eq!(&magic, b"MGR");
let reloaded = MemoryManager::init(backing);
let first = reloaded.get(MemoryId::new(1));
let second = reloaded.get(MemoryId::new(2));
let mut first_bytes = [0_u8; 3];
let mut second_bytes = [0_u8; 3];
first.read(STABLE_PAGE_SIZE - 1, &mut first_bytes);
second.read(0, &mut second_bytes);
assert_eq!(first_bytes, [1, 2, 3]);
assert_eq!(second_bytes, [4, 5, 6]);
}
#[test]
fn memory_manager_matches_upstream_layout_for_valid_operations() {
assert_matches_upstream_layout(1);
assert_matches_upstream_layout(128);
}
#[test]
fn pbt_memory_manager_matches_upstream_layout_for_random_operations() {
let mut runner = TestRunner::new(Config {
cases: 32,
..Config::default()
});
runner
.run(
&(
prop::sample::select(vec![1_u16, 2, 7, 128]),
operation_sequence(),
),
|(bucket_size, operations)| {
assert_random_ops_match_upstream(bucket_size, &operations);
Ok(())
},
)
.unwrap();
}
#[test]
#[should_panic(expected = "MemoryId(0): read out of bounds")]
fn cached_read_checks_logical_bounds() {
let backing = DefaultMemoryImpl::default();
let manager = MemoryManager::init(backing);
let memory = manager.get(MemoryId::new(0));
assert_eq!(memory.grow(1), 0);
memory.write(0, &[42]);
let mut byte = [0_u8; 1];
memory.read(STABLE_PAGE_SIZE, &mut byte);
}
#[test]
#[should_panic(expected = "MemoryId(0): write out of bounds")]
fn cached_write_checks_logical_bounds() {
let backing = DefaultMemoryImpl::default();
let manager = MemoryManager::init(backing);
let memory = manager.get(MemoryId::new(0));
assert_eq!(memory.grow(1), 0);
memory.write(0, &[42]);
memory.write(STABLE_PAGE_SIZE, &[1]);
}
#[test]
#[should_panic(expected = "bucket size must be greater than zero")]
fn init_with_bucket_size_rejects_zero_bucket_size() {
let backing = DefaultMemoryImpl::default();
let _manager = MemoryManager::init_with_bucket_size(backing, 0);
}
#[test]
fn grow_failure_returns_minus_one_without_metadata_changes() {
let backing = FailingGrowMemory::new(1);
let manager = MemoryManager::init_with_bucket_size(backing.clone(), 1);
let memory = manager.get(MemoryId::new(0));
assert_eq!(memory.grow(1), -1);
assert_eq!(memory.size(), 0);
assert_eq!(backing.allocation_owner(0), 255);
backing.set_max_pages(2);
let reloaded = MemoryManager::init(backing.clone());
let reloaded_memory = reloaded.get(MemoryId::new(0));
assert_eq!(reloaded_memory.size(), 0);
assert_eq!(backing.allocation_owner(0), 255);
}
#[test]
fn strict_init_rejects_non_memory_manager_layout() {
let backing = DefaultMemoryImpl::default();
assert_eq!(Memory::grow(&backing, 1), 0);
Memory::write(&backing, 0, b"not a memory manager");
let error = match MemoryManager::init_strict(backing) {
Ok(_) => panic!("strict init accepted non-memory-manager layout"),
Err(error) => error,
};
assert!(matches!(
error,
ic_sqlite_vfs::MemoryManagerInitError::NonMemoryManagerLayout
));
}
#[test]
fn strict_init_returns_error_for_invalid_memory_manager_layout() {
let backing = DefaultMemoryImpl::default();
assert_eq!(Memory::grow(&backing, 1), 0);
Memory::write(&backing, 0, &[b'M', b'G', b'R', 1]);
let error = match MemoryManager::init_strict(backing) {
Ok(_) => panic!("strict init accepted invalid layout"),
Err(error) => error,
};
assert!(matches!(
error,
ic_sqlite_vfs::MemoryManagerInitError::InvalidLayout(message)
if message.contains("bucket size is zero")
));
}
#[test]
fn grow_metadata_write_panic_rolls_back_allocation_table() {
for fail_on_write in [1_u64, 2] {
let backing = PanickingWriteMemory::default();
let manager = MemoryManager::init_with_bucket_size(backing.clone(), 1);
let memory = manager.get(MemoryId::new(0));
backing.reset_write_count();
backing.set_fail_on_write(Some(fail_on_write));
let result = catch_unwind(AssertUnwindSafe(|| {
let _ = memory.grow(1);
}));
assert!(result.is_err());
backing.set_fail_on_write(None);
let reloaded = MemoryManager::init(backing.clone());
let reloaded_memory = reloaded.get(MemoryId::new(0));
assert_eq!(reloaded_memory.size(), 0);
assert_eq!(backing.allocation_owner(0), 255);
}
}
fn assert_matches_upstream_layout(bucket_size_in_pages: u16) {
let local_backing = DefaultMemoryImpl::default();
let upstream_backing = upstream_ic_stable_structures::VectorMemory::default();
let local_manager =
MemoryManager::init_with_bucket_size(local_backing.clone(), bucket_size_in_pages);
let upstream_manager = UpstreamMemoryManager::init_with_bucket_size(
upstream_backing.clone(),
bucket_size_in_pages,
);
let local_memories: Vec<_> = (0_u8..5)
.map(|id| local_manager.get(MemoryId::new(id)))
.collect();
let upstream_memories: Vec<_> = (0_u8..5)
.map(|id| upstream_manager.get(UpstreamMemoryId::new(id)))
.collect();
let mut sizes = [0_u64; 5];
assert_eq!(
local_backing.borrow().as_slice(),
upstream_backing.borrow().as_slice()
);
for step in 0..384_u64 {
let id = usize::try_from(step % sizes.len() as u64).unwrap();
match step % 4 {
0 => {
let pages = (step.wrapping_mul(7) % 5) + 1;
let local_old = Memory::grow(&local_memories[id], pages);
let upstream_old = UpstreamMemory::grow(&upstream_memories[id], pages);
assert_eq!(
local_old, upstream_old,
"grow old size mismatch at step {step}"
);
sizes[id] += pages;
}
1 | 2 if sizes[id] > 0 => {
let capacity = sizes[id] * STABLE_PAGE_SIZE;
let len = ((step.wrapping_mul(13) % 4096) + 1).min(capacity) as usize;
let offset = step.wrapping_mul(7_919) % (capacity - len as u64 + 1);
let bytes = deterministic_bytes(step, len);
Memory::write(&local_memories[id], offset, &bytes);
UpstreamMemory::write(&upstream_memories[id], offset, &bytes);
}
_ if sizes[id] > 0 => {
let capacity = sizes[id] * STABLE_PAGE_SIZE;
let len = ((step.wrapping_mul(17) % 2048) + 1).min(capacity) as usize;
let offset = step.wrapping_mul(4_099) % (capacity - len as u64 + 1);
let mut local = vec![0_u8; len];
let mut upstream = vec![0_u8; len];
Memory::read(&local_memories[id], offset, &mut local);
UpstreamMemory::read(&upstream_memories[id], offset, &mut upstream);
assert_eq!(local, upstream, "read mismatch at step {step}");
}
_ => {}
}
for memory_id in 0..sizes.len() {
assert_eq!(
Memory::size(&local_memories[memory_id]),
UpstreamMemory::size(&upstream_memories[memory_id]),
"memory size mismatch at step {step}"
);
}
assert_eq!(
local_backing.borrow().as_slice(),
upstream_backing.borrow().as_slice(),
"stable layout diverged at step {step}"
);
}
assert_can_reload_with_either_manager(local_backing.clone(), upstream_backing.clone(), &sizes);
}
fn assert_can_reload_with_either_manager(
local_backing: DefaultMemoryImpl,
upstream_backing: upstream_ic_stable_structures::VectorMemory,
sizes: &[u64; 5],
) {
let local_from_upstream = MemoryManager::init(upstream_backing.clone());
let upstream_from_local = UpstreamMemoryManager::init(local_backing.clone());
for id in 0_u8..5 {
let local = local_from_upstream.get(MemoryId::new(id));
let upstream = upstream_from_local.get(UpstreamMemoryId::new(id));
assert_eq!(Memory::size(&local), sizes[id as usize]);
assert_eq!(UpstreamMemory::size(&upstream), sizes[id as usize]);
}
}
#[derive(Clone, Debug)]
enum Operation {
Grow {
id: usize,
pages_seed: u64,
},
Write {
id: usize,
offset_seed: u64,
len_seed: u64,
byte_seed: u64,
},
Read {
id: usize,
offset_seed: u64,
len_seed: u64,
},
}
fn operation_sequence() -> impl Strategy<Value = Vec<Operation>> {
let grow =
(0_usize..5, any::<u64>()).prop_map(|(id, pages_seed)| Operation::Grow { id, pages_seed });
let write = (0_usize..5, any::<u64>(), any::<u64>(), any::<u64>()).prop_map(
|(id, offset_seed, len_seed, byte_seed)| Operation::Write {
id,
offset_seed,
len_seed,
byte_seed,
},
);
let read = (0_usize..5, any::<u64>(), any::<u64>()).prop_map(|(id, offset_seed, len_seed)| {
Operation::Read {
id,
offset_seed,
len_seed,
}
});
proptest::collection::vec(prop_oneof![grow, write, read], 0..160)
}
fn assert_random_ops_match_upstream(bucket_size_in_pages: u16, operations: &[Operation]) {
let local_backing = DefaultMemoryImpl::default();
let upstream_backing = upstream_ic_stable_structures::VectorMemory::default();
let local_manager =
MemoryManager::init_with_bucket_size(local_backing.clone(), bucket_size_in_pages);
let upstream_manager = UpstreamMemoryManager::init_with_bucket_size(
upstream_backing.clone(),
bucket_size_in_pages,
);
let local_memories: Vec<_> = (0_u8..5)
.map(|id| local_manager.get(MemoryId::new(id)))
.collect();
let upstream_memories: Vec<_> = (0_u8..5)
.map(|id| upstream_manager.get(UpstreamMemoryId::new(id)))
.collect();
let mut sizes = [0_u64; 5];
for (step, operation) in operations.iter().enumerate() {
match *operation {
Operation::Grow { id, pages_seed } => {
let pages = projected_grow_pages(pages_seed, bucket_size_in_pages);
let local_old = Memory::grow(&local_memories[id], pages);
let upstream_old = UpstreamMemory::grow(&upstream_memories[id], pages);
assert_eq!(
local_old, upstream_old,
"grow old size mismatch at step {step}"
);
if local_old >= 0 {
sizes[id] = sizes[id].checked_add(pages).expect("test grow stays small");
}
}
Operation::Write {
id,
offset_seed,
len_seed,
byte_seed,
} if sizes[id] > 0 => {
let capacity = sizes[id] * STABLE_PAGE_SIZE;
let len = projected_len(len_seed, capacity, bucket_size_in_pages);
let offset = projected_offset(offset_seed, capacity, len, bucket_size_in_pages);
let bytes = deterministic_bytes(byte_seed, len);
Memory::write(&local_memories[id], offset, &bytes);
UpstreamMemory::write(&upstream_memories[id], offset, &bytes);
}
Operation::Read {
id,
offset_seed,
len_seed,
} if sizes[id] > 0 => {
let capacity = sizes[id] * STABLE_PAGE_SIZE;
let len = projected_len(len_seed, capacity, bucket_size_in_pages);
let offset = projected_offset(offset_seed, capacity, len, bucket_size_in_pages);
let mut local = vec![0_u8; len];
let mut upstream = vec![0_u8; len];
Memory::read(&local_memories[id], offset, &mut local);
UpstreamMemory::read(&upstream_memories[id], offset, &mut upstream);
assert_eq!(local, upstream, "read mismatch at step {step}");
}
_ => {}
}
for memory_id in 0..sizes.len() {
assert_eq!(
Memory::size(&local_memories[memory_id]),
UpstreamMemory::size(&upstream_memories[memory_id]),
"memory size mismatch at step {step}"
);
}
assert_eq!(
local_backing.borrow().as_slice(),
upstream_backing.borrow().as_slice(),
"stable layout diverged at step {step}"
);
}
assert_can_reload_with_either_manager(local_backing, upstream_backing, &sizes);
}
fn projected_grow_pages(seed: u64, bucket_size_in_pages: u16) -> u64 {
let bucket = u64::from(bucket_size_in_pages);
let candidates = [
0,
1,
bucket.saturating_sub(1),
bucket,
bucket.saturating_add(1),
seed % 5,
];
candidates[usize::try_from(seed % candidates.len() as u64).unwrap()]
}
fn projected_len(seed: u64, capacity: u64, bucket_size_in_pages: u16) -> usize {
let candidates = [
1,
2,
STABLE_PAGE_SIZE.saturating_sub(1),
STABLE_PAGE_SIZE,
STABLE_PAGE_SIZE.saturating_add(1),
u64::from(bucket_size_in_pages).saturating_add(1),
(seed % 4096).saturating_add(1),
];
let index = usize::try_from(seed % candidates.len() as u64).unwrap();
usize::try_from(candidates[index].min(capacity)).unwrap()
}
fn projected_offset(seed: u64, capacity: u64, len: usize, bucket_size_in_pages: u16) -> u64 {
let max = capacity - u64::try_from(len).unwrap();
let bucket_bytes = u64::from(bucket_size_in_pages) * STABLE_PAGE_SIZE;
let candidates = [
0,
1,
STABLE_PAGE_SIZE.saturating_sub(1),
STABLE_PAGE_SIZE,
STABLE_PAGE_SIZE.saturating_add(1),
bucket_bytes.saturating_sub(1),
bucket_bytes,
bucket_bytes.saturating_add(1),
max.saturating_sub(1),
max,
seed % (max + 1),
];
let index = usize::try_from(seed % candidates.len() as u64).unwrap();
candidates[index].min(max)
}
fn deterministic_bytes(seed: u64, len: usize) -> Vec<u8> {
(0..len)
.map(|index| seed.wrapping_add(index as u64).wrapping_mul(31) as u8)
.collect()
}
#[derive(Clone)]
struct FailingGrowMemory {
bytes: Rc<RefCell<Vec<u8>>>,
max_pages: Rc<Cell<u64>>,
}
impl FailingGrowMemory {
fn new(max_pages: u64) -> Self {
Self {
bytes: Rc::new(RefCell::new(Vec::new())),
max_pages: Rc::new(Cell::new(max_pages)),
}
}
fn set_max_pages(&self, max_pages: u64) {
self.max_pages.set(max_pages);
}
fn allocation_owner(&self, bucket: u64) -> u8 {
const HEADER_SIZE: u64 = 3 + 1 + 2 + 2 + 32 + 255 * 8;
self.bytes.borrow()[(HEADER_SIZE + bucket) as usize]
}
}
impl Memory for FailingGrowMemory {
fn size(&self) -> u64 {
self.bytes.borrow().len() as u64 / STABLE_PAGE_SIZE
}
fn grow(&self, pages: u64) -> i64 {
let size = self.size();
let Some(next_size) = size.checked_add(pages) else {
return -1;
};
if next_size > self.max_pages.get() {
return -1;
}
let Some(next_bytes) = next_size.checked_mul(STABLE_PAGE_SIZE) else {
return -1;
};
if next_bytes > usize::MAX as u64 {
return -1;
}
self.bytes.borrow_mut().resize(next_bytes as usize, 0);
size as i64
}
fn read(&self, offset: u64, dst: &mut [u8]) {
let end = checked_end(offset, dst.len(), "read");
dst.copy_from_slice(&self.bytes.borrow()[offset as usize..end]);
}
fn write(&self, offset: u64, src: &[u8]) {
let end = checked_end(offset, src.len(), "write");
self.bytes.borrow_mut()[offset as usize..end].copy_from_slice(src);
}
}
#[derive(Clone, Default)]
struct PanickingWriteMemory {
bytes: Rc<RefCell<Vec<u8>>>,
fail_on_write: Rc<Cell<Option<u64>>>,
write_count: Rc<Cell<u64>>,
}
impl PanickingWriteMemory {
fn set_fail_on_write(&self, ordinal: Option<u64>) {
self.fail_on_write.set(ordinal);
}
fn reset_write_count(&self) {
self.write_count.set(0);
}
fn allocation_owner(&self, bucket: u64) -> u8 {
const HEADER_SIZE: u64 = 3 + 1 + 2 + 2 + 32 + 255 * 8;
self.bytes.borrow()[(HEADER_SIZE + bucket) as usize]
}
}
impl Memory for PanickingWriteMemory {
fn size(&self) -> u64 {
self.bytes.borrow().len() as u64 / STABLE_PAGE_SIZE
}
fn grow(&self, pages: u64) -> i64 {
let size = self.size();
let Some(next_size) = size.checked_add(pages) else {
return -1;
};
let Some(next_bytes) = next_size.checked_mul(STABLE_PAGE_SIZE) else {
return -1;
};
if next_bytes > usize::MAX as u64 {
return -1;
}
self.bytes.borrow_mut().resize(next_bytes as usize, 0);
size as i64
}
fn read(&self, offset: u64, dst: &mut [u8]) {
let end = checked_end(offset, dst.len(), "read");
dst.copy_from_slice(&self.bytes.borrow()[offset as usize..end]);
}
fn write(&self, offset: u64, src: &[u8]) {
let count = self.write_count.get().saturating_add(1);
self.write_count.set(count);
if self.fail_on_write.get() == Some(count) {
panic!("write failpoint");
}
let end = checked_end(offset, src.len(), "write");
self.bytes.borrow_mut()[offset as usize..end].copy_from_slice(src);
}
}
fn checked_end(offset: u64, len: usize, operation: &str) -> usize {
let end = offset
.checked_add(len as u64)
.unwrap_or_else(|| panic!("{operation}: out of bounds"));
assert!(end <= usize::MAX as u64, "{operation}: out of bounds");
end as usize
}