use crate::io;
use crate::transaction_tracker::{TransactionId, TransactionTracker};
#[cfg(not(redb_no_std))]
use crate::tree_store::ReadOnlyBackend;
use crate::tree_store::{
AllocationPolicy, BtreeHeader, InternalTableDefinition, PAGE_SIZE, PageHint, PageNumber,
PageResolver, ShrinkPolicy, TableTree, TableType, TransactionalMemory,
};
use crate::types::{Key, Value};
use crate::{
CompactionError, DatabaseError, Error, ReadOnlyTable, ReadableTable, SavepointError,
StorageError, TableError,
};
use crate::{ReadTransaction, Result, WriteTransaction};
use alloc::boxed::Box;
use alloc::format;
use alloc::string::String;
use alloc::string::ToString;
use core::fmt::{Debug, Display, Formatter};
use alloc::sync::Arc;
use core::marker::PhantomData;
#[cfg(not(redb_no_std))]
use std::fs::{File, OpenOptions};
#[cfg(not(redb_no_std))]
use std::path::Path;
use crate::error::TransactionError;
use crate::sealed::{Sealed, SealedInApi5};
use crate::transactions::{
ALLOCATOR_STATE_TABLE_NAME, AllocatorStateKey, AllocatorStateTree, DATA_ALLOCATED_TABLE,
DATA_FREED_TABLE, PageList, SYSTEM_FREED_TABLE, SystemTableDefinition,
TransactionIdWithPagination,
};
#[cfg(not(redb_no_std))]
use crate::tree_store::file_backend::FileBackend;
#[cfg(feature = "logging")]
use log::{debug, warn};
#[allow(clippy::len_without_is_empty)]
pub trait StorageBackend: 'static + Debug + Send + Sync {
fn len(&self) -> core::result::Result<u64, io::Error>;
fn read(&self, offset: u64, out: &mut [u8]) -> core::result::Result<(), io::Error>;
fn set_len(&self, len: u64) -> core::result::Result<(), io::Error>;
fn sync_data(&self) -> core::result::Result<(), io::Error>;
fn write(&self, offset: u64, data: &[u8]) -> core::result::Result<(), io::Error>;
fn close(&self) -> core::result::Result<(), io::Error> {
Ok(())
}
}
pub trait TableHandle: Sealed {
fn name(&self) -> &str;
}
#[derive(Clone)]
pub struct UntypedTableHandle {
name: String,
}
impl UntypedTableHandle {
pub(crate) fn new(name: String) -> Self {
Self { name }
}
}
impl TableHandle for UntypedTableHandle {
fn name(&self) -> &str {
&self.name
}
}
impl Sealed for UntypedTableHandle {}
pub trait MultimapTableHandle: Sealed {
fn name(&self) -> &str;
}
#[derive(Clone)]
pub struct UntypedMultimapTableHandle {
name: String,
}
impl UntypedMultimapTableHandle {
pub(crate) fn new(name: String) -> Self {
Self { name }
}
}
impl MultimapTableHandle for UntypedMultimapTableHandle {
fn name(&self) -> &str {
&self.name
}
}
impl Sealed for UntypedMultimapTableHandle {}
pub struct TableDefinition<'a, K: Key + 'static, V: Value + 'static> {
name: &'a str,
_key_type: PhantomData<K>,
_value_type: PhantomData<V>,
}
impl<'a, K: Key + 'static, V: Value + 'static> TableDefinition<'a, K, V> {
pub const fn new(name: &'a str) -> Self {
assert!(!name.is_empty());
Self {
name,
_key_type: PhantomData,
_value_type: PhantomData,
}
}
}
impl<K: Key + 'static, V: Value + 'static> TableHandle for TableDefinition<'_, K, V> {
fn name(&self) -> &str {
self.name
}
}
impl<K: Key, V: Value> Sealed for TableDefinition<'_, K, V> {}
impl<K: Key + 'static, V: Value + 'static> Clone for TableDefinition<'_, K, V> {
fn clone(&self) -> Self {
*self
}
}
impl<K: Key + 'static, V: Value + 'static> Copy for TableDefinition<'_, K, V> {}
impl<K: Key + 'static, V: Value + 'static> Display for TableDefinition<'_, K, V> {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
write!(
f,
"{}<{}, {}>",
self.name,
K::type_name().name(),
V::type_name().name()
)
}
}
pub struct MultimapTableDefinition<'a, K: Key + 'static, V: Key + 'static> {
name: &'a str,
_key_type: PhantomData<K>,
_value_type: PhantomData<V>,
}
impl<'a, K: Key + 'static, V: Key + 'static> MultimapTableDefinition<'a, K, V> {
pub const fn new(name: &'a str) -> Self {
assert!(!name.is_empty());
Self {
name,
_key_type: PhantomData,
_value_type: PhantomData,
}
}
}
impl<K: Key + 'static, V: Key + 'static> MultimapTableHandle for MultimapTableDefinition<'_, K, V> {
fn name(&self) -> &str {
self.name
}
}
impl<K: Key, V: Key> Sealed for MultimapTableDefinition<'_, K, V> {}
impl<K: Key + 'static, V: Key + 'static> Clone for MultimapTableDefinition<'_, K, V> {
fn clone(&self) -> Self {
*self
}
}
impl<K: Key + 'static, V: Key + 'static> Copy for MultimapTableDefinition<'_, K, V> {}
impl<K: Key + 'static, V: Key + 'static> Display for MultimapTableDefinition<'_, K, V> {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
write!(
f,
"{}<{}, {}>",
self.name,
K::type_name().name(),
V::type_name().name()
)
}
}
#[derive(Debug)]
pub struct CacheStats {
pub(crate) evictions: u64,
pub(crate) read_hits: u64,
pub(crate) read_misses: u64,
pub(crate) write_hits: u64,
pub(crate) write_misses: u64,
pub(crate) used_bytes: usize,
}
impl CacheStats {
pub fn evictions(&self) -> u64 {
self.evictions
}
pub fn read_hits(&self) -> u64 {
self.read_hits
}
pub fn read_misses(&self) -> u64 {
self.read_misses
}
pub fn write_hits(&self) -> u64 {
self.write_hits
}
pub fn write_misses(&self) -> u64 {
self.write_misses
}
pub fn used_bytes(&self) -> usize {
self.used_bytes
}
}
pub(crate) enum TransactionGuard {
Read {
tracker: Arc<TransactionTracker>,
transaction_id: TransactionId,
},
Write {
tracker: Arc<TransactionTracker>,
transaction_id: TransactionId,
},
Untracked,
}
impl TransactionGuard {
pub(crate) fn new_read(
transaction_id: TransactionId,
tracker: Arc<TransactionTracker>,
) -> Self {
Self::Read {
tracker,
transaction_id,
}
}
pub(crate) fn allocate_read(
tracker: Arc<TransactionTracker>,
mem: &TransactionalMemory,
) -> Result<Self> {
let id = tracker.register_read_transaction(mem)?;
Ok(Self::new_read(id, tracker))
}
pub(crate) fn new_write(
transaction_id: TransactionId,
tracker: Arc<TransactionTracker>,
) -> Self {
Self::Write {
tracker,
transaction_id,
}
}
pub(crate) fn untracked() -> Self {
Self::Untracked
}
pub(crate) fn id(&self) -> TransactionId {
match self {
Self::Read { transaction_id, .. } | Self::Write { transaction_id, .. } => {
*transaction_id
}
Self::Untracked => {
panic!("TransactionGuard::id() called on an untracked guard")
}
}
}
}
impl Drop for TransactionGuard {
fn drop(&mut self) {
match self {
Self::Read {
tracker,
transaction_id,
} => tracker.deallocate_read_transaction(*transaction_id),
Self::Write {
tracker,
transaction_id,
} => {
if let Some(mem) = tracker.end_write_transaction(*transaction_id) {
close_database(tracker, &mem);
}
}
Self::Untracked => {}
}
}
}
pub trait ReadableDatabase: SealedInApi5 {
fn begin_read(&self) -> Result<ReadTransaction, TransactionError>;
fn cache_stats(&self) -> CacheStats;
}
#[cfg(not(redb_no_std))]
pub struct ReadOnlyDatabase {
mem: Arc<TransactionalMemory>,
transaction_tracker: Arc<TransactionTracker>,
}
#[cfg(not(redb_no_std))]
impl Sealed for ReadOnlyDatabase {}
#[cfg(not(redb_no_std))]
impl ReadableDatabase for ReadOnlyDatabase {
fn begin_read(&self) -> Result<ReadTransaction, TransactionError> {
let id = self
.transaction_tracker
.register_read_transaction(&self.mem)?;
#[cfg(feature = "logging")]
debug!("Beginning read transaction id={id:?}");
let guard = TransactionGuard::new_read(id, self.transaction_tracker.clone());
ReadTransaction::new(self.mem.clone(), guard)
}
fn cache_stats(&self) -> CacheStats {
self.mem.cache_stats()
}
}
#[cfg(not(redb_no_std))]
impl ReadOnlyDatabase {
#[cfg(not(redb_no_std))]
pub fn open(path: impl AsRef<Path>) -> Result<ReadOnlyDatabase, DatabaseError> {
Builder::new().open_read_only(path)
}
fn new(
file: Box<dyn StorageBackend>,
page_size: usize,
region_size: Option<u64>,
cache_size: usize,
) -> Result<Self, DatabaseError> {
#[cfg(feature = "logging")]
let file_path = format!("{:?}", &file);
#[cfg(feature = "logging")]
debug!("Opening database in read-only {:?}", &file_path);
let mem = TransactionalMemory::new(
Box::new(ReadOnlyBackend::new(file)),
false,
page_size,
region_size,
cache_size,
true,
)?;
let mem = Arc::new(mem);
if let Some(tree) = Database::get_allocator_state_table(&mem)? {
mem.load_allocator_state(&tree)?;
} else {
#[cfg(feature = "logging")]
warn!(
"Database {:?} not shutdown cleanly. Repair required",
&file_path
);
return Err(DatabaseError::RepairAborted);
}
let next_transaction_id = mem.get_last_committed_transaction_id()?.next();
let db = Self {
mem,
transaction_tracker: Arc::new(TransactionTracker::new(next_transaction_id)),
};
Ok(db)
}
}
pub struct Database {
mem: Arc<TransactionalMemory>,
transaction_tracker: Arc<TransactionTracker>,
}
impl Sealed for Database {}
impl ReadableDatabase for Database {
fn begin_read(&self) -> Result<ReadTransaction, TransactionError> {
let guard = TransactionGuard::allocate_read(self.transaction_tracker.clone(), &self.mem)?;
#[cfg(feature = "logging")]
debug!("Beginning read transaction id={:?}", guard.id());
ReadTransaction::new(self.get_memory(), guard)
}
fn cache_stats(&self) -> CacheStats {
self.mem.cache_stats()
}
}
impl Database {
#[cfg(not(redb_no_std))]
pub fn create(path: impl AsRef<Path>) -> Result<Database, DatabaseError> {
Self::builder().create(path)
}
#[cfg(not(redb_no_std))]
pub fn open(path: impl AsRef<Path>) -> Result<Database, DatabaseError> {
Self::builder().open(path)
}
pub(crate) fn get_memory(&self) -> Arc<TransactionalMemory> {
self.mem.clone()
}
pub(crate) fn verify_primary_checksums(mem: Arc<TransactionalMemory>) -> Result<bool> {
let data_root = mem.get_data_root();
let system_root = mem.get_system_root();
Self::verify_checksums(mem, data_root, system_root)
}
fn verify_checksums(
mem: Arc<TransactionalMemory>,
data_root: Option<BtreeHeader>,
system_root: Option<BtreeHeader>,
) -> Result<bool> {
let resolver = PageResolver::new(mem.clone());
let table_tree = TableTree::new(
data_root,
PageHint::None,
Arc::new(TransactionGuard::untracked()),
resolver.clone(),
)?;
if !table_tree.verify_checksums()? {
return Ok(false);
}
let system_table_tree = TableTree::new(
system_root,
PageHint::None,
Arc::new(TransactionGuard::untracked()),
resolver,
)?;
if !system_table_tree.verify_checksums()? {
return Ok(false);
}
Ok(true)
}
pub fn check_integrity(&mut self) -> Result<bool, DatabaseError> {
if Arc::get_mut(&mut self.mem).is_none() {
return Err(DatabaseError::TransactionInProgress);
}
if self.transaction_tracker.any_ephemeral_savepoint_exists() {
return Err(DatabaseError::TransactionInProgress);
}
self.mem.check_io_errors()?;
if !self.mem.allocator_state_loaded() {
return Err(StorageError::Corrupted(
"Allocator state was discarded by a failed integrity check or commit; reopen the database to repair it".to_string(),
)
.into());
}
let result = self.check_integrity_inner();
if result.is_err() {
self.mem.invalidate_allocator_state();
}
result
}
fn check_integrity_inner(&mut self) -> Result<bool, DatabaseError> {
let mut rolling_back_non_durable = false;
if self.mem.pending_non_durable_commit() {
self.mem.clear_read_cache();
if self.mem.file_len_matches_layout()?
&& let Some(live_allocator_clean) = self.repair_live_state()?
{
let durable_clean = self.durable_state_clean()?;
let mut txn = self
.begin_write()
.map_err(|e| DatabaseError::Storage(e.into_storage_error()))?;
txn.disable_post_commit_free();
txn.commit()
.map_err(|e| DatabaseError::Storage(e.into_storage_error()))?;
return Ok(live_allocator_clean && durable_clean);
}
rolling_back_non_durable = true;
}
let allocator_hash = self.mem.allocator_hash();
let mem = Arc::get_mut(&mut self.mem).unwrap();
let mut was_clean = mem.clear_cache_and_reload()?;
let old_roots = [self.mem.get_data_root(), self.mem.get_system_root()];
let new_roots = Self::do_repair(&mut self.mem, &|_| {}).map_err(|err| match err {
DatabaseError::Storage(storage_err) => storage_err,
_ => unreachable!(),
})?;
if old_roots != new_roots
|| allocator_hash != self.mem.allocator_hash()
|| rolling_back_non_durable
{
was_clean = false;
}
if !was_clean {
let next_transaction_id = self.mem.get_last_committed_transaction_id()?.next();
let [data_root, system_root] = new_roots;
self.mem.commit(
data_root,
system_root,
next_transaction_id,
true,
ShrinkPolicy::Never,
)?;
}
self.mem.begin_writable()?;
Ok(was_clean)
}
fn repair_live_state(&mut self) -> Result<Option<bool>, DatabaseError> {
match Self::verify_primary_checksums(self.mem.clone()) {
Ok(true) => {
let live_allocator_hash = self.mem.allocator_hash();
let live_roots = [self.mem.get_data_root(), self.mem.get_system_root()];
match Self::rebuild_allocator_state(&mut self.mem, &|_| {}) {
Ok(roots) if roots != live_roots => Ok(None),
Ok(_) => Ok(Some(live_allocator_hash == self.mem.allocator_hash())),
Err(DatabaseError::Storage(StorageError::Corrupted(_))) => Ok(None),
Err(err) => Err(err),
}
}
Ok(false) | Err(StorageError::Corrupted(_)) => Ok(None),
Err(err) => Err(err.into()),
}
}
fn durable_state_clean(&self) -> Result<bool, DatabaseError> {
match self.verify_durable_state() {
Ok(clean) => Ok(clean),
Err(DatabaseError::Storage(StorageError::Corrupted(_))) => Ok(false),
Err(err) => Err(err),
}
}
fn verify_durable_state(&self) -> Result<bool, DatabaseError> {
if self.mem.durable_primary_slot_corrupt()? {
return Ok(false);
}
let data_root = self.mem.get_durable_data_root();
let system_root = self.mem.get_durable_system_root();
Ok(Self::verify_checksums(
self.mem.clone(),
data_root,
system_root,
)?)
}
pub fn compact(&mut self) -> Result<bool, CompactionError> {
if self.transaction_tracker.any_persistent_savepoint_exists() {
return Err(CompactionError::PersistentSavepointExists);
}
if self.transaction_tracker.any_savepoint_exists() {
return Err(CompactionError::EphemeralSavepointExists);
}
if self.transaction_tracker.any_user_read_reference_exists() {
return Err(CompactionError::TransactionInProgress);
}
let txn = self.begin_write().map_err(|e| e.into_storage_error())?;
if txn.list_persistent_savepoints()?.next().is_some() {
return Err(CompactionError::PersistentSavepointExists);
}
if self.transaction_tracker.any_savepoint_exists() {
return Err(CompactionError::EphemeralSavepointExists);
}
if self.transaction_tracker.any_user_read_reference_exists() {
return Err(CompactionError::TransactionInProgress);
}
txn.abort()?;
self.drain_pending_free_pages(ShrinkPolicy::Maximum)?;
let mut compacted = false;
loop {
let mut progress = false;
let mut txn = self.begin_write().map_err(|e| e.into_storage_error())?;
if txn.compact_pages()? {
progress = true;
txn.commit().map_err(|e| e.into_storage_error())?;
} else {
txn.abort()?;
}
self.drain_pending_free_pages(ShrinkPolicy::Maximum)?;
if !progress {
break;
}
compacted = true;
}
Ok(compacted)
}
fn drain_pending_free_pages(&self, shrink_policy: ShrinkPolicy) -> Result {
let mut force_commit = true;
loop {
let mut txn = self.begin_write().map_err(|e| e.into_storage_error())?;
if !force_commit && !txn.pending_free_pages()? {
txn.abort()?;
return Ok(());
}
force_commit = false;
txn.set_two_phase_commit(true);
txn.set_shrink_policy(shrink_policy);
txn.commit().map_err(|e| e.into_storage_error())?;
}
}
#[cfg_attr(not(debug_assertions), expect(dead_code))]
fn check_repaired_allocated_pages_table(
system_root: Option<BtreeHeader>,
mem: Arc<TransactionalMemory>,
) -> Result {
let resolver = PageResolver::new(mem.clone());
let table_tree = TableTree::new(
system_root,
PageHint::None,
Arc::new(TransactionGuard::untracked()),
resolver.clone(),
)?;
if let Some(table_def) = table_tree
.get_table::<TransactionIdWithPagination, PageList>(
DATA_ALLOCATED_TABLE.name(),
TableType::Normal,
)
.map_err(|e| e.into_storage_error_or_corrupted("Allocated pages table corrupted"))?
{
let InternalTableDefinition::Normal { table_root, .. } = table_def else {
unreachable!()
};
let table: ReadOnlyTable<TransactionIdWithPagination, PageList> = ReadOnlyTable::new(
DATA_ALLOCATED_TABLE.name().to_string(),
table_root,
PageHint::None,
Arc::new(TransactionGuard::untracked()),
resolver,
)?;
for result in ReadableTable::iter(&table)? {
let (_, pages) = result?;
for i in 0..pages.value().len() {
assert!(mem.is_allocated(pages.value().get(i)));
}
}
}
Ok(())
}
fn visit_freed_tree<K: Key, V: Value, F>(
system_root: Option<BtreeHeader>,
table_def: SystemTableDefinition<K, V>,
mem: Arc<TransactionalMemory>,
mut visitor: F,
) -> Result
where
F: FnMut(PageNumber) -> Result,
{
let untracked_guard = Arc::new(TransactionGuard::untracked());
let resolver = PageResolver::new(mem.clone());
let system_tree = TableTree::new(
system_root,
PageHint::None,
untracked_guard,
resolver.clone(),
)?;
let table_name = table_def.name();
let result = match system_tree.get_table::<K, V>(table_name, TableType::Normal) {
Ok(result) => result,
Err(TableError::Storage(err)) => {
return Err(err);
}
Err(TableError::TableDoesNotExist(_)) => {
return Ok(());
}
Err(_) => {
return Err(StorageError::Corrupted(format!(
"Unable to open {table_name}"
)));
}
};
if let Some(definition) = result {
let table_root = match definition {
InternalTableDefinition::Normal { table_root, .. } => table_root,
InternalTableDefinition::Multimap { .. } => unreachable!(),
};
let table: ReadOnlyTable<TransactionIdWithPagination, PageList<'static>> =
ReadOnlyTable::new(
table_name.to_string(),
table_root,
PageHint::None,
Arc::new(TransactionGuard::untracked()),
resolver,
)?;
for result in ReadableTable::iter(&table)? {
let (_, page_list) = result?;
for i in 0..page_list.value().len() {
visitor(page_list.value().get(i))?;
}
}
}
Ok(())
}
#[cfg(debug_assertions)]
fn mark_allocated_page_for_debug(
mem: &mut Arc<TransactionalMemory>, ) -> Result {
let data_root = mem.get_data_root();
{
let untracked = Arc::new(TransactionGuard::untracked());
let tables = TableTree::new(
data_root,
PageHint::None,
untracked,
PageResolver::new(mem.clone()),
)?;
tables.visit_all_pages(|path| {
mem.mark_debug_allocated_page(path.page_number());
Ok(())
})?;
}
let system_root = mem.get_system_root();
{
let untracked = Arc::new(TransactionGuard::untracked());
let system_tables = TableTree::new(
system_root,
PageHint::None,
untracked,
PageResolver::new(mem.clone()),
)?;
system_tables.visit_all_pages(|path| {
mem.mark_debug_allocated_page(path.page_number());
Ok(())
})?;
}
Self::visit_freed_tree(system_root, DATA_FREED_TABLE, mem.clone(), |page| {
mem.mark_debug_allocated_page(page);
Ok(())
})?;
Self::visit_freed_tree(system_root, SYSTEM_FREED_TABLE, mem.clone(), |page| {
mem.mark_debug_allocated_page(page);
Ok(())
})?;
Ok(())
}
fn primary_verifies(mem: &Arc<TransactionalMemory>) -> Result<bool> {
match Self::verify_primary_checksums(mem.clone()) {
Ok(verified) => Ok(verified),
Err(StorageError::Corrupted(_)) => Ok(false),
Err(err) => Err(err),
}
}
fn do_repair(
mem: &mut Arc<TransactionalMemory>, repair_callback: &(dyn Fn(&mut RepairSession) + 'static),
) -> Result<[Option<BtreeHeader>; 2], DatabaseError> {
if !Self::primary_verifies(mem)? {
if mem.used_two_phase_commit() {
return Err(DatabaseError::Storage(StorageError::Corrupted(
"Primary is corrupted despite 2-phase commit".to_string(),
)));
}
let mut handle = RepairSession::new(0.3);
repair_callback(&mut handle);
if handle.aborted() {
return Err(DatabaseError::RepairAborted);
}
mem.repair_primary_corrupted();
mem.clear_read_cache();
if !Self::primary_verifies(mem)? {
return Err(DatabaseError::Storage(StorageError::Corrupted(
"Failed to repair database. All roots are corrupted".to_string(),
)));
}
}
let mut handle = RepairSession::new(0.6);
repair_callback(&mut handle);
if handle.aborted() {
return Err(DatabaseError::RepairAborted);
}
let [data_root, system_root] = Self::rebuild_allocator_state(mem, repair_callback)?;
mem.clear_recovery_required()?;
mem.clear_read_cache();
Ok([data_root, system_root])
}
fn rebuild_allocator_state(
mem: &mut Arc<TransactionalMemory>, repair_callback: &(dyn Fn(&mut RepairSession) + 'static),
) -> Result<[Option<BtreeHeader>; 2], DatabaseError> {
mem.reset_allocator_state()?;
let data_root = {
let root = mem.get_data_root();
let untracked = Arc::new(TransactionGuard::untracked());
let tables = TableTree::new(
root,
PageHint::None,
untracked,
PageResolver::new(mem.clone()),
)?;
tables.visit_all_pages(|path| mem.mark_page_allocated(path.page_number()))?;
Self::with_recounted_length(root, tables.count_tables()?)
};
let mut handle = RepairSession::new(0.9);
repair_callback(&mut handle);
if handle.aborted() {
return Err(DatabaseError::RepairAborted);
}
let system_root = {
let root = mem.get_system_root();
let untracked = Arc::new(TransactionGuard::untracked());
let system_tables = TableTree::new(
root,
PageHint::None,
untracked,
PageResolver::new(mem.clone()),
)?;
system_tables.visit_all_pages(|path| mem.mark_page_allocated(path.page_number()))?;
Self::with_recounted_length(root, system_tables.count_tables()?)
};
Self::visit_freed_tree(system_root, DATA_FREED_TABLE, mem.clone(), |page| {
mem.mark_page_allocated(page)
})?;
Self::visit_freed_tree(system_root, SYSTEM_FREED_TABLE, mem.clone(), |page| {
mem.mark_page_allocated(page)
})?;
for page in mem.unpersisted_data_freed_pages() {
mem.mark_page_allocated(page)?;
}
#[cfg(debug_assertions)]
{
Self::check_repaired_allocated_pages_table(system_root, mem.clone())?;
}
Ok([data_root, system_root])
}
fn with_recounted_length(root: Option<BtreeHeader>, length: u64) -> Option<BtreeHeader> {
root.map(|header| BtreeHeader::new(header.root, header.checksum, length))
}
fn new(
file: Box<dyn StorageBackend>,
allow_initialize: bool,
page_size: usize,
region_size: Option<u64>,
cache_size: usize,
repair_callback: &(dyn Fn(&mut RepairSession) + 'static),
) -> Result<Self, DatabaseError> {
#[cfg(feature = "logging")]
let file_path = format!("{:?}", &file);
#[cfg(feature = "logging")]
debug!("Opening database {:?}", &file_path);
let mem = TransactionalMemory::new(
file,
allow_initialize,
page_size,
region_size,
cache_size,
false,
)?;
let mut mem = Arc::new(mem);
if let Some(tree) = Self::get_allocator_state_table(&mem)? {
#[cfg(feature = "logging")]
debug!("Found valid allocator state, full repair not needed");
mem.load_allocator_state(&tree)?;
#[cfg(debug_assertions)]
Self::mark_allocated_page_for_debug(&mut mem)?;
} else {
#[cfg(feature = "logging")]
warn!("Database {:?} not shutdown cleanly. Repairing", &file_path);
let mut handle = RepairSession::new(0.0);
repair_callback(&mut handle);
if handle.aborted() {
return Err(DatabaseError::RepairAborted);
}
let [data_root, system_root] = Self::do_repair(&mut mem, repair_callback)?;
let next_transaction_id = mem.get_last_committed_transaction_id()?.next();
mem.commit(
data_root,
system_root,
next_transaction_id,
true,
ShrinkPolicy::Never,
)?;
}
mem.begin_writable()?;
let next_transaction_id = mem.get_last_committed_transaction_id()?.next();
let db = Database {
mem,
transaction_tracker: Arc::new(TransactionTracker::new(next_transaction_id)),
};
let txn = db.begin_write().map_err(|e| e.into_storage_error())?;
if let Some(next_id) = txn.next_persistent_savepoint_id()? {
db.transaction_tracker
.restore_savepoint_counter_state(next_id);
}
for id in txn.list_persistent_savepoints()? {
let savepoint = match txn.get_persistent_savepoint(id) {
Ok(savepoint) => savepoint,
Err(err) => match err {
SavepointError::InvalidSavepoint
| SavepointError::ImmediateDurabilityRequired => unreachable!(),
SavepointError::Storage(storage) => {
return Err(storage.into());
}
},
};
db.transaction_tracker
.register_persistent_savepoint(&savepoint);
}
txn.abort()?;
Ok(db)
}
fn get_allocator_state_table(
mem: &Arc<TransactionalMemory>,
) -> Result<Option<AllocatorStateTree>> {
if !mem.used_two_phase_commit() {
return Ok(None);
}
let resolver = PageResolver::new(mem.clone());
let system_table_tree = TableTree::new(
mem.get_system_root(),
PageHint::None,
Arc::new(TransactionGuard::untracked()),
resolver.clone(),
)?;
let Some(allocator_state_table) = system_table_tree
.get_table::<AllocatorStateKey, &[u8]>(ALLOCATOR_STATE_TABLE_NAME, TableType::Normal)
.map_err(|e| e.into_storage_error_or_corrupted("Unexpected TableError"))?
else {
return Ok(None);
};
let InternalTableDefinition::Normal { table_root, .. } = allocator_state_table else {
unreachable!();
};
let tree = AllocatorStateTree::new(
table_root,
PageHint::None,
Arc::new(TransactionGuard::untracked()),
resolver,
)?;
if !mem.is_valid_allocator_state(&tree)? {
return Ok(None);
}
Ok(Some(tree))
}
pub fn builder() -> Builder {
Builder::new()
}
pub fn begin_write(&self) -> Result<WriteTransaction, TransactionError> {
begin_write_with_allocation_policy(
&self.transaction_tracker,
&self.mem,
AllocationPolicy::Default,
)
}
}
fn begin_write_with_allocation_policy(
transaction_tracker: &Arc<TransactionTracker>,
mem: &Arc<TransactionalMemory>,
allocation_policy: AllocationPolicy,
) -> Result<WriteTransaction, TransactionError> {
mem.check_io_errors()?;
let guard = TransactionGuard::new_write(
transaction_tracker.start_write_transaction(),
transaction_tracker.clone(),
);
mem.check_io_errors()?;
if !mem.allocator_state_loaded() {
return Err(StorageError::Corrupted(
"Allocator state was discarded by a failed integrity check or commit; reopen the database to repair it".to_string(),
)
.into());
}
WriteTransaction::new(
guard,
transaction_tracker.clone(),
mem.clone(),
allocation_policy,
)
.map_err(|e| e.into())
}
fn ensure_allocator_state_table_and_trim(
transaction_tracker: &Arc<TransactionTracker>,
mem: &Arc<TransactionalMemory>,
) -> Result<(), Error> {
#[cfg(feature = "logging")]
debug!("Writing allocator state table");
let mut tx =
begin_write_with_allocation_policy(transaction_tracker, mem, AllocationPolicy::Lowest)?;
tx.set_quick_repair(true);
tx.disable_post_commit_free();
tx.set_shrink_policy(ShrinkPolicy::Maximum);
tx.commit()?;
Ok(())
}
fn close_database(transaction_tracker: &Arc<TransactionTracker>, mem: &Arc<TransactionalMemory>) {
if !crate::panicking()
&& ensure_allocator_state_table_and_trim(transaction_tracker, mem).is_err()
{
#[cfg(feature = "logging")]
warn!("Failed to write allocator state table. Repair may be required at restart.");
}
if mem.close().is_err() {
#[cfg(feature = "logging")]
warn!("Failed to flush database file. Repair may be required at restart.");
}
}
impl Drop for Database {
fn drop(&mut self) {
if self
.transaction_tracker
.defer_close_if_write_transaction_live(&self.mem)
{
#[cfg(feature = "logging")]
warn!(
"Database dropped while a write transaction is in progress. The database will remain open until the write transaction completes."
);
return;
}
close_database(&self.transaction_tracker, &self.mem);
}
}
pub struct RepairSession {
progress: f64,
aborted: bool,
}
impl RepairSession {
pub(crate) fn new(progress: f64) -> Self {
Self {
progress,
aborted: false,
}
}
pub(crate) fn aborted(&self) -> bool {
self.aborted
}
pub fn abort(&mut self) {
self.aborted = true;
}
pub fn progress(&self) -> f64 {
self.progress
}
}
pub struct Builder {
page_size: usize,
region_size: Option<u64>,
cache_size: usize,
repair_callback: Box<dyn Fn(&mut RepairSession)>,
}
impl Builder {
#[allow(clippy::new_without_default)]
pub fn new() -> Self {
Self {
page_size: PAGE_SIZE,
region_size: None,
cache_size: 1024 * 1024 * 1024,
repair_callback: Box::new(|_| {}),
}
}
pub fn set_repair_callback(
&mut self,
callback: impl Fn(&mut RepairSession) + 'static,
) -> &mut Self {
self.repair_callback = Box::new(callback);
self
}
#[cfg(any(fuzzing, test))]
pub fn set_page_size(&mut self, size: usize) -> &mut Self {
assert!(size.is_power_of_two());
self.page_size = core::cmp::max(size, 512);
self
}
pub fn set_cache_size(&mut self, bytes: usize) -> &mut Self {
self.cache_size = bytes;
self
}
#[cfg(any(test, fuzzing))]
pub fn set_region_size(&mut self, size: u64) -> &mut Self {
assert!(size.is_power_of_two());
self.region_size = Some(size);
self
}
#[cfg(not(redb_no_std))]
pub fn create(&self, path: impl AsRef<Path>) -> Result<Database, DatabaseError> {
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(path)?;
Database::new(
Box::new(FileBackend::new(file)?),
true,
self.page_size,
self.region_size,
self.cache_size,
&self.repair_callback,
)
}
#[cfg(not(redb_no_std))]
pub fn open(&self, path: impl AsRef<Path>) -> Result<Database, DatabaseError> {
let file = OpenOptions::new().read(true).write(true).open(path)?;
Database::new(
Box::new(FileBackend::new(file)?),
false,
self.page_size,
None,
self.cache_size,
&self.repair_callback,
)
}
#[cfg(not(redb_no_std))]
pub fn open_read_only(
&self,
path: impl AsRef<Path>,
) -> Result<ReadOnlyDatabase, DatabaseError> {
let file = OpenOptions::new().read(true).open(path)?;
ReadOnlyDatabase::new(
Box::new(FileBackend::new_internal(file, true)?),
self.page_size,
None,
self.cache_size,
)
}
#[cfg(not(redb_no_std))]
pub fn create_file(&self, file: File) -> Result<Database, DatabaseError> {
Database::new(
Box::new(FileBackend::new(file)?),
true,
self.page_size,
self.region_size,
self.cache_size,
&self.repair_callback,
)
}
pub fn create_with_backend(
&self,
backend: impl StorageBackend,
) -> Result<Database, DatabaseError> {
Database::new(
Box::new(backend),
true,
self.page_size,
self.region_size,
self.cache_size,
&self.repair_callback,
)
}
}
impl core::fmt::Debug for Database {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Database").finish()
}
}
#[cfg(test)]
mod test {
use crate::backends::FileBackend;
use crate::{
CommitError, Database, DatabaseError, Durability, ReadableTable, StorageBackend,
StorageError, TableDefinition, TransactionError,
};
use alloc::sync::Arc;
use core::sync::atomic::{AtomicU64, Ordering};
use std::fs::File;
use std::io::{ErrorKind, Read, Seek, SeekFrom};
#[derive(Debug)]
struct FailingBackend {
inner: FileBackend,
countdown: Arc<AtomicU64>,
}
impl FailingBackend {
fn new(backend: FileBackend, countdown: u64) -> Self {
Self {
inner: backend,
countdown: Arc::new(AtomicU64::new(countdown)),
}
}
fn check_countdown(&self) -> Result<(), std::io::Error> {
if self.countdown.load(Ordering::SeqCst) == 0 {
return Err(std::io::Error::from(ErrorKind::Other));
}
Ok(())
}
fn decrement_countdown(&self) -> Result<(), std::io::Error> {
if self
.countdown
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |x| {
if x > 0 { Some(x - 1) } else { None }
})
.is_err()
{
return Err(std::io::Error::from(ErrorKind::Other));
}
Ok(())
}
}
impl StorageBackend for FailingBackend {
fn len(&self) -> Result<u64, std::io::Error> {
self.inner.len()
}
fn read(&self, offset: u64, out: &mut [u8]) -> Result<(), std::io::Error> {
self.check_countdown()?;
self.inner.read(offset, out)
}
fn set_len(&self, len: u64) -> Result<(), std::io::Error> {
self.inner.set_len(len)
}
fn sync_data(&self) -> Result<(), std::io::Error> {
self.check_countdown()?;
self.inner.sync_data()
}
fn write(&self, offset: u64, data: &[u8]) -> Result<(), std::io::Error> {
self.decrement_countdown()?;
self.inner.write(offset, data)
}
}
#[test]
fn crash_regression4() {
let tmpfile = crate::create_tempfile();
let (file, path) = tmpfile.into_parts();
let backend = FailingBackend::new(FileBackend::new(file).unwrap(), 20);
let db = Database::builder()
.set_cache_size(12686)
.set_page_size(8 * 1024)
.set_region_size(32 * 4096)
.create_with_backend(backend)
.unwrap();
let table_def: TableDefinition<u64, &[u8]> = TableDefinition::new("x");
let tx = db.begin_write().unwrap();
let _savepoint = tx.ephemeral_savepoint().unwrap();
let _persistent_savepoint = tx.persistent_savepoint().unwrap();
tx.commit().unwrap();
let tx = db.begin_write().unwrap();
{
let mut table = tx.open_table(table_def).unwrap();
let _ = table.insert_reserve(118821, 360).unwrap();
}
let result = tx.commit();
assert!(result.is_err());
drop(db);
Database::builder()
.set_cache_size(1024 * 1024)
.set_page_size(8 * 1024)
.set_region_size(32 * 4096)
.create(&path)
.unwrap();
}
#[test]
fn transient_io_error() {
let tmpfile = crate::create_tempfile();
let (file, path) = tmpfile.into_parts();
let backend = FailingBackend::new(FileBackend::new(file).unwrap(), u64::MAX);
let countdown = backend.countdown.clone();
let db = Database::builder()
.set_cache_size(0)
.create_with_backend(backend)
.unwrap();
let table_def: TableDefinition<u64, u64> = TableDefinition::new("x");
let tx = db.begin_write().unwrap();
{
let mut table = tx.open_table(table_def).unwrap();
table.insert(0, 0).unwrap();
}
tx.commit().unwrap();
let tx = db.begin_write().unwrap();
{
let mut table = tx.open_table(table_def).unwrap();
table.insert(0, 1).unwrap();
}
tx.commit().unwrap();
let tx = db.begin_write().unwrap();
countdown.store(0, Ordering::SeqCst);
let result = tx.commit().err().unwrap();
assert!(matches!(result, CommitError::Storage(StorageError::Io(_))));
let result = db.begin_write().err().unwrap();
assert!(matches!(
result,
TransactionError::Storage(StorageError::PreviousIo)
));
countdown.store(u64::MAX, Ordering::SeqCst);
drop(db);
let mut file = File::open(&path).unwrap();
file.seek(SeekFrom::Start(9)).unwrap();
let mut god_byte = vec![0u8];
assert_eq!(file.read(&mut god_byte).unwrap(), 1);
assert_ne!(god_byte[0] & 2, 0);
}
#[test]
fn small_pages() {
let tmpfile = crate::create_tempfile();
let db = Database::builder()
.set_page_size(512)
.create(tmpfile.path())
.unwrap();
let table_definition: TableDefinition<u64, &[u8]> = TableDefinition::new("x");
let txn = db.begin_write().unwrap();
{
txn.open_table(table_definition).unwrap();
}
txn.commit().unwrap();
}
#[test]
fn small_pages2() {
let tmpfile = crate::create_tempfile();
let db = Database::builder()
.set_page_size(512)
.create(tmpfile.path())
.unwrap();
let table_def: TableDefinition<u64, &[u8]> = TableDefinition::new("x");
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
let savepoint0 = tx.ephemeral_savepoint().unwrap();
{
tx.open_table(table_def).unwrap();
}
tx.commit().unwrap();
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
let savepoint1 = tx.ephemeral_savepoint().unwrap();
tx.restore_savepoint(&savepoint0).unwrap();
tx.set_durability(Durability::None).unwrap();
{
let mut t = tx.open_table(table_def).unwrap();
t.insert_reserve(&660503, 489).unwrap().as_mut().fill(0xFF);
assert!(t.remove(&291295).unwrap().is_none());
}
tx.commit().unwrap();
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
tx.restore_savepoint(&savepoint0).unwrap();
{
tx.open_table(table_def).unwrap();
}
tx.commit().unwrap();
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
let savepoint2 = tx.ephemeral_savepoint().unwrap();
drop(savepoint0);
tx.restore_savepoint(&savepoint2).unwrap();
{
let mut t = tx.open_table(table_def).unwrap();
assert!(t.get(&2059).unwrap().is_none());
assert!(t.remove(&145227).unwrap().is_none());
assert!(t.remove(&145227).unwrap().is_none());
}
tx.commit().unwrap();
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
let savepoint3 = tx.ephemeral_savepoint().unwrap();
drop(savepoint1);
tx.restore_savepoint(&savepoint3).unwrap();
{
tx.open_table(table_def).unwrap();
}
tx.commit().unwrap();
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
let savepoint4 = tx.ephemeral_savepoint().unwrap();
drop(savepoint2);
tx.restore_savepoint(&savepoint3).unwrap();
tx.set_durability(Durability::None).unwrap();
{
let mut t = tx.open_table(table_def).unwrap();
assert!(t.remove(&207936).unwrap().is_none());
}
tx.abort().unwrap();
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
let _savepoint5 = tx.ephemeral_savepoint().unwrap();
drop(savepoint3);
tx.restore_savepoint(&savepoint4).unwrap();
{
tx.open_table(table_def).unwrap();
}
tx.commit().unwrap();
let mut tx = db.begin_write().unwrap();
tx.set_two_phase_commit(true);
tx.restore_savepoint(&savepoint4).unwrap();
tx.set_durability(Durability::None).unwrap();
{
tx.open_table(table_def).unwrap();
}
tx.commit().unwrap();
}
#[test]
fn small_pages3() {
let tmpfile = crate::create_tempfile();
let db = Database::builder()
.set_page_size(1024)
.create(tmpfile.path())
.unwrap();
let table_def: TableDefinition<u64, &[u8]> = TableDefinition::new("x");
let mut tx = db.begin_write().unwrap();
let _savepoint0 = tx.ephemeral_savepoint().unwrap();
tx.set_durability(Durability::None).unwrap();
{
let mut t = tx.open_table(table_def).unwrap();
let value = vec![0; 306];
t.insert(&539717, value.as_slice()).unwrap();
}
tx.abort().unwrap();
let mut tx = db.begin_write().unwrap();
let savepoint1 = tx.ephemeral_savepoint().unwrap();
tx.restore_savepoint(&savepoint1).unwrap();
tx.set_durability(Durability::None).unwrap();
{
let mut t = tx.open_table(table_def).unwrap();
let value = vec![0; 2008];
t.insert(&784384, value.as_slice()).unwrap();
}
tx.abort().unwrap();
}
#[test]
fn small_pages4() {
let tmpfile = crate::create_tempfile();
let db = Database::builder()
.set_cache_size(1024 * 1024)
.set_page_size(1024)
.create(tmpfile.path())
.unwrap();
let table_def: TableDefinition<u64, &[u8]> = TableDefinition::new("x");
let tx = db.begin_write().unwrap();
{
tx.open_table(table_def).unwrap();
}
tx.commit().unwrap();
let tx = db.begin_write().unwrap();
{
let mut t = tx.open_table(table_def).unwrap();
assert!(t.get(&131072).unwrap().is_none());
let value = vec![0xFF; 1130];
t.insert(&42394, value.as_slice()).unwrap();
t.insert_reserve(&744037, 3645).unwrap().as_mut().fill(0xFF);
assert!(t.get(&0).unwrap().is_none());
}
tx.abort().unwrap();
let tx = db.begin_write().unwrap();
{
let mut t = tx.open_table(table_def).unwrap();
t.insert_reserve(&118749, 734).unwrap().as_mut().fill(0xFF);
}
tx.abort().unwrap();
}
#[test]
fn dynamic_shrink() {
let tmpfile = crate::create_tempfile();
let table_definition: TableDefinition<u64, &[u8]> = TableDefinition::new("x");
let big_value = vec![0u8; 1024];
let db = Database::builder()
.set_region_size(1024 * 1024)
.create(tmpfile.path())
.unwrap();
let txn = db.begin_write().unwrap();
{
let mut table = txn.open_table(table_definition).unwrap();
for i in 0..2048 {
table.insert(&i, big_value.as_slice()).unwrap();
}
}
txn.commit().unwrap();
let file_size = tmpfile.as_file().metadata().unwrap().len();
let txn = db.begin_write().unwrap();
{
let mut table = txn.open_table(table_definition).unwrap();
for i in 0..2048 {
table.remove(&i).unwrap();
}
}
txn.commit().unwrap();
let txn = db.begin_write().unwrap();
{
let mut table = txn.open_table(table_definition).unwrap();
table.insert(0, [].as_slice()).unwrap();
}
txn.commit().unwrap();
let txn = db.begin_write().unwrap();
{
let mut table = txn.open_table(table_definition).unwrap();
table.remove(0).unwrap();
}
txn.commit().unwrap();
let txn = db.begin_write().unwrap();
txn.commit().unwrap();
let final_file_size = tmpfile.as_file().metadata().unwrap().len();
assert!(final_file_size < file_size);
}
#[test]
fn create_new_db_in_empty_file() {
let tmpfile = crate::create_tempfile();
let _db = Database::builder()
.create_file(tmpfile.into_file())
.unwrap();
}
#[test]
fn open_missing_file() {
let tmpfile = crate::create_tempfile();
let err = Database::builder()
.open(tmpfile.path().with_extension("missing"))
.unwrap_err();
match err {
DatabaseError::Storage(StorageError::Io(err)) if err.kind() == ErrorKind::NotFound => {}
err => panic!("Unexpected error for empty file: {err}"),
}
}
#[test]
fn open_empty_file() {
let tmpfile = crate::create_tempfile();
let err = Database::builder().open(tmpfile.path()).unwrap_err();
match err {
DatabaseError::Storage(StorageError::Io(err))
if err.kind() == ErrorKind::InvalidData => {}
err => panic!("Unexpected error for empty file: {err}"),
}
}
}