use std::{
fs,
hint::spin_loop,
io::Read,
iter::repeat_n,
path::{Path, PathBuf},
str,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
thread::yield_now,
};
use parking_lot::RwLock;
use whasher::{GxPapayaMap, fast_hash, hash128, new_papaya_map};
use crate::{
error::{Error, Result},
service::BfTreeService,
stub::{RANGE_INDEX_STUB_SIZE, RangeIndexStub},
types::{BfTreeConfig, StorageBackend, StorageBackendType, TreeTuning},
};
const SPIN_LIMIT: usize = 32;
const PREFIX_SEED_1: u64 = 0x27bb_2ee6_87b0_b0fd;
const PREFIX_SEED_2: u64 = 0x517c_c1b7_2722_0a95;
#[inline]
fn hex_padded_128(value: u128, width: usize) -> String {
const HEX: &[u8; 16] = b"0123456789abcdef";
let mut buf = [0u8; 32];
let mut v = value;
let mut pos = 32;
while v > 0 {
pos -= 1;
buf[pos] = HEX[(v & 0xf) as usize];
v >>= 4;
}
let hex_len = 32 - pos;
let pad = width.saturating_sub(hex_len);
let mut s = String::with_capacity(pad + hex_len);
s.extend(repeat_n('0', pad));
s.push_str(unsafe { str::from_utf8_unchecked(&buf[pos..]) });
s
}
const HASH_PREFIX_LEN: usize = 32;
const ADDR_HEX_LEN: usize = 16;
const DATA_FILE_SUFFIX: &str = ".data.bftree";
const FLUSH_FILE_SUFFIX: &str = ".flush.bftree";
const TREE_FILE_SUFFIX: &str = ".bftree";
const CPR_MAGIC: &[u8; 16] = b"BF-TREE-V0-BEGIN";
const CPR_MAGIC_LEN: usize = CPR_MAGIC.len();
pub const NUM_LOCK_STRIPES: usize = 128;
const STRIPE_MASK: usize = NUM_LOCK_STRIPES - 1;
#[repr(align(128))]
pub struct CacheAlignedLock(pub RwLock<()>);
impl CacheAlignedLock {
#[inline]
pub fn new() -> Self {
Self(RwLock::new(()))
}
}
impl Default for CacheAlignedLock {
#[inline]
fn default() -> Self {
Self::new()
}
}
pub struct RangeIndexLocks {
stripes: Box<[CacheAlignedLock]>,
}
impl Default for RangeIndexLocks {
fn default() -> Self {
Self::new()
}
}
impl RangeIndexLocks {
pub fn new() -> Self {
let mut stripes = Vec::with_capacity(NUM_LOCK_STRIPES);
stripes.extend((0..NUM_LOCK_STRIPES).map(|_| CacheAlignedLock::new()));
Self {
stripes: stripes.into_boxed_slice(),
}
}
#[inline]
pub fn read(&self, key_hash: u64) -> parking_lot::RwLockReadGuard<'_, ()> {
let idx = (key_hash as usize) & STRIPE_MASK;
unsafe { self.stripes.get_unchecked(idx) }.0.read()
}
#[inline]
pub fn write(&self, key_hash: u64) -> parking_lot::RwLockWriteGuard<'_, ()> {
let idx = (key_hash as usize) & STRIPE_MASK;
unsafe { self.stripes.get_unchecked(idx) }.0.write()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RangeIndexFileEntry {
pub path: PathBuf,
pub key_hash: String,
pub address: i64,
pub is_flush_file: bool,
}
pub struct TreeEntry {
pub tree: RwLock<Option<Arc<BfTreeService>>>,
pub key_hash: u64,
pub key_id: u128,
pub hash_prefix: String,
pub snapshot_pending: AtomicBool,
pub snapshot_in_progress: AtomicBool,
}
impl TreeEntry {
pub fn new(
tree: Option<Arc<BfTreeService>>,
key_hash: u64,
key_id: u128,
hash_prefix: String,
) -> Self {
Self {
tree: RwLock::new(tree),
key_hash,
key_id,
hash_prefix,
snapshot_pending: AtomicBool::new(false),
snapshot_in_progress: AtomicBool::new(false),
}
}
#[inline]
pub fn try_claim_snapshot(&self) -> bool {
self
.snapshot_in_progress
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
}
#[inline]
pub fn release_snapshot(&self) {
self.snapshot_in_progress.store(false, Ordering::Release);
}
pub fn snapshot_under_claim(&self, tree: &BfTreeService, destination_path: &Path) -> Result<()> {
let mut spins = 0usize;
while !self.try_claim_snapshot() {
if spins < SPIN_LIMIT {
spin_loop();
} else {
yield_now();
}
spins = spins.wrapping_add(1);
}
let _guard = SnapshotClaimGuard(self);
tree.cpr_snapshot(destination_path)
}
}
struct SnapshotClaimGuard<'a>(&'a TreeEntry);
impl Drop for SnapshotClaimGuard<'_> {
fn drop(&mut self) {
self.0.release_snapshot();
}
}
struct CheckpointGuard<'a>(&'a RangeIndexManager);
impl Drop for CheckpointGuard<'_> {
fn drop(&mut self) {
self.0.clear_checkpoint_barrier();
}
}
struct SnapshotPendingGuard<'a>(&'a TreeEntry);
impl Drop for SnapshotPendingGuard<'_> {
fn drop(&mut self) {
self.0.snapshot_pending.store(false, Ordering::SeqCst);
}
}
pub const DEFAULT_MIGRATION_CHUNK_SIZE: usize = 256 * 1024;
pub const INDEX_SIZE_BYTES: usize = RANGE_INDEX_STUB_SIZE;
pub struct RangeIndexManager {
ri_log_root: PathBuf,
cpr_dir: PathBuf,
migration_temp_dir: PathBuf,
live_indexes: GxPapayaMap<u128, Arc<TreeEntry>>,
checkpoint_in_progress: AtomicBool,
locks: RangeIndexLocks,
}
impl RangeIndexManager {
pub fn from_root(ri_log_root: impl Into<PathBuf>) -> Self {
let root = ri_log_root.into();
let cpr = root.join("cpr");
Self::new(root, cpr)
}
pub fn new(ri_log_root: impl Into<PathBuf>, cpr_dir: impl Into<PathBuf>) -> Self {
let ri_log_root = ri_log_root.into();
let cpr_dir = cpr_dir.into();
let _ = fs::create_dir_all(&ri_log_root);
let _ = fs::create_dir_all(&cpr_dir);
let migration_temp_dir = ri_log_root.join("migration-tmp");
if migration_temp_dir.exists() {
let _ = fs::remove_dir_all(&migration_temp_dir);
}
let _ = fs::create_dir_all(&migration_temp_dir);
Self {
ri_log_root,
cpr_dir,
migration_temp_dir,
live_indexes: new_papaya_map(),
checkpoint_in_progress: AtomicBool::new(false),
locks: RangeIndexLocks::new(),
}
}
#[inline]
pub fn derive_temp_migration_path(&self) -> PathBuf {
let rand_id = fastrand::u128(..);
let mut s = hex_padded_128(rand_id, 32);
s.push_str(TREE_FILE_SUFFIX);
self.migration_temp_dir.join(s)
}
pub fn dispose(&self) {
let pin = self.live_indexes.pin();
for entry in pin.values() {
if let Some(tree) = entry.tree.write().take() {
tree.dispose();
}
}
pin.clear();
}
#[inline]
pub fn locks(&self) -> &RangeIndexLocks {
&self.locks
}
#[inline]
pub fn key_id_of(key: &[u8]) -> u128 {
hash128(key, PREFIX_SEED_1, PREFIX_SEED_2)
}
#[inline]
pub fn hash_prefix_of(key: &[u8]) -> String {
let id = Self::key_id_of(key);
hex_padded_128(id, HASH_PREFIX_LEN)
}
#[inline]
pub fn key_hash_of(key: &[u8]) -> u64 {
fast_hash(key)
}
#[inline]
pub fn compute_leaf_page_size(max_record_size: usize) -> usize {
if max_record_size <= 2048 {
return 4096;
}
(max_record_size * 5 / 2).min(32768).next_power_of_two()
}
pub fn data_file_path(&self, hash_prefix: &str) -> PathBuf {
let mut file_name = String::with_capacity(hash_prefix.len() + DATA_FILE_SUFFIX.len());
file_name.push_str(hash_prefix);
file_name.push_str(DATA_FILE_SUFFIX);
self.ri_log_root.join(file_name)
}
#[inline]
pub fn data_file_path_for_key(&self, key: &[u8]) -> PathBuf {
let hash_prefix = Self::hash_prefix_of(key);
self.data_file_path(&hash_prefix)
}
pub fn log_flush_path(&self, hash_prefix: &str, logical_address: i64) -> PathBuf {
let mut s = String::from(hash_prefix);
s.push('.');
s.push_str(&hex_padded_128(logical_address as u64 as u128, 16));
s.push_str(FLUSH_FILE_SUFFIX);
self.ri_log_root.join(s)
}
#[inline]
pub fn live_indexes(&self) -> &GxPapayaMap<u128, Arc<TreeEntry>> {
&self.live_indexes
}
fn live_entries(&self) -> Vec<Arc<TreeEntry>> {
let pin = self.live_indexes.pin();
pin.values().cloned().collect()
}
#[inline]
pub fn get_tree(&self, key: &[u8]) -> Option<Arc<BfTreeService>> {
let key_id = Self::key_id_of(key);
let pin = self.live_indexes.pin();
pin
.get(&key_id)
.and_then(|e| e.tree.read().as_ref().cloned())
}
#[inline]
fn live_tree_of(&self, key_id: u128) -> Option<(Arc<TreeEntry>, Arc<BfTreeService>)> {
let pin = self.live_indexes.pin();
pin
.get(&key_id)
.and_then(|e| e.tree.read().as_ref().cloned().map(|t| (Arc::clone(e), t)))
}
pub fn checkpoint_snapshot_path(&self, token: &str, hash_prefix: &str) -> PathBuf {
let mut file_name = String::with_capacity(hash_prefix.len() + TREE_FILE_SUFFIX.len());
file_name.push_str(hash_prefix);
file_name.push_str(TREE_FILE_SUFFIX);
self.cpr_dir.join(token).join("rangeindex").join(file_name)
}
#[inline]
pub fn live_index_count(&self) -> usize {
self.live_indexes.pin().len()
}
#[inline]
pub fn is_checkpoint_in_progress(&self) -> bool {
self.checkpoint_in_progress.load(Ordering::Acquire)
}
pub fn set_checkpoint_barrier(&self) {
let pin = self.live_indexes.pin();
for entry in pin.values() {
entry.snapshot_pending.store(true, Ordering::SeqCst);
}
self.checkpoint_in_progress.store(true, Ordering::SeqCst);
}
pub fn clear_checkpoint_barrier(&self) {
self.checkpoint_in_progress.store(false, Ordering::SeqCst);
let pin = self.live_indexes.pin();
for entry in pin.values() {
entry.snapshot_pending.store(false, Ordering::SeqCst);
}
}
pub fn wait_for_tree_checkpoint(&self, key: &[u8]) -> bool {
if !self.checkpoint_in_progress.load(Ordering::Acquire) {
return false;
}
let key_id = Self::key_id_of(key);
let entry_opt = {
let pin = self.live_indexes.pin();
pin.get(&key_id).cloned()
};
if let Some(entry) = entry_opt
&& entry.snapshot_pending.load(Ordering::Acquire)
{
let mut spins = 0usize;
while entry.snapshot_pending.load(Ordering::Acquire) {
if spins < SPIN_LIMIT {
spin_loop();
} else {
yield_now();
}
spins = spins.wrapping_add(1);
}
return true;
}
false
}
#[inline]
pub fn wait_for_global_checkpoint(&self) {
let mut spins = 0usize;
while self.is_checkpoint_in_progress() {
if spins < SPIN_LIMIT {
spin_loop();
} else {
yield_now();
}
spins = spins.wrapping_add(1);
}
}
pub fn create_bftree(
&self,
key: &[u8],
storage_backend: StorageBackend,
tuning: TreeTuning,
) -> Result<Arc<BfTreeService>> {
let key_hash = Self::key_hash_of(key);
let _stripe_lock = self.locks.write(key_hash);
let key_id = Self::key_id_of(key);
let hash_prefix = Self::hash_prefix_of(key);
self.create_bftree_internal(key_id, key_hash, &hash_prefix, storage_backend, tuning)
}
fn instantiate_tree(
&self,
hash_prefix: &str,
storage_backend: StorageBackend,
tuning: TreeTuning,
) -> Result<Arc<BfTreeService>> {
let mut config = BfTreeConfig::default();
let (file_path_str, backend_type) = if storage_backend == StorageBackend::Memory {
config.cache_only(true);
(None, StorageBackendType::Memory)
} else {
let data_path = self.data_file_path(hash_prefix);
config.file_path(&data_path);
(
Some(data_path.to_string_lossy().into_owned()),
StorageBackendType::Disk,
)
};
if tuning.cache_size > 0 {
config.cb_size_byte(tuning.cache_size);
}
if tuning.min_record_size > 0 {
config.cb_min_record_size(tuning.min_record_size);
}
if tuning.max_record_size > 0 {
config.cb_max_record_size(tuning.max_record_size);
}
if tuning.max_key_len > 0 {
config.cb_max_key_len(tuning.max_key_len);
}
let actual_leaf_page_size = if tuning.leaf_page_size > 0 {
tuning.leaf_page_size
} else if tuning.max_record_size > 0 {
Self::compute_leaf_page_size(tuning.max_record_size)
} else {
0
};
if actual_leaf_page_size > 0 {
config.leaf_page_size(actual_leaf_page_size);
}
config.use_snapshot(true);
Ok(Arc::new(BfTreeService::new_with_backend(
config,
backend_type,
file_path_str,
)?))
}
fn create_bftree_internal(
&self,
key_id: u128,
key_hash: u64,
hash_prefix: &str,
storage_backend: StorageBackend,
tuning: TreeTuning,
) -> Result<Arc<BfTreeService>> {
let pin = self.live_indexes.pin();
if pin.contains_key(&key_id) {
return Err(Error::IndexExists);
}
let tree = self.instantiate_tree(hash_prefix, storage_backend, tuning)?;
let entry = Arc::new(TreeEntry::new(
Some(Arc::clone(&tree)),
key_hash,
key_id,
hash_prefix.to_string(),
));
pin.insert(key_id, entry);
Ok(tree)
}
pub fn get_or_open_tree(&self, key: &[u8], stub: &RangeIndexStub) -> Result<Arc<BfTreeService>> {
let key_hash = Self::key_hash_of(key);
let _stripe_lock = self.locks.write(key_hash);
let key_id = Self::key_id_of(key);
{
let pin = self.live_indexes.pin();
if let Some(entry) = pin.get(&key_id) {
let tree_guard = entry.tree.read();
if let Some(t) = tree_guard.as_ref() {
return Ok(Arc::clone(t));
}
}
}
let hash_prefix = Self::hash_prefix_of(key);
let backend = StorageBackendType::from_u8(stub.storage_backend);
let data_path = self.data_file_path(&hash_prefix);
let mut flush_name = String::with_capacity(hash_prefix.len() + FLUSH_FILE_SUFFIX.len());
flush_name.push_str(&hash_prefix);
flush_name.push_str(FLUSH_FILE_SUFFIX);
let flush_path = self.ri_log_root.join(flush_name);
if backend == StorageBackendType::Disk {
if flush_path.exists() {
fs::copy(&flush_path, &data_path)?;
} else if let Ok(entries) = fs::read_dir(&self.ri_log_root) {
let mut latest_candidate: Option<(i64, PathBuf)> = None;
for entry in entries.flatten() {
let path = entry.path();
if let Some(name_str) = path.file_name().and_then(|n| n.to_str())
&& let Some((prefix, addr)) = Self::parse_flush_file_name(name_str)
&& prefix == hash_prefix
&& latest_candidate
.as_ref()
.is_none_or(|(max_addr, _)| addr > *max_addr)
{
latest_candidate = Some((addr, path));
}
}
if let Some((_, path)) = latest_candidate {
fs::copy(path, &data_path)?;
}
}
if !data_path.exists() {
let mut msg = String::from("数据文件缺失且无可用刷盘快照: ");
msg.push_str(&data_path.display().to_string());
return Err(Error::Recovery(msg));
}
}
let is_cpr = backend == StorageBackendType::Disk
&& data_path
.metadata()
.map(|m| m.len() >= CPR_MAGIC_LEN as u64)
.unwrap_or(false)
&& fs::File::open(&data_path).is_ok_and(|mut f| {
let mut magic = [0u8; CPR_MAGIC_LEN];
f.read_exact(&mut magic).is_ok() && magic == *CPR_MAGIC
});
let tree = if is_cpr {
Arc::new(BfTreeService::recover_from_cpr_snapshot(
&data_path,
true,
StorageBackendType::Disk,
)?)
} else {
self.instantiate_tree(&hash_prefix, backend.into(), TreeTuning::from(stub))?
};
let pin = self.live_indexes.pin();
if let Some(entry) = pin.get(&key_id) {
*entry.tree.write() = Some(Arc::clone(&tree));
} else {
let entry = Arc::new(TreeEntry::new(
Some(Arc::clone(&tree)),
key_hash,
key_id,
hash_prefix,
));
pin.insert(key_id, entry);
}
Ok(tree)
}
pub fn pre_stage_and_register_pending(&self, key: &[u8], src_flush_address: i64) -> Result<()> {
let key_hash = Self::key_hash_of(key);
let _stripe_lock = self.locks.write(key_hash);
let key_id = Self::key_id_of(key);
let hash_prefix = Self::hash_prefix_of(key);
let snapshot_path = self.log_flush_path(&hash_prefix, src_flush_address);
if !snapshot_path.exists() {
return Ok(());
}
let data_path = self.data_file_path(&hash_prefix);
fs::copy(&snapshot_path, &data_path)?;
let entry = Arc::new(TreeEntry::new(None, key_hash, key_id, hash_prefix));
self.live_indexes.pin().insert(key_id, entry);
Ok(())
}
#[inline]
fn remove_and_dispose_entry(&self, key_id: u128) -> bool {
if let Some(entry) = self.live_indexes.pin().remove(&key_id).cloned() {
if let Some(tree) = entry.tree.write().take() {
tree.dispose();
}
true
} else {
false
}
}
pub fn unregister_index(&self, key: &[u8]) -> bool {
let key_hash = Self::key_hash_of(key);
let _stripe_lock = self.locks.write(key_hash);
let key_id = Self::key_id_of(key);
self.remove_and_dispose_entry(key_id)
}
pub fn delete_index(&self, key: &[u8]) -> bool {
let key_hash = Self::key_hash_of(key);
let _stripe_lock = self.locks.write(key_hash);
let removed = self.remove_and_dispose_entry(Self::key_id_of(key));
let data_path = self.data_file_path(&Self::hash_prefix_of(key));
if data_path.exists() {
let _ = fs::remove_file(data_path);
}
removed
}
pub fn dispose_tree(&self, key: &[u8], delete_file: bool) -> bool {
if delete_file {
self.delete_index(key)
} else {
self.unregister_index(key)
}
}
pub fn dispose_tree_under_lock(
&self,
key: &[u8],
stub: &RangeIndexStub,
delete_files: bool,
) -> bool {
if !delete_files && stub.is_transferred() {
return false;
}
self.dispose_tree(key, delete_files)
}
pub fn register_tree(&self, key: &[u8], tree: Arc<BfTreeService>) {
let key_hash = Self::key_hash_of(key);
let _stripe_lock = self.locks.write(key_hash);
let key_id = Self::key_id_of(key);
let hash_prefix = Self::hash_prefix_of(key);
let pin = self.live_indexes.pin();
if let Some(existing) = pin.get(&key_id) {
*existing.tree.write() = Some(tree);
} else {
let entry = Arc::new(TreeEntry::new(Some(tree), key_hash, key_id, hash_prefix));
pin.insert(key_id, entry);
}
}
pub fn on_flush(&self, key: &[u8], stub: &mut RangeIndexStub) -> Result<()> {
self.on_flush_internal(key, stub, None)
}
pub fn on_flush_address(
&self,
key: &[u8],
stub: &mut RangeIndexStub,
logical_address: i64,
) -> Result<()> {
self.on_flush_internal(key, stub, Some(logical_address))
}
fn on_flush_internal(
&self,
key: &[u8],
stub: &mut RangeIndexStub,
logical_address: Option<i64>,
) -> Result<()> {
if stub.is_transferred() {
return Ok(());
}
let key_id = Self::key_id_of(key);
let hash_prefix = Self::hash_prefix_of(key);
let flush_path = match logical_address {
Some(addr) => self.log_flush_path(&hash_prefix, addr),
None => {
let mut flush_name = String::with_capacity(hash_prefix.len() + FLUSH_FILE_SUFFIX.len());
flush_name.push_str(&hash_prefix);
flush_name.push_str(FLUSH_FILE_SUFFIX);
self.ri_log_root.join(flush_name)
}
};
if let Some((entry, tree)) = self.live_tree_of(key_id) {
entry.snapshot_under_claim(&tree, &flush_path)?;
stub.set_flushed(true);
return Ok(());
}
let key_hash = Self::key_hash_of(key);
let _stripe_lock = self.locks.write(key_hash);
if let Some((entry, tree)) = self.live_tree_of(key_id) {
entry.snapshot_under_claim(&tree, &flush_path)?;
stub.set_flushed(true);
return Ok(());
}
let data_path = self.data_file_path(&hash_prefix);
if data_path.exists() {
fs::copy(&data_path, &flush_path)?;
stub.set_flushed(true);
}
Ok(())
}
pub fn snapshot_all_trees_for_checkpoint(&self, checkpoint_token: &str) -> Result<()> {
self.snapshot_all_trees_to_dir(&self.cpr_dir, checkpoint_token)?;
Ok(())
}
pub fn snapshot_all_trees_to_dir(
&self,
target_dir: &Path,
checkpoint_token: &str,
) -> Result<usize> {
if !self.is_checkpoint_in_progress() {
self.set_checkpoint_barrier();
}
let _guard = CheckpointGuard(self);
let entries = self.live_entries();
let count = entries.len();
let token_snapshot_dir = target_dir.join(checkpoint_token).join("rangeindex");
let _ = fs::create_dir_all(&token_snapshot_dir);
for entry in &entries {
if !entry.snapshot_pending.load(Ordering::Acquire) {
continue;
}
let _pending_guard = SnapshotPendingGuard(entry);
let tree_opt = entry.tree.read().as_ref().cloned();
let mut dest_file_name =
String::with_capacity(entry.hash_prefix.len() + TREE_FILE_SUFFIX.len());
dest_file_name.push_str(&entry.hash_prefix);
dest_file_name.push_str(TREE_FILE_SUFFIX);
let token_dest = token_snapshot_dir.join(&dest_file_name);
if let Some(tree) = tree_opt {
entry.snapshot_under_claim(&tree, &token_dest)?;
} else {
let data_path = self.data_file_path(&entry.hash_prefix);
if data_path.exists() {
fs::copy(&data_path, &token_dest)?;
}
}
}
Ok(count)
}
#[inline]
fn parse_flush_file_name(file_name: &str) -> Option<(&str, i64)> {
let rest = file_name.strip_suffix(".flush.bftree")?;
let (prefix, addr_str) = rest.rsplit_once('.')?;
if prefix.len() != HASH_PREFIX_LEN || !prefix.as_bytes().iter().all(|b| b.is_ascii_hexdigit()) {
return None;
}
if addr_str.len() != ADDR_HEX_LEN || !addr_str.as_bytes().iter().all(|b| b.is_ascii_hexdigit())
{
return None;
}
let addr = i64::from_str_radix(addr_str, 16).ok()?;
Some((prefix, addr))
}
pub fn on_truncate(&self, new_begin_address: i64) -> Result<()> {
if !self.ri_log_root.exists() {
return Ok(());
}
for entry in fs::read_dir(&self.ri_log_root)? {
let entry = entry?;
let path = entry.path();
if let Some(file_name) = path.file_name().and_then(|n| n.to_str())
&& let Some((_prefix, addr)) = Self::parse_flush_file_name(file_name)
&& addr < new_begin_address
{
let _ = fs::remove_file(&path);
}
}
Ok(())
}
pub fn enumerate_files_for_replication(
&self,
checkpoint_token: &str,
hlog_start_address: i64,
hlog_end_address: i64,
) -> Result<Vec<RangeIndexFileEntry>> {
let mut result = Vec::new();
if self.ri_log_root.exists() {
for entry in fs::read_dir(&self.ri_log_root)? {
let entry = entry?;
let file_name = entry.file_name();
let Some(name) = file_name.to_str() else {
continue;
};
if let Some((prefix, addr)) = Self::parse_flush_file_name(name)
&& addr >= hlog_start_address
&& addr < hlog_end_address
{
result.push(RangeIndexFileEntry {
path: entry.path(),
key_hash: prefix.to_string(),
address: addr,
is_flush_file: true,
});
}
}
}
let snapshot_dir = self.cpr_dir.join(checkpoint_token).join("rangeindex");
if snapshot_dir.exists() {
for entry in fs::read_dir(&snapshot_dir)? {
let entry = entry?;
let file_name = entry.file_name();
let Some(name) = file_name.to_str() else {
continue;
};
if let Some(stem) = name.strip_suffix(".bftree")
&& stem.len() == HASH_PREFIX_LEN
&& stem.as_bytes().iter().all(|b| b.is_ascii_hexdigit())
{
result.push(RangeIndexFileEntry {
path: entry.path(),
key_hash: stem.to_string(),
address: 0,
is_flush_file: false,
});
}
}
}
Ok(result)
}
#[inline]
pub fn get_replication_file_names(
&self,
checkpoint_token: &str,
hlog_start_address: i64,
hlog_end_address: i64,
) -> Result<Vec<PathBuf>> {
let entries = self.enumerate_files_for_replication(
checkpoint_token,
hlog_start_address,
hlog_end_address,
)?;
Ok(entries.into_iter().map(|e| e.path).collect())
}
pub fn recover_all_trees_from_checkpoint(&self, checkpoint_token: &str) -> Result<()> {
self.recover_all_trees_from_dir(&self.cpr_dir, checkpoint_token)?;
Ok(())
}
pub fn recover_all_trees_from_dir(
&self,
target_dir: &Path,
checkpoint_token: &str,
) -> Result<usize> {
let candidate_dirs = [
target_dir.join(checkpoint_token).join("rangeindex"),
target_dir.join("rangeindex"),
self.cpr_dir.join(checkpoint_token).join("rangeindex"),
self.cpr_dir.join("rangeindex"),
];
let mut recovered_count = 0;
for snapshot_dir in &candidate_dirs {
if !snapshot_dir.exists() {
continue;
}
if let Ok(entries) = fs::read_dir(snapshot_dir) {
for entry in entries.flatten() {
let path = entry.path();
if path.extension().is_some_and(|ext| ext == "bftree")
&& let Some(stem) = path.file_stem().and_then(|s| s.to_str())
&& stem.len() == HASH_PREFIX_LEN
&& stem.as_bytes().iter().all(|b| b.is_ascii_hexdigit())
{
let target_data_path = self.data_file_path(stem);
if !target_data_path.exists() || target_data_path != path {
let _ = fs::copy(&path, &target_data_path);
}
let Ok(key_id) = u128::from_str_radix(stem, 16) else {
continue;
};
if let Ok(recovered_tree) = BfTreeService::recover_from_cpr_snapshot(
&target_data_path,
true,
StorageBackendType::Disk,
) {
let key_hash = Self::key_hash_of(stem.as_bytes());
let tree_entry = Arc::new(TreeEntry::new(
Some(Arc::new(recovered_tree)),
key_hash,
key_id,
stem.to_string(),
));
self.live_indexes.pin().insert(key_id, tree_entry);
recovered_count += 1;
}
}
}
}
if recovered_count > 0 {
break;
}
}
Ok(recovered_count)
}
}
impl Drop for RangeIndexManager {
fn drop(&mut self) {
self.dispose();
}
}