use std::hash::Hash;
use std::mem::size_of;
use std::ops::{Bound, ControlFlow};
use std::sync::{Arc, OnceLock, Weak};
use crate::Key;
use crate::codec::Codec;
use crate::compaction::{CompactionIndex, compact_shard};
use crate::config::Config;
use crate::ctx::{Ctx, Family, Held, Owned, ShardGuards};
use crate::disk_loc::DiskLoc;
use crate::engine::Engine;
use crate::error::{DbError, DbResult};
use crate::hook::{NoHook, TypedWriteHook};
use crate::map_index::{
Dir, KeyPage, MapIter, MapIterView, OrdKey, Page, ShardIndex, ShardSource, TryMapIter,
dir_from, key_index_batch,
};
use crate::projection::IndexKeyed;
use crate::recovery::recover_typed_map;
use crate::shard::ShardInner;
use crate::skiplist::node::TypedData;
use crate::sync::{self, Mutex};
use crate::ttl::{SweepOutcome, TtlSlots};
use crate::typed_tree::TypedRef;
use crate::upsert::Upsert;
use crate::watch::{self, CodecWatchDecoder, WatchOptions, WatchReceiver, WatchSpec};
use crate::watch::{WatchEventInner, WatchWakeBatch};
enum MapSweepAction<K: Key> {
Delete(K, u64),
Reschedule(K, u64),
Drop,
Retry(K),
}
pub(crate) struct TypedMapEntry<T> {
pub(crate) ptr: *mut TypedData<T>,
}
unsafe impl<T: Send> Send for TypedMapEntry<T> {}
unsafe impl<T: Sync> Sync for TypedMapEntry<T> {}
struct UpdatePtrs<T> {
old: *const TypedData<T>,
new: *const TypedData<T>,
retire: *mut TypedData<T>,
new_disk: DiskLoc,
}
unsafe impl<T: Send> Send for UpdatePtrs<T> {}
unsafe impl<T: Sync> Sync for UpdatePtrs<T> {}
pub(crate) struct MutationResult<K, T> {
ref_ptr: Option<*const TypedData<T>>,
hook: Option<(Option<*const T>, Option<*const T>)>,
retire: Option<*mut TypedData<T>>,
new_disk: Option<DiskLoc>,
watch: Option<Arc<WatchEventInner<K>>>,
}
pub(crate) type MutationEvents<K, T> = Vec<(K, MutationResult<K, T>)>;
pub struct TypedMap<
K: Key + Send + Sync + Hash + Eq,
T: Send + Sync,
C: Codec<T>,
H: TypedWriteHook<K, T> = NoHook,
> {
indexes: Vec<Mutex<ShardIndex<K, TypedMapEntry<T>>>>,
collector: seize::Collector,
engine: Engine,
codec: C,
compaction_threshold: f64,
shard_prefix_bits: usize,
pub(crate) hook: H,
iterable: bool,
reversed: bool,
len: std::sync::atomic::AtomicUsize,
encode_bufs: Vec<Mutex<Vec<u8>>>,
tx_self: OnceLock<Weak<Self>>,
ttl: TtlSlots<K>,
pub(crate) counter_dirty: OnceLock<crate::counter_dirty::CounterDirty>,
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>>
watch::sealed::Spec for TypedMap<K, T, C, H>
{
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Clone, H: TypedWriteHook<K, T>>
WatchSpec for TypedMap<K, T, C, H>
{
type Key = K;
type Value = T;
type Decoder = CodecWatchDecoder<C, T>;
}
impl<
K: Key + Send + Sync + Hash + Eq,
T: Send + Sync,
C: Codec<T> + Clone + Sync,
H: TypedWriteHook<K, T> + watch::sealed::Enabled,
> TypedMap<K, T, C, H>
{
pub fn watch(&self, options: WatchOptions) -> WatchReceiver<Self> {
WatchReceiver::subscribe(
self.hook.watch_channel().expect("enabled watch hook"),
CodecWatchDecoder::new(self.codec.clone()),
options,
)
}
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync> TypedMap<K, T, C> {
#[doc = include_str!("../docs/exclusive-access.md")]
pub fn open(path: impl AsRef<std::path::Path>, config: Config, codec: C) -> DbResult<Self> {
Self::open_hooked(path, config, codec, NoHook)
}
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
TypedMap<K, T, C, H>
{
pub(crate) fn floor_recovered_gsns(&self, gsn_high: u64, shard_gsn_high: &[u64]) {
self.engine.gsn().fetch_max(
gsn_high.saturating_add(1),
std::sync::atomic::Ordering::Relaxed,
);
for (shard_id, high) in shard_gsn_high.iter().copied().enumerate() {
if let Some(shard) = self.engine.shards().get(shard_id) {
shard.set_durable_recovered_gsn(shard.durable_recovered_gsn().max(high));
}
}
}
#[doc = include_str!("../docs/exclusive-access.md")]
pub fn open_hooked(
path: impl AsRef<std::path::Path>,
config: Config,
codec: C,
hook: H,
) -> DbResult<Self> {
crate::counter_dirty::check_no_counters(&hook)?;
crate::unique::check_no_claims(&hook)?;
#[cfg(feature = "armour")]
if H::EMITS {
return Err(DbError::ChangeStreamRequiresManagedOpen);
}
let config = config.with_resolved_hints(false);
Self::open_inner(path, config, codec, hook)
}
pub(crate) fn open_inner(
path: impl AsRef<std::path::Path>,
config: Config,
codec: C,
hook: H,
) -> DbResult<Self> {
if H::WATCHES && !config.iterable {
return Err(DbError::WatchRequiresIterable);
}
let compaction_threshold = config.compaction_threshold;
let shard_prefix_bits = config.shard_prefix_bits;
let iterable = config.iterable;
let reversed = config.reversed;
let engine = Engine::open(path, config)?;
let shard_count = engine.shards().len();
let mut indexes = Vec::with_capacity(shard_count);
for _ in 0..shard_count {
indexes.push(Mutex::new(ShardIndex::new(false)));
}
let encode_bufs: Vec<Mutex<Vec<u8>>> =
(0..shard_count).map(|_| Mutex::new(Vec::new())).collect();
let map = Self {
indexes,
collector: seize::Collector::new(),
engine,
codec,
compaction_threshold,
shard_prefix_bits,
hook,
iterable,
reversed,
len: std::sync::atomic::AtomicUsize::new(0),
encode_bufs,
tx_self: OnceLock::new(),
ttl: TtlSlots::new(),
counter_dirty: OnceLock::new(),
};
let shard_dirs = map.engine.shard_dirs();
let shard_dir_refs = Engine::shard_dir_refs(&shard_dirs);
let shard_ids = map.engine.shard_ids();
let hints = map.engine.hints();
let outcome = recover_typed_map::<K, T, C>(
&shard_dir_refs,
&shard_ids,
map.indexes(),
&map.codec,
hints,
#[cfg(feature = "encryption")]
map.engine.shard_ciphers(),
)?;
for tail in &outcome.active_tails {
map.engine.shards()[tail.shard_idx].apply_recovery_tail(tail)?;
}
for (shard_idx, dead) in outcome.shard_dead_bytes {
map.engine.shards()[shard_idx].install_dead_bytes(dead);
}
for (shard_idx, file_seqs) in outcome.shard_file_seqs {
map.engine.shards()[shard_idx].install_recovery_file_seqs(file_seqs);
}
for (shard_idx, max_gsn) in &outcome.shard_max_gsns {
map.engine.shards()[*shard_idx].set_durable_recovered_gsn(*max_gsn);
}
let max_gsn = outcome.max_gsn;
map.engine
.gsn()
.fetch_max(max_gsn + 1, std::sync::atomic::Ordering::Relaxed);
if hints {
for shard in map.engine.shards().iter() {
shard.set_key_len(size_of::<K>());
}
}
if iterable {
for idx in &map.indexes {
sync::lock(idx).set_iterable_and_rebuild();
}
}
let initial_len: usize = map.indexes.iter().map(|m| sync::lock(m).len()).sum();
map.len
.store(initial_len, std::sync::atomic::Ordering::Relaxed);
tracing::info!(
key_size = size_of::<K>(),
entries = initial_len,
"typed_map recovered"
);
Ok(map)
}
pub fn clean_shutdown(&self) -> DbResult<()> {
if self.engine.hints() {
self.sync_hints()?;
}
self.engine.close()
}
pub fn close(self) -> DbResult<()> {
self.clean_shutdown()
}
pub fn flush_buffers(&self) -> DbResult<()> {
self.engine.flush_buffers()
}
pub(crate) fn durable_flush(&self) -> DbResult<()> {
self.engine.flush()
}
pub fn config(&self) -> &Config {
self.engine.config()
}
}
impl<
K: Key + Send + Sync + Hash + Eq,
T: Clone + Send + Sync,
C: Codec<T> + Sync,
H: TypedWriteHook<K, T>,
> CompactionIndex<K> for TypedMap<K, T, C, H>
{
fn update_if_match(&self, key: &K, old_loc: DiskLoc, new_loc: DiskLoc) -> bool {
let mut index = sync::lock(&self.indexes[self.shard_for(key)]);
if let Some(entry) = index.get_mut(key) {
let data = unsafe { &*entry.ptr };
if data.disk == old_loc {
let new_data = Box::into_raw(Box::new(TypedData {
disk: new_loc,
value: data.value.clone(),
}));
let old_ptr = entry.ptr;
entry.ptr = new_data;
unsafe {
self.collector
.retire(old_ptr, seize::reclaim::boxed::<TypedData<T>>);
}
return true;
}
}
false
}
fn contains_key(&self, key: &K) -> bool {
self.contains(key)
}
fn is_live(&self, shard_id: u8, key: &K, loc: DiskLoc) -> bool {
let index = sync::lock(&self.indexes[shard_id as usize]);
match index.get(key) {
Some(entry) => unsafe { &*entry.ptr }.disk == loc,
None => false,
}
}
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
Family for TypedMap<K, T, C, H>
{
type Locked<'a>
= (
sync::MutexGuard<'a, ShardInner>,
sync::MutexGuard<'a, ShardIndex<K, TypedMapEntry<T>>>,
)
where
Self: 'a;
type Event<'v> = (K, MutationResult<K, T>);
type Owned = (K, MutationResult<K, T>);
#[inline]
fn lock_shard(&self, sid: usize) -> Self::Locked<'_> {
let inner = self.engine.shards()[sid].lock();
let index = sync::lock(&self.indexes[sid]);
(inner, index)
}
fn should_sync<'a>(locked: &Self::Locked<'a>) -> bool
where
Self: 'a,
{
locked.0.should_sync()
}
fn sync(&self, sid: usize) -> DbResult<()> {
#[cfg(feature = "ctx-gates")]
crate::ctx_gates_sync::record_sync_call();
self.engine.shards()[sid].flush()
}
fn fire(&self, ev: Self::Event<'_>) {
let (k, mr) = ev;
self.fire_hook_and_retire(&k, &mr);
}
fn own(ev: Self::Event<'_>) -> Self::Owned {
ev
}
fn fire_owned(&self, ev: &Self::Owned) {
self.fire_hook_and_retire(&ev.0, &ev.1);
}
fn materialize_watch<'a, 'v>(
&self,
locked: &mut Self::Locked<'a>,
ev: &mut Self::Event<'v>,
active: bool,
wakes: &mut WatchWakeBatch,
) where
Self: 'a,
{
if !H::WATCHES {
return;
}
if !active {
return;
}
let (key, mutation) = ev;
let Some((old_ptr, _)) = mutation.hook else {
return;
};
let old_existed = old_ptr.is_some();
let (inner, _index) = &*locked;
let old = if self.hook.watch_old() {
match old_ptr.map(|ptr| unsafe { &*ptr }) {
Some(value) => {
let mut buf = Vec::new();
match self.codec.encode_to(value, &mut buf) {
Ok(()) => Some(buf.into_boxed_slice()),
Err(e) => {
tracing::warn!(
key = ?key.as_bytes(),
error = %e,
"watch event lost its old value; reporting the write without it"
);
None
}
}
}
None => None,
}
} else {
None
};
let new = match mutation.new_disk.as_ref() {
Some(loc) => match inner.read_value_bytes_locked(loc) {
Ok(bytes) => Some(bytes.into_boxed_slice()),
Err(e) => {
tracing::error!(
key = ?key.as_bytes(),
error = %e,
"watch event could not be materialized; reporting a gap"
);
if let Some(channel) = self.hook.watch_channel() {
channel.publish_lost(wakes);
}
return;
}
},
None => None,
};
mutation.watch = Some(Arc::new(WatchEventInner::from_write(
*key,
old,
new,
self.hook.watch_old(),
old_existed,
)));
}
fn publish_watch(&self, ev: &Self::Event<'_>, active: bool, wakes: &mut WatchWakeBatch) {
if !H::WATCHES || !active {
return;
}
let Some(channel) = self.hook.watch_channel() else {
return;
};
if let Some(event) = ev.1.watch.as_ref() {
channel.publish(Arc::clone(event), wakes);
}
}
const WATCHES: bool = H::WATCHES;
fn watch_enter(&self) -> crate::watch::WatchProducer<'_> {
if H::WATCHES {
self.hook
.watch_channel()
.and_then(|channel| channel.enter_producer())
.unwrap_or_else(crate::watch::WatchProducer::inactive)
} else {
crate::watch::WatchProducer::inactive()
}
}
fn apply_projection_for_event(&self, ev: &Self::Event<'_>) -> DbResult<()> {
self.apply_projection_for_mutation(ev)
}
fn apply_projection_for_owned(&self, ev: &Self::Owned) -> DbResult<()> {
self.apply_projection_for_mutation(ev)
}
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
TypedMap<K, T, C, H>
{
#[doc(hidden)]
pub fn build_ttl_index(&self) -> DbResult<()> {
if !H::EXPIRES {
return Ok(());
}
let n = self.engine.shards().len();
let slots = self.ttl.install_and_clear(n);
let _guard = self.collector.enter();
for (sid, slot) in slots.iter().enumerate() {
let _shard = self.engine.shards()[sid].lock();
let index = sync::lock(&self.indexes[sid]);
for (key, entry) in index.iter() {
let data = unsafe { &*entry.ptr };
let deadline = self.hook.deadline(key, &data.value);
sync::lock(slot).record(key, deadline);
}
}
Ok(())
}
fn record_ttl(&self, shard_id: usize, key: &K, expires_at: Option<u64>) {
if !H::EXPIRES {
return;
}
self.ttl.record(shard_id, key, expires_at);
}
#[doc(hidden)]
pub fn ttl_debug_len(&self) -> usize {
self.ttl.debug_len()
}
#[doc(hidden)]
pub fn ttl_debug_deadline(&self, key: &K) -> Option<u64> {
self.ttl.debug_deadline(self.shard_for(key), key)
}
pub fn sweep_expired(&self, now_ms: u64, budget: usize) -> DbResult<SweepOutcome> {
if !H::EXPIRES {
return Ok(SweepOutcome::default());
}
let Some(slots) = self.ttl.get() else {
return Ok(SweepOutcome::default());
};
let Some(_pass) = self.ttl.begin_pass() else {
return Ok(SweepOutcome::default());
};
let shard_count = slots.len();
let mut remaining = budget;
let mut outcome = SweepOutcome::default();
let guard = self.collector.enter();
while remaining > 0 && self.ttl.any_due(now_ms) {
let sid = self.ttl.next_shard(shard_count);
let mut shard_events = Vec::new();
let mut wakes = WatchWakeBatch::default();
let mut locked = self.lock_shard(sid);
let producer = self.watch_enter();
let active = producer.active();
let mut idx = sync::lock(&slots[sid]);
while remaining > 0 {
let Some((candidate_dl, key)) = idx.pop_due(now_ms) else {
break;
};
remaining -= 1;
let action = {
let index = &*locked.1;
match index.get(&key) {
Some(entry) => {
let value = &unsafe { &*entry.ptr }.value;
match crate::ttl::catch_deadline(|| self.hook.deadline(&key, value)) {
Some(Some(dl)) if dl <= now_ms => {
MapSweepAction::Delete(key, candidate_dl)
}
Some(Some(dl)) => MapSweepAction::Reschedule(key, dl),
Some(None) => MapSweepAction::Drop,
None => MapSweepAction::Retry(key),
}
}
None => MapSweepAction::Drop,
}
};
match action {
MapSweepAction::Delete(key, candidate_dl) => {
drop(idx);
let (inner, index) = &mut locked;
match self.delete_locked_inner::<true>(sid, inner, index, &guard, &key) {
Ok(mr) => {
let mut event = (key, mr);
self.materialize_watch(&mut locked, &mut event, active, &mut wakes);
self.publish_watch(&event, active, &mut wakes);
idx = sync::lock(&slots[sid]);
if event.1.hook.is_some() {
shard_events.push(event);
}
outcome.deleted += 1;
}
Err(e) => {
sync::lock(&slots[sid]).record(&key, Some(candidate_dl));
let needs_sync = Self::should_sync(&locked);
locked.release();
crate::ttl::log_sweep_failures(None, &outcome);
let _ = crate::projection::replay_owned_events(
self,
&shard_events,
needs_sync,
sid,
);
<Self as Family>::wake_watches(wakes);
return Err(e);
}
}
}
MapSweepAction::Reschedule(key, dl) => {
idx.record(&key, Some(dl));
}
MapSweepAction::Retry(key) => {
outcome.failed += 1;
idx.record(&key, Some(now_ms + crate::ttl::TTL_RETRY_BACKOFF_MS));
}
MapSweepAction::Drop => {}
}
}
drop(idx);
let needs_sync = Self::should_sync(&locked);
locked.release();
let replay =
crate::projection::replay_owned_events(self, &shard_events, needs_sync, sid);
<Self as Family>::wake_watches(wakes);
replay?;
}
drop(guard);
outcome.more = self.ttl.any_due(now_ms);
crate::ttl::log_sweep_failures(None, &outcome);
Ok(outcome)
}
pub fn set_ttl_sweep(&self, enabled: bool) -> DbResult<()> {
if H::EXPIRES {
self.ttl.set_enabled(enabled)
} else {
Ok(())
}
}
pub fn mark_replication_follower(&self) {
if H::EXPIRES {
self.ttl.mark_replication_follower();
}
}
pub fn rebuild_ttl_index(&self) -> DbResult<()> {
if !H::EXPIRES {
return Err(DbError::Client("collection has no TTL policy"));
}
self.build_ttl_index()?;
self.ttl.clear_rebuild_requirement();
Ok(())
}
pub fn compact(&self) -> DbResult<usize>
where
T: Clone,
{
let mut total_compacted = 0;
for shard in self.engine.shards().iter() {
total_compacted += compact_shard(shard, self, self.compaction_threshold)?;
}
Ok(total_compacted)
}
pub fn get(&self, key: &K) -> Option<TypedRef<'_, T>> {
metrics::counter!("armdb.ops", "op" => "get", "tree" => "typed_map").increment(1);
let guard = self.collector.enter();
let data_ptr = {
let index = sync::lock(&self.indexes[self.shard_for(key)]);
let entry = index.get(key)?;
entry.ptr as *const TypedData<T>
};
Some(TypedRef::new(guard, data_ptr))
}
pub fn get_or_err(&self, key: &K) -> DbResult<TypedRef<'_, T>> {
self.get(key).ok_or(DbError::KeyNotFound)
}
pub fn contains(&self, key: &K) -> bool {
let index = sync::lock(&self.indexes[self.shard_for(key)]);
index.contains_key(key)
}
#[inline]
fn defer_or_retire_event(
&self,
key: &K,
mr: MutationResult<K, T>,
) -> Option<(K, MutationResult<K, T>)> {
if H::NEEDS_WRITE || (H::WATCHES && crate::watch::context_active()) {
Some((*key, mr))
} else {
if let Some(ptr) = mr.retire {
unsafe {
self.collector
.retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
}
}
None
}
}
pub fn put(&self, key: &K, value: T) -> DbResult<Option<TypedRef<'_, T>>> {
metrics::counter!("armdb.ops", "op" => "put", "tree" => "typed_map").increment(1);
let expires_at = if H::EXPIRES {
self.hook.deadline(key, &value)
} else {
None
};
let shard_id = self.shard_for(key);
let guard = self.collector.enter();
let mut ctx = Owned;
let ref_ptr = Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let mr = self.put_locked_inner::<true>(
shard_id,
inner,
index,
&guard,
key,
value,
crate::ttl::TtlAt::Precomputed(expires_at),
)?;
let ref_ptr = mr.ref_ptr;
*ev_out = Some((*key, mr));
Ok(ref_ptr)
})?;
Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
}
pub fn insert(&self, key: &K, value: T) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "insert", "tree" => "typed_map").increment(1);
let expires_at = if H::EXPIRES {
self.hook.deadline(key, &value)
} else {
None
};
let shard_id = self.shard_for(key);
let guard = self.collector.enter();
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let mr =
self.insert_locked::<true>(shard_id, inner, index, &guard, key, value, expires_at)?;
*ev_out = Some((*key, mr));
Ok(())
})
}
pub fn delete(&self, key: &K) -> DbResult<Option<TypedRef<'_, T>>> {
metrics::counter!("armdb.ops", "op" => "delete", "tree" => "typed_map").increment(1);
let shard_id = self.shard_for(key);
let guard = self.collector.enter();
let mut ctx = Owned;
let ref_ptr = Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let mr = self.delete_locked_inner::<true>(shard_id, inner, index, &guard, key)?;
let ref_ptr = mr.ref_ptr;
*ev_out = Some((*key, mr));
Ok(ref_ptr)
})?;
Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
}
pub fn cas(&self, key: &K, expected: &T, new_value: T) -> DbResult<()>
where
T: PartialEq,
{
metrics::counter!("armdb.ops", "op" => "cas", "tree" => "typed_map").increment(1);
let expires_at = if H::EXPIRES {
self.hook.deadline(key, &new_value)
} else {
None
};
let shard_id = self.shard_for(key);
let _guard = self.collector.enter();
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let entry = index.get(key).ok_or(DbError::KeyNotFound)?;
let current_data = unsafe { &*entry.ptr };
if current_data.value != *expected {
return Err(DbError::CasMismatch);
}
self.claim_guard(key, Some(¤t_data.value), Some(&new_value))?;
let mut buf = sync::lock(&self.encode_bufs[shard_id]);
self.codec.encode_to(&new_value, &mut buf)?;
let (disk_loc, _gsn) =
inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;
let new_data_ptr = Box::into_raw(Box::new(TypedData {
disk: disk_loc,
value: new_value,
}));
let old_ptr = entry.ptr;
index.get_mut(key).expect("key exists").ptr = new_data_ptr;
let old_disk = unsafe { (*old_ptr).disk };
inner.add_dead_bytes(
old_disk.file_id,
crate::entry::entry_size(size_of::<K>(), old_disk.len),
);
self.record_ttl(shard_id, key, expires_at);
let mr = MutationResult {
ref_ptr: None,
hook: Some((
Some(¤t_data.value as *const T),
Some(unsafe { &(*new_data_ptr).value as *const T }),
)),
retire: Some(old_ptr),
new_disk: Some(disk_loc),
watch: None,
};
*ev_out = Some((*key, mr));
Ok(())
})
}
pub fn compare_delete(&self, key: &K, expected: &T) -> DbResult<()>
where
T: PartialEq,
{
metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "typed_map")
.increment(1);
let shard_id = self.shard_for(key);
let guard = self.collector.enter();
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let entry = index.get(key).ok_or(DbError::KeyNotFound)?;
let current_data = unsafe { &*entry.ptr };
if current_data.value != *expected {
return Err(DbError::CasMismatch);
}
let mr = self.delete_locked_inner::<true>(shard_id, inner, index, &guard, key)?;
*ev_out = Some((*key, mr));
Ok(())
})
}
pub(crate) fn compare_delete_if_ptr(
&self,
key: &K,
expected: *const TypedData<T>,
) -> DbResult<()> {
metrics::counter!("armdb.ops", "op" => "compare_delete", "tree" => "typed_map")
.increment(1);
let shard_id = self.shard_for(key);
let guard = self.collector.enter();
let mut ctx = Owned;
Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let entry = index.get(key).ok_or(DbError::KeyNotFound)?;
if !std::ptr::eq(entry.ptr, expected as *mut TypedData<T>) {
return Err(DbError::CasMismatch);
}
let mr = self.delete_locked_inner::<true>(shard_id, inner, index, &guard, key)?;
*ev_out = Some((*key, mr));
Ok(())
})
}
pub fn update(&self, key: &K, f: impl FnOnce(&T) -> T) -> DbResult<Option<TypedRef<'_, T>>> {
self.update_inner(key, f, false)
}
pub fn fetch_update(
&self,
key: &K,
f: impl FnOnce(&T) -> T,
) -> DbResult<Option<TypedRef<'_, T>>> {
self.update_inner(key, f, true)
}
pub fn upsert<U: crate::IntoUpsert<T>>(
&self,
key: &K,
f: impl FnOnce(Option<&T>) -> U,
) -> crate::DbResult<U::Out> {
metrics::counter!("armdb.ops", "op" => "upsert", "tree" => "typed_map").increment(1);
self.atomic(key, |s| crate::upsert::upsert_via_held(s, key, f))
}
fn update_locked_inner(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, TypedMapEntry<T>>,
key: &K,
f: impl FnOnce(&T) -> T,
) -> DbResult<Option<UpdatePtrs<T>>> {
let entry = match index.get(key) {
Some(e) => e,
None => return Ok(None),
};
let old_ptr = entry.ptr;
let old_data = unsafe { &*old_ptr };
let new_value = f(&old_data.value);
self.claim_guard(key, Some(&old_data.value), Some(&new_value))?;
let mut buf = sync::lock(&self.encode_bufs[shard_id]);
self.codec.encode_to(&new_value, &mut buf)?;
let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;
drop(buf);
let new_data_ptr = Box::into_raw(Box::new(TypedData {
disk: disk_loc,
value: new_value,
}));
index.get_mut(key).expect("key exists").ptr = new_data_ptr;
let old_disk = unsafe { (*old_ptr).disk };
inner.add_dead_bytes(
old_disk.file_id,
crate::entry::entry_size(size_of::<K>(), old_disk.len),
);
let expires_at = if H::EXPIRES {
self.hook.deadline(key, unsafe { &(*new_data_ptr).value })
} else {
None
};
self.record_ttl(shard_id, key, expires_at);
Ok(Some(UpdatePtrs {
old: old_ptr as *const TypedData<T>,
new: new_data_ptr as *const TypedData<T>,
retire: old_ptr,
new_disk: disk_loc,
}))
}
fn update_inner(
&self,
key: &K,
f: impl FnOnce(&T) -> T,
return_old: bool,
) -> DbResult<Option<TypedRef<'_, T>>> {
metrics::counter!("armdb.ops", "op" => "update", "tree" => "typed_map").increment(1);
let shard_id = self.shard_for(key);
let guard = self.collector.enter();
let mut ctx = Owned;
let ref_ptr = Ctx::with_shard(&mut ctx, self, shard_id, |locked, ev_out| {
let (inner, index) = locked;
let ptrs = match self.update_locked_inner(shard_id, inner, index, key, f)? {
Some(p) => p,
None => return Ok(None),
};
let ref_ptr = if return_old { ptrs.old } else { ptrs.new };
let mr = MutationResult {
ref_ptr: None,
hook: Some((
Some(unsafe { &(*ptrs.old).value as *const T }),
Some(unsafe { &(*ptrs.new).value as *const T }),
)),
retire: Some(ptrs.retire),
new_disk: Some(ptrs.new_disk),
watch: None,
};
*ev_out = Some((*key, mr));
Ok(Some(ref_ptr))
})?;
Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
}
pub fn atomic<R>(
&self,
shard_key: &K,
f: impl FnOnce(&mut TypedMapShard<'_, K, T, C, H>) -> DbResult<R>,
) -> DbResult<R> {
let _atomic_depth = crate::AtomicDepthGuard::try_enter()?;
let shard_id = self.shard_for(shard_key);
let mut shard = self.tx_shard(shard_id);
let result = f(&mut shard);
let needs_sync = Self::should_sync(&shard.held.locked);
let (events, wakes, guard) = Self::tx_release_shard(shard);
let result =
crate::projection::finish_atomic_teardown(self, shard_id, result, needs_sync, &events);
drop(guard);
<Self as Family>::wake_watches(wakes);
result
}
pub fn len(&self) -> usize {
self.len.load(std::sync::atomic::Ordering::Relaxed)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub(crate) fn reversed(&self) -> bool {
self.reversed
}
pub fn sync_hints(&self) -> DbResult<()> {
for shard in self.engine.shards().iter() {
shard.write_active_hint(size_of::<K>())?;
}
Ok(())
}
pub(crate) fn shard_count(&self) -> usize {
self.engine.shards().len()
}
pub(crate) fn codec(&self) -> &C {
&self.codec
}
#[cfg_attr(not(test), allow(dead_code))]
pub(crate) fn bind_tx_self(arc: &Arc<Self>) {
let _ = arc.tx_self.set(Arc::downgrade(arc));
}
pub(crate) fn tx_arc(&self) -> DbResult<Arc<Self>> {
self.tx_self
.get()
.and_then(Weak::upgrade)
.ok_or(DbError::Internal(
"TypedMap::bind_tx_self must be called before transaction use",
))
}
pub(crate) fn tx_install_projection_metadata(&self, db: &crate::armour::Db) -> DbResult<()> {
if H::PROJECTS
&& let Some(targets) = self.hook.projection_targets()
{
db.tx_register_projection_targets(targets)?;
}
if H::COUNTS
&& let Some(targets) = self.hook.counter_targets()
{
db.tx_register_counter_targets(targets)?;
}
if H::CLAIMS
&& let Some(targets) = self.hook.claim_targets()
{
db.tx_register_claim_targets(targets)?;
}
if H::EMITS
&& let Some(targets) = self.hook.stream_targets()
{
db.tx_register_stream_targets(targets)?;
}
Ok(())
}
pub(crate) fn tx_preplan_projection_targets(
&self,
) -> DbResult<Vec<Arc<dyn crate::txn::participant::TxParticipant>>> {
let mut out = Vec::new();
if H::PROJECTS
&& let Some(targets) = self.hook.projection_targets()
{
out.extend(
targets
.erased_targets()
.iter()
.filter_map(|h| targets.resolve_tx_participant(h.name)),
);
}
if H::COUNTS
&& let Some(targets) = self.hook.counter_targets()
{
out.extend(targets.all().map(|m| m.as_tx_participant()));
}
if H::CLAIMS
&& let Some(targets) = self.hook.claim_targets()
{
out.extend(targets.all().map(|t| t.as_tx_participant()));
}
if H::EMITS
&& let Some(targets) = self.hook.stream_targets()
{
out.extend(targets.tx_participants());
}
Ok(out)
}
pub(crate) fn tx_preplan_stream_targets(
&self,
) -> Vec<Arc<dyn crate::txn::participant::TxParticipant>> {
self.hook
.stream_targets()
.map(|targets| targets.tx_participants())
.unwrap_or_default()
}
pub(crate) fn tx_prepare_streams(
&self,
) -> DbResult<Vec<crate::queue::retention::ProducerPermit>> {
self.hook
.stream_targets()
.map_or(Ok(Vec::new()), crate::QueueTargets::prepare_streams)
}
pub(crate) fn tx_stream_writes(
&self,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<crate::txn::participant::TxStreamWrites> {
if !H::EMITS {
return Ok(crate::txn::participant::TxStreamWrites::empty());
}
let Some(targets) = self.hook.stream_targets() else {
return Ok(crate::txn::participant::TxStreamWrites::empty());
};
let key = K::from_bytes(key);
let old = old.map(|bytes| self.codec.decode_from(bytes)).transpose()?;
let new = new.map(|bytes| self.codec.decode_from(bytes)).transpose()?;
let (max_rows, max_bytes) = self.effective_projection_limits();
let mut sink = crate::EmitSink::new(targets, max_rows, max_bytes);
self.hook
.stream_project(&key, old.as_ref(), new.as_ref(), &mut sink)?;
targets.lower(&sink)
}
pub(crate) fn tx_project_writes(
&self,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<Vec<crate::txn::participant::TxDerivedWrite>> {
if !H::PROJECTS {
return Ok(Vec::new());
}
let Some(targets) = self.hook.projection_targets() else {
return Ok(Vec::new());
};
let key_typed = K::from_bytes(key);
let old_t = match old {
Some(b) => Some(self.codec.decode_from(b)?),
None => None,
};
let new_t = match new {
Some(b) => Some(self.codec.decode_from(b)?),
None => None,
};
let (max_rows, max_bytes) = self.effective_projection_limits();
let plan = crate::projection::diff_projection(
&self.hook,
targets,
max_rows,
max_bytes,
&key_typed,
old_t.as_ref(),
new_t.as_ref(),
)?;
plan.into_tx_derived(targets)
}
pub(crate) fn tx_counter_deltas_for(
&self,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<Vec<crate::txn::participant::TxCounterDelta>> {
if !H::COUNTS {
return Ok(Vec::new());
}
let Some(targets) = self.hook.counter_targets() else {
return Ok(Vec::new());
};
let key_typed = K::from_bytes(key);
let old_t = match old {
Some(b) => Some(self.codec().decode_from(b)?),
None => None,
};
let new_t = match new {
Some(b) => Some(self.codec().decode_from(b)?),
None => None,
};
let (max_rows, max_bytes) = self.effective_projection_limits();
let deltas = crate::counter::diff_counters(
&self.hook,
targets,
max_rows,
max_bytes,
&key_typed,
old_t.as_ref(),
new_t.as_ref(),
)?;
crate::counter::into_tx_counter_deltas(targets, deltas)
}
pub(crate) fn tx_claims_for(
&self,
name: &'static str,
key: &[u8],
old: Option<&[u8]>,
new: Option<&[u8]>,
) -> DbResult<Vec<crate::txn::participant::ClaimOp>> {
if !H::CLAIMS {
return Ok(Vec::new());
}
let Some(targets) = self.hook.claim_targets() else {
return Ok(Vec::new());
};
let key_typed = K::from_bytes(key);
let old_t = match old {
Some(b) => Some(self.codec().decode_from(b)?),
None => None,
};
let new_t = match new {
Some(b) => Some(self.codec().decode_from(b)?),
None => None,
};
let (max_rows, max_bytes) = self.effective_projection_limits();
let delta = crate::unique::claim_delta(
&self.hook,
targets,
max_rows,
max_bytes,
&key_typed,
old_t.as_ref(),
new_t.as_ref(),
)?;
crate::unique::into_tx_claim_ops(targets, name, key, delta)
}
pub(crate) fn tx_shard(&self, sid: usize) -> TypedMapShard<'_, K, T, C, H> {
let locked = self.lock_shard(sid);
let guard = self.collector.enter();
TypedMapShard {
map: self,
held: Held {
locked,
sid,
events: Vec::new(),
watch_wakes: WatchWakeBatch::default(),
publish_watches: true,
},
guard,
}
}
pub(crate) fn tx_release_shard(
shard: TypedMapShard<'_, K, T, C, H>,
) -> (MutationEvents<K, T>, WatchWakeBatch, seize::LocalGuard<'_>) {
let TypedMapShard { held, guard, .. } = shard;
let Held {
locked,
events,
watch_wakes,
..
} = held;
locked.release();
(events, watch_wakes, guard)
}
pub(crate) fn tx_reserve_gsns(&self, count: usize) -> u64 {
self.engine
.gsn()
.fetch_add(count as u64, std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn tx_encrypted(&self) -> bool {
#[cfg(feature = "encryption")]
{
self.engine.config().encryption_key.is_some()
}
#[cfg(not(feature = "encryption"))]
{
false
}
}
pub(crate) fn tx_durable_recovered_gsn(&self, sid: usize) -> u64 {
self.engine.shards()[sid].durable_recovered_gsn()
}
pub fn shard_for(&self, key: &K) -> usize {
if self.shard_prefix_bits == 0 || self.shard_prefix_bits >= size_of::<K>() * 8 {
let hash = xxhash_rust::xxh3::xxh3_64(key.as_bytes());
return (hash as usize) % self.engine.shards().len();
}
let full_bytes = self.shard_prefix_bits / 8;
let extra_bits = self.shard_prefix_bits % 8;
let hash = if extra_bits == 0 {
xxhash_rust::xxh3::xxh3_64(&key.as_bytes()[..full_bytes])
} else {
let mut buf = K::zeroed();
buf.as_bytes_mut()[..full_bytes].copy_from_slice(&key.as_bytes()[..full_bytes]);
let mask = !((1u8 << (8 - extra_bits)) - 1);
buf.as_bytes_mut()[full_bytes] = key.as_bytes()[full_bytes] & mask;
xxhash_rust::xxh3::xxh3_64(&buf.as_bytes()[..full_bytes + 1])
};
(hash as usize) % self.engine.shards().len()
}
pub fn iter_view(&self) -> Option<TypedMapIterView<'_, K, T, C, H>> {
self.iterable.then(|| TypedMapIterView {
map: self,
inner: MapIterView::new(self, dir_from(self.reversed)),
})
}
pub fn scope<const N: usize>(
&self,
key: &K,
) -> Option<crate::scope::Scope<'_, crate::scope::MapScan<&Self>, N>> {
crate::scope::map_scope::<_, _, N>(self, key, dir_from(self.reversed()), self.iterable)
}
#[cfg(feature = "armour")]
pub fn tx_scope<const N: usize>(
self: &Arc<Self>,
key: &K,
) -> Option<crate::TxScope<Self, K, N>> {
self.iterable.then(|| {
crate::txn::scope::TxScope::from_parts(
self.clone(),
key,
dir_from(self.reversed()),
Self::tx_scope_keys,
)
})
}
fn tx_scope_keys<'a>(
collection: &'a Self,
start: Bound<&'a K>,
end: Bound<&'a K>,
dir: Dir,
) -> Box<dyn Iterator<Item = K> + 'a> {
Box::new(MapIterView::new(collection, dir).keys_range_bounds(start, end))
}
pub fn first(&self) -> DbResult<Option<(K, TypedRef<'_, T>)>> {
let view = self.iter_view().ok_or(DbError::NotImplemented)?;
Ok(view.iter().next())
}
pub fn last(&self) -> DbResult<Option<(K, TypedRef<'_, T>)>> {
self.iter_view().ok_or(DbError::NotImplemented)?;
let view = TypedMapIterView {
map: self,
inner: MapIterView::new(self, dir_from(self.reversed).flip()),
};
Ok(view.iter().next())
}
pub(crate) fn inspection_shard_count(&self) -> usize {
self.engine.shards().len()
}
pub(crate) fn inspection_guard(&self) -> seize::LocalGuard<'_> {
self.collector.enter()
}
pub(crate) fn inspection_range_bounds(
&self,
start: Bound<&K>,
end: Bound<&K>,
batch: usize,
) -> TryMapIter<'_, TypedMap<K, T, C, H>> {
MapIterView::with_batch(self, dir_from(self.reversed), batch).try_range_bounds(start, end)
}
pub fn for_each(&self, mut f: impl FnMut(&K, &T)) {
let _guard = self.collector.enter();
for i in 0..self.engine.shards().len() {
let pairs: Vec<(K, *const TypedData<T>)> = {
let index = sync::lock(&self.indexes[i]);
index
.iter()
.map(|(k, entry)| (*k, entry.ptr as *const TypedData<T>))
.collect()
};
for (key, ptr) in pairs {
let data = unsafe { &*ptr };
f(&key, &data.value);
}
}
}
pub(crate) fn for_each_while(
&self,
mut f: impl FnMut(&K, &T) -> ControlFlow<()>,
) -> ControlFlow<()> {
let _guard = self.collector.enter();
for i in 0..self.engine.shards().len() {
let pairs: Vec<(K, *const TypedData<T>)> = {
let index = sync::lock(&self.indexes[i]);
index
.iter()
.map(|(k, entry)| (*k, entry.ptr as *const TypedData<T>))
.collect()
};
for (key, ptr) in pairs {
let data = unsafe { &*ptr };
if let ControlFlow::Break(()) = f(&key, &data.value) {
return ControlFlow::Break(());
}
}
}
ControlFlow::Continue(())
}
pub fn migrate(&self, f: impl Fn(&K, &T) -> crate::MigrateAction<T>) -> DbResult<usize> {
self.migrate_inner(f, true)
}
pub(crate) fn migrate_inner(
&self,
f: impl Fn(&K, &T) -> crate::MigrateAction<T>,
fire_init: bool,
) -> DbResult<usize> {
use crate::MigrateAction;
let mut count = 0;
for i in 0..self.engine.shards().len() {
let keys: Vec<K> = {
let index = sync::lock(&self.indexes[i]);
index.keys().copied().collect()
};
for key in keys {
let value_ref = match self.get(&key) {
Some(v) => v,
None => continue,
};
let action = f(&key, &*value_ref);
drop(value_ref);
match action {
MigrateAction::Keep => {
if fire_init
&& H::NEEDS_INIT
&& let Some(v) = self.get(&key)
{
self.hook.on_init(&key, &*v);
}
}
MigrateAction::Update(value) => {
let shard_id = self.shard_for(&key);
{
let mut inner = self.engine.shards()[shard_id].lock();
let mut index = sync::lock(&self.indexes[shard_id]);
let guard = self.collector.enter();
self.put_locked_inner::<false>(
shard_id,
&mut inner,
&mut index,
&guard,
&key,
value,
crate::ttl::TtlAt::Derive,
)?;
}
if fire_init
&& H::NEEDS_INIT
&& let Some(v) = self.get(&key)
{
self.hook.on_init(&key, &*v);
}
count += 1;
}
MigrateAction::Delete => {
let shard_id = self.shard_for(&key);
let mut inner = self.engine.shards()[shard_id].lock();
let mut index = sync::lock(&self.indexes[shard_id]);
let guard = self.collector.enter();
self.delete_locked_inner::<false>(
shard_id, &mut inner, &mut index, &guard, &key,
)?;
count += 1;
}
}
}
}
tracing::info!(mutations = count, "typed_map migration complete");
Ok(count)
}
pub(crate) fn replay_init(&self) {
if !H::NEEDS_INIT {
return;
}
let _guard = self.collector.enter();
for shard in &self.indexes {
let pairs: Vec<(K, *const TypedData<T>)> = {
let index = sync::lock(shard);
index
.iter()
.map(|(k, entry)| (*k, entry.ptr as *const TypedData<T>))
.collect()
};
for (key, ptr) in pairs {
let data = unsafe { &*ptr };
self.hook.on_init(&key, &data.value);
}
}
}
pub(crate) fn effective_projection_limits(&self) -> (usize, usize) {
let cfg = self.config();
let (hook_rows, hook_bytes) = self.hook.projection_limits();
(
hook_rows.min(cfg.max_projected_rows),
hook_bytes.min(cfg.max_projected_bytes),
)
}
#[allow(clippy::collapsible_if)]
#[inline]
fn claim_guard(&self, key: &K, old: Option<&T>, new: Option<&T>) -> DbResult<()> {
#[cfg(feature = "armour")]
if H::EMITS && !crate::txn_depth::in_stream_apply() {
if let Some(targets) = self.hook.stream_targets() {
let collection = self
.counter_dirty
.get()
.map(crate::counter_dirty::CounterDirty::name)
.unwrap_or("");
return Err(DbError::ChangeStreamRequiresTx {
collection,
queue: targets.first_name(),
});
}
}
if !H::CLAIMS {
return Ok(());
}
if self.hook.claim_targets().is_none() {
return Ok(());
}
let collection = self
.counter_dirty
.get()
.map(crate::counter_dirty::CounterDirty::name)
.unwrap_or("");
crate::unique::guard_nontx_write(
collection,
&self.hook,
self.effective_projection_limits(),
key,
old,
new,
)
}
pub(crate) fn reconcile_projections(&self) -> DbResult<()> {
if !H::PROJECTS {
return Ok(());
}
let Some(targets) = self.hook.projection_targets() else {
return Ok(());
};
let (max_rows, max_bytes) = self.effective_projection_limits();
crate::projection::reconcile_carrier_projections(
targets,
|emit| {
let mut err: Option<DbError> = None;
self.for_each(|k, v| {
if err.is_some() {
return;
}
let mut sink =
crate::projection::ProjectionSink::new(targets, max_rows, max_bytes);
match self.hook.project(k, v, &mut sink).and_then(|()| {
for row in sink.rows() {
emit(row.target, &row.key, &row.value)?;
}
Ok(())
}) {
Ok(()) => {}
Err(e) => err = Some(e),
}
});
err.map_or(Ok(()), Err)
},
true,
)
}
pub(crate) fn reconcile_counters(&self) -> DbResult<u64> {
if !H::COUNTS {
return Ok(0);
}
let Some(targets) = self.hook.counter_targets() else {
return Ok(0);
};
let limits = self.effective_projection_limits();
crate::counter::recount(
|k: &K, v: &T, sink| self.hook.counter_project(k, v, sink),
targets,
limits,
&mut |visit| {
let mut err: Option<DbError> = None;
self.for_each(|k, v| {
if err.is_some() {
return;
}
if let Err(e) = visit(k, v) {
err = Some(e);
}
});
err.map_or(Ok(()), Err)
},
)
}
pub(crate) fn counter_target_names(&self) -> Vec<&'static str> {
self.hook
.counter_targets()
.map(|t| t.all().map(|m| m.collection_name()).collect())
.unwrap_or_default()
}
pub(crate) fn indexes(&self) -> &[Mutex<ShardIndex<K, TypedMapEntry<T>>>] {
&self.indexes
}
pub(crate) fn fire_hook_and_retire(&self, key: &K, mr: &MutationResult<K, T>) {
if let Some((old, new)) = mr.hook {
unsafe {
self.hook.on_write(key, old.map(|p| &*p), new.map(|p| &*p));
}
}
if let Some(ptr) = mr.retire {
unsafe {
self.collector
.retire(ptr, seize::reclaim::boxed::<TypedData<T>>);
}
}
}
pub(crate) fn mark_counter_dirty(&self) -> DbResult<()> {
match self.counter_dirty.get() {
Some(d) => d.mark(),
None => Ok(()),
}
}
fn apply_projection_for_mutation(&self, ev: &(K, MutationResult<K, T>)) -> DbResult<()> {
if !H::PROJECTS && !H::COUNTS {
return Ok(());
}
let (key, mr) = ev;
let Some((old_ptr, new_ptr)) = mr.hook else {
return Ok(());
};
let old = old_ptr.map(|p| unsafe { &*p });
let new = new_ptr.map(|p| unsafe { &*p });
let (max_rows, max_bytes) = self.effective_projection_limits();
let mut plan = match (H::PROJECTS, self.hook.projection_targets()) {
(true, Some(targets)) => crate::projection::diff_projection(
&self.hook, targets, max_rows, max_bytes, key, old, new,
)?,
_ => crate::projection::PendingProjection::default(),
};
if H::COUNTS
&& let Some(counter_targets) = self.hook.counter_targets()
{
let deltas = crate::counter::diff_counters(
&self.hook,
counter_targets,
max_rows,
max_bytes,
key,
old,
new,
)?;
if !deltas.is_empty() {
plan.set_deltas(deltas, counter_targets.clone());
}
}
plan.apply_or_incomplete()
}
#[allow(clippy::too_many_arguments)]
fn put_locked_inner<const HOOKS: bool>(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, TypedMapEntry<T>>,
_guard: &seize::LocalGuard<'_>,
key: &K,
value: T,
expires_at: crate::ttl::TtlAt,
) -> DbResult<MutationResult<K, T>> {
let expires_at = match expires_at {
crate::ttl::TtlAt::Precomputed(at) => at,
crate::ttl::TtlAt::Derive if H::EXPIRES => self.hook.deadline(key, &value),
crate::ttl::TtlAt::Derive => None,
};
if H::EMITS || (H::CLAIMS && self.hook.claim_targets().is_some()) {
#[cfg(feature = "ctx-gates")]
crate::ctx_gates_sync::record_claim_pre_read();
let old = index.get(key).map(|e| unsafe { &(*e.ptr).value });
self.claim_guard(key, old, Some(&value))?;
}
let mut buf = sync::lock(&self.encode_bufs[shard_id]);
self.codec.encode_to(&value, &mut buf)?;
let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;
let new_data_ptr = Box::into_raw(Box::new(TypedData {
disk: disk_loc,
value,
}));
if let Some(old_entry) = index.upsert(*key, TypedMapEntry { ptr: new_data_ptr }) {
let old_data = old_entry.ptr as *const TypedData<T>;
let old_disk = unsafe { (*old_entry.ptr).disk };
inner.add_dead_bytes(
old_disk.file_id,
crate::entry::entry_size(size_of::<K>(), old_disk.len),
);
if HOOKS {
self.record_ttl(shard_id, key, expires_at);
return Ok(MutationResult {
ref_ptr: Some(old_data),
hook: Some((
Some(unsafe { &(*old_data).value as *const T }),
Some(unsafe { &(*new_data_ptr).value as *const T }),
)),
retire: Some(old_entry.ptr),
new_disk: Some(disk_loc),
watch: None,
});
}
self.record_ttl(shard_id, key, expires_at);
unsafe {
self.collector
.retire(old_entry.ptr, seize::reclaim::boxed::<TypedData<T>>);
}
Ok(MutationResult {
ref_ptr: Some(old_data),
hook: None,
retire: None,
new_disk: None,
watch: None,
})
} else {
self.len.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.record_ttl(shard_id, key, expires_at);
let hook = if HOOKS {
Some((None, Some(unsafe { &(*new_data_ptr).value as *const T })))
} else {
None
};
Ok(MutationResult {
ref_ptr: None,
hook,
retire: None,
new_disk: Some(disk_loc),
watch: None,
})
}
}
#[allow(clippy::too_many_arguments)]
fn insert_locked<const HOOKS: bool>(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, TypedMapEntry<T>>,
_guard: &seize::LocalGuard<'_>,
key: &K,
value: T,
expires_at: Option<u64>,
) -> DbResult<MutationResult<K, T>> {
if index.contains_key(key) {
return Err(DbError::KeyExists);
}
self.claim_guard(key, None, Some(&value))?;
let mut buf = sync::lock(&self.encode_bufs[shard_id]);
self.codec.encode_to(&value, &mut buf)?;
let (disk_loc, _gsn) = inner.append_entry(shard_id as u8, key.as_bytes(), &buf, false)?;
let new_data_ptr = Box::into_raw(Box::new(TypedData {
disk: disk_loc,
value,
}));
index.upsert(*key, TypedMapEntry { ptr: new_data_ptr });
self.len.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.record_ttl(shard_id, key, expires_at);
let hook = if HOOKS {
Some((None, Some(unsafe { &(*new_data_ptr).value as *const T })))
} else {
None
};
Ok(MutationResult {
ref_ptr: None,
hook,
retire: None,
new_disk: Some(disk_loc),
watch: None,
})
}
fn delete_locked_inner<const HOOKS: bool>(
&self,
shard_id: usize,
inner: &mut ShardInner,
index: &mut ShardIndex<K, TypedMapEntry<T>>,
_guard: &seize::LocalGuard<'_>,
key: &K,
) -> DbResult<MutationResult<K, T>> {
#[cfg(test)]
if crate::durability::ttl_delete_should_fail() {
return Err(DbError::Internal("test: TTL delete injected failure"));
}
let entry = match index.get(key) {
Some(e) => e,
None => {
return Ok(MutationResult {
ref_ptr: None,
hook: None,
retire: None,
new_disk: None,
watch: None,
});
}
};
let old_data = entry.ptr as *const TypedData<T>;
self.claim_guard(key, Some(unsafe { &(*old_data).value }), None)?;
inner.append_entry(shard_id as u8, key.as_bytes(), &[], true)?;
let removed = index.remove(key).expect("key exists after probe");
self.len.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
let old_disk = unsafe { (*removed.ptr).disk };
inner.add_dead_bytes(
old_disk.file_id,
crate::entry::entry_size(size_of::<K>(), old_disk.len),
);
if HOOKS {
self.record_ttl(shard_id, key, None);
Ok(MutationResult {
ref_ptr: Some(old_data),
hook: Some((Some(unsafe { &(*old_data).value as *const T }), None)),
retire: Some(removed.ptr),
new_disk: None,
watch: None,
})
} else {
self.record_ttl(shard_id, key, None);
unsafe {
self.collector
.retire(removed.ptr, seize::reclaim::boxed::<TypedData<T>>);
}
Ok(MutationResult {
ref_ptr: Some(old_data),
hook: None,
retire: None,
new_disk: None,
watch: None,
})
}
}
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T>, H: TypedWriteHook<K, T>> Drop
for TypedMap<K, T, C, H>
{
fn drop(&mut self) {
for index_mutex in &self.indexes {
let map = sync::lock(index_mutex);
for entry in map.values() {
unsafe {
drop(Box::from_raw(entry.ptr));
}
}
}
}
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
ShardSource for TypedMap<K, T, C, H>
{
type Key = K;
type Val = *const TypedData<T>;
fn shard_count(&self) -> usize {
self.engine.shards().len()
}
fn next_batch(
&self,
shard: usize,
after: Option<&K>,
dir: Dir,
n: usize,
range_start: Bound<OrdKey<K>>,
range_end: Bound<OrdKey<K>>,
) -> Vec<(K, *const TypedData<T>)> {
let index = sync::lock(&self.indexes[shard]);
let Some(ki) = index.key_index() else {
return Vec::new();
};
let keys = key_index_batch(ki, after, dir, n, range_start, range_end);
let mut out = Vec::with_capacity(keys.len());
for k in keys {
if let Some(e) = index.get(&k) {
out.push((k, e.ptr as *const TypedData<T>));
}
}
out
}
fn next_keys_batch(
&self,
shard: usize,
after: Option<&K>,
dir: Dir,
n: usize,
range_start: Bound<OrdKey<K>>,
range_end: Bound<OrdKey<K>>,
) -> Vec<K> {
let index = sync::lock(&self.indexes[shard]);
let Some(ki) = index.key_index() else {
return Vec::new();
};
key_index_batch(ki, after, dir, n, range_start, range_end)
}
}
pub struct TypedMapIterView<
'a,
K: Key + Send + Sync + Hash + Eq,
T: Send + Sync,
C: Codec<T> + Sync,
H: TypedWriteHook<K, T>,
> {
map: &'a TypedMap<K, T, C, H>,
inner: MapIterView<'a, TypedMap<K, T, C, H>>,
}
impl<'a, K, T, C, H> TypedMapIterView<'a, K, T, C, H>
where
K: Key + Send + Sync + Hash + Eq,
T: Send + Sync,
C: Codec<T> + Sync,
H: TypedWriteHook<K, T>,
{
pub fn iter(&self) -> TypedMapIter<'a, K, T, C, H> {
let guard = self.map.collector.enter();
TypedMapIter {
map: self.map,
inner: self.inner.iter(),
_guard: guard,
}
}
pub fn range(&self, start: &K, end: &K) -> TypedMapIter<'a, K, T, C, H> {
let guard = self.map.collector.enter();
TypedMapIter {
map: self.map,
inner: self.inner.range(start, end),
_guard: guard,
}
}
pub fn range_bounds(&self, start: Bound<&K>, end: Bound<&K>) -> TypedMapIter<'a, K, T, C, H> {
let guard = self.map.collector.enter();
TypedMapIter {
map: self.map,
inner: self.inner.range_bounds(start, end),
_guard: guard,
}
}
pub fn prefix_iter(&self, prefix: &[u8]) -> TypedMapIter<'a, K, T, C, H> {
let guard = self.map.collector.enter();
TypedMapIter {
map: self.map,
inner: self.inner.prefix_iter(prefix),
_guard: guard,
}
}
pub fn paginate(&self, after: Option<&K>, limit: usize) -> Page<K, T>
where
T: Clone,
{
let iter = match (after, dir_from(self.map.reversed)) {
(None, _) => self.iter(),
(Some(k), Dir::Asc) => self.range_bounds(Bound::Excluded(k), Bound::Unbounded),
(Some(k), Dir::Desc) => self.range_bounds(Bound::Unbounded, Bound::Excluded(k)),
};
let items: Vec<_> = iter
.take(limit)
.map(|(k, r)| (k, r.clone_value()))
.collect();
let next = if items.len() == limit {
items.last().map(|(k, _)| *k)
} else {
None
};
Page { items, next }
}
pub fn keys(&self) -> impl Iterator<Item = K> + '_ {
self.inner.keys()
}
pub fn keys_range(&self, start: &K, end: &K) -> impl Iterator<Item = K> + '_ {
self.inner.keys_range(start, end)
}
pub fn keys_range_bounds(
&self,
start: Bound<&K>,
end: Bound<&K>,
) -> impl Iterator<Item = K> + '_ {
self.inner.keys_range_bounds(start, end)
}
pub fn keys_prefix(&self, prefix: &[u8]) -> impl Iterator<Item = K> + '_ {
self.inner.keys_prefix(prefix)
}
pub fn keys_paginate(&self, after: Option<&K>, limit: usize) -> KeyPage<K> {
self.inner.keys_paginate(after, limit)
}
pub fn retain(&self, mut f: impl FnMut(&K, &T) -> bool) -> DbResult<()> {
let _guard = self.map.collector.enter();
for (key, ptr) in self.inner.iter() {
let val = unsafe { &(*ptr).value };
if f(&key, val) {
continue;
}
match self.map.compare_delete_if_ptr(&key, ptr) {
Ok(()) => {}
Err(DbError::CasMismatch) | Err(DbError::KeyNotFound) => {
tracing::debug!("retain: compare_delete skipped (concurrent change)");
}
Err(e) => return Err(e),
}
}
Ok(())
}
}
pub struct TypedMapIter<
'a,
K: Key + Send + Sync + Hash + Eq,
T: Send + Sync,
C: Codec<T> + Sync,
H: TypedWriteHook<K, T>,
> {
map: &'a TypedMap<K, T, C, H>,
inner: MapIter<'a, TypedMap<K, T, C, H>>,
_guard: seize::LocalGuard<'a>,
}
impl<'a, K, T, C, H> Iterator for TypedMapIter<'a, K, T, C, H>
where
K: Key + Send + Sync + Hash + Eq,
T: Send + Sync,
C: Codec<T> + Sync,
H: TypedWriteHook<K, T>,
{
type Item = (K, TypedRef<'a, T>);
fn next(&mut self) -> Option<Self::Item> {
let (k, ptr) = self.inner.next()?;
let guard = self.map.collector.enter();
Some((k, TypedRef::new(guard, ptr)))
}
}
pub struct TypedMapShard<
'a,
K: Key + Send + Sync + Hash + Eq,
T: Send + Sync,
C: Codec<T>,
H: TypedWriteHook<K, T> = NoHook,
> {
map: &'a TypedMap<K, T, C, H>,
pub(crate) held: Held<'a, TypedMap<K, T, C, H>>,
guard: seize::LocalGuard<'a>,
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
TypedMapShard<'_, K, T, C, H>
{
pub fn put(&mut self, key: &K, value: T) -> DbResult<Option<TypedRef<'_, T>>> {
self.check_shard(key)?;
let map = self.map;
let sid = self.held.sid;
let ref_ptr = {
let guard = &self.guard;
Ctx::with_shard(&mut self.held, map, sid, |locked, ev_out| {
let (inner, index) = locked;
let mr = map.put_locked_inner::<true>(
sid,
inner,
index,
guard,
key,
value,
crate::ttl::TtlAt::Derive,
)?;
let ref_ptr = mr.ref_ptr;
*ev_out = map.defer_or_retire_event(key, mr);
Ok(ref_ptr)
})?
};
let guard = self.map.collector.enter();
Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
}
pub fn insert(&mut self, key: &K, value: T) -> DbResult<()> {
self.check_shard(key)?;
let map = self.map;
let sid = self.held.sid;
let expires_at = if H::EXPIRES {
map.hook.deadline(key, &value)
} else {
None
};
let guard = &self.guard;
Ctx::with_shard(&mut self.held, map, sid, |locked, ev_out| {
let (inner, index) = locked;
let mr = map.insert_locked::<true>(sid, inner, index, guard, key, value, expires_at)?;
*ev_out = map.defer_or_retire_event(key, mr);
Ok(())
})
}
pub fn delete(&mut self, key: &K) -> DbResult<Option<TypedRef<'_, T>>> {
self.check_shard(key)?;
let map = self.map;
let sid = self.held.sid;
let ref_ptr = {
let guard = &self.guard;
Ctx::with_shard(&mut self.held, map, sid, |locked, ev_out| {
let (inner, index) = locked;
let mr = map.delete_locked_inner::<true>(sid, inner, index, guard, key)?;
let ref_ptr = mr.ref_ptr;
*ev_out = map.defer_or_retire_event(key, mr);
Ok(ref_ptr)
})?
};
let guard = self.map.collector.enter();
Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
}
pub fn try_get(&self, key: &K) -> DbResult<Option<&T>> {
self.check_shard(key)?;
Ok(match self.held.locked.1.get(key) {
Some(entry) => Some(unsafe { &(*entry.ptr).value }),
None => None,
})
}
pub fn get_or_err(&self, key: &K) -> DbResult<&T> {
self.try_get(key)?.ok_or(DbError::KeyNotFound)
}
pub fn try_contains(&self, key: &K) -> DbResult<bool> {
self.check_shard(key)?;
Ok(self.held.locked.1.contains_key(key))
}
pub fn update(
&mut self,
key: &K,
f: impl FnOnce(&T) -> T,
) -> DbResult<Option<TypedRef<'_, T>>> {
self.update_inner(key, f, false)
}
pub fn fetch_update(
&mut self,
key: &K,
f: impl FnOnce(&T) -> T,
) -> DbResult<Option<TypedRef<'_, T>>> {
self.update_inner(key, f, true)
}
pub fn upsert<U: crate::IntoUpsert<T>>(
&mut self,
key: &K,
f: impl FnOnce(Option<&T>) -> U,
) -> crate::DbResult<U::Out> {
crate::upsert::upsert_via_held(self, key, f)
}
fn update_inner(
&mut self,
key: &K,
f: impl FnOnce(&T) -> T,
return_old: bool,
) -> DbResult<Option<TypedRef<'_, T>>> {
self.check_shard(key)?;
let map = self.map;
let sid = self.held.sid;
let ref_ptr = {
Ctx::with_shard(&mut self.held, map, sid, |locked, ev_out| {
let (inner, index) = locked;
let ptrs = match map.update_locked_inner(sid, inner, index, key, f)? {
Some(p) => p,
None => return Ok(None),
};
let ref_ptr = if return_old { ptrs.old } else { ptrs.new };
let mr = MutationResult {
ref_ptr: None,
hook: Some((
Some(unsafe { &(*ptrs.old).value as *const T }),
Some(unsafe { &(*ptrs.new).value as *const T }),
)),
retire: Some(ptrs.retire),
new_disk: Some(ptrs.new_disk),
watch: None,
};
*ev_out = map.defer_or_retire_event(key, mr);
Ok(Some(ref_ptr))
})?
};
let guard = self.map.collector.enter();
Ok(ref_ptr.map(|p| TypedRef::new(guard, p)))
}
fn check_shard(&self, key: &K) -> DbResult<()> {
if self.map.shard_for(key) != self.held.sid {
return Err(DbError::ShardMismatch);
}
Ok(())
}
}
impl<K: Key + Send + Sync + Hash + Eq, T: Send + Sync, C: Codec<T> + Sync, H: TypedWriteHook<K, T>>
TypedMapShard<'_, K, T, C, H>
{
pub(crate) fn tx_arm_gsn(&mut self, gsn: u64) {
self.held.locked.0.arm_gsn(gsn);
}
pub(crate) fn tx_clear_armed_gsn(&mut self) {
self.held.locked.0.clear_armed_gsn();
}
pub(crate) fn tx_fence(&mut self, gsn: u64) -> u64 {
self.held.locked.0.fence(gsn)
}
pub(crate) fn tx_buffered_events(&self) -> usize {
self.held.events.len()
}
}
#[cfg(feature = "armour")]
impl<T, C, H> crate::armour::collection::Collection for TypedMap<T::SelfId, T, C, H>
where
T: crate::CollectionMeta + Clone + Send + Sync,
C: crate::Codec<T> + Sync,
H: crate::hook::TypedWriteHook<T::SelfId, T>,
T::SelfId: crate::Key + Send + Sync + std::hash::Hash + Eq,
{
fn name(&self) -> &str {
T::NAME
}
fn len(&self) -> usize {
self.len()
}
fn compact(&self) -> crate::DbResult<usize> {
self.compact()
}
fn flush(&self) -> crate::DbResult<()> {
crate::projection::collection_flush_fault()?;
self.flush_buffers()?;
if self.engine.hints() {
self.sync_hints()?;
}
Ok(())
}
fn dirty_shards(&self) -> Option<usize> {
Some(self.engine.dirty_shards())
}
fn periodic_flush(&self) -> crate::DbResult<()> {
TypedMap::durable_flush(self)
}
fn durable_flush(&self) -> crate::DbResult<()> {
TypedMap::durable_flush(self)
}
fn clean_shutdown(&self) -> crate::DbResult<()> {
TypedMap::clean_shutdown(self)
}
fn sweep_expired(
&self,
now_ms: u64,
budget: usize,
) -> crate::DbResult<crate::ttl::SweepOutcome> {
self.sweep_expired(now_ms, budget)
}
fn set_ttl_sweep(&self, enabled: bool) -> crate::DbResult<()> {
self.set_ttl_sweep(enabled)
}
fn rebuild_ttl_index(&self) -> crate::DbResult<()> {
self.rebuild_ttl_index()
}
fn reset_counter_dirty_marker(&self) -> bool {
self.counter_dirty.get().is_some_and(|d| d.take_marked())
}
fn reconcile_counters(&self) -> crate::DbResult<u64> {
TypedMap::reconcile_counters(self)
}
fn counter_target_names(&self) -> Vec<&'static str> {
TypedMap::counter_target_names(self)
}
fn close_watches(&self) {
if let Some(channel) = self.hook.watch_channel() {
channel.close_terminal();
}
}
}
#[cfg(feature = "typed-tree")]
impl<
K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
T: armour_core::GetType
+ crate::tree_meta::CollectionMeta
+ Clone
+ PartialEq
+ Send
+ Sync
+ 'static,
C: Codec<T> + Sync + 'static,
H: TypedWriteHook<K, T> + 'static,
> crate::projection::sealed::Sealed for TypedMap<K, T, C, H>
{
}
impl<
K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
T: armour_core::GetType
+ crate::tree_meta::CollectionMeta
+ Clone
+ PartialEq
+ Send
+ Sync
+ 'static,
C: Codec<T> + Sync + 'static,
H: TypedWriteHook<K, T> + 'static,
> crate::projection::ProjectableIndex for TypedMap<K, T, C, H>
{
fn collection_name(&self) -> &'static str {
T::NAME
}
fn put_bytes(&self, key: &[u8], value: &[u8]) -> DbResult<()> {
if key.len() != size_of::<K>() {
return Err(DbError::CorruptedEntry { offset: 0 });
}
let k = K::from_bytes(key);
let v = self.codec().decode_from(value)?;
self.put(&k, v)?;
Ok(())
}
fn delete_bytes(&self, key: &[u8]) -> DbResult<()> {
if key.len() != size_of::<K>() {
return Err(DbError::CorruptedEntry { offset: 0 });
}
let k = K::from_bytes(key);
self.delete(&k)?;
Ok(())
}
fn values_equal(&self, a: &[u8], b: &[u8]) -> DbResult<bool> {
if a == b {
return Ok(true);
}
let va = self.codec().decode_from(a)?;
let vb = self.codec().decode_from(b)?;
Ok(va == vb)
}
fn for_each_bytes(&self, f: &mut dyn FnMut(&[u8], &[u8]) -> DbResult<()>) -> DbResult<()> {
let mut err: Option<DbError> = None;
self.for_each(|k, v| {
if err.is_some() {
return;
}
let mut buf = Vec::new();
match self
.codec()
.encode_to(v, &mut buf)
.and_then(|()| f(k.as_bytes(), &buf))
{
Ok(()) => {}
Err(e) => err = Some(e),
}
});
match err {
Some(e) => Err(e),
None => Ok(()),
}
}
fn durable_flush(&self) -> DbResult<()> {
TypedMap::durable_flush(self)
}
}
#[cfg(feature = "typed-tree")]
impl<
K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
T: armour_core::GetType
+ crate::tree_meta::CollectionMeta
+ Clone
+ PartialEq
+ Send
+ Sync
+ 'static,
C: Codec<T> + Sync + 'static,
H: TypedWriteHook<K, T> + 'static,
> crate::projection::ProjectableIndexTx for TypedMap<K, T, C, H>
{
fn as_tx_participant(self: Arc<Self>) -> Arc<dyn crate::txn::participant::TxParticipant> {
self
}
}
#[cfg(feature = "typed-tree")]
impl<
K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
T: armour_core::GetType
+ crate::tree_meta::CollectionMeta
+ Clone
+ PartialEq
+ Send
+ Sync
+ 'static,
C: Codec<T> + Sync + 'static,
H: TypedWriteHook<K, T> + 'static,
> crate::projection::IndexKeyed for TypedMap<K, T, C, H>
{
type K = K;
type V = T;
fn encode_value(&self, v: &T) -> DbResult<Vec<u8>> {
let mut buf = Vec::new();
self.codec().encode_to(v, &mut buf)?;
Ok(buf)
}
}
#[cfg(feature = "typed-tree")]
impl<
K: Key + Send + Sync + std::hash::Hash + Eq + 'static,
T: armour_core::GetType
+ crate::tree_meta::CollectionMeta
+ Clone
+ PartialEq
+ Send
+ Sync
+ 'static,
C: Codec<T> + Sync + 'static,
H: TypedWriteHook<K, T> + 'static,
> crate::counter::CounterCollection for TypedMap<K, T, C, H>
{
fn decode_value(&self, bytes: &[u8]) -> DbResult<T> {
self.codec().decode_from(bytes)
}
fn read_bytes(&self, key: &[u8]) -> DbResult<Option<Vec<u8>>> {
if key.len() != size_of::<K>() {
return Err(DbError::CorruptedEntry { offset: 0 });
}
let k = K::from_bytes(key);
self.get(&k).map(|v| self.encode_value(&v)).transpose()
}
fn contains_bytes(&self, key: &[u8]) -> DbResult<bool> {
if key.len() != size_of::<K>() {
return Err(DbError::CorruptedEntry { offset: 0 });
}
let k = K::from_bytes(key);
Ok(self.contains(&k))
}
fn upsert_counter(
&self,
key: &[u8],
f: &mut dyn FnMut(Option<&T>) -> Upsert<T>,
) -> DbResult<()> {
if key.len() != size_of::<K>() {
return Err(DbError::CorruptedEntry { offset: 0 });
}
let k = K::from_bytes(key);
self.upsert(&k, f)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::codec::Codec;
use crate::config::Config;
use crate::hook::TypedWriteHook;
use std::collections::HashMap as StdHashMap;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use tempfile::tempdir;
struct U64Codec;
impl Codec<u64> for U64Codec {
fn encode_to(&self, value: &u64, buf: &mut Vec<u8>) -> DbResult<()> {
buf.clear();
buf.extend_from_slice(&value.to_be_bytes());
Ok(())
}
fn decode_from(&self, bytes: &[u8]) -> DbResult<u64> {
let arr: [u8; 8] = bytes
.try_into()
.map_err(|_| DbError::CorruptedEntry { offset: 0 })?;
Ok(u64::from_be_bytes(arr))
}
}
type Map = TypedMap<[u8; 8], u64, U64Codec>;
fn open(dir: &std::path::Path) -> Map {
TypedMap::<[u8; 8], u64, U64Codec>::open(dir, Config::test(), U64Codec).unwrap()
}
fn key_on_other_shard(map: &Map, anchor: &[u8; 8]) -> [u8; 8] {
let anchor_shard = map.shard_for(anchor);
(2u64..1000)
.map(|i| i.to_be_bytes())
.find(|k| map.shard_for(k) != anchor_shard)
.expect("key on different shard")
}
#[derive(Default)]
struct TRecHook {
writes: AtomicUsize,
#[allow(clippy::type_complexity)]
seq: crate::sync::Mutex<Vec<(u64, Option<u64>, Option<u64>)>>,
}
impl TypedWriteHook<[u8; 8], u64> for TRecHook {
fn on_write(&self, key: &[u8; 8], old: Option<&u64>, new: Option<&u64>) {
self.writes.fetch_add(1, AtomicOrdering::Relaxed);
crate::sync::lock(&self.seq).push((
u64::from_be_bytes(*key),
old.copied(),
new.copied(),
));
}
}
fn open_typed_map_hooked<H: TypedWriteHook<[u8; 8], u64>>(
dir: &std::path::Path,
cfg: Config,
hook: H,
) -> TypedMap<[u8; 8], u64, U64Codec, H> {
TypedMap::open_hooked(dir, cfg, U64Codec, hook).expect("open hooked")
}
#[test]
fn typed_map_shard_update_and_fetch_update() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
let map: TypedMap<[u8; 8], u64, U64Codec> =
TypedMap::open(dir.path(), cfg, U64Codec).unwrap();
let key = 1u64.to_be_bytes();
map.put(&key, 10).unwrap();
let (new_v, old_v, missing) = map
.atomic(&key, |s| {
let new_v = s.update(&key, |old| old + 1)?.as_deref().copied();
let old_v = s.fetch_update(&key, |old| old + 1)?.as_deref().copied();
let missing = s.update(&2u64.to_be_bytes(), |old| *old)?.is_none();
Ok((new_v, old_v, missing))
})
.expect("atomic");
assert_eq!(new_v, Some(11));
assert_eq!(old_v, Some(11));
assert!(missing);
assert_eq!(map.get(&key).as_deref().copied(), Some(12));
}
#[test]
fn typed_map_atomic_update_fires_hook() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
let map = open_typed_map_hooked(dir.path(), cfg, TRecHook::default());
let k = 7u64.to_be_bytes();
map.put(&k, 100).unwrap();
map.hook.writes.store(0, AtomicOrdering::Relaxed);
crate::sync::lock(&map.hook.seq).clear();
map.atomic(&k, |s| {
s.update(&k, |old| old + 1)?;
Ok(())
})
.expect("atomic");
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 1);
let seq = crate::sync::lock(&map.hook.seq).clone();
assert_eq!(seq[0], (7, Some(100), Some(101)));
assert_eq!(map.get(&k).as_deref().copied(), Some(101));
}
#[test]
fn typed_map_atomic_fires_hooks_in_order() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
let map = open_typed_map_hooked(dir.path(), cfg, TRecHook::default());
let k = 5u64.to_be_bytes();
map.atomic(&k, |s| {
s.put(&k, 10)?; s.put(&k, 20)?; s.delete(&k)?; Ok(())
})
.expect("atomic");
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 3);
let seq = crate::sync::lock(&map.hook.seq).clone();
assert_eq!(seq[0], (5, None, Some(10)));
assert_eq!(seq[1], (5, Some(10), Some(20)));
assert_eq!(seq[2], (5, Some(20), None));
}
#[test]
fn typed_map_atomic_fires_for_applied_on_err() {
let dir = tempdir().unwrap();
let mut cfg = Config::test();
cfg.shard_count = 1;
let map = open_typed_map_hooked(dir.path(), cfg, TRecHook::default());
let k = 1u64.to_be_bytes();
let r: DbResult<()> = map.atomic(&k, |s| {
s.put(&k, 99)?;
Err(DbError::KeyNotFound)
});
assert!(r.is_err());
assert_eq!(map.hook.writes.load(AtomicOrdering::Relaxed), 1);
}
#[test]
fn typed_map_first_last_contract() {
let dir = tempdir().unwrap();
let cfg = Config::balanced()
.shard_count(2)
.hints(true)
.iterable(true)
.reversed(false)
.build();
let map = TypedMap::<[u8; 8], u64, U64Codec>::open(dir.path(), cfg, U64Codec).unwrap();
assert!(map.first().unwrap().is_none());
assert!(map.last().unwrap().is_none());
for id in [3_u64, 1, 2] {
map.put(&id.to_be_bytes(), id * 10).unwrap();
}
assert_eq!(
map.first()
.unwrap()
.map(|(k, v)| (u64::from_be_bytes(k), *v)),
Some((1, 10))
);
assert_eq!(
map.last()
.unwrap()
.map(|(k, v)| (u64::from_be_bytes(k), *v)),
Some((3, 30))
);
let no_index_dir = tempdir().unwrap();
let no_index =
TypedMap::<[u8; 8], u64, U64Codec>::open(no_index_dir.path(), Config::test(), U64Codec)
.unwrap();
assert!(matches!(no_index.first(), Err(DbError::NotImplemented)));
assert!(matches!(no_index.last(), Err(DbError::NotImplemented)));
}
#[cfg(feature = "rapira-codec")]
#[test]
fn typed_map_iter_no_clone_bound() {
#[derive(rapira::Rapira, PartialEq, Debug)]
struct NoClone {
n: u32,
s: String,
}
let dir = tempdir().unwrap();
let cfg = Config::balanced()
.shard_count(2)
.hints(true)
.iterable(true)
.build();
let m = TypedMap::<[u8; 2], NoClone, crate::RapiraCodec>::open(
dir.path(),
cfg,
crate::RapiraCodec,
)
.unwrap();
m.put(
&[1, 0],
NoClone {
n: 1,
s: "a".into(),
},
)
.unwrap();
m.put(
&[2, 0],
NoClone {
n: 2,
s: "b".into(),
},
)
.unwrap();
let v = m.iter_view().unwrap();
let ns: Vec<u32> = v.iter().map(|(_, r)| r.n).collect();
assert_eq!(ns, [2, 1]);
let keys: Vec<[u8; 2]> = v.keys().collect();
assert_eq!(keys, [[2, 0], [1, 0]]);
v.retain(|_k, val| val.n % 2 == 0).unwrap();
assert_eq!(m.iter_view().unwrap().iter().count(), 1);
}
#[test]
fn for_each_visits_all_entries() {
let dir = tempdir().unwrap();
let map = open(dir.path());
let key_a = 1u64.to_be_bytes();
let key_b = key_on_other_shard(&map, &key_a);
let key_c = key_on_other_shard(&map, &key_b);
let expected = StdHashMap::from([(key_a, 10u64), (key_b, 20u64), (key_c, 30u64)]);
for (key, value) in &expected {
map.put(key, *value).unwrap();
}
let mut seen = StdHashMap::new();
map.for_each(|key, value| {
seen.insert(*key, *value);
});
assert_eq!(seen, expected);
assert_ne!(map.shard_for(&key_a), map.shard_for(&key_b));
}
#[test]
fn for_each_while_stops_after_break_without_visiting_later_shards() {
use std::ops::ControlFlow;
use std::sync::atomic::{AtomicUsize, Ordering};
let dir = tempdir().unwrap();
let map = open(dir.path());
let key_a = 1u64.to_be_bytes();
let key_b = key_on_other_shard(&map, &key_a);
assert_ne!(map.shard_for(&key_a), map.shard_for(&key_b));
map.put(&key_a, 10).unwrap();
map.put(&key_b, 20).unwrap();
let seen = AtomicUsize::new(0);
let flow = map.for_each_while(|_k, _v| {
seen.fetch_add(1, Ordering::Relaxed);
ControlFlow::Break(())
});
assert_eq!(seen.load(Ordering::Relaxed), 1);
assert!(matches!(flow, ControlFlow::Break(())));
}
}