Github-Mirror/seaweedfs
1
0
Fork 0
mirror of https://github.com/seaweedfs/seaweedfs.git synced 2026-09-08 15:41:15 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
b58d52ac1683c5f77d88e2d959007c37b45c8594
seaweedfs/seaweed-volume/src/storage/needle_map.rs
T
Chris Lu b58d52ac16 rust volume: search .sdx for read-only volumes instead of holding the index (#10951) 
* rust volume: search .sdx for read-only volumes instead of holding the index

The Go volume server loads every read-only volume through SortedFileNeedleMap:
the index lives on disk as a sorted .sdx, a lookup is a binary search, and
since #10950 no descriptor is held between lookups. The Rust server had no
counterpart. Read-only volumes built a full in-memory CompactNeedleMap, and
cloud-tiered ones — noWriteCanDelete, so not the read-only branch — went
through the writable path and pinned an .idx append handle on top of it. At the
hundreds of thousands of tiered volumes a real server carries, that is an index
in RAM and a descriptor each, for volumes nobody reads.

Port the sorted map and the bounded handle pool. A tiered volume now costs zero
descriptors and zero index bytes when idle; the pool keeps the hot handles open
so a busy volume does not pay an open() per needle. Handles are Arc<File>, so
an eviction cannot close one a reader still holds.

The generated .sdx is byte-identical to Go's — same sort, same last-write-wins,
same dropped tombstones — so a volume moved between a Go and a Rust server reads
whichever copy is already on disk. A test pins the bytes against a Go-generated
fixture.

* rust volume: fail compaction on an unreadable .sdx, and rebuild the map on tier-down

Two ways the sorted map could lose data.

iter_entries swallowed read errors and returned however many entries it managed
to collect. Compaction takes that vector for the complete live set, so a
truncated .sdx or a mid-scan I/O fault would commit a volume missing every
needle past the failure. Return a Result instead and abort. redb's
collect_entries dropped errors the same way on the same path, so it goes with
it.

Tier-down clears the remote mode and publishes the volume as writable, but the
map it booted with is the read-only sorted one. Its put always fails, so the
first write would append to the local .dat and then fail to index it, leaving
bytes nothing references — and a non-fsync write repeats it. Fold the
reopen_idx_for_write swap into refresh_remote_write_mode so the map always
matches the mode it just published; a rebuild that fails pins the volume
read-only rather than letting it take writes it cannot record.

Go reaches neither: its tier-down leaves noWriteCanDelete set, so the volume
stays read-only until a reload or an explicit mark-writable, which already goes
through reopenIdxForWrite.

* rust volume: keep read-only volumes mountable on a read-only index dir, and batch the .sdx scan

Building .sdx writes to the index directory, and load_index_sorted_file also
created a missing .idx there. A volume whose index sits on a read-only mount
took both paths and failed to load, where before it mounted read-only off an
in-memory index and served reads. Create the .idx only where deletes are
allowed, and fall back to the in-memory map when the sorted one cannot be
built, so a directory nobody can write costs memory rather than availability.

The end-to-end scan behind iter_entries, ascending_visit and save_to_idx read
one entry per syscall. Read 1024 at a time instead, the batch size
idx::walk_index_file uses. Positional reads, not a cursor: the handle is shared
with any other borrower.

Also gate the Go byte-parity fixture on the 5bytes feature it describes, which
is otherwise dead code in a 4-byte-offset build.

* rust volume: roll back a failed writable mark, and rebuild a torn .sdx

set_writable clears the read-only flags before it can know the rest will
succeed, but only the map rebuild rolled them back. An .idx writer that fails to
attach left the volume advertising writable over a needle map with no writer, so
puts landed in memory and were gone after a restart — the exact failure the
function exists to prevent. The read-only-mount fallback made it reachable: that
path loads an in-memory map with no writer attached. All three steps now run
behind one rollback point.

A .sdx whose length is not a whole number of entries was accepted as long as it
looked fresh, and truncation is what makes it look fresh. The entry count then
floored, hiding the last needle from lookups and from compaction, which would
commit the shorter set. Treat a torn file like a stale one and rebuild it from
.idx. Go writes .sdx in place rather than through a temporary, so a crash
mid-generation is a real way to produce one.

Appends now start at the last whole .idx entry too, so a torn tail there is
overwritten by the next tombstone instead of misaligning every row after it.

* rust volume: trim a torn .idx before writing to it, keep delete-only volumes online, count sorted-map deletes

Three from review.

Flooring the sorted map's append offset only protected its own positional
writes. Every writable path appends at EOF instead, so a partial row left by a
short write pushed the next row off alignment and the following load parsed the
rest of the file as garbage. Drop the partial row before attaching any writable
index writer — it is unrecoverable anyway, and every loader already skips it.
Go refuses to load such a volume at all; trimming keeps it mountable with the
rows before the tear intact.

The unwritable-index-dir fallback stopped one step short for volumes that allow
deletes, which is every tiered one: the in-memory loader opens .idx read-write
there and fails on the same directory that just refused the .sdx, so the volume
stayed offline. Give up the deletes instead — without a writer no tombstone
could be recorded anyway — and a remount on a writable directory restores them.

Sorted-map deletes left the counters untouched, so a tiered volume reported
itself garbage-free until it restarted. They now land where a reload would put
them: the tombstone is another .idx row, and both it and the row it supersedes
count as deletions under the rule the load-time metric applies. Go skips this
too, and should not.
2026-08-25 15:21:37 -07:00

1550 lines
56 KiB
Rust
Raw Blame History

//! NeedleMapper: index mapping NeedleId -> (Offset, Size).
//!
//! Two implementations:
//! - `CompactNeedleMap`: in-memory segmented sorted arrays (~10 bytes/entry)
//! - `RedbNeedleMap`: disk-backed via redb (low RAM, slightly slower)
//!
//! The `NeedleMap` enum wraps both and provides a uniform interface.
//! Loaded from .idx file on volume mount. Supports Get, Put, Delete with
//! metrics tracking (file count, byte count, deleted count, deleted bytes).
use std::collections::HashMap;
use std::io::{self, Read, Seek, Write};
use std::path::Path;
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
mod compact_map;
pub mod file_pool;
pub mod sorted_file;
use compact_map::CompactMap;
use sorted_file::SortedFileNeedleMap;
use redb::{Database, Durability, ReadableDatabase, ReadableTable, TableDefinition};
use crate::storage::idx;
use crate::storage::needle::needle::get_actual_size;
use crate::storage::types::*;
// ============================================================================
// NeedleValue
// ============================================================================
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NeedleValue {
pub offset: Offset,
pub size: Size,
}
/// Packed size of a NeedleValue in redb storage: OFFSET_SIZE + SIZE_SIZE.
const PACKED_NEEDLE_VALUE_SIZE: usize = OFFSET_SIZE + SIZE_SIZE;
/// Pack an (Offset, Size) pair into bytes for redb storage.
/// Layout: [offset OFFSET_SIZE bytes] [size 4 bytes big-endian]
fn pack_needle_value(nv: &NeedleValue) -> [u8; PACKED_NEEDLE_VALUE_SIZE] {
let mut buf = [0u8; PACKED_NEEDLE_VALUE_SIZE];
nv.offset.to_bytes(&mut buf[..OFFSET_SIZE]);
nv.size.to_bytes(&mut buf[OFFSET_SIZE..]);
buf
}
/// Unpack bytes from redb storage into (Offset, Size).
fn unpack_needle_value(bytes: &[u8; PACKED_NEEDLE_VALUE_SIZE]) -> NeedleValue {
NeedleValue {
offset: Offset::from_bytes(&bytes[..OFFSET_SIZE]),
size: Size::from_bytes(&bytes[OFFSET_SIZE..]),
}
}
// ============================================================================
// NeedleMapMetric
// ============================================================================
/// Metrics tracking for needle map operations.
#[derive(Debug, Default)]
pub struct NeedleMapMetric {
pub file_count: AtomicI64,
pub file_byte_count: AtomicU64,
pub deletion_count: AtomicI64,
pub deletion_byte_count: AtomicU64,
pub max_file_key: AtomicU64,
/// Largest (offset.to_actual_offset() + get_actual_size(size, version))
/// observed during the load walk. Used at volume load to verify that no
/// .idx entry references bytes past the end of .dat (issue #8928)
/// without paying for a second linear scan.
pub max_needle_end: AtomicI64,
}
impl NeedleMapMetric {
/// Update metrics based on a Put operation (additive-only, matching Go's logPut).
fn on_put(&self, key: NeedleId, old: Option<&NeedleValue>, new_size: Size) {
self.maybe_set_max_file_key(key);
// Go: always LogFileCounter(newSize) which does FileCounter++ and FileByteCounter += newSize
self.file_count.fetch_add(1, Ordering::Relaxed);
self.file_byte_count
.fetch_add(new_size.0 as u64, Ordering::Relaxed);
// Go: if oldSize > 0 && oldSize.IsValid() { LogDeletionCounter(oldSize) }
if let Some(old_val) = old {
if old_val.size.0 > 0 && old_val.size.is_valid() {
self.deletion_count.fetch_add(1, Ordering::Relaxed);
self.deletion_byte_count
.fetch_add(old_val.size.0 as u64, Ordering::Relaxed);
}
}
}
/// Update metrics based on a Delete operation (additive-only, matching Go's logDelete).
fn on_delete(&self, old: &NeedleValue) {
if old.size.0 > 0 {
self.deletion_count.fetch_add(1, Ordering::Relaxed);
self.deletion_byte_count
.fetch_add(old.size.0 as u64, Ordering::Relaxed);
}
}
fn maybe_set_max_file_key(&self, key: NeedleId) {
let key_val: u64 = key.into();
loop {
let current = self.max_file_key.load(Ordering::Relaxed);
if key_val <= current {
break;
}
if self
.max_file_key
.compare_exchange(current, key_val, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
break;
}
}
}
/// Update `max_needle_end` if this entry's (offset + actual size) exceeds
/// the running maximum. Skips deleted/zero-offset entries because they
/// don't reserve space in .dat.
fn maybe_set_max_needle_end(&self, offset: Offset, size: Size, version: Version) {
if offset.is_zero() || !size.is_valid() {
return;
}
let end = offset.to_actual_offset() + get_actual_size(size, version);
loop {
let current = self.max_needle_end.load(Ordering::Relaxed);
if end <= current {
break;
}
if self
.max_needle_end
.compare_exchange(current, end, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
break;
}
}
}
}
// ============================================================================
// NeedleMapKind
// ============================================================================
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NeedleMapKind {
InMemory,
Redb,
RedbMedium,
RedbLarge,
}
// ============================================================================
// IdxFileWriter trait
// ============================================================================
/// Trait for appending to an index file.
pub trait IdxFileWriter: Write + Send + Sync {
fn sync_all(&self) -> io::Result<()>;
}
impl IdxFileWriter for std::fs::File {
fn sync_all(&self) -> io::Result<()> {
std::fs::File::sync_all(self)
}
}
// ============================================================================
// CompactNeedleMap (in-memory)
// ============================================================================
/// In-memory needle map backed by a CompactMap (segmented sorted arrays).
/// Uses ~10 bytes per entry instead of ~40-48 bytes with HashMap.
/// The .idx file is kept open for append-only writes.
pub struct CompactNeedleMap {
map: CompactMap,
metric: NeedleMapMetric,
idx_file: Option<Box<dyn IdxFileWriter>>,
idx_file_offset: u64,
}
impl CompactNeedleMap {
/// Create a new empty in-memory map.
pub fn new() -> Self {
CompactNeedleMap {
map: CompactMap::new(),
metric: NeedleMapMetric::default(),
idx_file: None,
idx_file_offset: 0,
}
}
/// Load from an .idx file, building the in-memory map.
pub fn load_from_idx<R: Read + Seek>(reader: &mut R, version: Version) -> io::Result<Self> {
let mut nm = CompactNeedleMap::new();
idx::walk_index_file(reader, 0, |key, offset, size| {
nm.metric.maybe_set_max_needle_end(offset, size, version);
if offset.is_zero() || size.is_deleted() {
nm.delete_from_map(key);
} else {
nm.set_internal(key, NeedleValue { offset, size });
}
Ok(())
})?;
Ok(nm)
}
/// Set the index file for append-only writes.
pub fn set_idx_file(&mut self, file: Box<dyn IdxFileWriter>, offset: u64) {
self.idx_file = Some(file);
self.idx_file_offset = offset;
}
/// True when an .idx file writer is attached. A read-only load leaves
/// this `false` — set_writable() must reattach a writer or subsequent
/// puts silently skip the disk append.
pub fn has_idx_writer(&self) -> bool {
self.idx_file.is_some()
}
// ---- Map operations ----
/// Insert or update an entry. Appends to .idx file if present.
pub fn put(&mut self, key: NeedleId, offset: Offset, size: Size) -> io::Result<()> {
// Persist to idx file BEFORE mutating in-memory state for crash consistency
if let Some(ref mut idx_file) = self.idx_file {
idx::write_index_entry(idx_file, key, offset, size)?;
self.idx_file_offset += NEEDLE_MAP_ENTRY_SIZE as u64;
}
let old = self.map.get(key);
self.metric.on_put(key, old.as_ref(), size);
self.map.set(key, offset, size);
Ok(())
}
/// Look up a needle.
pub fn get(&self, key: NeedleId) -> Option<NeedleValue> {
self.map.get(key)
}
/// Mark a needle as deleted. Appends tombstone to .idx file.
/// Matches Go's NeedleMap.Delete: ALWAYS writes tombstone to idx and
/// increments deletion counter, even if needle doesn't exist or is
/// already deleted (important for replication).
pub fn delete(&mut self, key: NeedleId, offset: Offset) -> io::Result<Option<Size>> {
// Go unconditionally calls nm.m.Delete(), nm.logDelete(), nm.appendToIndexFile()
let deleted_bytes = self.map.delete(key).unwrap_or(Size(0));
// Match Go's logDelete -> LogDeletionCounter: only increment when oldSize > 0.
// Go does NOT decrement FileCounter/FileByteCounter in Delete;
// live counts are computed as FileCounter - DeletionCounter.
if deleted_bytes.0 > 0 {
self.metric.deletion_count.fetch_add(1, Ordering::Relaxed);
self.metric
.deletion_byte_count
.fetch_add(deleted_bytes.0 as u64, Ordering::Relaxed);
}
// Always write tombstone to idx file (matching Go)
if let Some(ref mut idx_file) = self.idx_file {
idx::write_index_entry(idx_file, key, offset, TOMBSTONE_FILE_SIZE)?;
self.idx_file_offset += NEEDLE_MAP_ENTRY_SIZE as u64;
}
if deleted_bytes.0 > 0 {
Ok(Some(deleted_bytes))
} else {
Ok(None)
}
}
// ---- Internal helpers ----
/// Insert into map during loading (no idx file write).
fn set_internal(&mut self, key: NeedleId, nv: NeedleValue) {
let old = self.map.get(key);
self.metric.on_put(key, old.as_ref(), nv.size);
self.map.set(key, nv.offset, nv.size);
}
/// Remove from map during loading (handle deletions in idx walk).
/// Matches Go's doLoading else branch: always increments DeletionCounter,
/// and adds old size bytes to DeletionByteCounter.
fn delete_from_map(&mut self, key: NeedleId) {
self.metric.maybe_set_max_file_key(key);
// Go's CompactMap.Delete returns old size (0 if not found or already deleted).
// Go's doLoading always does DeletionCounter++ and DeletionByteCounter += uint64(oldSize).
let old_size = self.map.get(key).map(|nv| nv.size).unwrap_or(Size(0));
// Go unconditionally increments DeletionCounter
self.metric.deletion_count.fetch_add(1, Ordering::Relaxed);
// Go adds uint64(oldSize) which for valid sizes adds the value, for 0/negative adds 0
if old_size.0 > 0 {
self.metric
.deletion_byte_count
.fetch_add(old_size.0 as u64, Ordering::Relaxed);
}
self.map.remove(key);
}
// ---- Metrics accessors ----
pub fn content_size(&self) -> u64 {
self.metric.file_byte_count.load(Ordering::Relaxed)
}
pub fn deleted_size(&self) -> u64 {
self.metric.deletion_byte_count.load(Ordering::Relaxed)
}
pub fn file_count(&self) -> i64 {
self.metric.file_count.load(Ordering::Relaxed)
}
pub fn deleted_count(&self) -> i64 {
self.metric.deletion_count.load(Ordering::Relaxed)
}
pub fn max_file_key(&self) -> NeedleId {
NeedleId(self.metric.max_file_key.load(Ordering::Relaxed))
}
/// Largest (offset + actual size) seen during the load walk; 0 if the
/// map is empty. See `NeedleMapMetric::maybe_set_max_needle_end`.
pub fn max_needle_end(&self) -> i64 {
self.metric.max_needle_end.load(Ordering::Relaxed)
}
pub fn index_file_size(&self) -> u64 {
self.idx_file_offset
}
/// Sync index file to disk.
pub fn sync(&self) -> io::Result<()> {
if let Some(ref idx_file) = self.idx_file {
idx_file.sync_all()?;
}
Ok(())
}
/// Close index file.
pub fn close(&mut self) {
let _ = self.sync();
self.idx_file = None;
}
/// Save the in-memory map to an index file, sorted by needle ID ascending.
pub fn save_to_idx(&self, path: &str) -> io::Result<()> {
let mut file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path)?;
self.map.ascending_visit(|id, nv| {
if nv.size.is_valid() {
idx::write_index_entry(&mut file, id, nv.offset, nv.size)
} else {
Ok(())
}
})?;
file.sync_all()?;
Ok(())
}
/// Visit all entries in ascending order by needle ID.
pub fn ascending_visit<F>(&self, f: F) -> Result<(), String>
where
F: FnMut(NeedleId, &NeedleValue) -> Result<(), String>,
{
self.map.ascending_visit(f)
}
}
// ============================================================================
// RedbNeedleMap (disk-backed via redb)
// ============================================================================
/// redb table: NeedleId (u64) -> packed [offset(4) + size(4)]
const NEEDLE_TABLE: TableDefinition<u64, &[u8]> = TableDefinition::new("needles");
/// Metadata table: stores the .idx file size that was used to build this redb.
/// Key "idx_size" -> u64 byte offset. Used to detect whether the .rdb can be
/// reused on restart or needs a full/incremental rebuild.
const META_TABLE: TableDefinition<&str, u64> = TableDefinition::new("meta");
const META_IDX_SIZE: &str = "idx_size";
/// Disk-backed needle map using redb.
/// Low memory usage — data lives on disk with redb's page cache.
pub struct RedbNeedleMap {
db: Database,
metric: NeedleMapMetric,
idx_file: Option<Box<dyn IdxFileWriter>>,
idx_file_offset: u64,
}
impl RedbNeedleMap {
/// Begin a write transaction with `Durability::None` (no fsync).
/// The .idx file is the source of truth for crash recovery, so redb
/// is always rebuilt from .idx on startup — fsync is unnecessary.
fn begin_write_no_fsync(db: &Database) -> io::Result<redb::WriteTransaction> {
let mut txn = db.begin_write().map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb begin_write: {}", e))
})?;
let _ = txn.set_durability(Durability::None);
Ok(txn)
}
/// Create a new redb-backed needle map at the given path.
/// The database file will be created if it does not exist.
pub fn new(db_path: &str) -> io::Result<Self> {
let db = Database::create(db_path).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb create error: {}", e))
})?;
// Ensure tables exist
let txn = Self::begin_write_no_fsync(&db)?;
{
let _table = txn.open_table(NEEDLE_TABLE).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e))
})?;
let _meta = txn.open_table(META_TABLE).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb open_table meta: {}", e))
})?;
}
txn.commit()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb commit: {}", e)))?;
Ok(RedbNeedleMap {
db,
metric: NeedleMapMetric::default(),
idx_file: None,
idx_file_offset: 0,
})
}
/// Save the .idx file size into redb metadata so we can detect whether
/// the .rdb is up-to-date on the next startup.
fn save_idx_size_meta(&self, idx_size: u64) -> io::Result<()> {
let txn = Self::begin_write_no_fsync(&self.db)?;
{
let mut meta = txn.open_table(META_TABLE).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb open meta: {}", e))
})?;
meta.insert(META_IDX_SIZE, idx_size).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb insert meta: {}", e))
})?;
}
txn.commit().map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb commit meta: {}", e))
})?;
Ok(())
}
/// Read the stored .idx file size from redb metadata.
fn read_idx_size_meta(&self) -> io::Result<Option<u64>> {
let txn = self
.db
.begin_read()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb begin_read: {}", e)))?;
let meta = txn
.open_table(META_TABLE)
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb open meta: {}", e)))?;
match meta.get(META_IDX_SIZE) {
Ok(Some(guard)) => Ok(Some(guard.value())),
Ok(None) => Ok(None),
Err(e) => Err(io::Error::new(
io::ErrorKind::Other,
format!("redb get meta: {}", e),
)),
}
}
/// Rebuild metrics by scanning all entries in the redb table.
/// Called when reusing an existing .rdb without a full rebuild.
fn rebuild_metrics_from_db(&self, version: Version) -> io::Result<()> {
let txn = self
.db
.begin_read()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb begin_read: {}", e)))?;
let table = txn
.open_table(NEEDLE_TABLE)
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e)))?;
let iter = table
.iter()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb iter: {}", e)))?;
for entry in iter {
let (key_guard, val_guard) = entry.map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb iter next: {}", e))
})?;
let key = NeedleId(key_guard.value());
let bytes: &[u8] = val_guard.value();
if bytes.len() == PACKED_NEEDLE_VALUE_SIZE {
let mut arr = [0u8; PACKED_NEEDLE_VALUE_SIZE];
arr.copy_from_slice(bytes);
let nv = unpack_needle_value(&arr);
self.metric.maybe_set_max_file_key(key);
self.metric
.maybe_set_max_needle_end(nv.offset, nv.size, version);
if nv.size.is_valid() {
self.metric.file_count.fetch_add(1, Ordering::Relaxed);
self.metric
.file_byte_count
.fetch_add(nv.size.0 as u64, Ordering::Relaxed);
} else {
// Deleted entry (negative size)
self.metric.deletion_count.fetch_add(1, Ordering::Relaxed);
self.metric
.deletion_byte_count
.fetch_add((-nv.size.0) as u64, Ordering::Relaxed);
}
}
}
Ok(())
}
/// Load from an .idx file, reusing an existing .rdb if it is consistent.
///
/// Strategy:
/// 1. Try to open existing .rdb and read its stored .idx size
/// 2. If .idx size matches → reuse .rdb, rebuild metrics from scan
/// 3. If .idx is larger → replay new entries incrementally
/// 4. Otherwise (missing, corrupted, .idx smaller) → full rebuild
pub fn load_from_idx<R: Read + Seek>(
db_path: &str,
reader: &mut R,
version: Version,
) -> io::Result<Self> {
let idx_size = reader.seek(io::SeekFrom::End(0))?;
reader.seek(io::SeekFrom::Start(0))?;
// Try to reuse existing .rdb
if Path::new(db_path).exists() {
if let Ok(nm) = Self::try_reuse_rdb(db_path, reader, idx_size, version) {
return Ok(nm);
}
// Reuse failed — fall through to full rebuild
reader.seek(io::SeekFrom::Start(0))?;
}
Self::full_rebuild(db_path, reader, idx_size, version)
}
/// Try to reuse an existing .rdb file. Returns Ok if successful,
/// Err if a full rebuild is needed.
fn try_reuse_rdb<R: Read + Seek>(
db_path: &str,
reader: &mut R,
idx_size: u64,
version: Version,
) -> io::Result<Self> {
let db = Database::open(db_path)
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb open: {}", e)))?;
let nm = RedbNeedleMap {
db,
metric: NeedleMapMetric::default(),
idx_file: None,
idx_file_offset: 0,
};
let stored_idx_size = nm
.read_idx_size_meta()?
.ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no idx_size in redb meta"))?;
if stored_idx_size > idx_size {
// .idx shrank — corrupted or truncated, need full rebuild
return Err(io::Error::new(
io::ErrorKind::Other,
"idx file smaller than stored size",
));
}
// Rebuild metrics from existing data
nm.rebuild_metrics_from_db(version)?;
if stored_idx_size < idx_size {
// .idx grew — replay new entries incrementally
let start_entry = stored_idx_size / NEEDLE_MAP_ENTRY_SIZE as u64;
let txn = Self::begin_write_no_fsync(&nm.db)?;
{
let mut table = txn.open_table(NEEDLE_TABLE).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e))
})?;
idx::walk_index_file(reader, start_entry, |key, offset, size| {
nm.metric.maybe_set_max_needle_end(offset, size, version);
let key_u64: u64 = key.into();
if offset.is_zero() || size.is_deleted() {
// Delete: look up old value for metric update, then
// store tombstone (negative size with original offset)
if let Ok(Some(old)) = nm.get_via_table(&table, key_u64) {
if old.size.is_valid() {
nm.metric.on_delete(&old);
let deleted_nv = NeedleValue {
offset: old.offset,
size: Size(-(old.size.0)),
};
let packed = pack_needle_value(&deleted_nv);
table.insert(key_u64, packed.as_slice()).map_err(|e| {
io::Error::new(
io::ErrorKind::Other,
format!("redb insert: {}", e),
)
})?;
}
}
} else {
// Put: look up old value for metric update
let old = nm.get_via_table(&table, key_u64).ok().flatten();
let nv = NeedleValue { offset, size };
let packed = pack_needle_value(&nv);
table.insert(key_u64, packed.as_slice()).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb insert: {}", e))
})?;
nm.metric.on_put(key, old.as_ref(), size);
}
Ok(())
})?;
}
txn.commit()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb commit: {}", e)))?;
nm.save_idx_size_meta(idx_size)?;
}
Ok(nm)
}
/// Look up a needle value using an already-open table reference.
/// Used during incremental replay to avoid opening separate read transactions.
fn get_via_table(
&self,
table: &redb::Table<u64, &[u8]>,
key_u64: u64,
) -> io::Result<Option<NeedleValue>> {
match table.get(key_u64) {
Ok(Some(guard)) => {
let bytes: &[u8] = guard.value();
if bytes.len() == PACKED_NEEDLE_VALUE_SIZE {
let mut arr = [0u8; PACKED_NEEDLE_VALUE_SIZE];
arr.copy_from_slice(bytes);
Ok(Some(unpack_needle_value(&arr)))
} else {
Ok(None)
}
}
Ok(None) => Ok(None),
Err(e) => Err(io::Error::new(
io::ErrorKind::Other,
format!("redb get: {}", e),
)),
}
}
/// Full rebuild: delete existing .rdb and rebuild from entire .idx file.
fn full_rebuild<R: Read + Seek>(
db_path: &str,
reader: &mut R,
idx_size: u64,
version: Version,
) -> io::Result<Self> {
let _ = std::fs::remove_file(db_path);
let nm = RedbNeedleMap::new(db_path)?;
// Collect entries from idx file, resolving duplicates/deletions
let mut entries: HashMap<NeedleId, Option<NeedleValue>> = HashMap::new();
idx::walk_index_file(reader, 0, |key, offset, size| {
nm.metric.maybe_set_max_needle_end(offset, size, version);
if offset.is_zero() || size.is_deleted() {
entries.insert(key, None);
} else {
entries.insert(key, Some(NeedleValue { offset, size }));
}
Ok(())
})?;
// Write all live entries to redb in a single transaction
let txn = Self::begin_write_no_fsync(&nm.db)?;
{
let mut table = txn.open_table(NEEDLE_TABLE).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e))
})?;
for (key, maybe_nv) in &entries {
let key_u64: u64 = (*key).into();
if let Some(nv) = maybe_nv {
let packed = pack_needle_value(nv);
table.insert(key_u64, packed.as_slice()).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb insert: {}", e))
})?;
nm.metric.on_put(*key, None, nv.size);
} else {
// Entry was deleted — remove from redb if present
table.remove(key_u64).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb remove: {}", e))
})?;
}
}
}
txn.commit()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb commit: {}", e)))?;
nm.save_idx_size_meta(idx_size)?;
Ok(nm)
}
/// Set the index file for append-only writes.
pub fn set_idx_file(&mut self, file: Box<dyn IdxFileWriter>, offset: u64) {
self.idx_file = Some(file);
self.idx_file_offset = offset;
}
/// True when an .idx file writer is attached. See CompactNeedleMap.
pub fn has_idx_writer(&self) -> bool {
self.idx_file.is_some()
}
// ---- Map operations ----
/// Insert or update an entry. Writes to idx file first, then redb.
pub fn put(&mut self, key: NeedleId, offset: Offset, size: Size) -> io::Result<()> {
// Persist to idx file BEFORE mutating redb state for crash consistency
if let Some(ref mut idx_file) = self.idx_file {
idx::write_index_entry(idx_file, key, offset, size)?;
self.idx_file_offset += NEEDLE_MAP_ENTRY_SIZE as u64;
}
let key_u64: u64 = key.into();
let nv = NeedleValue { offset, size };
let packed = pack_needle_value(&nv);
// Read old value for metric update
let old = self.get_internal(key_u64)?;
let txn = Self::begin_write_no_fsync(&self.db)?;
{
let mut table = txn.open_table(NEEDLE_TABLE).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e))
})?;
table
.insert(key_u64, packed.as_slice())
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb insert: {}", e)))?;
}
txn.commit()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb commit: {}", e)))?;
self.metric.on_put(key, old.as_ref(), size);
Ok(())
}
/// Look up a needle.
pub fn get(&self, key: NeedleId) -> Option<NeedleValue> {
let key_u64: u64 = key.into();
self.get_internal(key_u64).ok().flatten()
}
/// Internal get that returns io::Result for error propagation.
fn get_internal(&self, key_u64: u64) -> io::Result<Option<NeedleValue>> {
let txn = self
.db
.begin_read()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb begin_read: {}", e)))?;
let table = txn
.open_table(NEEDLE_TABLE)
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e)))?;
match table.get(key_u64) {
Ok(Some(guard)) => {
let bytes: &[u8] = guard.value();
if bytes.len() == PACKED_NEEDLE_VALUE_SIZE {
let mut arr = [0u8; PACKED_NEEDLE_VALUE_SIZE];
arr.copy_from_slice(bytes);
Ok(Some(unpack_needle_value(&arr)))
} else {
Ok(None)
}
}
Ok(None) => Ok(None),
Err(e) => Err(io::Error::new(
io::ErrorKind::Other,
format!("redb get: {}", e),
)),
}
}
/// Mark a needle as deleted. Appends tombstone to .idx file, negates size in redb.
pub fn delete(&mut self, key: NeedleId, offset: Offset) -> io::Result<Option<Size>> {
let key_u64: u64 = key.into();
if let Some(old) = self.get_internal(key_u64)? {
if old.size.is_valid() {
// Persist tombstone to idx file BEFORE mutating redb
if let Some(ref mut idx_file) = self.idx_file {
idx::write_index_entry(idx_file, key, offset, TOMBSTONE_FILE_SIZE)?;
self.idx_file_offset += NEEDLE_MAP_ENTRY_SIZE as u64;
}
self.metric.on_delete(&old);
let deleted_size = Size(-(old.size.0));
// Keep original offset so readDeleted can find original data (matching Go behavior)
let deleted_nv = NeedleValue {
offset: old.offset,
size: deleted_size,
};
let packed = pack_needle_value(&deleted_nv);
let txn = Self::begin_write_no_fsync(&self.db)?;
{
let mut table = txn.open_table(NEEDLE_TABLE).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e))
})?;
table.insert(key_u64, packed.as_slice()).map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb insert: {}", e))
})?;
}
txn.commit().map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb commit: {}", e))
})?;
return Ok(Some(old.size));
}
}
Ok(None)
}
// ---- Metrics accessors ----
pub fn content_size(&self) -> u64 {
self.metric.file_byte_count.load(Ordering::Relaxed)
}
pub fn deleted_size(&self) -> u64 {
self.metric.deletion_byte_count.load(Ordering::Relaxed)
}
pub fn file_count(&self) -> i64 {
self.metric.file_count.load(Ordering::Relaxed)
}
pub fn deleted_count(&self) -> i64 {
self.metric.deletion_count.load(Ordering::Relaxed)
}
pub fn max_file_key(&self) -> NeedleId {
NeedleId(self.metric.max_file_key.load(Ordering::Relaxed))
}
/// Largest (offset + actual size) seen during the load walk; 0 if the
/// map is empty. See `NeedleMapMetric::maybe_set_max_needle_end`.
pub fn max_needle_end(&self) -> i64 {
self.metric.max_needle_end.load(Ordering::Relaxed)
}
pub fn index_file_size(&self) -> u64 {
self.idx_file_offset
}
/// Sync index file to disk.
pub fn sync(&self) -> io::Result<()> {
if let Some(ref idx_file) = self.idx_file {
idx_file.sync_all()?;
}
Ok(())
}
/// Close index file.
pub fn close(&mut self) {
let _ = self.sync();
self.idx_file = None;
}
/// Save the redb contents to an index file, sorted by needle ID ascending.
pub fn save_to_idx(&self, path: &str) -> io::Result<()> {
let txn = self
.db
.begin_read()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb begin_read: {}", e)))?;
let table = txn
.open_table(NEEDLE_TABLE)
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb open_table: {}", e)))?;
let mut file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path)?;
// redb iterates in key order (u64 ascending)
let iter = table
.iter()
.map_err(|e| io::Error::new(io::ErrorKind::Other, format!("redb iter: {}", e)))?;
for entry in iter {
let (key_guard, val_guard) = entry.map_err(|e| {
io::Error::new(io::ErrorKind::Other, format!("redb iter next: {}", e))
})?;
let key_u64: u64 = key_guard.value();
let bytes: &[u8] = val_guard.value();
if bytes.len() == PACKED_NEEDLE_VALUE_SIZE {
let mut arr = [0u8; PACKED_NEEDLE_VALUE_SIZE];
arr.copy_from_slice(bytes);
let nv = unpack_needle_value(&arr);
if nv.size.is_valid() {
idx::write_index_entry(&mut file, NeedleId(key_u64), nv.offset, nv.size)?;
}
}
}
file.sync_all()?;
Ok(())
}
/// Visit all entries in ascending order by needle ID.
pub fn ascending_visit<F>(&self, mut f: F) -> Result<(), String>
where
F: FnMut(NeedleId, &NeedleValue) -> Result<(), String>,
{
let txn = self
.db
.begin_read()
.map_err(|e| format!("redb begin_read: {}", e))?;
let table = txn
.open_table(NEEDLE_TABLE)
.map_err(|e| format!("redb open_table: {}", e))?;
let iter = table.iter().map_err(|e| format!("redb iter: {}", e))?;
for entry in iter {
let (key_guard, val_guard) = entry.map_err(|e| format!("redb iter next: {}", e))?;
let key_u64: u64 = key_guard.value();
let bytes: &[u8] = val_guard.value();
if bytes.len() == PACKED_NEEDLE_VALUE_SIZE {
let mut arr = [0u8; PACKED_NEEDLE_VALUE_SIZE];
arr.copy_from_slice(bytes);
let nv = unpack_needle_value(&arr);
f(NeedleId(key_u64), &nv)?;
}
}
Ok(())
}
/// Collect all entries as a Vec for iteration (used by volume.rs iter patterns).
pub fn collect_entries(&self) -> io::Result<Vec<(NeedleId, NeedleValue)>> {
let mut result = Vec::new();
let txn: redb::ReadTransaction = self
.db
.begin_read()
.map_err(|e| io::Error::other(format!("redb begin_read: {e}")))?;
let table = txn
.open_table(NEEDLE_TABLE)
.map_err(|e| io::Error::other(format!("redb open_table: {e}")))?;
let iter = table
.iter()
.map_err(|e| io::Error::other(format!("redb iter: {e}")))?;
for entry in iter {
let (key_guard, val_guard) =
entry.map_err(|e| io::Error::other(format!("redb entry: {e}")))?;
let key_u64: u64 = key_guard.value();
let bytes: &[u8] = val_guard.value();
if bytes.len() == PACKED_NEEDLE_VALUE_SIZE {
let mut arr = [0u8; PACKED_NEEDLE_VALUE_SIZE];
arr.copy_from_slice(bytes);
let nv = unpack_needle_value(&arr);
result.push((NeedleId(key_u64), nv));
}
}
Ok(result)
}
}
// ============================================================================
// NeedleMap enum — unified interface over both implementations
// ============================================================================
/// Unified needle map wrapping either in-memory or redb-backed storage.
pub enum NeedleMap {
InMemory(CompactNeedleMap),
Redb(RedbNeedleMap),
/// Read-only volumes — including every cloud-tiered one — search the sorted
/// `.sdx` on disk instead of holding an index in RAM. Mirrors Go's
/// `SortedFileNeedleMap`.
SortedFile(SortedFileNeedleMap),
}
impl NeedleMap {
/// Insert or update an entry.
pub fn put(&mut self, key: NeedleId, offset: Offset, size: Size) -> io::Result<()> {
match self {
NeedleMap::InMemory(nm) => nm.put(key, offset, size),
NeedleMap::Redb(nm) => nm.put(key, offset, size),
NeedleMap::SortedFile(nm) => nm.put(key, offset, size),
}
}
/// Look up a needle.
pub fn get(&self, key: NeedleId) -> Option<NeedleValue> {
match self {
NeedleMap::InMemory(nm) => nm.get(key),
NeedleMap::Redb(nm) => nm.get(key),
NeedleMap::SortedFile(nm) => nm.get(key),
}
}
/// Mark a needle as deleted.
pub fn delete(&mut self, key: NeedleId, offset: Offset) -> io::Result<Option<Size>> {
match self {
NeedleMap::InMemory(nm) => nm.delete(key, offset),
NeedleMap::Redb(nm) => nm.delete(key, offset),
NeedleMap::SortedFile(nm) => nm.delete(key, offset),
}
}
/// Set the index file for append-only writes.
pub fn set_idx_file(&mut self, file: Box<dyn IdxFileWriter>, offset: u64) {
match self {
NeedleMap::InMemory(nm) => nm.set_idx_file(file, offset),
NeedleMap::Redb(nm) => nm.set_idx_file(file, offset),
// The sorted map borrows its .idx per append, so there is no
// long-lived writer to install.
NeedleMap::SortedFile(_) => {}
}
}
/// True when an .idx file writer is attached.
pub fn has_idx_writer(&self) -> bool {
match self {
NeedleMap::InMemory(nm) => nm.has_idx_writer(),
NeedleMap::Redb(nm) => nm.has_idx_writer(),
// Appends open the .idx on demand, so one is always available.
NeedleMap::SortedFile(_) => true,
}
}
/// Content byte count.
pub fn content_size(&self) -> u64 {
match self {
NeedleMap::InMemory(nm) => nm.content_size(),
NeedleMap::Redb(nm) => nm.content_size(),
NeedleMap::SortedFile(nm) => nm.content_size(),
}
}
/// Deleted byte count.
pub fn deleted_size(&self) -> u64 {
match self {
NeedleMap::InMemory(nm) => nm.deleted_size(),
NeedleMap::Redb(nm) => nm.deleted_size(),
NeedleMap::SortedFile(nm) => nm.deleted_size(),
}
}
/// Live file count.
pub fn file_count(&self) -> i64 {
match self {
NeedleMap::InMemory(nm) => nm.file_count(),
NeedleMap::Redb(nm) => nm.file_count(),
NeedleMap::SortedFile(nm) => nm.file_count(),
}
}
/// Deleted file count.
pub fn deleted_count(&self) -> i64 {
match self {
NeedleMap::InMemory(nm) => nm.deleted_count(),
NeedleMap::Redb(nm) => nm.deleted_count(),
NeedleMap::SortedFile(nm) => nm.deleted_count(),
}
}
/// Maximum needle ID seen.
pub fn max_file_key(&self) -> NeedleId {
match self {
NeedleMap::InMemory(nm) => nm.max_file_key(),
NeedleMap::Redb(nm) => nm.max_file_key(),
NeedleMap::SortedFile(nm) => nm.max_file_key(),
}
}
/// Largest (offset + actual size) seen during the load walk; 0 if the
/// map is empty. Used at volume load to detect .idx entries that
/// reference past the end of .dat (issue #8928) without a second scan.
pub fn max_needle_end(&self) -> i64 {
match self {
NeedleMap::InMemory(nm) => nm.max_needle_end(),
NeedleMap::Redb(nm) => nm.max_needle_end(),
NeedleMap::SortedFile(nm) => nm.max_needle_end(),
}
}
/// Index file size in bytes.
pub fn index_file_size(&self) -> u64 {
match self {
NeedleMap::InMemory(nm) => nm.index_file_size(),
NeedleMap::Redb(nm) => nm.index_file_size(),
NeedleMap::SortedFile(nm) => nm.index_file_size(),
}
}
/// Sync index file to disk.
pub fn sync(&self) -> io::Result<()> {
match self {
NeedleMap::InMemory(nm) => nm.sync(),
NeedleMap::Redb(nm) => nm.sync(),
NeedleMap::SortedFile(nm) => nm.sync(),
}
}
/// Close index file.
pub fn close(&mut self) {
match self {
NeedleMap::InMemory(nm) => nm.close(),
NeedleMap::Redb(nm) => nm.close(),
NeedleMap::SortedFile(nm) => nm.close(),
}
}
/// Save to an index file.
pub fn save_to_idx(&self, path: &str) -> io::Result<()> {
match self {
NeedleMap::InMemory(nm) => nm.save_to_idx(path),
NeedleMap::Redb(nm) => nm.save_to_idx(path),
NeedleMap::SortedFile(nm) => nm.save_to_idx(path),
}
}
/// Visit all entries in ascending order by needle ID.
pub fn ascending_visit<F>(&self, f: F) -> Result<(), String>
where
F: FnMut(NeedleId, &NeedleValue) -> Result<(), String>,
{
match self {
NeedleMap::InMemory(nm) => nm.ascending_visit(f),
NeedleMap::Redb(nm) => nm.ascending_visit(f),
NeedleMap::SortedFile(nm) => nm.ascending_visit(f),
}
}
/// Iterate all entries. Returns a Vec of (NeedleId, NeedleValue) pairs.
/// For InMemory this collects via ascending visit; the disk-backed maps read
/// it back off disk, so a truncated .sdx or a redb read fault surfaces here
/// as an error. Compaction treats the result as the complete live set, so a
/// partial scan must never be mistaken for an empty tail.
pub fn iter_entries(&self) -> io::Result<Vec<(NeedleId, NeedleValue)>> {
match self {
NeedleMap::InMemory(nm) => {
let mut entries = Vec::new();
// The visitor never fails, so neither can this.
let _ = nm.ascending_visit(|id, nv| {
entries.push((id, *nv));
Ok(())
});
Ok(entries)
}
NeedleMap::Redb(nm) => nm.collect_entries(),
NeedleMap::SortedFile(nm) => nm.iter_entries(),
}
}
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
#[test]
fn test_needle_map_put_get() {
let mut nm = CompactNeedleMap::new();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
nm.put(NeedleId(2), Offset::from_actual_offset(128), Size(200))
.unwrap();
let v1 = nm.get(NeedleId(1)).unwrap();
assert_eq!(v1.size, Size(100));
let v2 = nm.get(NeedleId(2)).unwrap();
assert_eq!(v2.size, Size(200));
assert!(nm.get(NeedleId(99)).is_none());
}
#[test]
fn test_needle_map_delete() {
let mut nm = CompactNeedleMap::new();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
assert_eq!(nm.file_count(), 1);
assert_eq!(nm.content_size(), 100);
let deleted = nm
.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(deleted, Some(Size(100)));
// Additive-only: file_count stays at 1 after delete
assert_eq!(nm.file_count(), 1);
assert_eq!(nm.deleted_count(), 1);
assert_eq!(nm.deleted_size(), 100);
}
#[test]
fn test_needle_map_metrics() {
let mut nm = CompactNeedleMap::new();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
nm.put(NeedleId(2), Offset::from_actual_offset(128), Size(200))
.unwrap();
nm.put(NeedleId(3), Offset::from_actual_offset(384), Size(300))
.unwrap();
assert_eq!(nm.file_count(), 3);
assert_eq!(nm.content_size(), 600);
assert_eq!(nm.max_file_key(), NeedleId(3));
// Update existing — additive-only: file_count increments, content_size adds
nm.put(NeedleId(2), Offset::from_actual_offset(700), Size(250))
.unwrap();
assert_eq!(nm.file_count(), 4); // 3 + 1 (always increments)
assert_eq!(nm.content_size(), 850); // 600 + 250 (always adds)
// Delete — additive-only: file_count unchanged
nm.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(nm.file_count(), 4); // unchanged
assert_eq!(nm.deleted_count(), 2); // 1 from overwrite + 1 from delete
}
#[test]
fn test_needle_map_load_from_idx() {
// Build an idx file in memory
// Note: offset 0 is reserved for the SuperBlock, so real needles start at offset >= 8
let mut idx_data = Vec::new();
idx::write_index_entry(
&mut idx_data,
NeedleId(1),
Offset::from_actual_offset(8),
Size(100),
)
.unwrap();
idx::write_index_entry(
&mut idx_data,
NeedleId(2),
Offset::from_actual_offset(128),
Size(200),
)
.unwrap();
idx::write_index_entry(
&mut idx_data,
NeedleId(3),
Offset::from_actual_offset(384),
Size(300),
)
.unwrap();
// Delete needle 2
idx::write_index_entry(
&mut idx_data,
NeedleId(2),
Offset::default(),
TOMBSTONE_FILE_SIZE,
)
.unwrap();
let mut cursor = Cursor::new(idx_data);
let nm = CompactNeedleMap::load_from_idx(&mut cursor, Version::current()).unwrap();
assert!(nm.get(NeedleId(1)).is_some());
assert!(nm.get(NeedleId(2)).is_none()); // deleted
assert!(nm.get(NeedleId(3)).is_some());
// Additive-only: put(1)+put(2)+put(3) = 3, delete doesn't decrement
assert_eq!(nm.file_count(), 3);
}
#[test]
fn test_needle_map_double_delete() {
let mut nm = CompactNeedleMap::new();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
let r1 = nm
.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(r1, Some(Size(100)));
// Second delete should return None (already deleted)
let r2 = nm
.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(r2, None);
assert_eq!(nm.deleted_count(), 1); // not double counted
}
// ---- RedbNeedleMap tests ----
#[test]
fn test_redb_needle_map_put_get() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let mut nm = RedbNeedleMap::new(db_path.to_str().unwrap()).unwrap();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
nm.put(NeedleId(2), Offset::from_actual_offset(128), Size(200))
.unwrap();
let v1 = nm.get(NeedleId(1)).unwrap();
assert_eq!(v1.size, Size(100));
let v2 = nm.get(NeedleId(2)).unwrap();
assert_eq!(v2.size, Size(200));
assert!(nm.get(NeedleId(99)).is_none());
}
#[test]
fn test_redb_needle_map_delete() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let mut nm = RedbNeedleMap::new(db_path.to_str().unwrap()).unwrap();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
assert_eq!(nm.file_count(), 1);
assert_eq!(nm.content_size(), 100);
let deleted = nm
.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(deleted, Some(Size(100)));
// Additive-only: file_count stays at 1 after delete
assert_eq!(nm.file_count(), 1);
assert_eq!(nm.deleted_count(), 1);
assert_eq!(nm.deleted_size(), 100);
// Deleted entry should have negated size
let nv = nm.get(NeedleId(1)).unwrap();
assert_eq!(nv.size, Size(-100));
}
#[test]
fn test_redb_needle_map_metrics() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let mut nm = RedbNeedleMap::new(db_path.to_str().unwrap()).unwrap();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
nm.put(NeedleId(2), Offset::from_actual_offset(128), Size(200))
.unwrap();
nm.put(NeedleId(3), Offset::from_actual_offset(384), Size(300))
.unwrap();
assert_eq!(nm.file_count(), 3);
assert_eq!(nm.content_size(), 600);
assert_eq!(nm.max_file_key(), NeedleId(3));
// Update existing — additive-only: file_count increments, content_size adds
nm.put(NeedleId(2), Offset::from_actual_offset(700), Size(250))
.unwrap();
assert_eq!(nm.file_count(), 4); // 3 + 1 (always increments)
assert_eq!(nm.content_size(), 850); // 600 + 250 (always adds)
// Delete — additive-only: file_count unchanged
nm.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(nm.file_count(), 4); // unchanged
assert_eq!(nm.deleted_count(), 2); // 1 from overwrite + 1 from delete
}
#[test]
fn test_redb_needle_map_load_from_idx() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let mut idx_data = Vec::new();
idx::write_index_entry(
&mut idx_data,
NeedleId(1),
Offset::from_actual_offset(8),
Size(100),
)
.unwrap();
idx::write_index_entry(
&mut idx_data,
NeedleId(2),
Offset::from_actual_offset(128),
Size(200),
)
.unwrap();
idx::write_index_entry(
&mut idx_data,
NeedleId(3),
Offset::from_actual_offset(384),
Size(300),
)
.unwrap();
// Delete needle 2
idx::write_index_entry(
&mut idx_data,
NeedleId(2),
Offset::default(),
TOMBSTONE_FILE_SIZE,
)
.unwrap();
let mut cursor = Cursor::new(idx_data);
let nm = RedbNeedleMap::load_from_idx(db_path.to_str().unwrap(), &mut cursor, Version::current()).unwrap();
assert!(nm.get(NeedleId(1)).is_some());
assert!(nm.get(NeedleId(2)).is_none()); // deleted and removed
assert!(nm.get(NeedleId(3)).is_some());
assert_eq!(nm.file_count(), 2);
}
#[test]
fn test_redb_needle_map_double_delete() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let mut nm = RedbNeedleMap::new(db_path.to_str().unwrap()).unwrap();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
let r1 = nm
.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(r1, Some(Size(100)));
// Second delete should return None (already deleted)
let r2 = nm
.delete(NeedleId(1), Offset::from_actual_offset(0))
.unwrap();
assert_eq!(r2, None);
assert_eq!(nm.deleted_count(), 1); // not double counted
}
#[test]
fn test_redb_needle_map_ascending_visit() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let mut nm = RedbNeedleMap::new(db_path.to_str().unwrap()).unwrap();
nm.put(NeedleId(3), Offset::from_actual_offset(384), Size(300))
.unwrap();
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
nm.put(NeedleId(2), Offset::from_actual_offset(128), Size(200))
.unwrap();
let mut visited = Vec::new();
nm.ascending_visit(|id, nv| {
visited.push((id, nv.size));
Ok(())
})
.unwrap();
assert_eq!(visited.len(), 3);
assert_eq!(visited[0], (NeedleId(1), Size(100)));
assert_eq!(visited[1], (NeedleId(2), Size(200)));
assert_eq!(visited[2], (NeedleId(3), Size(300)));
}
#[test]
fn test_redb_needle_map_save_to_idx() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let idx_path = dir.path().join("test.idx");
let mut nm = RedbNeedleMap::new(db_path.to_str().unwrap()).unwrap();
nm.put(NeedleId(1), Offset::from_actual_offset(8), Size(100))
.unwrap();
nm.put(NeedleId(2), Offset::from_actual_offset(128), Size(200))
.unwrap();
nm.put(NeedleId(3), Offset::from_actual_offset(384), Size(300))
.unwrap();
// Delete needle 2
nm.delete(NeedleId(2), Offset::from_actual_offset(128))
.unwrap();
nm.save_to_idx(idx_path.to_str().unwrap()).unwrap();
// Load back with CompactNeedleMap to verify
let mut idx_file = std::fs::File::open(&idx_path).unwrap();
let loaded = CompactNeedleMap::load_from_idx(&mut idx_file, Version::current()).unwrap();
assert_eq!(loaded.file_count(), 2); // only live entries
assert!(loaded.get(NeedleId(1)).is_some());
assert!(loaded.get(NeedleId(2)).is_none()); // deleted, not saved
assert!(loaded.get(NeedleId(3)).is_some());
}
#[test]
fn test_pack_unpack_needle_value() {
let nv = NeedleValue {
offset: Offset::from_actual_offset(8 * 1000),
size: Size(4096),
};
let packed = pack_needle_value(&nv);
let unpacked = unpack_needle_value(&packed);
assert_eq!(
nv.offset.to_actual_offset(),
unpacked.offset.to_actual_offset()
);
assert_eq!(nv.size, unpacked.size);
}
#[test]
fn test_pack_unpack_negative_size() {
let nv = NeedleValue {
offset: Offset::from_actual_offset(8 * 500),
size: Size(-100),
};
let packed = pack_needle_value(&nv);
let unpacked = unpack_needle_value(&packed);
assert_eq!(
nv.offset.to_actual_offset(),
unpacked.offset.to_actual_offset()
);
assert_eq!(nv.size, unpacked.size);
}
// ---- NeedleMap enum tests ----
#[test]
fn test_needle_map_enum_inmemory() {
let mut nm = NeedleMap::InMemory(CompactNeedleMap::new());
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
assert_eq!(nm.get(NeedleId(1)).unwrap().size, Size(100));
assert_eq!(nm.file_count(), 1);
}
#[test]
fn test_needle_map_enum_redb() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("test.rdb");
let mut nm = NeedleMap::Redb(RedbNeedleMap::new(db_path.to_str().unwrap()).unwrap());
nm.put(NeedleId(1), Offset::from_actual_offset(0), Size(100))
.unwrap();
assert_eq!(nm.get(NeedleId(1)).unwrap().size, Size(100));
assert_eq!(nm.file_count(), 1);
}
}
Reference in New Issue View Git Blame Copy Permalink