mod graph_index_merge;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Weak};
use super::entity::{EntityId, EntityKind, UnifiedEntity};
use super::metadata::{Metadata, MetadataFilter};
use super::segment::{
GraphEntityKind, GrowingSegment, SegmentConfig, SegmentError, SegmentId, SegmentState,
SegmentStats, UnifiedSegment, ZoneColPred, ZoneColPredKind,
};
use super::visibility_map::VisibilityMap;
use crate::runtime::mvcc::{entity_visible_with_context, SnapshotContext};
pub const CONSOLIDATION_TOMBSTONE_RATIO: f64 = 0.20;
pub const CONSOLIDATION_FRAGMENTATION_RATIO: f64 = 0.30;
pub const CONSOLIDATION_ENTITIES_PER_TICK: usize = 4_096;
#[derive(Debug, Clone)]
pub struct ManagerConfig {
pub segment_config: SegmentConfig,
pub max_sealed_segments: usize,
pub idle_seal_secs: u64,
pub enable_consolidation: bool,
pub consolidation_entities_per_tick: usize,
pub enable_archival: bool,
pub archive_age_secs: u64,
}
impl Default for ManagerConfig {
fn default() -> Self {
Self {
segment_config: SegmentConfig::default(),
max_sealed_segments: 10,
idle_seal_secs: 300, enable_consolidation: true,
consolidation_entities_per_tick: CONSOLIDATION_ENTITIES_PER_TICK,
enable_archival: true,
archive_age_secs: 86400 * 7, }
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ConsolidationStats {
pub runs_started: u64,
pub runs_completed: u64,
pub segments_merged: u64,
pub tombstones_reclaimed: u64,
pub bytes_reclaimed: u64,
}
#[derive(Debug, Clone, Default)]
pub struct ManagerStats {
pub total_entities: usize,
pub growing_count: usize,
pub sealed_count: usize,
pub archived_count: usize,
pub total_memory_bytes: usize,
pub seal_ops: u64,
pub consolidation: ConsolidationStats,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SegmentScanStats {
pub segments_total: u64,
pub segments_scanned: u64,
pub segments_pruned: u64,
}
#[derive(Debug, Clone)]
pub enum LifecycleEvent {
SegmentCreated(SegmentId),
SegmentSealed(SegmentId),
SegmentConsolidated {
source: Vec<SegmentId>,
target: SegmentId,
},
SegmentArchived(SegmentId),
EntityInserted(EntityId, SegmentId),
EntityDeleted(EntityId, SegmentId),
}
struct Consolidation {
sources: Vec<SourceProgress>,
merged_id: SegmentId,
merged: GrowingSegment,
cursor: usize,
pending_ids: Vec<EntityId>,
pending_cursor: usize,
entities_copied: u64,
}
struct SourceProgress {
id: SegmentId,
copied: Vec<EntityId>,
tombstones_at_open: Option<usize>,
}
struct PhysicalRootCacheEntry {
segments: Vec<(u64, u64)>,
seed: u64,
root: u64,
}
pub struct SegmentManager {
row_constraint_lock: Arc<parking_lot::ReentrantMutex<()>>,
collection: String,
config: ManagerConfig,
next_segment_id: AtomicU64,
next_entity_id: AtomicU64,
next_row_id: AtomicU64,
total_entities_atomic: AtomicU64,
growing: RwLock<Option<Arc<RwLock<GrowingSegment>>>>,
sealed: RwLock<Vec<Arc<RwLock<GrowingSegment>>>>,
retired: RwLock<Vec<Weak<RwLock<GrowingSegment>>>>,
archived: RwLock<Vec<SegmentId>>,
entity_segment: RwLock<HashMap<EntityId, SegmentId>>,
column_schema: RwLock<Option<Arc<Vec<String>>>>,
stats: RwLock<ManagerStats>,
consolidation: RwLock<Option<Consolidation>>,
events: RwLock<Vec<LifecycleEvent>>,
visibility_map: VisibilityMap,
physical_root_cache: parking_lot::Mutex<Option<PhysicalRootCacheEntry>>,
}
impl SegmentManager {
pub fn new(collection: impl Into<String>) -> Self {
Self::with_config(collection, ManagerConfig::default())
}
pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
Self {
row_constraint_lock: Arc::new(parking_lot::ReentrantMutex::new(())),
physical_root_cache: parking_lot::Mutex::new(None),
collection: collection.into(),
config,
next_segment_id: AtomicU64::new(1),
next_entity_id: AtomicU64::new(1),
next_row_id: AtomicU64::new(1),
total_entities_atomic: AtomicU64::new(0),
growing: RwLock::new(None),
sealed: RwLock::new(Vec::new()),
retired: RwLock::new(Vec::new()),
archived: RwLock::new(Vec::new()),
entity_segment: RwLock::new(HashMap::new()),
column_schema: RwLock::new(None),
stats: RwLock::new(ManagerStats::default()),
consolidation: RwLock::new(None),
events: RwLock::new(Vec::new()),
visibility_map: VisibilityMap::new(),
}
}
pub(crate) fn row_constraint_lock(&self) -> Arc<parking_lot::ReentrantMutex<()>> {
Arc::clone(&self.row_constraint_lock)
}
pub fn get_or_init_schema(
&self,
named: &HashMap<String, reddb_types::Value>,
) -> Arc<Vec<String>> {
{
let schema = self.column_schema.read();
if let Some(ref s) = *schema {
return Arc::clone(s);
}
}
let cols: Vec<String> = named.keys().cloned().collect();
let arc = Arc::new(cols);
*self.column_schema.write() = Some(Arc::clone(&arc));
arc
}
pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
self.column_schema.read().clone()
}
pub(crate) fn set_column_schema_if_empty(&self, columns: Vec<String>) {
if columns.is_empty() {
return;
}
let mut schema = self.column_schema.write();
if schema.is_none() {
*schema = Some(Arc::new(columns));
}
}
pub fn collection(&self) -> &str {
&self.collection
}
pub fn config(&self) -> &ManagerConfig {
&self.config
}
pub fn stats(&self) -> ManagerStats {
let mut s = self.stats.read().clone();
s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
s.total_memory_bytes = self.resident_bytes() as usize;
s
}
pub fn resident_bytes(&self) -> u64 {
self.memory_usage().1
}
pub fn memory_bytes(&self) -> u64 {
self.memory_usage().0
}
fn memory_usage(&self) -> (u64, u64) {
let consolidation = self.consolidation.read();
let growing = self.growing.read();
let sealed = self.sealed.read();
let mut retired = self.retired.write();
let mut payload = 0u64;
let mut resident = 0u64;
{
let mut add = |segment: &GrowingSegment| {
payload = payload.saturating_add(segment.memory_bytes());
resident = resident.saturating_add(segment.resident_bytes());
};
for segment in growing.iter().chain(sealed.iter()) {
add(&segment.read());
}
retired.retain(|source| {
if let Some(source) = source.upgrade() {
add(&source.read());
true
} else {
false
}
});
if let Some(run) = consolidation.as_ref() {
add(&run.merged);
}
}
let mut bookkeeping =
(retired.capacity() * std::mem::size_of::<Weak<RwLock<GrowingSegment>>>()) as u64;
if let Some(run) = consolidation.as_ref() {
bookkeeping = bookkeeping
.saturating_add(
(run.sources.capacity() * std::mem::size_of::<SourceProgress>()) as u64,
)
.saturating_add(
(run.pending_ids.capacity() * std::mem::size_of::<EntityId>()) as u64,
);
for source in &run.sources {
bookkeeping = bookkeeping.saturating_add(
(source.copied.capacity() * std::mem::size_of::<EntityId>()) as u64,
);
}
}
(
payload.saturating_add(bookkeeping),
resident.saturating_add(bookkeeping),
)
}
pub fn reclaimable_bytes(&self) -> u64 {
self.sealed
.read()
.iter()
.map(|segment| segment.read().reclaimable_bytes())
.fold(0, u64::saturating_add)
}
pub fn next_entity_id(&self) -> EntityId {
EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
}
pub fn next_row_id(&self) -> u64 {
self.next_row_id.fetch_add(1, Ordering::SeqCst)
}
pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
start..start + n
}
pub fn register_row_id(&self, id: u64) {
let candidate = id.saturating_add(1);
let mut current = self.next_row_id.load(Ordering::SeqCst);
while candidate > current {
match self.next_row_id.compare_exchange(
current,
candidate,
Ordering::SeqCst,
Ordering::SeqCst,
) {
Ok(_) => break,
Err(updated) => current = updated,
}
}
}
fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
{
let growing = self.growing.read();
if let Some(segment) = growing.as_ref() {
return Arc::clone(segment);
}
}
let mut growing = self.growing.write();
if let Some(segment) = growing.as_ref() {
return Arc::clone(segment);
}
let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
let segment = GrowingSegment::new(id, &self.collection);
let segment_arc = Arc::new(RwLock::new(segment));
*growing = Some(Arc::clone(&segment_arc));
self.emit(LifecycleEvent::SegmentCreated(id));
self.stats.write().growing_count += 1;
segment_arc
}
pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
self.maybe_seal_growing()?;
let segment_arc = self.get_or_create_growing();
let mut segment = segment_arc.write();
if entity.id.raw() == 0 {
entity.id = self.next_entity_id();
}
let entity_id = entity.id;
let segment_id = segment.id();
segment.insert(entity)?;
self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
Ok(entity_id)
}
pub(crate) fn restore_snapshot_entity(
&self,
entity: UnifiedEntity,
) -> Result<EntityId, SegmentError> {
self.maybe_seal_growing()?;
if self.get(entity.id).is_some() {
return Err(SegmentError::AlreadyExists(entity.id));
}
let segment_arc = self.get_or_create_growing();
let mut segment = segment_arc.write();
let entity_id = entity.id;
let segment_id = segment.id();
segment.adopt_entity(entity, None);
self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
Ok(entity_id)
}
pub fn insert_batch(
&self,
entities: Vec<UnifiedEntity>,
) -> Result<Vec<EntityId>, SegmentError> {
let mut ids = Vec::with_capacity(entities.len());
for entity in entities {
ids.push(self.insert(entity)?);
}
Ok(ids)
}
pub fn bulk_insert(
&self,
mut entities: Vec<UnifiedEntity>,
) -> Result<Vec<EntityId>, SegmentError> {
for entity in &mut entities {
if entity.id.raw() == 0 {
entity.id = self.next_entity_id();
}
if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
if *row_id == 0 {
*row_id = self.next_row_id();
} else {
self.register_row_id(*row_id);
}
}
}
if let Some(first_row) = entities.first() {
if let super::entity::EntityData::Row(ref row) = first_row.data {
if let Some(ref named) = row.named {
let schema = self.get_or_init_schema(named);
for entity in &mut entities {
if let super::entity::EntityData::Row(ref mut row) = entity.data {
if let Some(named) = row.named.take() {
let mut cols = Vec::with_capacity(schema.len());
for col_name in schema.iter() {
cols.push(
named
.get(col_name)
.cloned()
.unwrap_or(reddb_types::Value::Null),
);
}
row.columns = cols;
row.schema = Some(Arc::clone(&schema));
}
}
}
}
}
}
let segment_arc = self.get_or_create_growing();
let mut segment = segment_arc.write();
let segment_id = segment.id();
let ids = segment.bulk_insert(entities)?;
self.total_entities_atomic
.fetch_add(ids.len() as u64, Ordering::Relaxed);
Ok(ids)
}
pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
self.get_with(id, Clone::clone)
}
pub(crate) fn get_with<T>(
&self,
id: EntityId,
visit: impl FnOnce(&UnifiedEntity) -> T,
) -> Option<T> {
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
if let Some(entity) = growing.get(id) {
return Some(visit(entity));
}
}
let sealed = self.sealed.read();
for segment in sealed.iter() {
let seg = segment.read();
if let Some(entity) = seg.get(id) {
return Some(visit(entity));
}
}
None
}
pub(crate) fn graph_node_logical_id(
&self,
id: EntityId,
mut before_segment: impl FnMut() -> bool,
) -> Option<EntityId> {
if let Some(growing_arc) = self.growing.read().as_ref() {
if !before_segment() {
return None;
}
let growing = growing_arc.read();
if let Some(entity) = growing.get(id) {
return matches!(entity.kind, EntityKind::GraphNode(_))
.then(|| entity.logical_id());
}
}
let sealed = self.sealed.read();
for segment in sealed.iter() {
if !before_segment() {
return None;
}
let segment = segment.read();
if let Some(entity) = segment.get(id) {
return matches!(entity.kind, EntityKind::GraphNode(_))
.then(|| entity.logical_id());
}
}
None
}
pub(crate) fn visit_graph_index_batches<E>(
&self,
kind: GraphEntityKind,
keys: &[EntityId],
mut before_work: impl FnMut() -> bool,
mut admit: impl FnMut(usize) -> Result<(), E>,
mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
) -> Result<bool, E> {
use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
let mut admitted = 0;
let sources = loop {
let count = {
let growing = self.growing.read();
let sealed = self.sealed.read();
sealed.len() + usize::from(growing.is_some())
};
let bytes = count.saturating_mul(std::mem::size_of::<(
Arc<RwLock<GrowingSegment>>,
GraphIndexCursor,
)>());
admit(bytes.saturating_sub(admitted))?;
admitted = admitted.max(bytes);
let mut sources = Vec::with_capacity(count);
let growing = self.growing.read();
let sealed = self.sealed.read();
if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
continue;
}
for segment in growing.iter().chain(sealed.iter()) {
if !before_work() {
return Ok(false);
}
let view = segment.read();
if !view.may_contain_graph_kind(kind) {
continue;
}
let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
return Ok(false);
};
sources.push((Arc::clone(segment), cursor));
}
break sources;
};
let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
for (source, mut cursor) in sources {
while !cursor.finished(keys) {
let mut count = 0;
{
let segment = source.read();
if !segment.may_contain_graph_kind(kind) {
break;
}
if !segment.graph_index_batch(
kind,
keys,
&mut cursor,
&mut before_work,
&mut |id| {
ids[count] = id;
count += 1;
},
) {
return Ok(false);
}
}
if !before_work() || !consume(&ids[..count])? {
return Ok(false);
}
}
}
Ok(true)
}
pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
let schema = self.column_schema();
let clone_row = |entity: &UnifiedEntity| {
let mut owned = entity.clone();
if let super::entity::EntityData::Row(row) = &mut owned.data {
if row.schema.is_none() {
row.schema = schema.clone();
}
}
owned
};
let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
let mut remaining: Vec<usize> = (0..ids.len()).collect();
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = if let Some(g) = growing_arc.try_read() {
g
} else {
growing_arc.read()
};
remaining.retain(|&i| {
if let Some(entity) = growing.get(ids[i]) {
out[i] = Some(clone_row(entity));
false } else {
true }
});
}
if remaining.is_empty() {
return out;
}
let sealed = self.sealed.read();
for segment in sealed.iter() {
if remaining.is_empty() {
break;
}
let seg = segment.read();
remaining.retain(|&i| {
if let Some(entity) = seg.get(ids[i]) {
out[i] = Some(clone_row(entity));
false
} else {
true
}
});
}
out
}
pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
where
F: FnMut(usize, &UnifiedEntity),
{
thread_local! {
static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
const { std::cell::Cell::new(Vec::new()) };
}
let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
remaining.clear();
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = if let Some(g) = growing_arc.try_read() {
g
} else {
growing_arc.read()
};
for (i, id) in ids.iter().enumerate() {
if let Some(entity) = growing.get(*id) {
f(i, entity);
} else {
remaining.push(i);
}
}
} else {
remaining.reserve(ids.len());
remaining.extend(0..ids.len());
}
if !remaining.is_empty() {
let sealed = self.sealed.read();
for segment in sealed.iter() {
if remaining.is_empty() {
break;
}
let seg = segment.read();
remaining.retain(|&i| {
if let Some(entity) = seg.get(ids[i]) {
f(i, entity);
false
} else {
true
}
});
}
}
REMAINING_SCRATCH.with(|cell| cell.set(remaining));
}
fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
if let Some(entity) = growing.get(id) {
return Some(entity.clone());
}
}
let sealed = self.sealed.read();
for segment in sealed.iter() {
if let Some(entity) = segment.get(id) {
return Some(entity.clone());
}
}
None
}
fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
let sealed = self.sealed.read();
sealed
.iter()
.find(|segment_arc| segment_arc.read().contains(id))
.map(Arc::clone)
}
fn rewrite_sealed_entity_into_growing(
&self,
entity: UnifiedEntity,
metadata: Option<&Metadata>,
) -> Result<(), SegmentError> {
let entity_id = entity.id;
let sealed_arc = self
.find_sealed_segment_arc(entity_id)
.ok_or(SegmentError::NotFound(entity_id))?;
let metadata_to_apply = {
let mut sealed = sealed_arc.write();
let existing_metadata = sealed.get_metadata(entity_id);
if !sealed.force_delete(entity_id) {
return Err(SegmentError::NotFound(entity_id));
}
metadata.cloned().or(existing_metadata)
};
let growing_arc = self.get_or_create_growing();
let growing_id = {
let mut growing = growing_arc.write();
growing.insert(entity)?;
if let Some(metadata) = metadata_to_apply {
growing.set_metadata(entity_id, metadata)?;
}
growing.id()
};
self.entity_segment.write().insert(entity_id, growing_id);
Ok(())
}
pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
let entity_id = entity.id;
let mut entity = Some(entity);
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.contains(entity_id) && growing.state().is_writable() {
return growing.update(entity.take().expect("entity already moved"));
}
}
let segment_id = self.entity_segment.read().get(&entity_id).copied();
if let Some(seg_id) = segment_id {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.id() == seg_id && growing.state().is_writable() {
return growing.update(entity.take().expect("entity already moved"));
}
}
}
if let Some(entity) = entity.take() {
return self.rewrite_sealed_entity_into_growing(entity, None);
}
Err(SegmentError::NotFound(entity_id))
}
pub fn update_with_metadata(
&self,
entity: UnifiedEntity,
metadata: Option<&Metadata>,
) -> Result<(), SegmentError> {
let entity_id = entity.id;
let mut entity = Some(entity);
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.contains(entity_id) && growing.state().is_writable() {
growing.update(entity.take().expect("entity already moved"))?;
if let Some(metadata) = metadata {
growing.set_metadata(entity_id, metadata.clone())?;
}
return Ok(());
}
}
let segment_id = self.entity_segment.read().get(&entity_id).copied();
if let Some(seg_id) = segment_id {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.id() == seg_id && growing.state().is_writable() {
growing.update(entity.take().expect("entity already moved"))?;
if let Some(metadata) = metadata {
growing.set_metadata(entity_id, metadata.clone())?;
}
return Ok(());
}
}
}
if let Some(entity) = entity.take() {
return self.rewrite_sealed_entity_into_growing(entity, metadata);
}
Err(SegmentError::NotFound(entity_id))
}
pub fn update_hot(
&self,
entity: UnifiedEntity,
modified_columns: &[String],
) -> Result<(), SegmentError> {
let entity_id = entity.id;
let mut entity = Some(entity);
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.contains(entity_id) && growing.state().is_writable() {
return growing.update_hot(
entity.take().expect("entity already moved"),
modified_columns,
);
}
}
let segment_id = self.entity_segment.read().get(&entity_id).copied();
if let Some(seg_id) = segment_id {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.id() == seg_id && growing.state().is_writable() {
return growing.update_hot(
entity.take().expect("entity already moved"),
modified_columns,
);
}
}
}
if let Some(entity) = entity.take() {
return self.rewrite_sealed_entity_into_growing(entity, None);
}
Err(SegmentError::NotFound(entity_id))
}
pub fn update_hot_with_metadata(
&self,
entity: UnifiedEntity,
modified_columns: &[String],
metadata: Option<&Metadata>,
) -> Result<(), SegmentError> {
let entity_id = entity.id;
let mut entity = Some(entity);
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.contains(entity_id) && growing.state().is_writable() {
growing.update_hot(
entity.take().expect("entity already moved"),
modified_columns,
)?;
if let Some(metadata) = metadata {
growing.set_metadata(entity_id, metadata.clone())?;
}
return Ok(());
}
}
let segment_id = self.entity_segment.read().get(&entity_id).copied();
if let Some(seg_id) = segment_id {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.id() == seg_id && growing.state().is_writable() {
growing.update_hot(
entity.take().expect("entity already moved"),
modified_columns,
)?;
if let Some(metadata) = metadata {
growing.set_metadata(entity_id, metadata.clone())?;
}
return Ok(());
}
}
}
if let Some(entity) = entity.take() {
return self.rewrite_sealed_entity_into_growing(entity, metadata);
}
Err(SegmentError::NotFound(entity_id))
}
pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
where
I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
{
let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
items.into_iter().collect();
if items.is_empty() {
return Ok(());
}
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.state().is_writable() {
match growing.update_hot_batch_with_metadata(items.iter().copied()) {
Ok(()) => return Ok(()),
Err(SegmentError::NotFound(_)) => {}
Err(other) => return Err(other),
}
}
}
for (entity, modified_columns, metadata) in items {
self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
}
Ok(())
}
pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.contains(id) && growing.state().is_writable() {
let seg_id = growing.id();
let deleted = growing.delete(id)?;
if deleted {
self.entity_segment.write().remove(&id);
self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
}
return Ok(deleted);
}
}
let segment_id = self.entity_segment.read().get(&id).copied();
if let Some(seg_id) = segment_id {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.id() == seg_id && growing.state().is_writable() {
let deleted = growing.delete(id)?;
if deleted {
self.entity_segment.write().remove(&id);
self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
}
return Ok(deleted);
}
}
}
{
let sealed = self.sealed.read();
for segment_arc in sealed.iter() {
let mut seg = segment_arc.write();
let seg_id = seg.id();
if seg.contains(id) {
let deleted = seg.force_delete(id);
drop(seg);
if deleted {
self.entity_segment.write().remove(&id);
self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
}
return Ok(deleted);
}
}
}
Ok(false)
}
pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
if ids.is_empty() {
return Ok(Vec::new());
}
let mut deleted_ids = Vec::with_capacity(ids.len());
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.state().is_writable() {
let seg_id = growing.id();
let deleted = growing.delete_batch(ids)?;
if !deleted.is_empty() {
{
let mut entity_segment = self.entity_segment.write();
for id in &deleted {
entity_segment.remove(id);
}
}
self.total_entities_atomic
.fetch_sub(deleted.len() as u64, Ordering::Relaxed);
for id in &deleted {
self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
}
deleted_ids.extend(deleted);
}
}
}
if deleted_ids.len() == ids.len() {
return Ok(deleted_ids);
}
let deleted_set: std::collections::HashSet<EntityId> =
deleted_ids.iter().copied().collect();
for &id in ids {
if deleted_set.contains(&id) {
continue;
}
if self.delete(id)? {
deleted_ids.push(id);
}
}
Ok(deleted_ids)
}
pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
if growing.contains(id) {
return growing.get_metadata(id);
}
}
let segment_id = self.entity_segment.read().get(&id).copied();
if let Some(seg_id) = segment_id {
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
if growing.id() == seg_id {
return growing.get_metadata(id);
}
}
let sealed = self.sealed.read();
for segment in sealed.iter() {
if segment.id() == seg_id {
return segment.get_metadata(id);
}
}
}
if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
return segment_arc.read().get_metadata(id);
}
None
}
pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.contains(id) && growing.state().is_writable() {
return growing.set_metadata(id, metadata);
}
}
let segment_id = self.entity_segment.read().get(&id).copied();
if let Some(seg_id) = segment_id {
if let Some(growing_arc) = self.growing.read().as_ref() {
let mut growing = growing_arc.write();
if growing.id() == seg_id && growing.state().is_writable() {
return growing.set_metadata(id, metadata);
}
}
}
if let Some(entity) = self.get(id) {
return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
}
Err(SegmentError::NotFound(id))
}
fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
let should_seal = {
let growing_opt = self.growing.read();
if let Some(growing_arc) = growing_opt.as_ref() {
let growing = growing_arc.read();
growing.should_seal(&self.config.segment_config)
|| growing.idle_secs() >= self.config.idle_seal_secs
} else {
false
}
};
if should_seal {
self.seal_current()?;
}
Ok(())
}
pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
let mut growing_slot = self.growing.write();
if let Some(growing_arc) = growing_slot.take() {
let mut growing = growing_arc.write();
let seg_id = growing.id();
let entity_count = growing.stats().entity_count as u64;
growing.seal()?;
drop(growing);
self.sealed.write().push(growing_arc);
drop(growing_slot);
self.mark_sealed_pages_visible(entity_count);
{
let mut stats = self.stats.write();
stats.growing_count = stats.growing_count.saturating_sub(1);
stats.sealed_count += 1;
stats.seal_ops += 1;
}
self.emit(LifecycleEvent::SegmentSealed(seg_id));
return Ok(seg_id);
}
Err(SegmentError::InvalidState(SegmentState::Sealed))
}
pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
let has_growing = self.growing.read().is_some();
if has_growing {
Ok(Some(self.seal_current()?))
} else {
Ok(None)
}
}
pub fn all_visible_fraction(&self) -> f64 {
const ROWS_PER_PAGE: u32 = 256;
let sealed = self.sealed.read();
if sealed.is_empty() {
return 0.0;
}
let mut total_pages: u64 = 0;
for seg_arc in sealed.iter() {
let seg = seg_arc.read();
let entity_count = seg.stats().entity_count as u64;
let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
total_pages += pages;
}
if total_pages == 0 {
return 0.0;
}
let visible = self.visibility_map.all_visible_count();
(visible as f64 / total_pages as f64).min(1.0)
}
fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
const ROWS_PER_PAGE: u32 = 256;
let existing_visible = self.visibility_map.all_visible_count();
let start_page = existing_visible as u32;
let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
let end_page = start_page + new_pages as u32;
self.visibility_map.mark_range_visible(start_page, end_page);
}
pub(crate) fn find_unique_key_conflict(
&self,
columns: &[Vec<String>],
key_index: usize,
signatures: &[String],
exclude_id: Option<EntityId>,
reserves_key: &impl Fn(&UnifiedEntity) -> bool,
) -> Option<EntityId> {
if let Some(growing) = self.growing.read().as_ref() {
if let Some(id) = growing.read().find_unique_key_conflict(
columns,
key_index,
signatures,
exclude_id,
reserves_key,
) {
return Some(id);
}
}
for segment in self.sealed.read().iter() {
if let Some(id) = segment.read().find_unique_key_conflict(
columns,
key_index,
signatures,
exclude_id,
reserves_key,
) {
return Some(id);
}
}
None
}
pub fn for_each_entity<F>(&self, callback: F)
where
F: FnMut(&UnifiedEntity) -> bool,
{
self.for_each_entity_in_segments(|_| true, || true, callback);
}
pub(crate) fn checkpoint_revisions(&self) -> Vec<(u64, u64)> {
let growing = self.growing.read();
let sealed = self.sealed.read();
growing
.iter()
.chain(sealed.iter())
.map(|segment| {
let segment = segment.read();
(segment.id(), segment.checkpoint_revision())
})
.collect()
}
pub(crate) fn cached_physical_root(&self, seed: u64, compute: impl FnOnce() -> u64) -> u64 {
let segments = self.checkpoint_revisions();
let mut cache = self.physical_root_cache.lock();
if let Some(previous) = cache.as_ref() {
if previous.seed == seed && previous.segments == segments {
return previous.root;
}
}
let root = compute();
*cache = Some(PhysicalRootCacheEntry {
segments,
seed,
root,
});
root
}
pub(crate) fn for_each_entity_in_id_order(
&self,
mut visit: impl FnMut(&UnifiedEntity) -> bool,
) {
self.for_each_entity_source_in_id_order(|_, entity| visit(entity));
}
pub(crate) fn for_each_entity_with_metadata_in_id_order(
&self,
mut visit: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
) {
self.for_each_entity_source_in_id_order(|segment, entity| {
visit(entity, segment.get_metadata(entity.id))
});
}
fn for_each_entity_source_in_id_order(
&self,
mut visit: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
) {
let segments: Vec<_> = {
let growing = self.growing.read();
let sealed = self.sealed.read();
growing.iter().chain(sealed.iter()).cloned().collect()
};
let mut order = Vec::new();
for (source, segment) in segments.iter().enumerate() {
segment.read().for_each_fast(|entity| {
order.push((entity.id, source));
true
});
}
order.sort_by_key(|(id, _)| id.raw());
for (id, source) in order {
let segment = segments[source].read();
if let Some(entity) = segment.get(id) {
if !visit(&segment, entity) {
break;
}
}
}
}
pub(crate) fn for_each_entity_with_metadata(
&self,
mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
) {
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
if !growing.for_each_fast(|entity| callback(entity, growing.get_metadata(entity.id))) {
return;
}
}
let sealed = self.sealed.read();
for segment_arc in sealed.iter() {
let segment = segment_arc.read();
if !segment.for_each_fast(|entity| callback(entity, segment.get_metadata(entity.id))) {
return;
}
}
}
fn for_each_entity_in_segments<F>(
&self,
segment_filter: impl Fn(&GrowingSegment) -> bool,
mut before_segment: impl FnMut() -> bool,
mut callback: F,
) where
F: FnMut(&UnifiedEntity) -> bool,
{
if let Some(growing_arc) = self.growing.read().as_ref() {
if !before_segment() {
return;
}
let growing = if let Some(g) = growing_arc.try_read() {
g
} else {
growing_arc.read()
};
if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
return;
}
}
let sealed = self.sealed.read();
for segment_arc in sealed.iter() {
if !before_segment() {
return;
}
let segment = segment_arc.read();
if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
return;
}
}
}
pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
&self,
init: FInit,
fold: FFold,
reduce: FReduce,
) -> T
where
T: Send,
FInit: Fn() -> T + Send + Sync,
FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
FReduce: Fn(T, T) -> T + Send + Sync,
{
use rayon::prelude::*;
let mut acc = init();
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = if let Some(g) = growing_arc.try_read() {
g
} else {
growing_arc.read()
};
growing.for_each_fast(|entity| {
acc = fold(std::mem::replace(&mut acc, init()), entity);
true
});
}
let segments: Vec<_> = {
let sealed = self.sealed.read();
sealed.iter().cloned().collect()
};
if segments.len() <= 1 {
for seg_arc in &segments {
let seg = seg_arc.read();
seg.for_each_fast(|entity| {
acc = fold(std::mem::replace(&mut acc, init()), entity);
true
});
}
return acc;
}
let sealed_acc = segments
.into_par_iter()
.map(|seg_arc| {
let mut local = init();
let seg = seg_arc.read();
seg.for_each_fast(|entity| {
local = fold(std::mem::replace(&mut local, init()), entity);
true
});
local
})
.reduce(&init, &reduce);
reduce(acc, sealed_acc)
}
pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
where
F: FnMut(&UnifiedEntity) -> bool,
{
let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
}
pub fn for_each_entity_zoned_with_stats<F>(
&self,
zone_preds: &[(&str, ZoneColPred<'_>)],
mut callback: F,
) -> SegmentScanStats
where
F: FnMut(&UnifiedEntity) -> bool,
{
let mut scan_stats = SegmentScanStats::default();
if let Some(growing_arc) = self.growing.read().as_ref() {
scan_stats.segments_total += 1;
scan_stats.segments_scanned += 1;
let growing = if let Some(g) = growing_arc.try_read() {
g
} else {
growing_arc.read()
};
if !growing.for_each_fast(&mut callback) {
return scan_stats;
}
}
let sealed = self.sealed.read();
for segment_arc in sealed.iter() {
scan_stats.segments_total += 1;
let segment = segment_arc.read();
if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
scan_stats.segments_pruned += 1;
continue; }
scan_stats.segments_scanned += 1;
if !segment.for_each_fast(&mut callback) {
return scan_stats;
}
}
scan_stats
}
pub fn query_all_zoned<F>(
&self,
zone_preds: &[(&str, ZoneColPred<'_>)],
filter: F,
) -> Vec<UnifiedEntity>
where
F: Fn(&UnifiedEntity) -> bool + Sync,
{
self.query_all_zoned_with_stats(zone_preds, filter).0
}
pub fn query_all_zoned_with_stats<F>(
&self,
zone_preds: &[(&str, ZoneColPred<'_>)],
filter: F,
) -> (Vec<UnifiedEntity>, SegmentScanStats)
where
F: Fn(&UnifiedEntity) -> bool + Sync,
{
let mut results = Vec::new();
let mut scan_stats = SegmentScanStats::default();
if let Some(growing_arc) = self.growing.read().as_ref() {
scan_stats.segments_total += 1;
scan_stats.segments_scanned += 1;
let growing = if let Some(g) = growing_arc.try_read() {
g
} else {
growing_arc.read()
};
results.extend(growing.iter().filter(|e| filter(e)).cloned());
}
let sealed = self.sealed.read();
let surviving: Vec<_> = sealed
.iter()
.filter(|seg_arc| {
scan_stats.segments_total += 1;
if zone_preds.is_empty() {
scan_stats.segments_scanned += 1;
return true;
}
let seg = seg_arc.read();
if seg.can_skip_zone_preds(zone_preds) {
scan_stats.segments_pruned += 1;
false
} else {
scan_stats.segments_scanned += 1;
true
}
})
.collect();
let use_parallel = surviving.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
if use_parallel {
let filter_ref = &filter;
let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
surviving
.iter()
.map(|segment| {
s.spawn(move || {
segment
.read()
.iter()
.filter(|e| filter_ref(e))
.cloned()
.collect::<Vec<_>>()
})
})
.collect::<Vec<_>>()
.into_iter()
.map(|handle| handle.join().unwrap_or_default())
.collect()
});
for batch in segment_results {
results.extend(batch);
}
} else {
for segment_arc in surviving {
let seg = segment_arc.read();
results.extend(seg.iter().filter(|e| filter(e)).cloned());
}
}
(results, scan_stats)
}
fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
match snapshot {
Some(ctx) => entity_visible_with_context(Some(ctx), entity),
None => {
!crate::runtime::ai::moderation::entity_moderation_hidden(entity)
&& entity.xmax == 0
}
}
}
pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
where
F: Fn(&UnifiedEntity) -> bool + Sync,
{
self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
}
pub(crate) fn project_key_entity<T>(
&self,
key: &str,
snapshot: Option<&SnapshotContext>,
versioned: bool,
mut project: impl FnMut(&UnifiedEntity) -> T,
) -> Option<T> {
let mut probe = |segment: &GrowingSegment, best: &mut Option<(u64, T)>| {
if let Some(entity) = segment.find_key_entity(key, versioned, &mut |entity| {
!versioned || Self::scan_entity_visible(snapshot, entity)
}) {
if best.as_ref().is_none_or(|(xmin, _)| entity.xmin >= *xmin) {
*best = Some((entity.xmin, project(entity)));
}
}
};
let mut best = None;
if let Some(growing) = self.growing.read().as_ref() {
probe(&growing.read(), &mut best);
}
for segment in self.sealed.read().iter() {
if !versioned && best.is_some() {
break;
}
probe(&segment.read(), &mut best);
}
best.map(|(_, projected)| projected)
}
pub(crate) fn scan_graph_kind(
&self,
snapshot: Option<&SnapshotContext>,
kind: GraphEntityKind,
) -> Vec<UnifiedEntity> {
self.query_all_in_segments(
|segment| segment.may_contain_graph_kind(kind),
|entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
)
}
pub(crate) fn scan_graph_candidates_for_each<F>(
&self,
snapshot: Option<&SnapshotContext>,
kind: GraphEntityKind,
before_segment: impl FnMut() -> bool,
mut callback: F,
) where
F: FnMut(&UnifiedEntity) -> bool,
{
self.for_each_entity_in_segments(
|segment| segment.may_contain_graph_kind(kind),
before_segment,
|entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
);
}
pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
where
F: FnMut(&UnifiedEntity) -> bool,
{
self.for_each_entity(|entity| {
!Self::scan_entity_visible(snapshot, entity) || callback(entity)
});
}
pub(crate) fn scan_batches(
&self,
snapshot: Option<&SnapshotContext>,
mut before_work: impl FnMut() -> bool,
mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
) -> bool {
let segments = {
let growing = self.growing.read();
let sealed = self.sealed.read();
let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
for segment in growing.iter().chain(sealed.iter()) {
if !before_work() {
return false;
}
let cursor = segment.read().scan_cursor();
segments.push((Arc::clone(segment), cursor));
}
segments
};
let mut ids = Vec::new();
let mut batch = Vec::new();
for (segment, mut cursor) in segments {
while !cursor.finished() {
ids.clear();
batch.clear();
let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
ids.reserve_exact(batch_size);
batch.reserve_exact(batch_size);
{
let segment = segment.read();
if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
return false;
}
}
if !before_work() {
return false;
}
self.for_each_id(&ids, |_, entity| {
if Self::scan_entity_visible(snapshot, entity) {
batch.push(entity.clone());
}
});
if !before_work() {
return false;
}
if !consume(&batch) {
return false;
}
}
}
true
}
pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
&self,
snapshot: Option<&SnapshotContext>,
init: FInit,
fold: FFold,
reduce: FReduce,
) -> T
where
T: Send,
FInit: Fn() -> T + Send + Sync,
FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
FReduce: Fn(T, T) -> T + Send + Sync,
{
self.fold_entities_parallel(
init,
|acc, entity| {
if Self::scan_entity_visible(snapshot, entity) {
fold(acc, entity)
} else {
acc
}
},
reduce,
)
}
pub fn scan_for_each_zoned_with_stats<F>(
&self,
snapshot: Option<&SnapshotContext>,
zone_preds: &[(&str, ZoneColPred<'_>)],
mut callback: F,
) -> SegmentScanStats
where
F: FnMut(&UnifiedEntity) -> bool,
{
self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
!Self::scan_entity_visible(snapshot, entity) || callback(entity)
})
}
pub fn scan_zoned_with_stats<F>(
&self,
snapshot: Option<&SnapshotContext>,
zone_preds: &[(&str, ZoneColPred<'_>)],
filter: F,
) -> (Vec<UnifiedEntity>, SegmentScanStats)
where
F: Fn(&UnifiedEntity) -> bool + Sync,
{
self.query_all_zoned_with_stats(zone_preds, |entity| {
Self::scan_entity_visible(snapshot, entity) && filter(entity)
})
}
pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
where
F: Fn(&UnifiedEntity) -> bool + Sync,
{
self.query_all_in_segments(|_| true, filter)
}
fn query_all_in_segments<F>(
&self,
segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
filter: F,
) -> Vec<UnifiedEntity>
where
F: Fn(&UnifiedEntity) -> bool + Sync,
{
let mut results = Vec::new();
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = if let Some(g) = growing_arc.try_read() {
g
} else {
growing_arc.read()
};
if segment_filter(&growing) {
results.extend(growing.iter().filter(|e| filter(e)).cloned());
}
}
let sealed = self.sealed.read();
let use_parallel = sealed.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
if use_parallel {
let filter_ref = &filter;
let segment_filter_ref = &segment_filter;
let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
sealed
.iter()
.map(|segment| {
s.spawn(move || {
let segment = segment.read();
if segment_filter_ref(&segment) {
segment
.iter()
.filter(|e| filter_ref(e))
.cloned()
.collect::<Vec<_>>()
} else {
Vec::new()
}
})
})
.collect::<Vec<_>>()
.into_iter()
.map(|handle| handle.join().unwrap_or_default())
.collect()
});
for batch in segment_results {
results.extend(batch);
}
} else {
for segment in sealed.iter() {
let seg = segment.read();
if segment_filter(&seg) {
results.extend(seg.iter().filter(|e| filter(e)).cloned());
}
}
}
results
}
pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
if growing.may_contain_exact_key(key) {
return true;
}
} else {
return true;
}
let sealed = self.sealed.read();
for segment_arc in sealed.iter() {
let segment = segment_arc.read();
if segment.may_contain_exact_key(key) {
return true;
}
}
false
}
pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
let mut results = Vec::new();
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
results.extend(growing.filter_metadata(filters));
}
let sealed = self.sealed.read();
for segment in sealed.iter() {
results.extend(segment.filter_metadata(filters));
}
results
}
pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
let mut results = Vec::new();
if let Some(growing_arc) = self.growing.read().as_ref() {
let growing = growing_arc.read();
for entity in growing.iter_kind(kind) {
results.push(entity.clone());
}
}
let sealed = self.sealed.read();
for segment_arc in sealed.iter() {
let segment = segment_arc.read();
for entity in segment.iter_kind(kind) {
results.push(entity.clone());
}
}
results
}
pub fn count(&self) -> usize {
self.total_entities_atomic.load(Ordering::Relaxed) as usize
}
pub fn segment_ids(&self) -> Vec<SegmentId> {
let mut ids = Vec::new();
if let Some(growing_arc) = self.growing.read().as_ref() {
ids.push(growing_arc.read().id());
}
let sealed = self.sealed.read();
for segment in sealed.iter() {
ids.push(segment.id());
}
ids.extend(self.archived.read().iter().copied());
ids
}
#[inline]
#[allow(clippy::unused_self)]
fn emit(&self, _event: LifecycleEvent) {}
pub fn drain_events(&self) -> Vec<LifecycleEvent> {
std::mem::take(&mut *self.events.write())
}
pub fn run_maintenance(&self) -> Result<(), SegmentError> {
self.maybe_seal_growing()?;
if self.config.enable_consolidation {
self.consolidation_tick();
}
Ok(())
}
fn consolidation_tick(&self) {
let running = self.consolidation.read().is_some();
if !running && !self.start_consolidation() {
return;
}
let mut guard = self.consolidation.write();
let Some(run) = guard.as_mut() else {
return;
};
self.copy_bounded(run);
if run.cursor >= run.sources.len() {
if let Some(run) = guard.take() {
self.finish_consolidation(run);
}
}
}
fn start_consolidation(&self) -> bool {
let sources = self.select_candidates();
self.start_consolidation_from_sources(sources)
}
fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
return false;
}
let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
let merged = GrowingSegment::new(merged_id, &self.collection);
*self.consolidation.write() = Some(Consolidation {
sources: sources
.into_iter()
.map(|id| SourceProgress {
id,
copied: Vec::new(),
tombstones_at_open: None,
})
.collect(),
merged_id,
merged,
cursor: 0,
pending_ids: Vec::new(),
pending_cursor: 0,
entities_copied: 0,
});
self.stats.write().consolidation.runs_started += 1;
true
}
pub fn pressure_consolidation_tick(&self) -> bool {
if !self.config.enable_consolidation {
return false;
}
let running = self.consolidation.read().is_some();
if !running {
let sources = self.dirty_sealed_segments();
if !self.start_consolidation_from_sources(sources) {
return false;
}
}
self.consolidation_tick();
true
}
fn select_candidates(&self) -> Vec<SegmentId> {
let sealed = self.sealed.read();
if sealed.is_empty() {
return Vec::new();
}
if sealed.len() > self.config.max_sealed_segments {
return sealed.iter().map(|segment| segment.read().id()).collect();
}
let mut live = 0usize;
let mut tombstones = 0usize;
let mut held = 0u64;
let mut reclaimable = 0u64;
let mut dirty = Vec::new();
for segment_arc in sealed.iter() {
let segment = segment_arc.read();
live += segment.live_entity_count();
tombstones += segment.tombstone_count();
held += segment.resident_bytes();
reclaimable += segment.reclaimable_bytes();
if segment.tombstone_count() > 0 {
dirty.push(segment.id());
}
}
let total = live + tombstones;
let tombstone_ratio = if total == 0 {
0.0
} else {
tombstones as f64 / total as f64
};
let fragmentation = if held == 0 {
0.0
} else {
reclaimable as f64 / held as f64
};
if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
|| fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
{
return dirty;
}
Vec::new()
}
fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
self.sealed
.read()
.iter()
.filter_map(|segment_arc| {
let segment = segment_arc.read();
(segment.reclaimable_bytes() > 0).then(|| segment.id())
})
.collect()
}
fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
if sources.is_empty() {
return false;
}
let sealed = self.sealed.read();
sealed.iter().any(|segment_arc| {
let segment = segment_arc.read();
sources.contains(&segment.id()) && segment.tombstone_count() > 0
})
}
fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
let sealed = self.sealed.read();
sealed
.iter()
.find(|segment| segment.read().id() == id)
.map(Arc::clone)
}
fn copy_bounded(&self, run: &mut Consolidation) {
let mut budget = self.config.consolidation_entities_per_tick;
while budget > 0 && run.cursor < run.sources.len() {
let source_id = run.sources[run.cursor].id;
let Some(source_arc) = self.sealed_segment(source_id) else {
run.cursor += 1;
run.pending_ids.clear();
run.pending_cursor = 0;
continue;
};
if run.sources[run.cursor].tombstones_at_open.is_none() {
let source = source_arc.read();
run.pending_ids = source.live_entity_ids();
run.pending_cursor = 0;
run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
}
{
let source = source_arc.read();
while budget > 0 && run.pending_cursor < run.pending_ids.len() {
let id = run.pending_ids[run.pending_cursor];
run.pending_cursor += 1;
budget -= 1;
let Some(entity) = source.get(id) else {
continue;
};
let metadata = source.get_metadata(id);
run.merged.adopt_entity(entity.clone(), metadata);
run.sources[run.cursor].copied.push(id);
run.entities_copied += 1;
}
}
if run.pending_cursor >= run.pending_ids.len() {
run.cursor += 1;
run.pending_ids.clear();
run.pending_cursor = 0;
}
}
}
fn finish_consolidation(&self, mut run: Consolidation) {
let mut sealed = self.sealed.write();
let mut positions = Vec::with_capacity(run.sources.len());
for source in &run.sources {
if let Some(index) = sealed
.iter()
.position(|segment| segment.read().id() == source.id)
{
positions.push(index);
}
}
if positions.is_empty() {
return;
}
positions.sort_unstable();
let mut stale = Vec::new();
for source in &run.sources {
let Some(source_arc) = sealed
.iter()
.find(|segment| segment.read().id() == source.id)
else {
continue;
};
let segment = source_arc.read();
if source.tombstones_at_open == Some(segment.tombstone_count()) {
continue;
}
stale.extend(
source
.copied
.iter()
.copied()
.filter(|id| !segment.contains(*id)),
);
}
for id in stale {
run.merged.evict_entity(id);
}
let reclaimed_bytes: u64 = positions
.iter()
.map(|&index| sealed[index].read().resident_bytes())
.sum();
let reclaimed_tombstones: u64 = positions
.iter()
.map(|&index| sealed[index].read().tombstone_count() as u64)
.sum();
if run.merged.seal().is_err() {
return;
}
let merged_bytes = run.merged.resident_bytes();
let merged_arc = Arc::new(RwLock::new(run.merged));
{
let mut retired = self.retired.write();
retired.retain(|source| source.strong_count() > 0);
for &index in positions.iter().rev() {
let source = sealed.remove(index);
if Arc::strong_count(&source) > 1 {
retired.push(Arc::downgrade(&source));
}
}
}
let insert_at = positions[0].min(sealed.len());
sealed.insert(insert_at, Arc::clone(&merged_arc));
drop(sealed);
let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
{
let mut entity_segment = self.entity_segment.write();
for segment_id in entity_segment.values_mut() {
if source_ids.contains(segment_id) {
*segment_id = run.merged_id;
}
}
}
{
let mut stats = self.stats.write();
stats.sealed_count = stats
.sealed_count
.saturating_sub(positions.len())
.saturating_add(1);
let consolidation = &mut stats.consolidation;
consolidation.runs_completed += 1;
consolidation.segments_merged += positions.len() as u64;
consolidation.tombstones_reclaimed += reclaimed_tombstones;
consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
}
self.emit(LifecycleEvent::SegmentConsolidated {
source: source_ids,
target: run.merged_id,
});
}
#[cfg(test)]
fn consolidation_progress(&self) -> Option<u64> {
self.consolidation
.read()
.as_ref()
.map(|run| run.entities_copied)
}
}
impl UnifiedSegment for Arc<RwLock<GrowingSegment>> {
fn id(&self) -> SegmentId {
self.read().id()
}
fn state(&self) -> SegmentState {
self.read().state()
}
fn collection(&self) -> &str {
"unknown"
}
fn stats(&self) -> SegmentStats {
self.read().stats()
}
fn entity_count(&self) -> usize {
self.read().entity_count()
}
fn contains(&self, id: EntityId) -> bool {
self.read().contains(id)
}
fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
None
}
fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
None
}
fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
self.write().insert(entity)
}
fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
self.write().update(entity)
}
fn update_hot(
&mut self,
entity: UnifiedEntity,
modified_columns: &[String],
) -> Result<(), SegmentError> {
self.write().update_hot(entity, modified_columns)
}
fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
self.write().delete(id)
}
fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
self.read().get_metadata(id)
}
fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
self.write().set_metadata(id, metadata)
}
fn seal(&mut self) -> Result<(), SegmentError> {
self.write().seal()
}
fn should_seal(&self, config: &SegmentConfig) -> bool {
self.read().should_seal(config)
}
fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
Box::new(std::iter::empty())
}
fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
Box::new(std::iter::empty())
}
fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
self.read().filter_metadata(filters)
}
}
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::*;
use crate::runtime::mvcc::SnapshotContext;
use crate::storage::transaction::snapshot::{Snapshot, SnapshotManager};
use crate::storage::unified::entity::{EntityData, EntityKind, RowData};
use reddb_types::Value;
fn lookup_version(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
entity.data = EntityData::Row(RowData::with_names(
vec![Value::text(key)],
vec!["key".into()],
));
entity.set_xmin(xmin);
entity.set_xmax(xmax);
entity
}
#[test]
fn key_lookup_matches_snapshot_scan_across_segments() {
let manager = Arc::new(SegmentManager::new("keys"));
for entity in [
lookup_version(1, "changed", 3, 12),
lookup_version(2, "deleted", 3, 7),
] {
manager.insert(entity).expect("old version");
}
manager.force_seal().expect("seal old versions");
for entity in [
lookup_version(3, "changed", 12, 0),
lookup_version(4, "pending", 4, 0),
lookup_version(5, "own", 15, 0),
lookup_version(6, "live", 3, 0),
] {
manager.insert(entity).expect("new version");
}
std::thread::spawn(move || {
let snapshot = SnapshotContext {
snapshot: Snapshot {
xid: 10,
in_progress: HashSet::from([4]),
},
manager: Arc::new(SnapshotManager::new()),
own_xids: HashSet::from([15]),
requires_index_fallback: false,
serializable_reader: None,
};
for snapshot in [Some(&snapshot), None] {
for key in ["changed", "deleted", "pending", "own", "live", "absent"] {
let expected = manager
.scan(snapshot, |entity| {
super::super::segment_key_lookup::SegmentKeyLookup::key(entity)
== Some(key)
})
.into_iter()
.max_by_key(|entity| entity.xmin)
.map(|entity| entity.id);
assert_eq!(
manager.project_key_entity(key, snapshot, true, |entity| entity.id),
expected,
"key={key}"
);
}
}
assert_eq!(
manager.project_key_entity("changed", Some(&snapshot), true, |entity| entity.id),
Some(EntityId::new(1))
);
assert_eq!(
manager.project_key_entity("changed", None, true, |entity| entity.id),
Some(EntityId::new(3))
);
assert!(manager
.project_key_entity("deleted", None, true, |entity| entity.id)
.is_none());
})
.join()
.expect("snapshot worker");
}
#[test]
fn key_lookup_survives_paced_consolidation() {
let manager = consolidating_manager(1);
manager
.bulk_insert(
(1..=10)
.map(|id| lookup_version(id, &format!("key-{id}"), 1, 0))
.collect(),
)
.expect("rows");
manager.force_seal().expect("seal");
for id in 1..=4 {
manager.delete(EntityId::new(id)).expect("delete");
}
assert!(manager.pressure_consolidation_tick());
drain_maintenance(&manager);
for id in 1..=10 {
assert_eq!(
manager.project_key_entity(&format!("key-{id}"), None, false, |entity| entity.id),
(id > 4).then_some(EntityId::new(id))
);
}
}
#[test]
#[ignore = "diagnostic of segment fragmentation and retained history; no timing assertions"]
fn key_lookup_fragmentation_history_diagnostic() {
use std::{hint::black_box, time::Instant};
for segments in [1, 16, 64] {
for versions in [1, 100, 1000] {
let manager = SegmentManager::new("keys");
let mut id = 1;
for segment in 0..segments {
let mut rows: Vec<_> = (0..128)
.map(|_| {
let entity = lookup_version(id, &format!("other-{id}"), 1, 0);
id += 1;
entity
})
.collect();
for version in (segment..versions).step_by(segments) {
rows.push(lookup_version(
id,
"history",
version as u64 + 1,
if version + 1 == versions {
0
} else {
version as u64 + 2
},
));
id += 1;
}
manager.bulk_insert(rows).expect("rows");
manager.force_seal().expect("seal");
}
let started = Instant::now();
for _ in 0..2000 {
let xmin =
manager.project_key_entity(black_box("history"), None, true, |entity| {
entity.xmin
});
assert_eq!(xmin, Some(versions as u64));
}
eprintln!("segments={segments} versions={versions} reads=2000 elapsed_us={} resident_bytes={}", started.elapsed().as_micros(), manager.resident_bytes());
}
}
}
#[test]
fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
let manager = SegmentManager::new("fingerprint");
let row = |id, label: &str| {
UnifiedEntity::table_row(
EntityId::new(id),
"fingerprint",
id,
vec![Value::text(label)],
)
};
manager
.bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
.expect("sealed rows");
manager.force_seal().expect("seal");
manager
.bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
.expect("growing rows");
let mut expected = manager.query_all(|_| true);
expected.sort_by_key(|entity| entity.id.raw());
let expected: Vec<_> = expected
.iter()
.map(|entity| format!("{entity:?}"))
.collect();
let mut stored_addresses = HashSet::new();
manager.for_each_entity(|entity| {
stored_addresses.insert(entity as *const UnifiedEntity as usize);
true
});
let mut actual = Vec::new();
manager.for_each_entity_in_id_order(|entity| {
assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
assert!(
manager.growing.try_write().is_some(),
"topology guard retained"
);
assert!(
manager.sealed.try_write().is_some(),
"topology guard retained"
);
actual.push(format!("{entity:?}"));
true
});
assert_eq!(
actual, expected,
"equal IDs retain their original segment order"
);
}
#[test]
fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
let manager = SegmentManager::new("fingerprint");
for id in [3, 1, 2] {
manager
.insert(UnifiedEntity::table_row(
EntityId::new(id),
"fingerprint",
id,
vec![Value::Integer(id as i64)],
))
.expect("row");
}
let source = manager.growing.read().as_ref().expect("source").clone();
let handles = Arc::strong_count(&source);
let mut seen = Vec::new();
manager.for_each_entity_in_id_order(|entity| {
seen.push(entity.id.raw());
false
});
assert_eq!(seen, vec![1]);
assert_eq!(Arc::strong_count(&source), handles);
assert!(
source.try_write().is_some(),
"segment guard retained after early stop"
);
}
#[test]
fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
let manager = SegmentManager::with_config(
"vectors",
ManagerConfig {
max_sealed_segments: 1,
consolidation_entities_per_tick: 10_000,
..Default::default()
},
);
for part in 0..3 {
manager
.bulk_insert(
(1..=300)
.map(|offset| {
UnifiedEntity::vector(
EntityId::new(part * 300 + offset),
"vectors",
vec![1.0, 0.0],
)
})
.collect(),
)
.expect("batch");
if part < 2 {
manager.force_seal().expect("seal");
}
}
let original_growing = manager.growing.read().as_ref().expect("growing").clone();
let mut changed = false;
let mut seen = Vec::new();
assert!(manager.scan_batches(
None,
|| true,
|batch| {
assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
assert!(
manager.growing.try_write().is_some(),
"consumer holds no topology guard"
);
assert!(
original_growing.try_write().is_some(),
"consumer holds no segment guard"
);
for entity in batch {
assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
seen.push(entity.id.raw());
}
if !changed {
changed = true;
manager.force_seal().expect("seal during scan");
manager
.insert(UnifiedEntity::vector(
EntityId::new(5000),
"vectors",
vec![1.0, 0.0],
))
.expect("new write");
drain_maintenance(&manager);
assert!(manager.stats().consolidation.runs_completed > 0);
assert!(manager
.delete(EntityId::new(1))
.expect("delete after source retirement"));
assert!(manager.get(EntityId::new(1)).is_none());
}
true
}
));
seen.sort_unstable();
assert_eq!(
seen,
(2..=900).collect::<Vec<_>>(),
"no duplicate, missing or late candidates"
);
let mut after = Vec::new();
assert!(manager.scan_batches(
None,
|| true,
|batch| {
after.extend(batch.iter().map(|entity| entity.id.raw()));
true
}
));
after.sort_unstable();
let mut expected: Vec<_> = (2..=900).collect();
expected.push(5000);
assert_eq!(after, expected, "adopted entities have scan positions");
}
#[test]
fn bounded_scan_applies_snapshot_to_relocated_versions() {
for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
let manager = SegmentManager::with_config(
"vectors",
ManagerConfig {
max_sealed_segments: 1,
..Default::default()
},
);
for ids in [[1, 2], [3, 4]] {
let entities = ids.map(|id| {
let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
entity.set_xmin(3);
entity
});
if bulk {
manager
.bulk_insert(entities.into())
.expect("initial versions");
} else {
for entity in entities {
manager.insert(entity).expect("initial version");
}
}
manager.force_seal().expect("seal source");
}
let snapshot = SnapshotContext {
snapshot: Snapshot {
xid: 10,
in_progress: HashSet::new(),
},
manager: Arc::new(SnapshotManager::new()),
own_xids: HashSet::new(),
requires_index_fallback: false,
serializable_reader: None,
};
let mut changed = false;
let mut seen = Vec::new();
assert!(manager.scan_batches(
Some(&snapshot),
|| true,
|batch| {
seen.extend(batch.iter().map(|entity| entity.id.raw()));
if !changed {
changed = true;
if consolidate {
drain_maintenance(&manager);
assert!(manager.stats().consolidation.runs_completed > 0);
}
for (id, xmax) in [(3, 8), (4, 12)] {
let mut entity =
manager.get(EntityId::new(id)).expect("retained version");
entity.set_xmax(xmax);
manager
.update(entity)
.expect("update after source retirement");
}
}
true
}
));
seen.sort_unstable();
assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
let mut again = Vec::new();
assert!(manager.scan_batches(
Some(&snapshot),
|| true,
|batch| {
again.extend(batch.iter().map(|entity| entity.id.raw()));
true
}
));
again.sort_unstable();
assert_eq!(
again,
vec![1, 2, 4],
"old slots do not duplicate relocated versions"
);
let mut current = Vec::new();
assert!(manager.scan_batches(
None,
|| true,
|batch| {
current.extend(batch.iter().map(|entity| entity.id.raw()));
true
}
));
current.sort_unstable();
assert_eq!(
current,
vec![1, 2],
"current view excludes both deleted versions"
);
}
}
#[test]
fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
let manager = SegmentManager::new("mixed");
let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
manager
.bulk_insert((1..=300).map(vector).collect())
.expect("flat");
manager.insert(vector(500)).expect("hash insert");
manager
.bulk_insert(vec![vector(502), vector(505)])
.expect("gapped bulk");
manager.delete(EntityId::new(2)).expect("flat delete");
manager.delete(EntityId::new(505)).expect("hash delete");
let mut changed = false;
let mut seen = Vec::new();
assert!(manager.scan_batches(
None,
|| true,
|batch| {
seen.extend(batch.iter().map(|entity| entity.id.raw()));
if !changed {
changed = true;
manager
.insert(vector(700))
.expect("append after cursor boundary");
manager
.update(UnifiedEntity::graph_node(
EntityId::new(500),
"mixed",
"Node",
HashMap::new(),
))
.expect("structural update");
}
true
}
));
seen.sort_unstable();
let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
expected.extend([500, 502]);
assert_eq!(seen, expected);
}
#[test]
fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
let manager = SegmentManager::new("vectors");
let mut entities = Vec::new();
for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
entity.set_xmin(xmin);
entity.set_xmax(xmax);
entities.push(entity);
}
manager.bulk_insert(entities).expect("versions");
manager.force_seal().expect("seal");
let snapshot = SnapshotContext {
snapshot: Snapshot {
xid: 10,
in_progress: HashSet::from([4]),
},
manager: Arc::new(SnapshotManager::new()),
own_xids: HashSet::new(),
requires_index_fallback: false,
serializable_reader: None,
};
let mut work = 0;
let mut callbacks = 0;
assert!(!manager.scan_batches(
Some(&snapshot),
|| {
work += 1;
work <= 2
},
|_| {
callbacks += 1;
true
}
));
assert_eq!(callbacks, 0, "a stopped batch is not consumed");
for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
let mut seen = Vec::new();
assert!(manager.scan_batches(
view,
|| true,
|batch| {
seen.extend(batch.iter().map(|entity| entity.id.raw()));
true
}
));
seen.sort_unstable();
assert_eq!(seen, expected);
}
}
#[test]
fn test_manager_basic() {
let manager = SegmentManager::new("test_collection");
let entity = UnifiedEntity::table_row(
manager.next_entity_id(),
"users",
1,
vec![Value::text("Alice".to_string())],
);
let id = manager.insert(entity).unwrap();
assert!(manager.get(id).is_some());
assert_eq!(manager.count(), 1);
}
#[test]
fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
let manager = Arc::new(SegmentManager::new("mixed"));
manager
.bulk_insert(
(1..=128)
.map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
.collect(),
)
.expect("vector-only segment");
manager.force_seal().expect("seal vectors");
let mut old =
UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
old.set_xmin(3);
old.set_xmax(12);
let mut new =
UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
new.set_xmin(12);
let edge = UnifiedEntity::graph_edge(
EntityId::new(202),
"mixed",
"a",
"b",
1.0,
Default::default(),
);
manager
.bulk_insert(vec![old, new, edge])
.expect("physical graph versions");
manager.force_seal().expect("seal graph");
manager
.insert(UnifiedEntity::graph_node(
EntityId::new(203),
"mixed",
"person",
Default::default(),
))
.expect("growing node");
let snapshot = SnapshotContext {
snapshot: Snapshot {
xid: 10,
in_progress: HashSet::new(),
},
manager: Arc::new(SnapshotManager::new()),
own_xids: HashSet::new(),
requires_index_fallback: false,
serializable_reader: None,
};
std::thread::spawn(move || {
for view in [None, Some(&snapshot)] {
for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
let expected: HashSet<_> = manager
.scan(view, |item| kind.matches(&item.kind))
.into_iter()
.map(|item| item.id)
.collect();
let actual: HashSet<_> = manager
.scan_graph_kind(view, kind)
.into_iter()
.map(|item| item.id)
.collect();
assert_eq!(actual, expected);
let mut visited = 0;
let mut candidates = HashSet::new();
manager.scan_graph_candidates_for_each(
view,
kind,
|| true,
|item| {
visited += 1;
if kind.matches(&item.kind) {
candidates.insert(item.id);
}
true
},
);
assert_eq!(candidates, expected);
assert_eq!(
visited,
if matches!(kind, GraphEntityKind::Node) {
3
} else {
2
}
);
}
}
let ids: HashSet<_> = manager
.scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
.into_iter()
.map(|item| item.id.raw())
.collect();
assert_eq!(ids, HashSet::from([200, 203]));
let ids: HashSet<_> = manager
.scan_graph_kind(None, GraphEntityKind::Node)
.into_iter()
.map(|item| item.id.raw())
.collect();
assert_eq!(ids, HashSet::from([201, 203]));
let mut visits = 0;
manager.scan_graph_candidates_for_each(
None,
GraphEntityKind::Node,
|| true,
|_| {
visits += 1;
false
},
);
assert_eq!(
visits, 1,
"early termination must also stop subsequent segments"
);
})
.join()
.expect("snapshot worker");
}
#[test]
fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
let manager = Arc::new(SegmentManager::new("accounts"));
let cases = [
("legacy row", 0, 0, true),
("committed creator", 3, 0, true),
("future creator", 11, 0, false),
("in-progress creator", 4, 0, false),
("committed deleter", 3, 7, false),
("future deleter", 3, 11, true),
("in-progress deleter", 3, 4, true),
];
for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
let row_id = index as u64 + 1;
let mut entity = UnifiedEntity::table_row(
EntityId::new(row_id),
"accounts",
row_id,
vec![Value::Integer(row_id as i64)],
);
entity.set_xmin(*xmin);
entity.set_xmax(*xmax);
manager.insert(entity).unwrap();
}
let snapshot = SnapshotContext {
snapshot: Snapshot {
xid: 10,
in_progress: HashSet::from([4]),
},
manager: Arc::new(SnapshotManager::new()),
own_xids: HashSet::new(),
requires_index_fallback: false,
serializable_reader: None,
};
let raw_visible_row_ids: HashSet<u64> = manager
.query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
.into_iter()
.filter_map(|entity| match entity.kind {
EntityKind::TableRow { row_id, .. } => Some(row_id),
_ => None,
})
.collect();
let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
.join()
.unwrap();
let visible_row_ids: HashSet<u64> = scanned
.into_iter()
.filter_map(|entity| match entity.kind {
EntityKind::TableRow { row_id, .. } => Some(row_id),
_ => None,
})
.collect();
assert_eq!(visible_row_ids, raw_visible_row_ids);
for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
let row_id = index as u64 + 1;
assert_eq!(
visible_row_ids.contains(&row_id),
*expected_visible,
"{name} visibility mismatch"
);
}
}
#[test]
fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
let manager = SegmentManager::new("accounts");
let mut live =
UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
live.set_xmin(2);
live.set_xmax(0);
manager.insert(live).unwrap();
let mut superseded =
UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
superseded.set_xmin(2);
superseded.set_xmax(9);
manager.insert(superseded).unwrap();
let rows: Vec<u64> = manager
.scan(None, |_| true)
.into_iter()
.filter_map(|entity| match entity.kind {
EntityKind::TableRow { row_id, .. } => Some(row_id),
_ => None,
})
.collect();
assert_eq!(
rows,
vec![1],
"frameless scan must hide superseded versions"
);
}
#[test]
fn query_scan_callers_use_snapshot_api_for_visibility() {
let table = include_str!("../../runtime/query_exec/table.rs");
let aggregate = include_str!("../../runtime/query_exec/aggregate.rs");
assert_eq!(
table.matches(".for_each_entity(|entity|").count(),
1,
"only the visibility-independent schema inspection may use the raw table iterator"
);
assert!(!table.contains(".query_all_zoned_with_stats("));
assert!(!table.contains(".for_each_entity_zoned_with_stats("));
assert!(!table.contains("TableRowMvccReadResolver::captured"));
assert!(!aggregate.contains("TableRowMvccReadResolver"));
assert!(!aggregate.contains(".for_each_entity("));
assert!(!aggregate.contains(".fold_entities_parallel("));
}
#[test]
fn remaining_visibility_scan_callers_use_snapshot_api() {
let aggregate_pushdown =
include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs");
let ask = include_str!("../../runtime/ask_pipeline.rs");
let dml = include_str!("../../runtime/impl_dml.rs");
let dml_targets = include_str!("../../runtime/dml_target_scan.rs");
let events = include_str!("../../runtime/impl_events.rs");
let graph_commands = include_str!("../../runtime/impl_graph_commands.rs");
let graph_dsl = include_str!("../../runtime/graph_dsl.rs");
let graph_tvf = include_str!("../../runtime/graph_tvf.rs");
let kv = include_str!("../../runtime/impl_kv.rs");
let queue = include_str!("../../runtime/impl_queue.rs");
let primary_queue = include_str!("../../runtime/primary_queue_store.rs");
let record_search = include_str!("../../runtime/record_search.rs");
let search = include_str!("../../runtime/impl_search.rs");
let vector = include_str!("../../runtime/query_exec/vector.rs");
assert!(!aggregate_pushdown.contains("TableRowMvccReadResolver"));
assert!(!aggregate_pushdown.contains(".for_each_entity("));
assert!(!ask.contains("for entity in manager.query_all(|_| true)"));
assert!(!dml.contains("let entities = manager.query_all(|_| true);"));
assert_eq!(
dml_targets
.matches("self.candidate_visible(snapshot.as_ref(), entity)")
.count(),
2,
"both raw DML iterations must route through candidate_visible"
);
assert_eq!(
events.matches("manager.query_all(").count(),
1,
"only the visibility-free backfill dedup may use query_all"
);
assert!(!graph_commands.contains("TableRowMvccReadResolver"));
assert!(!graph_dsl.contains("entity_visible_with_context(snap_ctx.as_ref(), e)"));
assert!(!graph_tvf.contains("entity_visible_with_context(snap_ctx.as_ref(), &entity)"));
assert!(!kv.contains("TableRowMvccReadResolver"));
assert!(!queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
assert!(!primary_queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
assert!(!record_search.contains("TableRowMvccReadResolver"));
assert!(!search.contains("for entity in manager.query_all(|_| true)"));
assert!(!vector.contains("manager.query_all("));
}
fn production_source(src: &str) -> String {
let mut kept: Vec<&str> = Vec::new();
let lines: Vec<&str> = src.lines().collect();
let mut idx = 0;
while idx < lines.len() {
let starts_test_mod = lines[idx].starts_with("#[cfg(test)]")
&& lines
.get(idx + 1)
.is_some_and(|next| next.starts_with("mod "));
if !starts_test_mod {
kept.push(lines[idx]);
idx += 1;
continue;
}
idx += 1;
let mut depth: isize = 0;
while idx < lines.len() {
depth += lines[idx].matches('{').count() as isize;
depth -= lines[idx].matches('}').count() as isize;
idx += 1;
if depth <= 0 {
break;
}
}
}
kept.join("\n")
}
fn raw_iteration_sites(src: &str) -> usize {
src.matches("query_all(").count()
+ src.matches(".for_each_entity").count()
+ src.matches("fold_entities_parallel(").count()
}
#[test]
fn raw_iteration_confined_to_maintenance_allowlist() {
let allowlist: &[(&str, &str, usize, &str)] = &[
(
"runtime/query_exec/table.rs",
include_str!("../../runtime/query_exec/table.rs"),
1,
"documented exception: schema inspection (does the collection ever \
observe these fields?) is visibility-independent",
),
(
"runtime/query_exec/aggregate.rs",
include_str!("../../runtime/query_exec/aggregate.rs"),
0,
"fully migrated by batch 1",
),
(
"runtime/query_exec/aggregate_pushdown_dispatch.rs",
include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs"),
0,
"fully migrated by batch 1",
),
(
"runtime/query_exec/vector.rs",
include_str!("../../runtime/query_exec/vector.rs"),
0,
"fully migrated by batch 2",
),
(
"runtime/ask_pipeline.rs",
include_str!("../../runtime/ask_pipeline.rs"),
0,
"fully migrated by batch 2",
),
(
"runtime/impl_dml.rs",
include_str!("../../runtime/impl_dml.rs"),
1,
"documented exception: ON CONFLICT resolves a unique key against \
physical versions on purpose. A snapshot scan cannot see a row an \
uncommitted concurrent writer already holds the key for, which is \
exactly the case that must raise a serialization conflict rather \
than insert a duplicate; the same pass reaps index entries left by \
aborted writers. \
TODO(#2263 follow-up): narrow the candidates through the unique \
index instead of walking the collection",
),
(
"runtime/dml_target_scan.rs",
include_str!("../../runtime/dml_target_scan.rs"),
2,
"documented exception: kind-conditional visibility the uniform scan \
API cannot express; both iterations gate on candidate_visible",
),
(
"runtime/impl_events.rs",
include_str!("../../runtime/impl_events.rs"),
1,
"documented exception: backfill dedup must include consumed messages, \
so it is deliberately visibility-free",
),
(
"runtime/impl_graph_commands.rs",
include_str!("../../runtime/impl_graph_commands.rs"),
0,
"native node properties resolve one snapshot-visible, RLS-admitted version",
),
(
"runtime/graph_dsl.rs",
include_str!("../../runtime/graph_dsl.rs"),
1,
"remaining legacy debt: unused edge-property helper; search expansion \
now uses snapshot visitors and RLS-gated hydration",
),
(
"runtime/graph_tvf.rs",
include_str!("../../runtime/graph_tvf.rs"),
0,
"fully migrated by batch 2",
),
(
"runtime/impl_kv.rs",
include_str!("../../runtime/impl_kv.rs"),
3,
"one non-versioned listing fallback (the versioned arm scans under \
the captured snapshot). Point reads use physical key candidates with \
versioned visibility rechecks. Vault reads use snapshot scans. Offline vault \
export preserves all \
physical versions across tenants; restore's emptiness check must reject \
every physical row, including versions invisible to a snapshot",
),
(
"runtime/impl_queue.rs",
include_str!("../../runtime/impl_queue.rs"),
10,
"maintenance: red_queue_meta bookkeeping (config, groups, acks, \
pending, dedup TTL reaper, gauge counts) — queue message bodies are \
read through scan. \
TODO(#2138 follow-up): QUEUE PENDING projects meta rows to clients",
),
(
"runtime/primary_queue_store.rs",
include_str!("../../runtime/primary_queue_store.rs"),
3,
"maintenance: generic red_queue_meta read/delete helpers and the \
next-position allocator, which must not collide with rows invisible \
to the writer's snapshot",
),
(
"runtime/record_search.rs",
include_str!("../../runtime/record_search.rs"),
0,
"fully migrated by batch 2",
),
(
"runtime/search_graph.rs",
include_str!("../../runtime/search_graph.rs"),
0,
"search graph preparation uses snapshot visitors and gated hydration",
),
(
"runtime/impl_search.rs",
include_str!("../../runtime/impl_search.rs"),
1,
"maintenance: ask-audit retention sweep must see every physical row",
),
];
for (path, src, allowed, reason) in allowlist {
let found = raw_iteration_sites(&production_source(src));
assert_eq!(
found, *allowed,
"{path}: {found} raw-iteration sites, allowlist permits {allowed} \
({reason}). Raw iteration bypasses MVCC visibility — use \
SegmentManager::scan / scan_for_each / scan_fold_parallel / \
scan_zoned_with_stats, or extend the allowlist with a maintenance \
justification."
);
}
}
#[test]
fn explicit_snapshot_scan_ignores_ambient_thread_local() {
let manager = Arc::new(SegmentManager::new("accounts"));
let versions = [(0, 0), (3, 0), (11, 0), (4, 0), (3, 7), (3, 11), (3, 4)];
for (index, (xmin, xmax)) in versions.iter().enumerate() {
let row_id = index as u64 + 1;
let mut entity = UnifiedEntity::table_row(
EntityId::new(row_id),
"accounts",
row_id,
vec![Value::Integer(row_id as i64)],
);
entity.set_xmin(*xmin);
entity.set_xmax(*xmax);
manager.insert(entity).unwrap();
}
let explicit = SnapshotContext {
snapshot: Snapshot {
xid: 10,
in_progress: HashSet::from([4]),
},
manager: Arc::new(SnapshotManager::new()),
own_xids: HashSet::new(),
requires_index_fallback: false,
serializable_reader: None,
};
let decoy = SnapshotContext {
snapshot: Snapshot {
xid: 1,
in_progress: HashSet::new(),
},
manager: Arc::new(SnapshotManager::new()),
own_xids: HashSet::new(),
requires_index_fallback: false,
serializable_reader: None,
};
let row_ids = |entities: Vec<UnifiedEntity>| -> HashSet<u64> {
entities
.into_iter()
.filter_map(|entity| match entity.kind {
EntityKind::TableRow { row_id, .. } => Some(row_id),
_ => None,
})
.collect()
};
let guard = crate::runtime::execution_context::CurrentSnapshotGuard::install(decoy);
let on_installing_thread = row_ids(manager.scan(Some(&explicit), |_| true));
drop(guard);
let spawned_manager = Arc::clone(&manager);
let spawned_snapshot = explicit.clone();
let on_fresh_thread = std::thread::spawn(move || {
row_ids(spawned_manager.scan(Some(&spawned_snapshot), |_| true))
})
.join()
.unwrap();
assert_eq!(
on_installing_thread, on_fresh_thread,
"an explicit snapshot must produce the same visible set on any thread"
);
assert_eq!(on_installing_thread, HashSet::from([1, 2, 6, 7]));
let raw_manager = Arc::clone(&manager);
let raw_count = std::thread::spawn(move || raw_manager.query_all(|_| true).len())
.join()
.unwrap();
assert_eq!(
raw_count,
versions.len(),
"raw iteration must expose every physical version"
);
}
#[test]
fn bloom_may_contain_key_probes_inserted_rid_bytes() {
let manager = SegmentManager::new("test_collection");
let entity = UnifiedEntity::table_row(
manager.next_entity_id(),
"users",
1,
vec![Value::text("Alice".to_string())],
);
let id = manager.insert(entity).unwrap();
assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
assert!(manager.delete(id).unwrap());
assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
}
#[test]
fn zoned_scan_reports_pruned_segments_and_preserves_results() {
let config = ManagerConfig {
segment_config: SegmentConfig {
max_entities: 2,
..Default::default()
},
..Default::default()
};
let manager = SegmentManager::with_config("events", config);
for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
let entity = UnifiedEntity::new(
manager.next_entity_id(),
EntityKind::TableRow {
table: "events".into(),
row_id: (idx + 1) as u64,
},
EntityData::Row(row),
);
manager.insert(entity).unwrap();
}
let full = manager.query_all(|entity| match &entity.data {
EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
_ => false,
});
let probe = Value::Integer(100);
let (pruned, stats) =
manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
match &entity.data {
EntityData::Row(row) => {
matches!(row.get_field("ts"), Some(Value::Integer(100)))
}
_ => false,
}
});
assert_eq!(pruned.len(), full.len());
assert_eq!(pruned[0].id, full[0].id);
assert_eq!(stats.segments_total, 3);
assert_eq!(stats.segments_scanned, 2);
assert_eq!(stats.segments_pruned, 1);
}
#[test]
fn test_manager_auto_seal() {
let config = ManagerConfig {
segment_config: SegmentConfig {
max_entities: 2,
..Default::default()
},
..Default::default()
};
let manager = SegmentManager::with_config("test", config);
manager
.insert(UnifiedEntity::vector(
manager.next_entity_id(),
"v",
vec![0.1],
))
.unwrap();
manager
.insert(UnifiedEntity::vector(
manager.next_entity_id(),
"v",
vec![0.2],
))
.unwrap();
manager
.insert(UnifiedEntity::vector(
manager.next_entity_id(),
"v",
vec![0.3],
))
.unwrap();
let stats = manager.stats();
assert_eq!(stats.total_entities, 3);
}
#[test]
fn test_manager_delete() {
let manager = SegmentManager::new("test");
let id = manager
.insert(UnifiedEntity::vector(
manager.next_entity_id(),
"v",
vec![0.1],
))
.unwrap();
assert!(manager.get(id).is_some());
assert!(manager.delete(id).unwrap());
assert!(manager.get(id).is_none());
}
#[test]
fn test_manager_metadata() {
let manager = SegmentManager::new("test");
let id = manager
.insert(UnifiedEntity::table_row(
manager.next_entity_id(),
"hosts",
1,
vec![Value::text("192.168.1.1".to_string())],
))
.unwrap();
let mut meta = Metadata::new();
meta.set(
"os",
super::super::metadata::MetadataValue::String("linux".to_string()),
);
manager.set_metadata(id, meta).unwrap();
let retrieved = manager.get_metadata(id).unwrap();
assert!(retrieved.has("os"));
}
#[test]
fn test_manager_query_by_kind() {
let manager = SegmentManager::new("test");
manager
.insert(UnifiedEntity::table_row(
manager.next_entity_id(),
"hosts",
1,
vec![],
))
.unwrap();
manager
.insert(UnifiedEntity::vector(
manager.next_entity_id(),
"embeddings",
vec![0.1],
))
.unwrap();
manager
.insert(UnifiedEntity::table_row(
manager.next_entity_id(),
"hosts",
2,
vec![],
))
.unwrap();
let rows = manager.get_by_kind("table");
assert_eq!(rows.len(), 2);
let vectors = manager.get_by_kind("vector");
assert_eq!(vectors.len(), 1);
}
fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
SegmentManager::with_config(
"test",
ManagerConfig {
max_sealed_segments: 1_000,
consolidation_entities_per_tick: entities_per_tick,
..Default::default()
},
)
}
fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
let entities: Vec<UnifiedEntity> = (0..count)
.map(|i| {
UnifiedEntity::table_row(
manager.next_entity_id(),
"users",
i as u64,
vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
)
})
.collect();
let ids = manager.bulk_insert(entities).expect("bulk insert");
manager.force_seal().expect("seal");
ids
}
fn live_ids(manager: &SegmentManager) -> Vec<u64> {
let mut ids: Vec<u64> = manager
.query_all(|_| true)
.into_iter()
.map(|entity| entity.id.raw())
.collect();
ids.sort_unstable();
ids
}
fn drain_maintenance(manager: &SegmentManager) -> usize {
let mut ticks = 0;
while manager.consolidation.read().is_some() || ticks == 0 {
manager.run_maintenance().expect("maintenance");
ticks += 1;
assert!(ticks < 1_000, "consolidation failed to converge");
}
ticks
}
#[test]
fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
let ids = manager
.bulk_insert(
(1..=1024)
.map(|id| {
UnifiedEntity::graph_edge(
EntityId::new(id),
"step",
"2000",
"3000",
1.0,
HashMap::new(),
)
})
.collect(),
)
.expect("edges");
manager.force_seal().expect("seal");
for id in ids.iter().take(300) {
manager.delete(*id).expect("fragment source");
}
let source = Arc::downgrade(&manager.sealed.read()[0]);
let mut seen = Vec::new();
let mut changed = false;
let mut metadata_bytes = 0;
assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
|| true,
|bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
|batch| {
assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
assert!(manager.growing.try_write().is_some(), "no growing topology guard");
assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
if !changed {
changed = true;
drain_maintenance(&manager);
assert!(manager.stats().consolidation.runs_completed > 0);
assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
manager.update(updated).expect("update authoritative merged copy");
}
for id in batch {
let entity = manager.get(*id).expect("hydrate current manager");
if id.raw() == 1024 {
assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
}
seen.push(id.raw());
}
Ok(true)
}).expect("cursor"));
assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
assert!(
source.upgrade().is_none(),
"finished cursor releases retired source"
);
assert!(
metadata_bytes < 4096,
"fixed ID buffer plus source handles: {metadata_bytes}"
);
}
#[test]
fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
let manager = SegmentManager::new("graph");
manager
.bulk_insert(
(1..=1024)
.map(|id| {
UnifiedEntity::graph_edge(
EntityId::new(id),
"step",
"2000",
"3000",
1.0,
HashMap::new(),
)
})
.collect(),
)
.expect("edges");
let mut delivered = 0;
let denied = manager.visit_graph_index_batches(
GraphEntityKind::Edge,
&[EntityId::new(2000)],
|| true,
|_| Err("out of memory"),
|ids| {
delivered += ids.len();
Ok(true)
},
);
assert_eq!(denied, Err("out of memory"));
assert_eq!(delivered, 0);
let mut work = 0;
assert!(!manager
.visit_graph_index_batches(
GraphEntityKind::Edge,
&[EntityId::new(2000)],
|| {
work += 1;
work < 100
},
|_| Ok::<_, ()>(()),
|ids| {
delivered += ids.len();
Ok(true)
}
)
.expect("cancelled cursor"));
assert_eq!(
delivered, 0,
"cancelled partial batch cannot reach consumer"
);
assert!(!manager
.visit_graph_index_batches(
GraphEntityKind::Edge,
&[EntityId::new(2000)],
|| true,
|_| Ok::<_, ()>(()),
|ids| {
delivered += ids.len();
Ok(false)
}
)
.expect("consumer stop"));
assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
}
#[test]
fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
let ids = seal_bulk_segment(&manager, 100);
for id in ids.iter().take(40) {
manager.delete(*id).expect("delete before scan");
}
let source_bytes = manager.resident_bytes();
let source = Arc::downgrade(&manager.sealed.read()[0]);
std::thread::scope(|scope| {
let (ready, waiting) = std::sync::mpsc::channel();
let (release_first, first_wait) = std::sync::mpsc::channel();
let (release_last, last_wait) = std::sync::mpsc::channel();
let reader = |resume: std::sync::mpsc::Receiver<()>,
ready: std::sync::mpsc::Sender<()>| {
assert!(manager.scan_batches(
None,
|| true,
|_| {
ready.send(()).expect("signal captured cursor");
let _ = resume.recv();
true
}
));
};
let ready_first = ready.clone();
let first = scope.spawn(move || reader(first_wait, ready_first));
let last = scope.spawn(move || reader(last_wait, ready));
for _ in 0..2 {
waiting
.recv_timeout(std::time::Duration::from_secs(10))
.expect("reader ready");
}
drain_maintenance(&manager);
assert!(
source.upgrade().is_some(),
"both cursors pin the retired source"
);
let active_bytes = manager.sealed.read()[0].read().resident_bytes();
let held = manager.resident_bytes();
assert!(
held >= active_bytes + source_bytes,
"retired source remains resident"
);
release_first.send(()).expect("release first reader");
first.join().expect("first reader");
assert_eq!(
manager.resident_bytes(),
held,
"one reader still pins the source"
);
release_last.send(()).expect("release last reader");
last.join().expect("last reader");
assert!(
source.upgrade().is_none(),
"accounting must not retain payloads"
);
assert_eq!(held - manager.resident_bytes(), source_bytes);
assert_eq!(
manager.stats().total_memory_bytes as u64,
manager.resident_bytes()
);
});
}
#[test]
fn retained_memory_gates_runtime_growth_until_the_cursor_is_cancelled() {
use crate::storage::memory_pools::MemoryPool;
use crate::{RedDBOptions, RedDBRuntime};
let budget = 128 * 1024;
let runtime =
RedDBRuntime::with_options(RedDBOptions::in_memory().with_memory_budget(budget))
.expect("runtime");
runtime
.db()
.store()
.create_collection("test")
.expect("collection");
let manager = runtime
.db()
.store()
.get_collection("test")
.expect("manager");
let ids = seal_bulk_segment(&manager, 100);
for id in ids.iter().take(40) {
manager.delete(*id).expect("delete before scan");
}
let source_bytes = manager.resident_bytes();
let mut growth = 0;
let mut held_used = 0;
assert!(!manager.scan_batches(
None,
|| true,
|_| {
drain_maintenance(&manager);
runtime.refresh_memory_accounting();
held_used = runtime.memory_accounting().total_used_bytes();
growth = budget
.checked_sub(held_used)
.expect("fixture fits while the cursor is open")
+ source_bytes / 2;
assert!(growth > 0);
assert!(
runtime
.admit_non_evictable_growth(
MemoryPool::SegmentArena,
"growth with pinned source",
growth,
)
.is_err(),
"retained source must participate in runtime admission"
);
false
}
));
runtime.refresh_memory_accounting();
assert_eq!(
held_used - runtime.memory_accounting().total_used_bytes(),
source_bytes
);
let _reservation = runtime
.admit_non_evictable_growth(
MemoryPool::SegmentArena,
"growth after cursor cancellation",
growth,
)
.expect("the same growth fits once the cursor releases its source");
}
#[test]
fn retained_memory_includes_an_inflight_consolidation_copy() {
let manager = consolidating_manager(1);
let ids = seal_bulk_segment(&manager, 100);
for id in ids.iter().take(40) {
manager.delete(*id).expect("delete");
}
let before = manager.resident_bytes();
assert!(manager.pressure_consolidation_tick());
let copied = manager
.consolidation
.read()
.as_ref()
.expect("paced merge")
.merged
.resident_bytes();
assert!(copied > 0);
assert!(
manager.resident_bytes() >= before + copied,
"unpublished copy consumes memory"
);
drain_maintenance(&manager);
assert!(
manager.resident_bytes() < before,
"completed merge releases the old storage"
);
}
#[test]
fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
let ids = seal_bulk_segment(&manager, 100);
let bytes_before_delete = manager.resident_bytes();
for id in ids.iter().take(40) {
assert!(manager.delete(*id).expect("delete"));
}
let bytes_with_tombstones = manager.resident_bytes();
assert!(bytes_with_tombstones > bytes_before_delete);
let survivors = live_ids(&manager);
assert_eq!(survivors.len(), 60);
drain_maintenance(&manager);
let stats = manager.stats();
assert_eq!(stats.consolidation.runs_started, 1);
assert_eq!(stats.consolidation.runs_completed, 1);
assert_eq!(stats.consolidation.segments_merged, 1);
assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
assert!(stats.consolidation.bytes_reclaimed > 0);
let bytes_after = manager.resident_bytes();
assert!(
bytes_after <= bytes_before_delete,
"consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
);
assert_eq!(
bytes_with_tombstones - bytes_after,
stats.consolidation.bytes_reclaimed
);
assert_eq!(stats.total_memory_bytes, bytes_after as usize);
assert_eq!(live_ids(&manager), survivors);
assert_eq!(manager.sealed.read().len(), 1);
assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
}
#[test]
fn query_results_are_identical_before_during_and_after_consolidation() {
let manager = consolidating_manager(3);
let ids = seal_bulk_segment(&manager, 20);
seal_bulk_segment(&manager, 20);
for id in ids.iter().take(10) {
assert!(manager.delete(*id).expect("delete"));
}
let before = live_ids(&manager);
let survivor = *ids.last().expect("at least one row");
let tombstoned = ids[0];
let mut mid_tick_observations = 0;
while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
manager.run_maintenance().expect("maintenance");
assert_eq!(live_ids(&manager), before);
assert!(manager.get(survivor).is_some());
assert!(manager.get(tombstoned).is_none());
assert_eq!(manager.get_by_kind("table").len(), before.len());
mid_tick_observations += 1;
assert!(mid_tick_observations < 1_000);
}
assert!(
mid_tick_observations > 2,
"expected the run to span several ticks, saw {mid_tick_observations}"
);
assert_eq!(live_ids(&manager), before);
assert!(manager.get(survivor).is_some());
assert!(manager.get(tombstoned).is_none());
assert_eq!(manager.get_by_kind("table").len(), before.len());
}
#[test]
fn get_by_kind_reaches_sealed_segments() {
let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
let ids = seal_bulk_segment(&manager, 5);
assert!(manager.delete(ids[0]).expect("delete"));
assert_eq!(manager.get_by_kind("table").len(), 4);
assert_eq!(manager.get_by_kind("vector").len(), 0);
}
#[test]
fn consolidation_is_paced_by_the_per_tick_entity_bound() {
const PER_TICK: usize = 4;
let manager = consolidating_manager(PER_TICK);
let ids = seal_bulk_segment(&manager, 60);
for id in ids.iter().take(30) {
assert!(manager.delete(*id).expect("delete"));
}
manager.run_maintenance().expect("maintenance");
assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
let mut previous = PER_TICK as u64;
let mut ticks = 1;
while let Some(copied) = manager.consolidation_progress() {
assert!(
copied - previous <= PER_TICK as u64,
"tick copied {} entities, bound is {PER_TICK}",
copied - previous
);
previous = copied;
manager.run_maintenance().expect("maintenance");
ticks += 1;
assert!(ticks < 100);
}
assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
assert_eq!(manager.stats().consolidation.runs_completed, 1);
assert_eq!(live_ids(&manager).len(), 30);
}
#[test]
fn sealed_segment_count_above_the_limit_triggers_a_merge() {
let manager = SegmentManager::with_config(
"test",
ManagerConfig {
max_sealed_segments: 3,
..Default::default()
},
);
for _ in 0..4 {
seal_bulk_segment(&manager, 5);
}
assert_eq!(manager.sealed.read().len(), 4);
let before = live_ids(&manager);
drain_maintenance(&manager);
assert_eq!(manager.sealed.read().len(), 1);
assert_eq!(manager.stats().consolidation.segments_merged, 4);
assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
assert_eq!(live_ids(&manager), before);
}
#[test]
fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
seal_bulk_segment(&manager, 10);
seal_bulk_segment(&manager, 10);
for _ in 0..5 {
manager.run_maintenance().expect("maintenance");
}
assert_eq!(manager.stats().consolidation.runs_started, 0);
assert_eq!(manager.sealed.read().len(), 2);
}
#[test]
fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
use std::sync::atomic::AtomicBool;
let manager = Arc::new(consolidating_manager(2));
let ids = seal_bulk_segment(&manager, 40);
seal_bulk_segment(&manager, 40);
for id in ids.iter().take(20) {
assert!(manager.delete(*id).expect("delete"));
}
let expected = live_ids(&manager);
let stop = Arc::new(AtomicBool::new(false));
let readers: Vec<_> = (0..4)
.map(|_| {
let manager = Arc::clone(&manager);
let stop = Arc::clone(&stop);
let expected = expected.clone();
std::thread::spawn(move || {
let mut scans = 0u64;
loop {
let mut seen: Vec<u64> = manager
.query_all(|_| true)
.into_iter()
.map(|entity| entity.id.raw())
.collect();
seen.sort_unstable();
assert_eq!(seen, expected, "reader saw a torn sealed set");
scans += 1;
if stop.load(Ordering::Relaxed) {
break;
}
}
scans
})
})
.collect();
drain_maintenance(&manager);
stop.store(true, Ordering::Relaxed);
for reader in readers {
let scans = reader.join().expect("reader thread");
assert!(scans > 0, "reader never got to scan");
}
assert_eq!(manager.stats().consolidation.runs_completed, 1);
assert_eq!(live_ids(&manager), expected);
}
#[test]
fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
let manager = consolidating_manager(4);
let ids = seal_bulk_segment(&manager, 20);
for id in ids.iter().take(6) {
assert!(manager.delete(*id).expect("delete"));
}
manager.run_maintenance().expect("maintenance");
assert_eq!(manager.consolidation_progress(), Some(4));
let doomed = ids[6];
assert!(manager.delete(doomed).expect("delete"));
let expected = live_ids(&manager);
drain_maintenance(&manager);
assert!(manager.get(doomed).is_none(), "deleted row came back");
assert_eq!(live_ids(&manager), expected);
}
#[test]
fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
let manager = consolidating_manager(3);
let ids = seal_bulk_segment(&manager, 30);
for id in ids.iter().take(12) {
assert!(manager.delete(*id).expect("delete"));
}
let before = live_ids(&manager);
let sealed_before: Vec<SegmentId> = manager
.sealed
.read()
.iter()
.map(|segment| segment.read().id())
.collect();
manager.run_maintenance().expect("maintenance");
manager.run_maintenance().expect("maintenance");
assert!(manager.consolidation.read().is_some());
*manager.consolidation.write() = None;
assert_eq!(live_ids(&manager), before);
let sealed_after: Vec<SegmentId> = manager
.sealed
.read()
.iter()
.map(|segment| segment.read().id())
.collect();
assert_eq!(sealed_after, sealed_before);
assert_eq!(manager.stats().consolidation.runs_completed, 0);
drain_maintenance(&manager);
assert_eq!(live_ids(&manager), before);
assert_eq!(manager.stats().consolidation.runs_completed, 1);
}
#[test]
fn consolidation_preserves_entity_payload_metadata_and_sequence() {
let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
let entities: Vec<UnifiedEntity> = (0..10)
.map(|i| {
UnifiedEntity::table_row(
manager.next_entity_id(),
"users",
i as u64,
vec![Value::Integer(i as i64)],
)
})
.collect();
let ids = manager.bulk_insert(entities).expect("bulk insert");
let mut meta = Metadata::new();
meta.set(
"os",
super::super::metadata::MetadataValue::String("linux".to_string()),
);
manager.set_metadata(ids[9], meta).expect("set metadata");
manager.force_seal().expect("seal");
assert!(manager.delete(ids[0]).expect("delete"));
assert!(manager.delete(ids[1]).expect("delete"));
assert!(manager.delete(ids[2]).expect("delete"));
let before: Vec<UnifiedEntity> = {
let mut entities = manager.query_all(|_| true);
entities.sort_by_key(|entity| entity.id.raw());
entities
};
drain_maintenance(&manager);
let after: Vec<UnifiedEntity> = {
let mut entities = manager.query_all(|_| true);
entities.sort_by_key(|entity| entity.id.raw());
entities
};
let columns = |entity: &UnifiedEntity| match &entity.data {
super::super::entity::EntityData::Row(row) => row.columns.clone(),
_ => Vec::new(),
};
assert_eq!(before.len(), after.len());
for (old, new) in before.iter().zip(after.iter()) {
assert_eq!(old.id, new.id);
assert_eq!(columns(old), columns(new));
assert_eq!(
old.sequence_id, new.sequence_id,
"merge must not re-issue sequence ids"
);
}
assert!(manager
.get_metadata(ids[9])
.expect("metadata survives the merge")
.has("os"));
}
#[test]
#[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
fn test_lifecycle_events() {
let manager = SegmentManager::new("test");
manager
.insert(UnifiedEntity::vector(
manager.next_entity_id(),
"v",
vec![0.1],
))
.unwrap();
let events = manager.drain_events();
assert!(events
.iter()
.any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
assert!(events
.iter()
.any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
}
}