use std::{
error::Error as StdError,
path::{Path, PathBuf},
sync::{Arc, atomic::Ordering},
};
use itoa::Buffer;
use wbftree::{
RANGE_INDEX_STUB_SIZE, RangeIndexManager, RangeIndexStub, StorageBackendType, TreeEntry,
};
use wdev::Device;
use windex::{HashBucket, HashBucketEntry};
use wval::{CollectionType, META_VALUE_SIZE, MetaValue, NamespaceDbCodec};
use crate::{
config::StoreConfig,
error::{Error, Result},
range_index::encode_meta_stub_record,
read_cache::is_read_cache_addr,
store::{KEY_ID_ASSIGN_MARGIN, WedbStore},
};
pub const BFTREE_SNAPSHOT_DIR: &str = "bftree";
pub const BFTREE_SNAPSHOT_FILE: &str = "shared.bftree";
fn bftree_snapshot_path(checkpoint_dir: &Path, token: u128) -> PathBuf {
let mut buf = Buffer::new();
checkpoint_dir
.join(buf.format(token))
.join(BFTREE_SNAPSHOT_DIR)
.join(BFTREE_SNAPSHOT_FILE)
}
fn cpr_err(e: Error) -> wcpr::Error {
match e {
Error::Device(e) => e.into(),
Error::Epoch(e) => e.into(),
Error::HLog(e) => e.into(),
Error::Index(e) => e.into(),
Error::Io(e) => e.into(),
Error::Cpr(e) => e,
other => (Box::new(other) as Box<dyn StdError + Send + Sync>).into(),
}
}
impl<D: Device> WedbStore<D> {
pub fn take_range_index_checkpoints(
&self,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<usize> {
let mut buf = Buffer::new();
let token_str = buf.format(token);
Ok(
self
.range_index
.snapshot_all_trees_to_dir(checkpoint_dir.as_ref(), token_str)?,
)
}
pub async fn recover_range_indexes(
&self,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<usize> {
let mut buf = Buffer::new();
let token_str = buf.format(token);
let dir = checkpoint_dir.as_ref();
let mut count = self
.range_index
.recover_all_trees_from_dir(dir, token_str)?;
let begin_addr = self.begin_address();
let participant = self.epoch.register()?;
let _guard = participant.enter();
for bucket in self.index.buckets.iter() {
let mut curr_bucket = bucket;
loop {
for item in curr_bucket.entries.iter().take(HashBucket::DATA_ENTRIES) {
let raw = item.load(Ordering::Acquire);
if raw == 0 {
continue;
}
let entry = HashBucketEntry::from_raw(raw);
if entry.is_tentative() {
continue;
}
let mut addr = entry.address();
if is_read_cache_addr(addr) {
addr = self.read_cache.skip_read_cache(addr);
}
if addr < begin_addr {
continue;
}
if let Ok(record) = self.hlog.read_record(addr).await
&& let Ok(key) = record.key()
&& let Some(user_key) = NamespaceDbCodec::decode_meta_user_key(key)
&& let Ok(val) = record.value()
&& val.len() >= META_VALUE_SIZE + RANGE_INDEX_STUB_SIZE
&& let Ok(meta) = MetaValue::from_slice(val)
&& meta.collection_type == CollectionType::RangeIndex
{
let stub_slice = &val[META_VALUE_SIZE..META_VALUE_SIZE + RANGE_INDEX_STUB_SIZE];
if let Ok(mut stub) = RangeIndexStub::decode(stub_slice) {
stub.mark_recovered_from_checkpoint();
let healed = stub.encode();
if stub_slice != healed.as_slice() {
let new_val = encode_meta_stub_record(&meta, &stub);
if !self.hlog.try_update_in_place(addr, key, &new_val)? {
let new_addr = self.hlog.append(key, &new_val, addr, false)?;
self.index.update_address(key, addr, new_addr);
}
}
let key_id = RangeIndexManager::key_id_of(user_key);
let pin = self.range_index.live_indexes().pin();
if !pin.contains_key(&key_id) {
let key_hash = RangeIndexManager::key_hash_of(user_key);
let hash_prefix = RangeIndexManager::base32_prefix_of(user_key);
pin.insert(
key_id,
Arc::new(TreeEntry::new(None, key_hash, key_id, hash_prefix)),
);
count += 1;
}
}
}
}
let overflow_idx = curr_bucket.overflow_index();
if overflow_idx == 0 {
break;
}
match self.index.overflow_pool.get(overflow_idx) {
Some(next) => curr_bucket = next,
None => break,
}
}
}
Ok(count)
}
pub fn take_bftree_checkpoint(
&self,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<usize> {
if self.config.bftree_path.is_none() {
return Ok(0);
}
if self.bftree.storage_backend() != StorageBackendType::Disk {
let mut msg =
String::from("共享 BfTree 引擎非磁盘后端,无法为配置的持久工作文件生成快照: token=");
let mut buf = itoa::Buffer::new();
msg.push_str(buf.format(token));
return Err(wbftree::Error::Snapshot(msg).into());
}
self
.bftree
.cpr_snapshot(bftree_snapshot_path(checkpoint_dir.as_ref(), token))?;
Ok(1)
}
pub fn recover_bftree_from_checkpoint(
&self,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<()> {
let snapshot_path = bftree_snapshot_path(checkpoint_dir.as_ref(), token);
let Some(work_path) = &self.config.bftree_path else {
if snapshot_path.exists() {
let mut msg = String::from(
"Checkpoint 含共享 BfTree 快照但 StoreMeta 未记录 bftree_path,拒绝静默丢弃可恢复数据,请配置 bftree_path 后重试: token=",
);
let mut buf = itoa::Buffer::new();
msg.push_str(buf.format(token));
return Err(wbftree::Error::Recovery(msg).into());
}
return Ok(());
};
if !snapshot_path.exists() {
return Ok(());
}
self.bftree.recover_in_place(&snapshot_path, work_path)?;
Ok(())
}
}
pub fn take_cpr_snapshots<D: Device>(
store: &WedbStore<D>,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<usize> {
store.take_range_index_checkpoints(checkpoint_dir, token)
}
pub async fn recover_cpr_snapshots<D: Device>(
store: &WedbStore<D>,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<usize> {
store.recover_range_indexes(checkpoint_dir, token).await
}
pub fn take_shared_bftree_snapshot<D: Device>(
store: &WedbStore<D>,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<usize> {
store.take_bftree_checkpoint(checkpoint_dir, token)
}
pub fn recover_shared_bftree<D: Device>(
store: &WedbStore<D>,
checkpoint_dir: impl AsRef<Path>,
token: u128,
) -> Result<()> {
store.recover_bftree_from_checkpoint(checkpoint_dir, token)
}
impl<D: Device> wcpr::CprStore for WedbStore<D> {
type Device = D;
#[inline]
fn hlog(&self) -> &whlog::HybridLog<D> {
&self.hlog
}
#[inline]
fn index(&self) -> &windex::HashIndex {
&self.index
}
#[inline]
fn epoch(&self) -> &wepoch::LightEpoch {
&self.epoch
}
#[inline]
fn tail_address(&self) -> u64 {
self.tail_address()
}
#[inline]
fn begin_address(&self) -> u64 {
self.begin_address()
}
#[inline]
fn head_address(&self) -> u64 {
self.head_address()
}
#[inline]
fn shift_read_only_address(&self, target: u64) {
self.shift_read_only_address(target);
}
#[inline]
async fn flush_all(&self) -> wcpr::Result<()> {
self.flush_all().await.map_err(cpr_err)
}
#[inline]
fn entry_count(&self) -> usize {
self.entry_count()
}
#[inline]
fn skip_read_cache(&self, addr: u64) -> u64 {
if self.read_cache.is_enabled {
self.read_cache.skip_read_cache(addr)
} else {
0
}
}
#[inline]
fn take_range_index_checkpoints(&self, dir: &Path, token: u128) -> wcpr::Result<usize> {
self
.take_range_index_checkpoints(dir, token)
.map_err(cpr_err)
}
#[inline]
fn take_bftree_checkpoint(&self, dir: &Path, token: u128) -> wcpr::Result<usize> {
self.take_bftree_checkpoint(dir, token).map_err(cpr_err)
}
#[inline]
fn checkpoint_store_meta(&self) -> wcpr::StoreMeta {
wcpr::StoreMeta {
index_size: self.config.index_size,
page_size: self.config.page_size,
num_pages: self.config.num_pages,
mutable_fraction: self.config.mutable_fraction,
max_sessions: self.config.max_sessions,
enable_revivification: self.config.enable_revivification,
enable_read_cache: self.config.enable_read_cache,
read_cache_num_pages: self.config.read_cache_num_pages,
range_index_dir: self
.config
.range_index_dir
.as_ref()
.map(|p| p.to_string_lossy().into_owned()),
bftree_path: self
.config
.bftree_path
.as_ref()
.map(|p| p.to_string_lossy().into_owned()),
next_key_id: self.next_key_id.load(Ordering::Relaxed),
}
}
}
impl<D: Device> wcpr::CprRecover for WedbStore<D> {
async fn from_recovered(
recovered: wcpr::RecoveredCheckpoint<D>,
checkpoint_dir: &Path,
device: Arc<D>,
) -> wcpr::Result<Self> {
let mut config = StoreConfig::new(
recovered.meta.store_meta.index_size,
recovered.meta.store_meta.page_size,
recovered.meta.store_meta.num_pages,
recovered.meta.store_meta.mutable_fraction,
)
.map_err(cpr_err)?
.with_max_sessions(recovered.meta.store_meta.max_sessions)
.map_err(cpr_err)?;
if recovered.meta.store_meta.enable_revivification {
config = config.with_revivification(true);
}
if recovered.meta.store_meta.enable_read_cache {
config = config
.with_read_cache(true)
.with_read_cache_pages(recovered.meta.store_meta.read_cache_num_pages)
.map_err(cpr_err)?;
}
if let Some(p) = &recovered.meta.store_meta.range_index_dir {
config = config.with_range_index_dir(p);
}
if let Some(p) = &recovered.meta.store_meta.bftree_path {
config = config.with_bftree_path(p);
}
let store = Self::from_components(
config,
recovered.index,
recovered.hlog,
recovered.epoch,
device,
)
.map_err(cpr_err)?;
store.raise_key_id_floor(
recovered
.meta
.store_meta
.next_key_id
.saturating_add(KEY_ID_ASSIGN_MARGIN),
);
store
.recover_range_indexes(checkpoint_dir, recovered.meta.token)
.await
.map_err(cpr_err)?;
store
.recover_bftree_from_checkpoint(checkpoint_dir, recovered.meta.token)
.map_err(cpr_err)?;
Ok(store)
}
}
impl<D: Device> WedbStore<D> {
#[inline]
pub async fn create_checkpoint(
&self,
checkpoint_dir: impl AsRef<Path>,
cp_type: wcpr::CheckpointType,
) -> Result<wcpr::CheckpointMeta> {
let mgr = wcpr::CheckpointManager::<D>::new();
mgr
.create_checkpoint(self, checkpoint_dir, cp_type)
.await
.map_err(Error::from)
}
#[inline]
pub async fn create_checkpoint_with_token(
&self,
checkpoint_dir: impl AsRef<Path>,
cp_type: wcpr::CheckpointType,
token: u128,
) -> Result<wcpr::CheckpointMeta> {
let mgr = wcpr::CheckpointManager::<D>::new();
mgr
.create_checkpoint_with_token(self, checkpoint_dir, cp_type, token)
.await
.map_err(Error::from)
}
#[inline]
pub async fn recover(
checkpoint_dir: impl AsRef<Path>,
token: u128,
device: Arc<D>,
) -> Result<Self> {
wcpr::CheckpointManager::<D>::recover(checkpoint_dir, token, device)
.await
.map_err(Error::from)
}
#[inline]
pub async fn recover_latest(checkpoint_dir: impl AsRef<Path>, device: Arc<D>) -> Result<Self> {
wcpr::CheckpointManager::<D>::recover_latest(checkpoint_dir, device)
.await
.map_err(Error::from)
}
}
#[derive(Debug, Default)]
pub struct CheckpointManager<D: Device = wdev::SegmentedDevice> {
inner: wcpr::CheckpointManager<D>,
}
impl<D: Device> CheckpointManager<D> {
#[inline]
pub const fn new() -> Self {
Self {
inner: wcpr::CheckpointManager::new(),
}
}
#[inline]
pub const fn with_device() -> Self {
Self::new()
}
#[inline]
pub async fn create_checkpoint<S: wcpr::CprStore<Device = D>>(
&self,
store: &S,
checkpoint_dir: impl AsRef<Path>,
cp_type: wcpr::CheckpointType,
) -> wcpr::Result<wcpr::CheckpointMeta> {
self
.inner
.create_checkpoint(store, checkpoint_dir, cp_type)
.await
}
#[inline]
pub async fn create_checkpoint_with_token<S: wcpr::CprStore<Device = D>>(
&self,
store: &S,
checkpoint_dir: impl AsRef<Path>,
cp_type: wcpr::CheckpointType,
token: u128,
) -> wcpr::Result<wcpr::CheckpointMeta> {
self
.inner
.create_checkpoint_with_token(store, checkpoint_dir, cp_type, token)
.await
}
#[inline]
pub async fn recover(
checkpoint_dir: impl AsRef<Path>,
token: u128,
device: Arc<D>,
) -> wcpr::Result<WedbStore<D>> {
wcpr::CheckpointManager::<D>::recover(checkpoint_dir, token, device).await
}
#[inline]
pub async fn recover_store(
&self,
checkpoint_dir: impl AsRef<Path>,
token: u128,
device: Arc<D>,
) -> wcpr::Result<WedbStore<D>> {
Self::recover(checkpoint_dir, token, device).await
}
#[inline]
pub async fn recover_latest(
checkpoint_dir: impl AsRef<Path>,
device: Arc<D>,
) -> wcpr::Result<WedbStore<D>> {
wcpr::CheckpointManager::<D>::recover_latest(checkpoint_dir, device).await
}
#[inline]
pub async fn recover_latest_store(
&self,
checkpoint_dir: impl AsRef<Path>,
device: Arc<D>,
) -> wcpr::Result<WedbStore<D>> {
Self::recover_latest(checkpoint_dir, device).await
}
#[inline]
pub fn list_checkpoints(checkpoint_dir: impl AsRef<Path>) -> wcpr::Result<Vec<u128>> {
wcpr::CheckpointManager::<D>::list_checkpoints(checkpoint_dir)
}
#[inline]
pub fn find_latest_checkpoint(checkpoint_dir: impl AsRef<Path>) -> wcpr::Result<Option<u128>> {
wcpr::CheckpointManager::<D>::find_latest_checkpoint(checkpoint_dir)
}
#[inline]
pub fn purge_checkpoint(checkpoint_dir: impl AsRef<Path>, token: u128) -> wcpr::Result<()> {
wcpr::CheckpointManager::<D>::purge_checkpoint(checkpoint_dir, token)
}
#[inline]
pub fn purge_all(checkpoint_dir: impl AsRef<Path>) -> wcpr::Result<()> {
wcpr::CheckpointManager::<D>::purge_all(checkpoint_dir)
}
#[inline]
pub fn purge_outdated(checkpoint_dir: impl AsRef<Path>, keep: usize) -> wcpr::Result<Vec<u128>> {
wcpr::CheckpointManager::<D>::purge_outdated(checkpoint_dir, keep)
}
#[inline]
pub fn purge(&self, checkpoint_dir: impl AsRef<Path>, token: u128) -> wcpr::Result<()> {
Self::purge_checkpoint(checkpoint_dir, token)
}
#[inline]
pub fn purge_all_checkpoints(&self, checkpoint_dir: impl AsRef<Path>) -> wcpr::Result<()> {
Self::purge_all(checkpoint_dir)
}
#[inline]
pub fn purge_outdated_checkpoints(
&self,
checkpoint_dir: impl AsRef<Path>,
keep: usize,
) -> wcpr::Result<Vec<u128>> {
Self::purge_outdated(checkpoint_dir, keep)
}
#[inline]
pub async fn take_index_checkpoint(
&self,
index: &windex::HashIndex,
entry_count: usize,
checkpoint_dir: impl AsRef<Path>,
token: u128,
rc_skip: impl Fn(u64) -> u64,
) -> wcpr::Result<wcpr::IndexMeta> {
self
.inner
.take_index_checkpoint(index, entry_count, checkpoint_dir, token, rc_skip)
.await
}
}