#![allow(clippy::type_complexity)]
use std::hash::Hash;
use std::sync::Arc;
use std::time::Duration;
use std::{collections::BTreeMap, path::PathBuf};
use crate::sync::{Mutex, lock};
use armour_core::persist::Persist;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use crate::armour::multi_tx::{MultiTx, unique_sorted_shards};
use crate::compaction::Compactor;
use crate::error::SchemaMismatchKind;
use crate::flusher::Flusher;
use crate::hook::TypedWriteHook;
use crate::shutdown::ShutdownSignal;
use crate::{Codec, CollectionMeta, Config, DbError, DbResult, FixedConfig, Key, TreeMeta};
use crate::{TypedMap, TypedTree, ZeroMap, ZeroTree};
use super::collection::Collection;
use super::migration::TypedMigration;
use super::seq::SeqGen;
#[derive(Serialize, Deserialize, Default, Debug, Clone)]
pub struct DbInfo {
#[serde(default)]
pub version: u32,
pub collections: BTreeMap<String, CollectionInfo>,
}
#[derive(Serialize, Deserialize, Default, Debug, Clone, Copy)]
pub struct CollectionInfo {
pub version: u16,
pub typ_hash: u64,
#[serde(default)]
pub seq: u64,
}
#[derive(Serialize, Deserialize, Default, Debug, Clone)]
struct UserMeta {
db: Option<serde_json::Value>,
collections: BTreeMap<String, serde_json::Value>,
}
const DEFAULT_COMPACTION_INTERVAL: Duration = Duration::from_secs(60);
const DEFAULT_FLUSH_INTERVAL: Duration = Duration::from_secs(1);
pub struct Db {
pub path: PathBuf,
pub db_info: Persist<DbInfo>,
user_meta: Persist<UserMeta>,
pub seq: Arc<SeqGen>,
pub(crate) shutdown: ShutdownSignal,
collections: Arc<Mutex<Vec<Arc<dyn Collection>>>>,
#[cfg(feature = "rpc")]
handlers: Arc<Mutex<Vec<(u64, Arc<dyn super::handler::RpcHandler>)>>>,
#[cfg(feature = "rpc")]
pub(crate) rpc_handles: Mutex<Vec<super::rpc::RpcHandle>>,
compactor: Option<Compactor>,
flusher: Option<Flusher>,
}
impl Db {
pub fn open(path: impl AsRef<std::path::Path>) -> DbResult<Self> {
Self::open_with_options(
path,
Some(DEFAULT_COMPACTION_INTERVAL),
Some(DEFAULT_FLUSH_INTERVAL),
true,
)
}
pub fn open_test(path: impl AsRef<std::path::Path>) -> DbResult<Self> {
Self::open_with_options(path, None, None, false)
}
pub fn open_with_compaction(
path: impl AsRef<std::path::Path>,
compaction_interval: Option<Duration>,
) -> DbResult<Self> {
Self::open_with_options(
path,
compaction_interval,
Some(DEFAULT_FLUSH_INTERVAL),
true,
)
}
pub fn open_with_intervals(
path: impl AsRef<std::path::Path>,
compaction_interval: Option<Duration>,
flush_interval: Option<Duration>,
) -> DbResult<Self> {
Self::open_with_options(path, compaction_interval, flush_interval, true)
}
fn open_with_options(
path: impl AsRef<std::path::Path>,
compaction_interval: Option<Duration>,
flush_interval: Option<Duration>,
sync_persist: bool,
) -> DbResult<Self> {
let path = path.as_ref().to_path_buf();
std::fs::create_dir_all(&path).map_err(crate::DbError::Io)?;
let db_info = if sync_persist {
Persist::open(path.join("db.info"))?
} else {
Persist::open_no_sync(path.join("db.info"))?
};
let user_meta = if sync_persist {
Persist::open(path.join("db.user"))?
} else {
Persist::open_no_sync(path.join("db.user"))?
};
let seq = SeqGen::open(path.join("__seq"))?;
let collections: Arc<Mutex<Vec<Arc<dyn Collection>>>> = Arc::new(Mutex::new(Vec::new()));
let shutdown = ShutdownSignal::new();
let compactor = compaction_interval.map(|interval| {
let cols = collections.clone();
Compactor::start_with_signal(
move || {
let snapshot = lock(&cols).clone();
let mut total = 0;
for c in &snapshot {
total += c.compact()?;
}
Ok(total)
},
interval,
shutdown.clone(),
)
});
let flusher = flush_interval.map(|interval| {
let cols = collections.clone();
Flusher::start_with_signal(
move || {
let snapshot = lock(&cols).clone();
let mut first_err = None;
for c in &snapshot {
if let Err(e) = c.periodic_flush()
&& first_err.is_none()
{
first_err = Some(e);
}
}
match first_err {
Some(e) => Err(e),
None => Ok(()),
}
},
interval,
shutdown.clone(),
)
});
Ok(Self {
path,
db_info,
user_meta,
seq,
shutdown,
collections,
#[cfg(feature = "rpc")]
handlers: Arc::new(Mutex::new(Vec::new())),
#[cfg(feature = "rpc")]
rpc_handles: Mutex::new(Vec::new()),
compactor,
flusher,
})
}
pub fn version(&self) -> u32 {
self.db_info.cloned().version
}
pub fn set_version(&self, f: impl FnOnce(u32) -> u32) -> DbResult<u32> {
let mut new_version = 0;
self.db_info.update(|info| {
new_version = f(info.version);
info.version = new_version;
})?;
Ok(new_version)
}
pub fn metadata<M: DeserializeOwned>(&self) -> Option<M> {
let um = self.user_meta.cloned();
um.db.and_then(|v| serde_json::from_value(v).ok())
}
pub fn set_metadata<M: Serialize + DeserializeOwned>(
&self,
f: impl FnOnce(Option<M>) -> M,
) -> DbResult<M> {
let mut result = None;
self.user_meta.update(|um| {
let current = um
.db
.as_ref()
.and_then(|v| serde_json::from_value(v.clone()).ok());
let new_val = f(current);
um.db = Some(serde_json::to_value(&new_val).expect("serialize metadata"));
result = Some(new_val);
})?;
Ok(result.expect("update executed"))
}
pub fn collection_metadata<M: DeserializeOwned>(&self, name: &str) -> Option<M> {
let um = self.user_meta.cloned();
um.collections
.get(name)
.and_then(|v| serde_json::from_value(v.clone()).ok())
}
pub fn set_collection_metadata<M: Serialize + DeserializeOwned>(
&self,
name: &str,
f: impl FnOnce(Option<M>) -> M,
) -> DbResult<M> {
let mut result = None;
self.user_meta.update(|um| {
let current = um
.collections
.get(name)
.and_then(|v| serde_json::from_value(v.clone()).ok());
let new_val = f(current);
um.collections.insert(
name.to_owned(),
serde_json::to_value(&new_val).expect("serialize metadata"),
);
result = Some(new_val);
})?;
Ok(result.expect("update executed"))
}
pub fn tree_path(&self, name: &str) -> PathBuf {
self.path.join(name)
}
pub fn next_id(&self, name: &str) -> DbResult<u64> {
self.seq.next_id(name)
}
pub fn collection_len<T: CollectionMeta>(&self) -> u64 {
self.stored_info(T::NAME).map_or(0, |ci| ci.seq)
}
pub fn collection_names(&self) -> Vec<String> {
lock(&self.collections)
.iter()
.map(|c| c.name().to_string())
.collect()
}
pub fn close(&self) -> DbResult<()> {
for c in lock(&self.collections).iter() {
c.flush()?;
self.save_collection_len(c.name(), c.len() as u64)?;
}
self.seq.flush()
}
pub fn shutdown(&mut self) -> DbResult<()> {
self.shutdown.shutdown();
#[cfg(feature = "rpc")]
for h in self.rpc_handles.get_mut().drain(..) {
drop(h); }
if let Some(ref mut c) = self.compactor {
c.stop();
}
if let Some(ref mut f) = self.flusher {
f.stop();
}
self.close()
}
pub fn shutdown_signal(&self) -> ShutdownSignal {
self.shutdown.clone()
}
pub fn compact(&self) -> DbResult<usize> {
let snapshot = lock(&self.collections).clone();
let mut total = 0;
for c in &snapshot {
total += c.compact()?;
}
Ok(total)
}
pub fn open_typed_tree<T, C, H>(
&self,
config: Config,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<TypedTree<T::SelfId, T, C, H>>>
where
T: CollectionMeta + Clone + Send + Sync + 'static,
T::SelfId: Key + Ord + Send + Sync,
C: Codec<T> + Default + 'static,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
let meta = TreeMeta::of::<T>();
let stored = self.stored_info(meta.name);
let tree = TypedTree::open_hooked(self.tree_path(meta.name), config, C::default(), hook)?;
let migrated = self.run_migration(
&meta,
stored.as_ref(),
migrations,
|mfn| tree.migrate(mfn),
|| tree.flush_buffers(),
)?;
if !migrated {
tree.replay_init();
}
self.save_info(&meta, tree.len() as u64)?;
let tree = Arc::new(tree);
lock(&self.collections).push(tree.clone());
#[cfg(feature = "rpc")]
self.register_typed_tree_handler::<T, C, H>(&meta, tree.clone());
Ok(tree)
}
pub fn open_typed_map<T, C, H>(
&self,
config: Config,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<TypedMap<T::SelfId, T, C, H>>>
where
T: CollectionMeta + Clone + Send + Sync + 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
C: Codec<T> + Default + 'static,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
let meta = TreeMeta::of::<T>();
let stored = self.stored_info(meta.name);
let map = TypedMap::open_hooked(self.tree_path(meta.name), config, C::default(), hook)?;
let migrated = self.run_migration(
&meta,
stored.as_ref(),
migrations,
|mfn| map.migrate(mfn),
|| map.flush_buffers(),
)?;
if !migrated {
map.replay_init();
}
self.save_info(&meta, map.len() as u64)?;
let map = Arc::new(map);
lock(&self.collections).push(map.clone());
#[cfg(feature = "rpc")]
self.register_typed_map_handler::<T, C, H>(&meta, map.clone());
Ok(map)
}
#[cfg(feature = "var-collections")]
pub fn open_var_typed_tree<T, C, H>(
&self,
config: Config,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<crate::VarTypedTree<T::SelfId, T, C, H>>>
where
T: CollectionMeta + Clone + Send + Sync + 'static,
T::SelfId: Key + Ord + Send + Sync,
C: Codec<T> + Clone + Default + 'static,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
let meta = TreeMeta::of::<T>();
let stored = self.stored_info(meta.name);
let tree = crate::VarTypedTree::open_hooked(
self.tree_path(meta.name),
config,
C::default(),
hook,
)?;
let migrated = self.run_migration(
&meta,
stored.as_ref(),
migrations,
|mfn| tree.migrate(mfn),
|| tree.flush_buffers(),
)?;
if !migrated {
tree.replay_init();
}
self.save_info(&meta, tree.len() as u64)?;
let tree = Arc::new(tree);
lock(&self.collections).push(tree.clone());
#[cfg(feature = "rpc")]
self.register_var_typed_tree_handler::<T, C, H>(&meta, tree.clone());
Ok(tree)
}
#[cfg(feature = "var-collections")]
pub fn open_var_typed_map<T, C, H>(
&self,
config: Config,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<crate::VarTypedMap<T::SelfId, T, C, H>>>
where
T: CollectionMeta + Clone + Send + Sync + 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
C: Codec<T> + Clone + Default + 'static,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
let meta = TreeMeta::of::<T>();
let stored = self.stored_info(meta.name);
let map =
crate::VarTypedMap::open_hooked(self.tree_path(meta.name), config, C::default(), hook)?;
let migrated = self.run_migration(
&meta,
stored.as_ref(),
migrations,
|mfn| map.migrate(mfn),
|| map.flush_buffers(),
)?;
if !migrated {
map.replay_init();
}
self.save_info(&meta, map.len() as u64)?;
let map = Arc::new(map);
lock(&self.collections).push(map.clone());
#[cfg(feature = "rpc")]
self.register_var_typed_map_handler::<T, C, H>(&meta, map.clone());
Ok(map)
}
pub fn open_zero_tree<T, const V: usize, H>(
&self,
config: Config,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<ZeroTree<T::SelfId, V, T, H, crate::durability::Bitcask>>>
where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Ord + Send + Sync,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
self.open_zero_tree_with(
|path, hook| ZeroTree::open_hooked(path, config, hook),
hook,
migrations,
)
}
fn open_zero_tree_with<T, const V: usize, H, D>(
&self,
open_fn: impl FnOnce(PathBuf, H) -> DbResult<ZeroTree<T::SelfId, V, T, H, D>>,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<ZeroTree<T::SelfId, V, T, H, D>>>
where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Ord + Send + Sync,
H: TypedWriteHook<T::SelfId, T> + 'static,
D: crate::durability::Durability + 'static,
ZeroTree<T::SelfId, V, T, H, D>: Collection,
{
let meta = TreeMeta::of::<T>();
let stored = self.stored_info(meta.name);
let tree = open_fn(self.tree_path(meta.name), hook)?;
let migrated = self.run_migration(
&meta,
stored.as_ref(),
migrations,
|mfn| tree.migrate(mfn),
|| Collection::flush(&tree),
)?;
if !migrated {
tree.replay_init();
}
self.save_info(&meta, tree.len() as u64)?;
let tree = Arc::new(tree);
lock(&self.collections).push(tree.clone());
#[cfg(feature = "rpc")]
self.register_zero_tree_handler::<T, V, H, D>(&meta, tree.clone());
Ok(tree)
}
pub fn open_zero_tree_fixed<T, const V: usize, H>(
&self,
config: FixedConfig,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<ZeroTree<T::SelfId, V, T, H, crate::durability::Fixed>>>
where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Ord + Send + Sync,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
self.open_zero_tree_with(
|path, hook| ZeroTree::open_with_hook(path, config, hook),
hook,
migrations,
)
}
pub fn open_zero_map<T, const V: usize, H>(
&self,
config: Config,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<ZeroMap<T::SelfId, V, T, H, crate::durability::Bitcask>>>
where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
self.open_zero_map_with(
|path, hook| ZeroMap::open_hooked(path, config, hook),
hook,
migrations,
)
}
fn open_zero_map_with<T, const V: usize, H, D>(
&self,
open_fn: impl FnOnce(PathBuf, H) -> DbResult<ZeroMap<T::SelfId, V, T, H, D>>,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<ZeroMap<T::SelfId, V, T, H, D>>>
where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
H: TypedWriteHook<T::SelfId, T> + 'static,
D: crate::durability::Durability + 'static,
ZeroMap<T::SelfId, V, T, H, D>: Collection,
{
let meta = TreeMeta::of::<T>();
let stored = self.stored_info(meta.name);
let map = open_fn(self.tree_path(meta.name), hook)?;
let migrated = self.run_migration(
&meta,
stored.as_ref(),
migrations,
|mfn| map.migrate(mfn),
|| Collection::flush(&map),
)?;
if !migrated {
map.replay_init();
}
self.save_info(&meta, map.len() as u64)?;
let map = Arc::new(map);
lock(&self.collections).push(map.clone());
#[cfg(feature = "rpc")]
self.register_zero_map_handler::<T, V, H, D>(&meta, map.clone());
Ok(map)
}
pub fn open_zero_map_fixed<T, const V: usize, H>(
&self,
config: FixedConfig,
hook: H,
migrations: &[TypedMigration<T::SelfId, T>],
) -> DbResult<Arc<ZeroMap<T::SelfId, V, T, H, crate::durability::Fixed>>>
where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
H: TypedWriteHook<T::SelfId, T> + 'static,
{
self.open_zero_map_with(
|path, hook| ZeroMap::open_with_hook(path, config, hook),
hook,
migrations,
)
}
fn stored_info(&self, name: &str) -> Option<CollectionInfo> {
self.db_info.cloned().collections.get(name).copied()
}
fn run_migration<K, T>(
&self,
meta: &TreeMeta,
stored: Option<&CollectionInfo>,
migrations: &[TypedMigration<K, T>],
mut migrate_fn: impl FnMut(&super::migration::TypedMigrationFn<K, T>) -> DbResult<usize>,
flush_fn: impl Fn() -> DbResult<()>,
) -> DbResult<bool> {
let Some(stored) = stored else {
return Ok(false);
};
if stored.version == meta.version && stored.typ_hash != meta.ty.h() {
return Err(DbError::SchemaMismatch {
name: meta.name.into(),
kind: SchemaMismatchKind::TypHash {
stored: stored.typ_hash,
expected: meta.ty.h(),
},
});
}
if stored.version > meta.version {
return Err(DbError::SchemaMismatch {
name: meta.name.into(),
kind: SchemaMismatchKind::Downgrade {
stored: stored.version,
requested: meta.version,
},
});
}
if stored.version == meta.version {
return Ok(false);
}
let mut current = stored.version;
let mut mutated_any = false;
while current < meta.version {
let (_, mfn) = migrations
.iter()
.find(|(v, _)| *v == current)
.ok_or_else(|| DbError::SchemaMismatch {
name: meta.name.into(),
kind: SchemaMismatchKind::MissingStep { from: current },
})?;
let mutated = migrate_fn(mfn)?;
flush_fn()?;
self.update_schema_version(meta.name, current + 1, meta.ty.h())?;
tracing::info!(mutated, from = current, to = current + 1, "migration step");
current += 1;
mutated_any = true;
}
Ok(mutated_any)
}
fn update_schema_version(&self, name: &str, version: u16, typ_hash: u64) -> DbResult<()> {
self.db_info.update(|info| {
let entry = info.collections.entry(name.to_owned()).or_default();
entry.version = version;
entry.typ_hash = typ_hash;
})?;
Ok(())
}
fn save_info(&self, meta: &TreeMeta, seq: u64) -> DbResult<()> {
let typ_hash = meta.ty.h();
self.db_info.update(|info| {
info.collections.insert(
meta.name.to_owned(),
CollectionInfo {
version: meta.version,
typ_hash,
seq,
},
);
})?;
Ok(())
}
fn save_collection_len(&self, name: &str, len: u64) -> DbResult<()> {
self.db_info.update(|info| {
if let Some(ci) = info.collections.get_mut(name) {
ci.seq = len;
}
})?;
Ok(())
}
}
#[cfg(feature = "rpc")]
impl Db {
fn register_handler(&self, name: &str, handler: Arc<dyn super::handler::RpcHandler>) {
let hashname = xxhash_rust::xxh3::xxh3_64(name.as_bytes());
lock(&self.handlers).push((hashname, handler));
}
fn register_typed_tree_handler<T, C, H>(
&self,
meta: &TreeMeta,
tree: Arc<TypedTree<T::SelfId, T, C, H>>,
) where
T: CollectionMeta + Send + Sync + 'static,
T::SelfId: Key + Ord + Send + Sync,
C: Codec<T> + Default + 'static,
H: crate::hook::TypedWriteHook<T::SelfId, T> + 'static,
{
self.register_handler(
meta.name,
Arc::new(super::handler::TypedTreeHandler {
name: meta.name.to_owned(),
typ_hash: meta.ty.h(),
ty: meta.ty,
key_scheme: <T::SelfId as crate::Key>::KEY_SCHEME,
version: meta.version,
tree,
codec: Arc::new(C::default()),
seq: self.seq.clone(),
}),
);
}
fn register_typed_map_handler<T, C, H>(
&self,
meta: &TreeMeta,
map: Arc<TypedMap<T::SelfId, T, C, H>>,
) where
T: CollectionMeta + Send + Sync + 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
C: Codec<T> + Default + 'static,
H: crate::hook::TypedWriteHook<T::SelfId, T> + 'static,
{
self.register_handler(
meta.name,
Arc::new(super::handler::TypedMapHandler {
name: meta.name.to_owned(),
typ_hash: meta.ty.h(),
ty: meta.ty,
key_scheme: <T::SelfId as crate::Key>::KEY_SCHEME,
version: meta.version,
map,
codec: Arc::new(C::default()),
seq: self.seq.clone(),
}),
);
}
fn register_zero_tree_handler<T, const V: usize, H, D>(
&self,
meta: &TreeMeta,
tree: Arc<ZeroTree<T::SelfId, V, T, H, D>>,
) where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Ord + Send + Sync,
H: crate::hook::TypedWriteHook<T::SelfId, T> + 'static,
D: crate::durability::Durability + 'static,
{
self.register_handler(
meta.name,
Arc::new(super::handler::ZeroTreeHandler {
name: meta.name.to_owned(),
typ_hash: meta.ty.h(),
ty: meta.ty,
key_scheme: <T::SelfId as crate::Key>::KEY_SCHEME,
version: meta.version,
tree,
seq: self.seq.clone(),
}),
);
}
fn register_zero_map_handler<T, const V: usize, H, D>(
&self,
meta: &TreeMeta,
map: Arc<ZeroMap<T::SelfId, V, T, H, D>>,
) where
T: CollectionMeta
+ Copy
+ zerocopy::IntoBytes
+ zerocopy::FromBytes
+ zerocopy::Immutable
+ Send
+ Sync
+ 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
H: crate::hook::TypedWriteHook<T::SelfId, T> + 'static,
D: crate::durability::Durability + 'static,
{
self.register_handler(
meta.name,
Arc::new(super::handler::ZeroMapHandler {
name: meta.name.to_owned(),
typ_hash: meta.ty.h(),
ty: meta.ty,
key_scheme: <T::SelfId as crate::Key>::KEY_SCHEME,
version: meta.version,
map,
seq: self.seq.clone(),
}),
);
}
#[cfg(feature = "var-collections")]
fn register_var_typed_tree_handler<T, C, H>(
&self,
meta: &TreeMeta,
tree: Arc<crate::VarTypedTree<T::SelfId, T, C, H>>,
) where
T: CollectionMeta + Send + Sync + 'static,
T::SelfId: Key + Ord + Send + Sync,
C: Codec<T> + Clone + Default + 'static,
H: crate::hook::TypedWriteHook<T::SelfId, T> + 'static,
{
self.register_handler(
meta.name,
Arc::new(super::handler::VarTypedTreeHandler {
name: meta.name.to_owned(),
typ_hash: meta.ty.h(),
ty: meta.ty,
key_scheme: <T::SelfId as crate::Key>::KEY_SCHEME,
version: meta.version,
tree,
codec: Arc::new(C::default()),
seq: self.seq.clone(),
}),
);
}
#[cfg(feature = "var-collections")]
fn register_var_typed_map_handler<T, C, H>(
&self,
meta: &TreeMeta,
map: Arc<crate::VarTypedMap<T::SelfId, T, C, H>>,
) where
T: CollectionMeta + Send + Sync + 'static,
T::SelfId: Key + Send + Sync + Hash + Eq,
C: Codec<T> + Clone + Default + 'static,
H: crate::hook::TypedWriteHook<T::SelfId, T> + 'static,
{
self.register_handler(
meta.name,
Arc::new(super::handler::VarTypedMapHandler {
name: meta.name.to_owned(),
typ_hash: meta.ty.h(),
ty: meta.ty,
key_scheme: <T::SelfId as crate::Key>::KEY_SCHEME,
version: meta.version,
map,
codec: Arc::new(C::default()),
seq: self.seq.clone(),
}),
);
}
pub fn build_tree_map(&self) -> super::rpc::TreeMap {
Arc::new(
lock(&self.handlers)
.iter()
.map(|(h, handler)| (*h, handler.clone()))
.collect(),
)
}
}
impl Drop for Db {
fn drop(&mut self) {
let _ = self.shutdown();
}
}
#[allow(clippy::unused_self)]
impl Db {
pub fn atomic2<A, B, R>(
&self,
a: &A,
a_keys: &[A::Key],
b: &B,
b_keys: &[B::Key],
f: impl FnOnce(&mut A::Tx<'_>, &mut B::Tx<'_>) -> DbResult<R>,
) -> DbResult<R>
where
A: MultiTx,
B: MultiTx,
{
let ca = a.collection_id();
let cb = b.collection_id();
if ca == cb {
return Err(DbError::DuplicateCollectionInTx);
}
let a_sh = unique_sorted_shards(a_keys, |k| a.shard_for_key(k));
let b_sh = unique_sorted_shards(b_keys, |k| b.shard_for_key(k));
let mut plan: Vec<(usize, usize, u8)> = Vec::with_capacity(a_sh.len() + b_sh.len());
plan.extend(a_sh.iter().map(|&s| (ca, s, 0u8)));
plan.extend(b_sh.iter().map(|&s| (cb, s, 1u8)));
plan.sort_unstable();
let mut ta = a.begin_tx();
let mut tb = b.begin_tx();
for (_, shard, which) in &plan {
if *which == 0 {
a.lock_shard_into(*shard, &mut ta);
} else {
b.lock_shard_into(*shard, &mut tb);
}
}
let result = f(&mut ta, &mut tb);
let sa = a.release_locks(&mut ta);
let sb = b.release_locks(&mut tb);
let r1 = a.run_sync(sa);
let r2 = b.run_sync(sb);
a.replay_hooks(ta);
b.replay_hooks(tb);
if result.is_err() {
result
} else {
r1?;
r2?;
result
}
}
#[allow(clippy::too_many_arguments)]
pub fn atomic3<A, B, C, R>(
&self,
a: &A,
a_keys: &[A::Key],
b: &B,
b_keys: &[B::Key],
c: &C,
c_keys: &[C::Key],
f: impl FnOnce(&mut A::Tx<'_>, &mut B::Tx<'_>, &mut C::Tx<'_>) -> DbResult<R>,
) -> DbResult<R>
where
A: MultiTx,
B: MultiTx,
C: MultiTx,
{
let (ca, cb, cc) = (a.collection_id(), b.collection_id(), c.collection_id());
if ca == cb || ca == cc || cb == cc {
return Err(DbError::DuplicateCollectionInTx);
}
let a_sh = unique_sorted_shards(a_keys, |k| a.shard_for_key(k));
let b_sh = unique_sorted_shards(b_keys, |k| b.shard_for_key(k));
let c_sh = unique_sorted_shards(c_keys, |k| c.shard_for_key(k));
let mut plan: Vec<(usize, usize, u8)> =
Vec::with_capacity(a_sh.len() + b_sh.len() + c_sh.len());
plan.extend(a_sh.iter().map(|&s| (ca, s, 0u8)));
plan.extend(b_sh.iter().map(|&s| (cb, s, 1u8)));
plan.extend(c_sh.iter().map(|&s| (cc, s, 2u8)));
plan.sort_unstable();
let mut ta = a.begin_tx();
let mut tb = b.begin_tx();
let mut tc = c.begin_tx();
for (_, shard, which) in &plan {
match which {
0 => a.lock_shard_into(*shard, &mut ta),
1 => b.lock_shard_into(*shard, &mut tb),
_ => c.lock_shard_into(*shard, &mut tc),
}
}
let result = f(&mut ta, &mut tb, &mut tc);
let sa = a.release_locks(&mut ta);
let sb = b.release_locks(&mut tb);
let sc = c.release_locks(&mut tc);
let r1 = a.run_sync(sa);
let r2 = b.run_sync(sb);
let r3 = c.run_sync(sc);
a.replay_hooks(ta);
b.replay_hooks(tb);
c.replay_hooks(tc);
if result.is_err() {
result
} else {
r1?;
r2?;
r3?;
result
}
}
#[allow(clippy::too_many_arguments)]
pub fn atomic4<A, B, C, D, R>(
&self,
a: &A,
a_keys: &[A::Key],
b: &B,
b_keys: &[B::Key],
c: &C,
c_keys: &[C::Key],
d: &D,
d_keys: &[D::Key],
f: impl FnOnce(&mut A::Tx<'_>, &mut B::Tx<'_>, &mut C::Tx<'_>, &mut D::Tx<'_>) -> DbResult<R>,
) -> DbResult<R>
where
A: MultiTx,
B: MultiTx,
C: MultiTx,
D: MultiTx,
{
let (ca, cb, cc, cd) = (
a.collection_id(),
b.collection_id(),
c.collection_id(),
d.collection_id(),
);
let ids = [ca, cb, cc, cd];
for i in 0..ids.len() {
for j in (i + 1)..ids.len() {
if ids[i] == ids[j] {
return Err(DbError::DuplicateCollectionInTx);
}
}
}
let a_sh = unique_sorted_shards(a_keys, |k| a.shard_for_key(k));
let b_sh = unique_sorted_shards(b_keys, |k| b.shard_for_key(k));
let c_sh = unique_sorted_shards(c_keys, |k| c.shard_for_key(k));
let d_sh = unique_sorted_shards(d_keys, |k| d.shard_for_key(k));
let mut plan: Vec<(usize, usize, u8)> =
Vec::with_capacity(a_sh.len() + b_sh.len() + c_sh.len() + d_sh.len());
plan.extend(a_sh.iter().map(|&s| (ca, s, 0u8)));
plan.extend(b_sh.iter().map(|&s| (cb, s, 1u8)));
plan.extend(c_sh.iter().map(|&s| (cc, s, 2u8)));
plan.extend(d_sh.iter().map(|&s| (cd, s, 3u8)));
plan.sort_unstable();
let mut ta = a.begin_tx();
let mut tb = b.begin_tx();
let mut tc = c.begin_tx();
let mut td = d.begin_tx();
for (_, shard, which) in &plan {
match which {
0 => a.lock_shard_into(*shard, &mut ta),
1 => b.lock_shard_into(*shard, &mut tb),
2 => c.lock_shard_into(*shard, &mut tc),
_ => d.lock_shard_into(*shard, &mut td),
}
}
let result = f(&mut ta, &mut tb, &mut tc, &mut td);
let sa = a.release_locks(&mut ta);
let sb = b.release_locks(&mut tb);
let sc = c.release_locks(&mut tc);
let sd = d.release_locks(&mut td);
let r1 = a.run_sync(sa);
let r2 = b.run_sync(sb);
let r3 = c.run_sync(sc);
let r4 = d.run_sync(sd);
a.replay_hooks(ta);
b.replay_hooks(tb);
c.replay_hooks(tc);
d.replay_hooks(td);
if result.is_err() {
result
} else {
r1?;
r2?;
r3?;
r4?;
result
}
}
}
#[cfg(all(test, feature = "rapira-codec"))]
mod tests {
use super::*;
use armour_core::GetType;
use rapira::Rapira;
use tempfile::tempdir;
use crate::{NoHook, RapiraCodec};
#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct User {
id: u64,
name: String,
}
impl CollectionMeta for User {
type SelfId = [u8; 8];
const NAME: &'static str = "iterable_toggle_users";
const VERSION: u16 = 1;
}
fn key(id: u64) -> [u8; 8] {
id.to_be_bytes()
}
fn demo_user() -> User {
User {
id: 1,
name: "alice".into(),
}
}
#[test]
fn db_map_iterable_toggles_across_reopen() {
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let users = db
.open_typed_map::<User, RapiraCodec, _>(Config::test(), NoHook, &[])
.unwrap();
assert!(users.iter_view().is_none());
users.put(&key(1), demo_user()).unwrap();
db.close().unwrap();
}
let db = Db::open_test(dir.path()).unwrap();
let cfg = Config::balanced()
.shard_count(2)
.hints(true)
.iterable(true)
.build();
let users = db
.open_typed_map::<User, RapiraCodec, _>(cfg, NoHook, &[])
.unwrap();
assert!(users.iter_view().is_some());
assert_eq!(users.iter_view().unwrap().iter().count(), 1);
}
#[derive(
Clone,
Copy,
Debug,
PartialEq,
GetType,
zerocopy::FromBytes,
zerocopy::IntoBytes,
zerocopy::Immutable,
)]
#[repr(C)]
struct FxCounter {
n: u32,
}
impl CollectionMeta for FxCounter {
type SelfId = [u8; 8];
const NAME: &'static str = "fixed_reversed_db";
const VERSION: u16 = 1;
}
#[test]
fn db_fixed_openers_honor_reversed() {
let dir = tempdir().unwrap();
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree_fixed::<FxCounter, 4, _>(FixedConfig::test(), NoHook, &[])
.unwrap();
for i in 1u64..=3 {
tree.put(&i.to_be_bytes(), &FxCounter { n: i as u32 })
.unwrap();
}
let tk: Vec<u64> = tree.iter().map(|(k, _)| u64::from_be_bytes(k)).collect();
assert_eq!(tk, [3, 2, 1], "Db fixed tree default DESC");
let map_cfg = FixedConfig {
iterable: true,
reversed: false,
..FixedConfig::test()
};
let map = db
.open_zero_map_fixed::<FxCounter, 4, _>(map_cfg, NoHook, &[])
.unwrap();
for i in 1u64..=3 {
map.put(&i.to_be_bytes(), &FxCounter { n: i as u32 })
.unwrap();
}
let mk: Vec<u64> = map
.iter_view()
.unwrap()
.keys()
.map(u64::from_be_bytes)
.collect();
assert_eq!(mk, [1, 2, 3], "Db fixed map reversed=false ASC");
}
}