use std::{
cell::RefCell,
fs,
io::Read as _,
panic::{self, AssertUnwindSafe},
path::{Path, PathBuf},
result::Result as StdResult,
sync::{
Arc,
atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering},
},
};
use bf_tree::{BfTree, ConfigError, LeafInsertResult, LeafReadResult, ScanIter, ScanIterError};
use parking_lot::RwLock;
use wbase::{
align::CachePadded64,
time::{Duration, Instant},
};
use crate::{
error::{Error, Result},
manager::CPR_MAGIC,
types::{
BfTreeConfig, BfTreeDeleteResult, BfTreeInsertResult, BfTreeReadResult, ScanRecord,
ScanReturnField, StorageBackendType,
},
};
const STACK_READ_BUF_SIZE: usize = 4096;
const STACK_SCAN_BUF_SIZE: usize = 8192;
const DRAIN_TIMEOUT: Duration = Duration::from_secs(30);
use wbase::backoff::backoff;
pub(crate) fn spin_until(cond: impl Fn() -> bool, deadline: Instant) -> bool {
let mut spins = 0u32;
while !cond() {
if Instant::now() >= deadline {
return false;
}
backoff(spins);
spins = spins.wrapping_add(1);
}
true
}
fn snapshot_missing(path: &Path) -> Error {
let mut msg = String::from(SNAPSHOT_MISSING_PREFIX);
msg.push_str(&path.display().to_string());
Error::Recovery(msg)
}
const PRESET_LEAF_PAGE_SIZE: usize = 16384;
const PRESET_MAX_RECORD_SIZE: usize = 4096;
const PRESET_MAX_KEY_LEN: usize = 512;
const PRESET_MIN_RECORD_SIZE: usize = 4;
const SNAPSHOT_MISSING_PREFIX: &str = "快照文件不存在: ";
const MIN_MAX_RECORD_SIZE: usize = STACK_READ_BUF_SIZE;
thread_local! {
static READ_SCRATCH: RefCell<Vec<u8>> = const { RefCell::new(Vec::new()) };
}
#[inline]
fn with_read_scratch<R>(max_record_size: usize, f: impl FnOnce(&mut [u8]) -> R) -> R {
READ_SCRATCH.with(|scratch| {
let mut scratch = scratch.borrow_mut();
if scratch.len() < max_record_size {
scratch.resize(max_record_size, 0);
}
f(&mut scratch[..max_record_size])
})
}
#[inline]
pub fn file_has_cpr_magic(path: &Path) -> bool {
let Ok(mut file) = fs::File::open(path) else {
return false;
};
let mut magic = [0u8; CPR_MAGIC.len()];
file.read_exact(&mut magic).is_ok() && magic == *CPR_MAGIC
}
#[inline]
fn config_error_to_string(e: ConfigError) -> String {
match e {
ConfigError::MinimumRecordSize(s) => format!("MinimumRecordSize: {s}"),
ConfigError::MaximumRecordSize(s) => format!("MaximumRecordSize: {s}"),
ConfigError::LeafPageSize(s) => format!("LeafPageSize: {s}"),
ConfigError::MaxKeyLen(s) => format!("MaxKeyLen: {s}"),
ConfigError::CircularBufferSize(s) => format!("CircularBufferSize: {s}"),
ConfigError::SnapshotFileInvalid(s) => format!("SnapshotFileInvalid: {s}"),
ConfigError::SnapshotDisabled => "SnapshotDisabled".to_string(),
}
}
#[inline]
fn scan_iter_error_to_string(e: ScanIterError) -> &'static str {
match e {
ScanIterError::CacheOnlyMode => "CacheOnlyMode",
ScanIterError::InvalidStartKey => "InvalidStartKey",
ScanIterError::InvalidEndKey => "InvalidEndKey",
ScanIterError::InvalidCount => "InvalidCount",
ScanIterError::InvalidKeyRange => "InvalidKeyRange",
}
}
pub struct BfTreeService {
tree: RwLock<Option<Arc<BfTree>>>,
storage_backend: AtomicU8,
file_path: RwLock<Option<String>>,
max_record_size: AtomicUsize,
disposed: AtomicBool,
barriers: AtomicUsize,
writers: CachePadded64<AtomicUsize>,
}
struct WriteGuard<'a> {
writers: &'a AtomicUsize,
}
impl<'a> WriteGuard<'a> {
#[inline]
fn acquire(service: &'a BfTreeService) -> Self {
let mut spins = 0u32;
loop {
if service.barriers.load(Ordering::Acquire) == 0 {
service.writers.fetch_add(1, Ordering::SeqCst);
if service.barriers.load(Ordering::SeqCst) == 0 {
return Self {
writers: &service.writers,
};
}
service.writers.fetch_sub(1, Ordering::Release);
}
backoff(spins);
spins = spins.wrapping_add(1);
}
}
}
impl Drop for WriteGuard<'_> {
#[inline]
fn drop(&mut self) {
self.writers.fetch_sub(1, Ordering::Release);
}
}
pub struct WriteBarrierGuard<'a> {
service: &'a BfTreeService,
}
impl Drop for WriteBarrierGuard<'_> {
#[inline]
fn drop(&mut self) {
self.service.barriers.fetch_sub(1, Ordering::Release);
}
}
impl BfTreeService {
pub(crate) fn new_with_backend(
config: impl Into<bf_tree::Config>,
storage_backend: StorageBackendType,
file_path: Option<String>,
) -> Result<Self> {
if storage_backend == StorageBackendType::Disk && file_path.is_none() {
return Err(Error::InvalidArgument(
"磁盘后端必须指定数据文件路径 (file_path)".into(),
));
}
let inner_cfg: bf_tree::Config = config.into();
let max_record_size = inner_cfg.get_cb_max_record_size().max(MIN_MAX_RECORD_SIZE);
let tree = match BfTree::with_config(inner_cfg, None) {
Ok(t) => Arc::new(t),
Err(e) => return Err(Error::InvalidConfig(config_error_to_string(e))),
};
Ok(Self {
tree: RwLock::new(Some(tree)),
storage_backend: AtomicU8::new(storage_backend as u8),
file_path: RwLock::new(file_path),
max_record_size: AtomicUsize::new(max_record_size),
disposed: AtomicBool::new(false),
barriers: AtomicUsize::new(0),
writers: CachePadded64::new(AtomicUsize::new(0)),
})
}
pub fn new(config: BfTreeConfig) -> Result<Self> {
let storage_backend = config.storage_backend;
let file_path = config.file_path;
Self::new_with_backend(config.inner, storage_backend, file_path)
}
fn preset_config(cb_min_record_size: usize) -> BfTreeConfig {
let mut config = BfTreeConfig::default();
config
.use_snapshot(true)
.leaf_page_size(PRESET_LEAF_PAGE_SIZE)
.cb_max_record_size(PRESET_MAX_RECORD_SIZE)
.cb_max_key_len(PRESET_MAX_KEY_LEN)
.cb_min_record_size(if cb_min_record_size > 0 {
cb_min_record_size
} else {
PRESET_MIN_RECORD_SIZE
});
config
}
pub fn open_disk(path: impl AsRef<Path>, cb_min_record_size: usize) -> Result<Self> {
let p = path.as_ref();
if p.as_os_str().is_empty() {
return Err(Error::InvalidArgument(
"磁盘后端必须指定有效的数据文件路径".into(),
));
}
if let Some(parent) = p.parent()
&& !parent.as_os_str().is_empty()
{
fs::create_dir_all(parent)?;
}
let mut config = Self::preset_config(cb_min_record_size);
config.file_path(p);
Self::new_with_backend(
config,
StorageBackendType::Disk,
Some(p.to_string_lossy().into_owned()),
)
}
pub fn open_memory(cb_min_record_size: usize) -> Result<Self> {
let mut config = Self::preset_config(cb_min_record_size);
config.cache_only(true);
Self::new_with_backend(config, StorageBackendType::Memory, None)
}
#[inline]
fn tree_arc(&self) -> Result<Arc<BfTree>> {
self.check_disposed()?;
self.tree.read().as_ref().cloned().ok_or(Error::Disposed)
}
#[inline]
fn with_tree<R>(&self, f: impl FnOnce(&BfTree) -> R) -> Result<R> {
self.check_disposed()?;
let guard = self.tree.read();
match guard.as_ref() {
Some(tree) => Ok(f(tree)),
None => Err(Error::Disposed),
}
}
#[inline]
pub fn native_ptr(&self) -> u64 {
if let Some(tree) = self.tree.read().as_ref() {
Arc::as_ptr(tree) as usize as u64
} else {
0
}
}
#[inline]
pub fn file_path(&self) -> Option<String> {
self.file_path.read().clone()
}
#[inline]
pub fn storage_backend(&self) -> StorageBackendType {
StorageBackendType::from_u8(self.storage_backend.load(Ordering::Acquire))
}
#[inline]
fn max_record_size(&self) -> usize {
self.max_record_size.load(Ordering::Relaxed)
}
#[inline]
pub fn is_disposed(&self) -> bool {
self.disposed.load(Ordering::Acquire)
}
#[inline]
fn check_disposed(&self) -> Result<()> {
if self.is_disposed() {
Err(Error::Disposed)
} else {
Ok(())
}
}
#[inline]
pub fn insert(&self, key: &[u8], value: &[u8]) -> BfTreeInsertResult {
if value.is_empty() {
return BfTreeInsertResult::InvalidKV;
}
let _guard = WriteGuard::acquire(self);
self
.with_tree(|tree| match tree.insert(key, value) {
LeafInsertResult::Success => BfTreeInsertResult::Success,
LeafInsertResult::InvalidKV(_) => BfTreeInsertResult::InvalidKV,
})
.unwrap_or(BfTreeInsertResult::InvalidArguments)
}
pub fn read(&self, key: &[u8]) -> (BfTreeReadResult, Option<Vec<u8>>) {
self
.with_tree(|tree| {
let max_record_size = self.max_record_size();
if max_record_size <= STACK_READ_BUF_SIZE {
let mut stack_buf = [0u8; STACK_READ_BUF_SIZE];
let (res, len) = Self::read_tree(tree, key, &mut stack_buf);
(
res,
(res == BfTreeReadResult::Found).then(|| stack_buf[..len].to_vec()),
)
} else {
with_read_scratch(max_record_size, |scratch| {
let (res, len) = Self::read_tree(tree, key, scratch);
(
res,
(res == BfTreeReadResult::Found).then(|| scratch[..len].to_vec()),
)
})
}
})
.unwrap_or((BfTreeReadResult::InvalidArguments, None))
}
#[inline]
fn read_tree(tree: &BfTree, key: &[u8], out_buf: &mut [u8]) -> (BfTreeReadResult, usize) {
match tree.read(key, out_buf) {
LeafReadResult::Found(n) => (BfTreeReadResult::Found, n as usize),
LeafReadResult::NotFound => (BfTreeReadResult::NotFound, 0),
LeafReadResult::Deleted => (BfTreeReadResult::Deleted, 0),
LeafReadResult::InvalidKey => (BfTreeReadResult::InvalidKey, 0),
}
}
#[inline]
fn read_via_scratch(
tree: &BfTree,
key: &[u8],
out_buf: &mut [u8],
scratch: &mut [u8],
) -> (BfTreeReadResult, usize) {
let (res, len) = Self::read_tree(tree, key, scratch);
match res {
BfTreeReadResult::Found if out_buf.len() >= len => {
out_buf[..len].copy_from_slice(&scratch[..len]);
(BfTreeReadResult::Found, len)
}
BfTreeReadResult::Found => (BfTreeReadResult::InvalidArguments, 0),
r => (r, 0),
}
}
pub fn read_into(&self, key: &[u8], out_buf: &mut [u8]) -> (BfTreeReadResult, usize) {
self
.with_tree(|tree| {
let max_record_size = self.max_record_size();
if out_buf.len() >= max_record_size {
return Self::read_tree(tree, key, out_buf);
}
if max_record_size <= STACK_READ_BUF_SIZE {
let mut stack_buf = [0u8; STACK_READ_BUF_SIZE];
Self::read_via_scratch(tree, key, out_buf, &mut stack_buf)
} else {
with_read_scratch(max_record_size, |scratch| {
Self::read_via_scratch(tree, key, out_buf, scratch)
})
}
})
.unwrap_or((BfTreeReadResult::InvalidArguments, 0))
}
#[inline]
pub fn delete(&self, key: &[u8]) -> BfTreeDeleteResult {
let _guard = WriteGuard::acquire(self);
self
.with_tree(|tree| {
tree.delete(key);
BfTreeDeleteResult::Success
})
.unwrap_or(BfTreeDeleteResult::InvalidArguments)
}
pub fn scan_with_count_callback<F>(
&self,
start_key: &[u8],
count: usize,
return_field: ScanReturnField,
on_record: F,
) -> Result<usize>
where
F: FnMut(&[u8], &[u8]) -> bool,
{
if count == 0 {
return Ok(0);
}
self.scan_callback(
|tree| tree.scan_with_count(start_key, count, return_field),
return_field,
on_record,
)
}
pub fn scan_with_count(
&self,
start_key: &[u8],
count: usize,
return_field: ScanReturnField,
) -> Result<Vec<ScanRecord>> {
let mut records = Vec::with_capacity(count.min(1024));
self.scan_with_count_callback(
start_key,
count,
return_field,
ScanRecord::sink(&mut records),
)?;
Ok(records)
}
pub fn scan_with_end_key_callback<F>(
&self,
start_key: &[u8],
end_key: &[u8],
return_field: ScanReturnField,
on_record: F,
) -> Result<usize>
where
F: FnMut(&[u8], &[u8]) -> bool,
{
if start_key > end_key {
return Ok(0);
}
self.scan_callback(
|tree| tree.scan_with_end_key(start_key, end_key, return_field),
return_field,
on_record,
)
}
pub fn scan_with_end_key(
&self,
start_key: &[u8],
end_key: &[u8],
return_field: ScanReturnField,
) -> Result<Vec<ScanRecord>> {
let mut records = Vec::with_capacity(32);
self.scan_with_end_key_callback(
start_key,
end_key,
return_field,
ScanRecord::sink(&mut records),
)?;
Ok(records)
}
pub fn scan_all_callback<F>(&self, return_field: ScanReturnField, on_record: F) -> Result<usize>
where
F: FnMut(&[u8], &[u8]) -> bool,
{
self.scan_with_count_callback(&[0], usize::MAX, return_field, on_record)
}
pub fn scan_all(&self, return_field: ScanReturnField) -> Result<Vec<ScanRecord>> {
let mut records = Vec::with_capacity(32);
self.scan_all_callback(return_field, ScanRecord::sink(&mut records))?;
Ok(records)
}
fn scan_callback<F>(
&self,
make_iter: impl FnOnce(&BfTree) -> StdResult<ScanIter<'_, '_>, ScanIterError>,
return_field: ScanReturnField,
mut on_record: F,
) -> Result<usize>
where
F: FnMut(&[u8], &[u8]) -> bool,
{
let tree = self.tree_arc()?;
let mut iter = make_iter(&tree)
.map_err(|e| Error::InvalidArgument(scan_iter_error_to_string(e).to_string()))?;
let mut stack_buf = [0u8; STACK_SCAN_BUF_SIZE];
let mut heap_buf;
let max_record_size = self.max_record_size();
let buf: &mut [u8] = if max_record_size <= STACK_SCAN_BUF_SIZE {
&mut stack_buf
} else {
heap_buf = vec![0u8; max_record_size];
&mut heap_buf
};
let mut scanned = 0;
while let Some((k_len, v_len)) = iter.next(buf) {
let k = if return_field != ScanReturnField::Value {
&buf[..k_len]
} else {
&[]
};
let v = if return_field != ScanReturnField::Key {
&buf[k_len..k_len + v_len]
} else {
&[]
};
scanned += 1;
if !on_record(k, v) {
break;
}
}
Ok(scanned)
}
fn drain_writers(&self) -> Result<()> {
self.drain_writers_within(DRAIN_TIMEOUT)
}
fn drain_writers_within(&self, timeout: Duration) -> Result<()> {
let deadline = Instant::now() + timeout;
if spin_until(|| self.writers.load(Ordering::SeqCst) == 0, deadline) {
Ok(())
} else {
Err(Error::Timeout)
}
}
pub fn write_barrier(&self) -> WriteBarrierGuard<'_> {
self.barriers.fetch_add(1, Ordering::SeqCst);
WriteBarrierGuard { service: self }
}
pub fn cpr_snapshot(&self, snapshot_path: impl AsRef<Path>) -> Result<()> {
let tree = self.tree_arc()?;
let p = snapshot_path.as_ref();
if let Some(parent) = p.parent()
&& !parent.as_os_str().is_empty()
{
fs::create_dir_all(parent)?;
}
panic::catch_unwind(AssertUnwindSafe(|| tree.cpr_snapshot(p)))
.map_err(|_| Error::Snapshot("底层引擎异常 (快照未启用或内部状态异常)".into()))
}
pub fn recover_in_place(&self, snapshot_path: &Path, work_path: &Path) -> Result<()> {
if !snapshot_path.exists() {
return Err(snapshot_missing(snapshot_path));
}
if let Some(parent) = work_path.parent()
&& !parent.as_os_str().is_empty()
{
fs::create_dir_all(parent)?;
}
let mut tmp_os = work_path.as_os_str().to_os_string();
tmp_os.push(".recovering");
let tmp_path = PathBuf::from(tmp_os);
fs::copy(snapshot_path, &tmp_path)?;
let recovered = match Self::recover_from_cpr_snapshot(&tmp_path, true, StorageBackendType::Disk)
{
Ok(tree) => tree,
Err(e) => {
let _ = fs::remove_file(&tmp_path);
return Err(e);
}
};
if self.disposed.load(Ordering::Acquire) {
let _ = fs::remove_file(&tmp_path);
return Err(Error::Disposed);
}
if let Err(e) = fs::rename(&tmp_path, work_path) {
let _ = fs::remove_file(&tmp_path);
return Err(e.into());
}
let old_tree = {
let _barrier = self.write_barrier();
self.drain_writers()?;
let mut guard = self.tree.write();
if self.disposed.load(Ordering::Acquire) {
return Err(Error::Disposed);
}
let old_tree = guard.take();
*guard = recovered.tree.write().take();
self
.max_record_size
.store(recovered.max_record_size(), Ordering::Release);
old_tree
};
drop(old_tree);
self
.storage_backend
.store(recovered.storage_backend() as u8, Ordering::Release);
*self.file_path.write() = Some(work_path.to_string_lossy().into_owned());
Ok(())
}
pub fn recover_from_cpr_snapshot(
recovery_path: impl AsRef<Path>,
enable_snapshots: bool,
storage_backend: impl Into<StorageBackendType>,
) -> Result<Self> {
let p = recovery_path.as_ref();
if !p.exists() {
return Err(snapshot_missing(p));
}
if !file_has_cpr_magic(p) {
let mut msg = String::from("快照文件损坏或格式非法 (魔数不匹配): ");
msg.push_str(&p.display().to_string());
return Err(Error::Recovery(msg));
}
let backend = storage_backend.into();
let use_snapshot = enable_snapshots;
match panic::catch_unwind(AssertUnwindSafe(|| {
BfTree::new_from_cpr_snapshot(p, use_snapshot, None, None, None)
})) {
Ok(Ok(tree)) => {
let max_record_size = tree
.config()
.get_cb_max_record_size()
.max(MIN_MAX_RECORD_SIZE);
Ok(Self {
tree: RwLock::new(Some(Arc::new(tree))),
storage_backend: AtomicU8::new(backend as u8),
file_path: RwLock::new(Some(p.to_string_lossy().into_owned())),
max_record_size: AtomicUsize::new(max_record_size),
disposed: AtomicBool::new(false),
barriers: AtomicUsize::new(0),
writers: CachePadded64::new(AtomicUsize::new(0)),
})
}
Ok(Err(e)) => Err(Error::Recovery(config_error_to_string(e))),
Err(_) => {
let mut msg = String::from("快照文件损坏或格式非法: ");
msg.push_str(&p.display().to_string());
Err(Error::Recovery(msg))
}
}
}
pub fn dispose(&self) {
if !self.disposed.swap(true, Ordering::SeqCst) {
self.tree.write().take();
}
}
pub fn dispose_quiesced(&self) -> Result<()> {
if self.disposed.swap(true, Ordering::SeqCst) {
return Ok(());
}
let _barrier = self.write_barrier();
self.drain_writers()?;
self.tree.write().take();
Ok(())
}
}
impl Drop for BfTreeService {
fn drop(&mut self) {
self.dispose();
}
}
#[cfg(test)]
mod tests {
use std::{
env, process,
thread::{sleep, spawn},
};
use super::*;
fn mem_service() -> BfTreeService {
BfTreeService::open_memory(0).unwrap()
}
#[test]
fn test_insert_empty_value_rejected_without_engine_panic() {
let service = mem_service();
assert_eq!(
service.insert(b"long_enough_key", b""),
BfTreeInsertResult::InvalidKV
);
assert_eq!(service.insert(b"k", b""), BfTreeInsertResult::InvalidKV);
let (res, v) = service.read(b"long_enough_key");
assert_eq!(res, BfTreeReadResult::NotFound);
assert_eq!(v, None);
}
#[test]
fn test_dispose_quiesced_waits_inflight_writer() {
let service = Arc::new(mem_service());
assert_eq!(
service.insert(b"key1", b"val1"),
BfTreeInsertResult::Success
);
service.writers.store(1, Ordering::SeqCst);
let done = Arc::new(AtomicBool::new(false));
let (service_t, done_t) = (Arc::clone(&service), Arc::clone(&done));
let handle = spawn(move || {
service_t.dispose_quiesced().unwrap();
done_t.store(true, Ordering::Release);
});
sleep(Duration::from_millis(100).into());
assert!(!done.load(Ordering::Acquire));
assert!(service.tree.read().is_some());
service.writers.store(0, Ordering::SeqCst);
handle.join().unwrap();
assert!(done.load(Ordering::Acquire));
assert!(service.tree.read().is_none());
assert!(service.is_disposed());
assert_eq!(
service.insert(b"key2", b"val2"),
BfTreeInsertResult::InvalidArguments
);
}
#[test]
fn test_drain_writers_timeout() {
let service = mem_service();
service.writers.store(1, Ordering::SeqCst);
let err = service
.drain_writers_within(Duration::from_millis(50))
.unwrap_err();
assert!(matches!(err, Error::Timeout));
assert!(!service.is_disposed());
assert!(service.tree.read().is_some());
service.writers.store(0, Ordering::SeqCst);
}
#[test]
fn test_cpr_snapshot_concurrent_with_writers() {
let service = Arc::new(mem_service());
for i in 0..100u32 {
let k = format!("base{i:04}");
assert_eq!(
service.insert(k.as_bytes(), b"base_value"),
BfTreeInsertResult::Success
);
}
let dir = env::temp_dir().join(format!(
"wbftree_cpr_concurrent_{}_{}",
process::id(),
fastrand::u64(..)
));
fs::create_dir_all(&dir).unwrap();
let snap = dir.join("snap.bftree");
let wsvc = Arc::clone(&service);
let writer = spawn(move || {
for i in 0..5000u32 {
let k = format!("live{i:05}");
assert_eq!(
wsvc.insert(k.as_bytes(), b"payload"),
BfTreeInsertResult::Success
);
}
});
service.cpr_snapshot(&snap).unwrap();
writer.join().unwrap();
let recovered =
BfTreeService::recover_from_cpr_snapshot(&snap, true, StorageBackendType::Disk).unwrap();
for i in 0..100u32 {
let k = format!("base{i:04}");
let (res, v) = recovered.read(k.as_bytes());
assert_eq!(res, BfTreeReadResult::Found, "基线键 {k} 必须在快照中");
assert_eq!(v.as_deref(), Some(&b"base_value"[..]));
}
fs::remove_dir_all(&dir).unwrap();
}
#[test]
fn test_recover_from_corrupt_snapshot_returns_err() {
let dir = env::temp_dir().join(format!(
"wbftree_corrupt_{}_{}",
process::id(),
fastrand::u64(..)
));
fs::create_dir_all(&dir).unwrap();
let bad_magic = dir.join("bad_magic.bftree");
fs::write(&bad_magic, b"garbage payload").unwrap();
let err = BfTreeService::recover_from_cpr_snapshot(&bad_magic, true, StorageBackendType::Disk);
assert!(matches!(err, Err(Error::Recovery(_))));
let truncated = dir.join("truncated.bftree");
fs::write(&truncated, b"BF-TREE").unwrap();
let err = BfTreeService::recover_from_cpr_snapshot(&truncated, true, StorageBackendType::Disk);
assert!(matches!(err, Err(Error::Recovery(_))));
let wrong_magic = dir.join("wrong_magic.bftree");
fs::write(&wrong_magic, b"XX-TREE-V0-BEGIN_PAYLOAD").unwrap();
let err =
BfTreeService::recover_from_cpr_snapshot(&wrong_magic, true, StorageBackendType::Disk);
assert!(matches!(err, Err(Error::Recovery(_))));
fs::remove_dir_all(&dir).unwrap();
}
#[test]
fn test_recover_in_place_republishes_max_record_size() {
let dir = env::temp_dir().join(format!(
"wbftree_swap_max_{}_{}",
process::id(),
fastrand::u64(..)
));
fs::create_dir_all(&dir).unwrap();
let src_work = dir.join("src.data.bftree");
let snap = dir.join("snap.bftree");
{
let mut config = BfTreeConfig::default();
config
.use_snapshot(true)
.leaf_page_size(32768)
.cb_max_record_size(8192)
.cb_max_key_len(PRESET_MAX_KEY_LEN)
.cb_min_record_size(8);
config.file_path(&src_work);
let src = BfTreeService::new(config).unwrap();
let big = [b'x'; 6000];
assert_eq!(src.insert(b"big_key", &big), BfTreeInsertResult::Success);
src.cpr_snapshot(&snap).unwrap();
}
let work = dir.join("work.data.bftree");
let target = BfTreeService::open_disk(&work, 0).unwrap();
assert_eq!(target.max_record_size(), PRESET_MAX_RECORD_SIZE);
let big = [b'x'; 6000];
assert_eq!(
target.insert(b"big_key", &big),
BfTreeInsertResult::InvalidKV
);
target.recover_in_place(&snap, &work).unwrap();
assert_eq!(target.max_record_size(), 8192);
let (res, v) = target.read(b"big_key");
assert_eq!(res, BfTreeReadResult::Found);
assert_eq!(v.as_deref(), Some(&big[..]));
assert_eq!(target.insert(b"fresh", &big), BfTreeInsertResult::Success);
let (res, v) = target.read(b"fresh");
assert_eq!(res, BfTreeReadResult::Found);
assert_eq!(v.as_deref(), Some(&big[..]));
let mut big_out = [0u8; 8192];
let (res, len) = target.read_into(b"big_key", &mut big_out);
assert_eq!(res, BfTreeReadResult::Found);
assert_eq!(&big_out[..len], &big[..]);
let mut small_out = [0u8; 64];
let (res, _) = target.read_into(b"big_key", &mut small_out);
assert_eq!(res, BfTreeReadResult::InvalidArguments);
fs::remove_dir_all(&dir).unwrap();
}
}