1mod graph_index_merge;
16
17use parking_lot::RwLock;
18use std::collections::HashMap;
19use std::sync::atomic::{AtomicU64, Ordering};
20use std::sync::{Arc, Weak};
21
22use super::entity::{EntityId, EntityKind, UnifiedEntity};
23use super::metadata::{Metadata, MetadataFilter};
24use super::segment::{
25 GraphEntityKind, GrowingSegment, SegmentConfig, SegmentError, SegmentId, SegmentScanCursor,
26 SegmentState, SegmentStats, UnifiedSegment, ZoneColPred, ZoneColPredKind,
27};
28use super::visibility_map::VisibilityMap;
29use crate::runtime::mvcc::{entity_visible_with_context, SnapshotContext};
30
31pub const CONSOLIDATION_TOMBSTONE_RATIO: f64 = 0.20;
38
39pub const CONSOLIDATION_FRAGMENTATION_RATIO: f64 = 0.30;
44
45pub const CONSOLIDATION_ENTITIES_PER_TICK: usize = 4_096;
51
52#[derive(Debug, Clone)]
54pub struct ManagerConfig {
55 pub segment_config: SegmentConfig,
57 pub max_sealed_segments: usize,
59 pub idle_seal_secs: u64,
61 pub enable_consolidation: bool,
63 pub consolidation_entities_per_tick: usize,
65 pub enable_archival: bool,
67 pub archive_age_secs: u64,
69}
70
71impl Default for ManagerConfig {
72 fn default() -> Self {
73 Self {
74 segment_config: SegmentConfig::default(),
75 max_sealed_segments: 10,
76 idle_seal_secs: 300, enable_consolidation: true,
78 consolidation_entities_per_tick: CONSOLIDATION_ENTITIES_PER_TICK,
79 enable_archival: true,
80 archive_age_secs: 86400 * 7, }
82 }
83}
84
85#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
90pub struct ConsolidationStats {
91 pub runs_started: u64,
93 pub runs_completed: u64,
95 pub segments_merged: u64,
97 pub tombstones_reclaimed: u64,
99 pub bytes_reclaimed: u64,
101}
102
103#[derive(Debug, Clone, Default)]
105pub struct ManagerStats {
106 pub total_entities: usize,
108 pub growing_count: usize,
110 pub sealed_count: usize,
112 pub archived_count: usize,
114 pub total_memory_bytes: usize,
116 pub seal_ops: u64,
118 pub consolidation: ConsolidationStats,
120}
121
122#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
123pub struct SegmentScanStats {
124 pub segments_total: u64,
125 pub segments_scanned: u64,
126 pub segments_pruned: u64,
127}
128
129#[derive(Debug, Clone)]
131pub enum LifecycleEvent {
132 SegmentCreated(SegmentId),
133 SegmentSealed(SegmentId),
134 SegmentConsolidated {
135 source: Vec<SegmentId>,
136 target: SegmentId,
137 },
138 SegmentArchived(SegmentId),
139 EntityInserted(EntityId, SegmentId),
140 EntityDeleted(EntityId, SegmentId),
141}
142
143struct Consolidation {
150 sources: Vec<SourceProgress>,
152 merged_id: SegmentId,
154 merged: GrowingSegment,
156 cursor: usize,
158 pending_ids: Vec<EntityId>,
160 pending_cursor: usize,
162 entities_copied: u64,
164}
165
166struct SourceProgress {
168 id: SegmentId,
169 copied: Vec<EntityId>,
171 tombstones_at_open: Option<usize>,
180}
181
182struct PhysicalRootCacheEntry {
183 segments: Vec<(u64, u64)>,
184 seed: u64,
185 root: u64,
186}
187
188pub(crate) struct PhysicalScanCursor {
191 sources: Vec<(Weak<RwLock<GrowingSegment>>, SegmentScanCursor)>,
192 source: usize,
193}
194
195impl PhysicalScanCursor {
196 pub(crate) fn finished(&self) -> bool {
197 self.source == self.sources.len()
198 }
199}
200
201pub struct SegmentManager {
203 row_constraint_lock: Arc<parking_lot::ReentrantMutex<()>>,
206 collection: String,
208 config: ManagerConfig,
210 next_segment_id: AtomicU64,
212 next_entity_id: AtomicU64,
214 next_row_id: AtomicU64,
216 total_entities_atomic: AtomicU64,
220 growing: RwLock<Option<Arc<RwLock<GrowingSegment>>>>,
222 sealed: RwLock<Vec<Arc<RwLock<GrowingSegment>>>>,
224 retired: RwLock<Vec<Weak<RwLock<GrowingSegment>>>>,
226 archived: RwLock<Vec<SegmentId>>,
228 entity_segment: RwLock<HashMap<EntityId, SegmentId>>,
232 column_schema: RwLock<Option<Arc<Vec<String>>>>,
235 stats: RwLock<ManagerStats>,
237 consolidation: RwLock<Option<Consolidation>>,
240 events: RwLock<Vec<LifecycleEvent>>,
242 visibility_map: VisibilityMap,
246 physical_root_cache: parking_lot::Mutex<Option<PhysicalRootCacheEntry>>,
247}
248
249impl SegmentManager {
250 pub fn new(collection: impl Into<String>) -> Self {
252 Self::with_config(collection, ManagerConfig::default())
253 }
254
255 pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
257 Self {
258 row_constraint_lock: Arc::new(parking_lot::ReentrantMutex::new(())),
259 physical_root_cache: parking_lot::Mutex::new(None),
260 collection: collection.into(),
261 config,
262 next_segment_id: AtomicU64::new(1),
263 next_entity_id: AtomicU64::new(1),
264 next_row_id: AtomicU64::new(1),
265 total_entities_atomic: AtomicU64::new(0),
266 growing: RwLock::new(None),
267 sealed: RwLock::new(Vec::new()),
268 retired: RwLock::new(Vec::new()),
269 archived: RwLock::new(Vec::new()),
270 entity_segment: RwLock::new(HashMap::new()),
271 column_schema: RwLock::new(None),
272 stats: RwLock::new(ManagerStats::default()),
273 consolidation: RwLock::new(None),
274 events: RwLock::new(Vec::new()),
275 visibility_map: VisibilityMap::new(),
276 }
277 }
278
279 pub(crate) fn row_constraint_lock(&self) -> Arc<parking_lot::ReentrantMutex<()>> {
280 Arc::clone(&self.row_constraint_lock)
281 }
282
283 pub fn get_or_init_schema(
285 &self,
286 named: &HashMap<String, reddb_types::Value>,
287 ) -> Arc<Vec<String>> {
288 {
289 let schema = self.column_schema.read();
290 if let Some(ref s) = *schema {
291 return Arc::clone(s);
292 }
293 }
294 let cols: Vec<String> = named.keys().cloned().collect();
295 let arc = Arc::new(cols);
296 *self.column_schema.write() = Some(Arc::clone(&arc));
297 arc
298 }
299
300 pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
302 self.column_schema.read().clone()
303 }
304
305 fn restore_owned_row_schemas<'a>(&self, entities: impl Iterator<Item = &'a mut UnifiedEntity>) {
309 let schema = self.column_schema();
310 for entity in entities {
311 if let super::entity::EntityData::Row(row) = &mut entity.data {
312 if row.named.is_none() && row.schema.is_none() {
313 row.schema = schema.clone();
314 }
315 }
316 }
317 }
318
319 pub(crate) fn set_column_schema_if_empty(&self, columns: Vec<String>) {
320 if columns.is_empty() {
321 return;
322 }
323 let mut schema = self.column_schema.write();
324 if schema.is_none() {
325 *schema = Some(Arc::new(columns));
326 }
327 }
328
329 pub fn collection(&self) -> &str {
331 &self.collection
332 }
333
334 pub fn config(&self) -> &ManagerConfig {
336 &self.config
337 }
338
339 pub fn stats(&self) -> ManagerStats {
342 let mut s = self.stats.read().clone();
343 s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
344 s.total_memory_bytes = self.resident_bytes() as usize;
345 s
346 }
347
348 pub fn resident_bytes(&self) -> u64 {
351 self.memory_usage().1
352 }
353
354 pub fn memory_bytes(&self) -> u64 {
357 self.memory_usage().0
358 }
359
360 fn memory_usage(&self) -> (u64, u64) {
364 let consolidation = self.consolidation.read();
365 let growing = self.growing.read();
366 let sealed = self.sealed.read();
367 let mut retired = self.retired.write();
368 let mut payload = 0u64;
369 let mut resident = 0u64;
370 {
371 let mut add = |segment: &GrowingSegment| {
372 payload = payload.saturating_add(segment.memory_bytes());
373 resident = resident.saturating_add(segment.resident_bytes());
374 };
375 for segment in growing.iter().chain(sealed.iter()) {
376 add(&segment.read());
377 }
378 retired.retain(|source| {
379 if let Some(source) = source.upgrade() {
380 add(&source.read());
381 true
382 } else {
383 false
384 }
385 });
386 if let Some(run) = consolidation.as_ref() {
387 add(&run.merged);
388 }
389 }
390 let mut bookkeeping =
391 (retired.capacity() * std::mem::size_of::<Weak<RwLock<GrowingSegment>>>()) as u64;
392 if let Some(run) = consolidation.as_ref() {
393 bookkeeping = bookkeeping
394 .saturating_add(
395 (run.sources.capacity() * std::mem::size_of::<SourceProgress>()) as u64,
396 )
397 .saturating_add(
398 (run.pending_ids.capacity() * std::mem::size_of::<EntityId>()) as u64,
399 );
400 for source in &run.sources {
401 bookkeeping = bookkeeping.saturating_add(
402 (source.copied.capacity() * std::mem::size_of::<EntityId>()) as u64,
403 );
404 }
405 }
406 (
407 payload.saturating_add(bookkeeping),
408 resident.saturating_add(bookkeeping),
409 )
410 }
411
412 pub fn reclaimable_bytes(&self) -> u64 {
414 self.sealed
415 .read()
416 .iter()
417 .map(|segment| segment.read().reclaimable_bytes())
418 .fold(0, u64::saturating_add)
419 }
420
421 pub fn next_entity_id(&self) -> EntityId {
423 EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
424 }
425
426 pub fn next_row_id(&self) -> u64 {
428 self.next_row_id.fetch_add(1, Ordering::SeqCst)
429 }
430
431 pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
435 let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
436 start..start + n
437 }
438
439 pub fn register_row_id(&self, id: u64) {
442 let candidate = id.saturating_add(1);
443 let mut current = self.next_row_id.load(Ordering::SeqCst);
444 while candidate > current {
445 match self.next_row_id.compare_exchange(
446 current,
447 candidate,
448 Ordering::SeqCst,
449 Ordering::SeqCst,
450 ) {
451 Ok(_) => break,
452 Err(updated) => current = updated,
453 }
454 }
455 }
456
457 fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
464 {
466 let growing = self.growing.read();
467 if let Some(segment) = growing.as_ref() {
468 return Arc::clone(segment);
469 }
470 }
471
472 let mut growing = self.growing.write();
474 if let Some(segment) = growing.as_ref() {
476 return Arc::clone(segment);
477 }
478
479 let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
480 let segment = GrowingSegment::new(id, &self.collection);
481 let segment_arc = Arc::new(RwLock::new(segment));
482 *growing = Some(Arc::clone(&segment_arc));
483
484 self.emit(LifecycleEvent::SegmentCreated(id));
485
486 self.stats.write().growing_count += 1;
489
490 segment_arc
491 }
492
493 pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
495 self.maybe_seal_growing()?;
497
498 let segment_arc = self.get_or_create_growing();
499 let mut segment = segment_arc.write();
500
501 if entity.id.raw() == 0 {
503 entity.id = self.next_entity_id();
504 }
505
506 let entity_id = entity.id;
507 let segment_id = segment.id();
508
509 segment.insert(entity)?;
510
511 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
513
514 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
523
524 Ok(entity_id)
525 }
526
527 pub(crate) fn restore_snapshot_entity(
529 &self,
530 entity: UnifiedEntity,
531 ) -> Result<EntityId, SegmentError> {
532 self.maybe_seal_growing()?;
533 if self.get(entity.id).is_some() {
534 return Err(SegmentError::AlreadyExists(entity.id));
535 }
536 let segment_arc = self.get_or_create_growing();
537 let mut segment = segment_arc.write();
538 let entity_id = entity.id;
539 let segment_id = segment.id();
540 segment.adopt_entity(entity, None);
543 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
544 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
545 Ok(entity_id)
546 }
547
548 pub fn insert_batch(
550 &self,
551 entities: Vec<UnifiedEntity>,
552 ) -> Result<Vec<EntityId>, SegmentError> {
553 let mut ids = Vec::with_capacity(entities.len());
554 for entity in entities {
555 ids.push(self.insert(entity)?);
556 }
557 Ok(ids)
558 }
559
560 pub fn bulk_insert(
563 &self,
564 mut entities: Vec<UnifiedEntity>,
565 ) -> Result<Vec<EntityId>, SegmentError> {
566 for entity in &mut entities {
568 if entity.id.raw() == 0 {
569 entity.id = self.next_entity_id();
570 }
571 if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
572 if *row_id == 0 {
573 *row_id = self.next_row_id();
574 } else {
575 self.register_row_id(*row_id);
576 }
577 }
578 }
579
580 if let Some(first_row) = entities.first() {
583 if let super::entity::EntityData::Row(ref row) = first_row.data {
584 if let Some(ref named) = row.named {
585 let schema = self.get_or_init_schema(named);
586 for entity in &mut entities {
587 if let super::entity::EntityData::Row(ref mut row) = entity.data {
588 if let Some(named) = row.named.take() {
589 let mut cols = Vec::with_capacity(schema.len());
590 for col_name in schema.iter() {
591 cols.push(
592 named
593 .get(col_name)
594 .cloned()
595 .unwrap_or(reddb_types::Value::Null),
596 );
597 }
598 row.columns = cols;
599 row.schema = Some(Arc::clone(&schema));
600 }
601 }
602 }
603 }
604 }
605 }
606
607 let segment_arc = self.get_or_create_growing();
608 let mut segment = segment_arc.write();
609 let segment_id = segment.id();
610
611 let ids = segment.bulk_insert(entities)?;
613
614 self.total_entities_atomic
619 .fetch_add(ids.len() as u64, Ordering::Relaxed);
620
621 Ok(ids)
622 }
623
624 pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
626 let mut entity = self.get_with(id, Clone::clone)?;
627 self.restore_owned_row_schemas(std::iter::once(&mut entity));
628 Some(entity)
629 }
630
631 pub(crate) fn get_with<T>(
634 &self,
635 id: EntityId,
636 visit: impl FnOnce(&UnifiedEntity) -> T,
637 ) -> Option<T> {
638 if let Some(growing_arc) = self.growing.read().as_ref() {
640 let growing = growing_arc.read();
641 if let Some(entity) = growing.get(id) {
642 return Some(visit(entity));
643 }
644 }
645
646 let sealed = self.sealed.read();
648 for segment in sealed.iter() {
649 let seg = segment.read();
650 if let Some(entity) = seg.get(id) {
651 return Some(visit(entity));
652 }
653 }
654
655 None
656 }
657
658 pub(crate) fn graph_node_logical_id(
662 &self,
663 id: EntityId,
664 mut before_segment: impl FnMut() -> bool,
665 ) -> Option<EntityId> {
666 if let Some(growing_arc) = self.growing.read().as_ref() {
667 if !before_segment() {
668 return None;
669 }
670 let growing = growing_arc.read();
671 if let Some(entity) = growing.get(id) {
672 return matches!(entity.kind, EntityKind::GraphNode(_))
673 .then(|| entity.logical_id());
674 }
675 }
676 let sealed = self.sealed.read();
677 for segment in sealed.iter() {
678 if !before_segment() {
679 return None;
680 }
681 let segment = segment.read();
682 if let Some(entity) = segment.get(id) {
683 return matches!(entity.kind, EntityKind::GraphNode(_))
684 .then(|| entity.logical_id());
685 }
686 }
687 None
688 }
689
690 pub(crate) fn visit_graph_index_batches<E>(
695 &self,
696 kind: GraphEntityKind,
697 keys: &[EntityId],
698 mut before_work: impl FnMut() -> bool,
699 mut admit: impl FnMut(usize) -> Result<(), E>,
700 mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
701 ) -> Result<bool, E> {
702 use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
703 admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
704 let mut admitted = 0;
705 let sources = loop {
706 let count = {
707 let growing = self.growing.read();
708 let sealed = self.sealed.read();
709 sealed.len() + usize::from(growing.is_some())
710 };
711 let bytes = count.saturating_mul(std::mem::size_of::<(
712 Arc<RwLock<GrowingSegment>>,
713 GraphIndexCursor,
714 )>());
715 admit(bytes.saturating_sub(admitted))?;
716 admitted = admitted.max(bytes);
717 let mut sources = Vec::with_capacity(count);
718 let growing = self.growing.read();
719 let sealed = self.sealed.read();
720 if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
721 continue;
722 }
723 for segment in growing.iter().chain(sealed.iter()) {
724 if !before_work() {
725 return Ok(false);
726 }
727 let view = segment.read();
728 if !view.may_contain_graph_kind(kind) {
729 continue;
730 }
731 let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
732 return Ok(false);
733 };
734 sources.push((Arc::clone(segment), cursor));
735 }
736 break sources;
737 };
738 let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
739 for (source, mut cursor) in sources {
740 while !cursor.finished(keys) {
741 let mut count = 0;
742 {
743 let segment = source.read();
744 if !segment.may_contain_graph_kind(kind) {
745 break;
746 }
747 if !segment.graph_index_batch(
748 kind,
749 keys,
750 &mut cursor,
751 &mut before_work,
752 &mut |id| {
753 ids[count] = id;
754 count += 1;
755 },
756 ) {
757 return Ok(false);
758 }
759 }
760 if !before_work() || !consume(&ids[..count])? {
761 return Ok(false);
762 }
763 }
764 }
765 Ok(true)
766 }
767
768 pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
773 let schema = self.column_schema();
774 let clone_row = |entity: &UnifiedEntity| {
775 let mut owned = entity.clone();
776 if let super::entity::EntityData::Row(row) = &mut owned.data {
777 if row.schema.is_none() {
778 row.schema = schema.clone();
779 }
780 }
781 owned
782 };
783 let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
784 let mut remaining: Vec<usize> = (0..ids.len()).collect(); if let Some(growing_arc) = self.growing.read().as_ref() {
789 let growing = if let Some(g) = growing_arc.try_read() {
790 g
791 } else {
792 growing_arc.read()
793 };
794 remaining.retain(|&i| {
795 if let Some(entity) = growing.get(ids[i]) {
796 out[i] = Some(clone_row(entity));
797 false } else {
799 true }
801 });
802 }
803
804 if remaining.is_empty() {
805 self.restore_owned_row_schemas(out.iter_mut().flatten());
806 return out;
807 }
808
809 let sealed = self.sealed.read();
811 for segment in sealed.iter() {
812 if remaining.is_empty() {
813 break;
814 }
815 let seg = segment.read();
816 remaining.retain(|&i| {
817 if let Some(entity) = seg.get(ids[i]) {
818 out[i] = Some(clone_row(entity));
819 false
820 } else {
821 true
822 }
823 });
824 }
825
826 drop(sealed);
827 self.restore_owned_row_schemas(out.iter_mut().flatten());
828 out
829 }
830
831 pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
844 where
845 F: FnMut(usize, &UnifiedEntity),
846 {
847 thread_local! {
856 static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
857 const { std::cell::Cell::new(Vec::new()) };
858 }
859
860 let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
861 remaining.clear();
862
863 if let Some(growing_arc) = self.growing.read().as_ref() {
864 let growing = if let Some(g) = growing_arc.try_read() {
865 g
866 } else {
867 growing_arc.read()
868 };
869 for (i, id) in ids.iter().enumerate() {
870 if let Some(entity) = growing.get(*id) {
871 f(i, entity);
872 } else {
873 remaining.push(i);
874 }
875 }
876 } else {
877 remaining.reserve(ids.len());
878 remaining.extend(0..ids.len());
879 }
880
881 if !remaining.is_empty() {
882 let sealed = self.sealed.read();
883 for segment in sealed.iter() {
884 if remaining.is_empty() {
885 break;
886 }
887 let seg = segment.read();
888 remaining.retain(|&i| {
889 if let Some(entity) = seg.get(ids[i]) {
890 f(i, entity);
891 false
892 } else {
893 true
894 }
895 });
896 }
897 }
898
899 REMAINING_SCRATCH.with(|cell| cell.set(remaining));
900 }
901
902 fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
904 if let Some(growing_arc) = self.growing.read().as_ref() {
906 let growing = growing_arc.read();
907 if let Some(entity) = growing.get(id) {
908 return Some(entity.clone());
909 }
910 }
911
912 let sealed = self.sealed.read();
914 for segment in sealed.iter() {
915 if let Some(entity) = segment.get(id) {
916 return Some(entity.clone());
917 }
918 }
919
920 None
921 }
922
923 fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
924 let sealed = self.sealed.read();
925 sealed
926 .iter()
927 .find(|segment_arc| segment_arc.read().contains(id))
928 .map(Arc::clone)
929 }
930
931 fn rewrite_sealed_entity_into_growing(
932 &self,
933 entity: UnifiedEntity,
934 metadata: Option<&Metadata>,
935 ) -> Result<(), SegmentError> {
936 let entity_id = entity.id;
937 let sealed_arc = self
938 .find_sealed_segment_arc(entity_id)
939 .ok_or(SegmentError::NotFound(entity_id))?;
940
941 let metadata_to_apply = {
942 let mut sealed = sealed_arc.write();
943 let existing_metadata = sealed.get_metadata(entity_id);
944 if !sealed.force_delete(entity_id) {
945 return Err(SegmentError::NotFound(entity_id));
946 }
947 metadata.cloned().or(existing_metadata)
948 };
949
950 let growing_arc = self.get_or_create_growing();
951 let growing_id = {
952 let mut growing = growing_arc.write();
953 growing.insert(entity)?;
954 if let Some(metadata) = metadata_to_apply {
955 growing.set_metadata(entity_id, metadata)?;
956 }
957 growing.id()
958 };
959
960 self.entity_segment.write().insert(entity_id, growing_id);
961 Ok(())
962 }
963
964 pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
966 let entity_id = entity.id;
967 let mut entity = Some(entity);
968
969 if let Some(growing_arc) = self.growing.read().as_ref() {
971 let mut growing = growing_arc.write();
972 if growing.contains(entity_id) && growing.state().is_writable() {
973 return growing.update(entity.take().expect("entity already moved"));
974 }
975 }
976
977 let segment_id = self.entity_segment.read().get(&entity_id).copied();
979 if let Some(seg_id) = segment_id {
980 if let Some(growing_arc) = self.growing.read().as_ref() {
981 let mut growing = growing_arc.write();
982 if growing.id() == seg_id && growing.state().is_writable() {
983 return growing.update(entity.take().expect("entity already moved"));
984 }
985 }
986 }
987
988 if let Some(entity) = entity.take() {
989 return self.rewrite_sealed_entity_into_growing(entity, None);
990 }
991
992 Err(SegmentError::NotFound(entity_id))
993 }
994
995 pub fn update_with_metadata(
998 &self,
999 entity: UnifiedEntity,
1000 metadata: Option<&Metadata>,
1001 ) -> Result<(), SegmentError> {
1002 let entity_id = entity.id;
1003 let mut entity = Some(entity);
1004
1005 if let Some(growing_arc) = self.growing.read().as_ref() {
1007 let mut growing = growing_arc.write();
1008 if growing.contains(entity_id) && growing.state().is_writable() {
1009 growing.update(entity.take().expect("entity already moved"))?;
1010 if let Some(metadata) = metadata {
1011 growing.set_metadata(entity_id, metadata.clone())?;
1012 }
1013 return Ok(());
1014 }
1015 }
1016
1017 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1019 if let Some(seg_id) = segment_id {
1020 if let Some(growing_arc) = self.growing.read().as_ref() {
1021 let mut growing = growing_arc.write();
1022 if growing.id() == seg_id && growing.state().is_writable() {
1023 growing.update(entity.take().expect("entity already moved"))?;
1024 if let Some(metadata) = metadata {
1025 growing.set_metadata(entity_id, metadata.clone())?;
1026 }
1027 return Ok(());
1028 }
1029 }
1030 }
1031
1032 if let Some(entity) = entity.take() {
1033 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1034 }
1035
1036 Err(SegmentError::NotFound(entity_id))
1037 }
1038
1039 pub fn update_hot(
1043 &self,
1044 entity: UnifiedEntity,
1045 modified_columns: &[String],
1046 ) -> Result<(), SegmentError> {
1047 let entity_id = entity.id;
1048 let mut entity = Some(entity);
1049
1050 if let Some(growing_arc) = self.growing.read().as_ref() {
1051 let mut growing = growing_arc.write();
1052 if growing.contains(entity_id) && growing.state().is_writable() {
1053 return growing.update_hot(
1054 entity.take().expect("entity already moved"),
1055 modified_columns,
1056 );
1057 }
1058 }
1059
1060 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1061 if let Some(seg_id) = segment_id {
1062 if let Some(growing_arc) = self.growing.read().as_ref() {
1063 let mut growing = growing_arc.write();
1064 if growing.id() == seg_id && growing.state().is_writable() {
1065 return growing.update_hot(
1066 entity.take().expect("entity already moved"),
1067 modified_columns,
1068 );
1069 }
1070 }
1071 }
1072
1073 if let Some(entity) = entity.take() {
1074 return self.rewrite_sealed_entity_into_growing(entity, None);
1075 }
1076
1077 Err(SegmentError::NotFound(entity_id))
1078 }
1079
1080 pub fn update_hot_with_metadata(
1083 &self,
1084 entity: UnifiedEntity,
1085 modified_columns: &[String],
1086 metadata: Option<&Metadata>,
1087 ) -> Result<(), SegmentError> {
1088 let entity_id = entity.id;
1089 let mut entity = Some(entity);
1090
1091 if let Some(growing_arc) = self.growing.read().as_ref() {
1092 let mut growing = growing_arc.write();
1093 if growing.contains(entity_id) && growing.state().is_writable() {
1094 growing.update_hot(
1095 entity.take().expect("entity already moved"),
1096 modified_columns,
1097 )?;
1098 if let Some(metadata) = metadata {
1099 growing.set_metadata(entity_id, metadata.clone())?;
1100 }
1101 return Ok(());
1102 }
1103 }
1104
1105 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1106 if let Some(seg_id) = segment_id {
1107 if let Some(growing_arc) = self.growing.read().as_ref() {
1108 let mut growing = growing_arc.write();
1109 if growing.id() == seg_id && growing.state().is_writable() {
1110 growing.update_hot(
1111 entity.take().expect("entity already moved"),
1112 modified_columns,
1113 )?;
1114 if let Some(metadata) = metadata {
1115 growing.set_metadata(entity_id, metadata.clone())?;
1116 }
1117 return Ok(());
1118 }
1119 }
1120 }
1121
1122 if let Some(entity) = entity.take() {
1123 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1124 }
1125
1126 Err(SegmentError::NotFound(entity_id))
1127 }
1128
1129 pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
1132 where
1133 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
1134 {
1135 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
1136 items.into_iter().collect();
1137 if items.is_empty() {
1138 return Ok(());
1139 }
1140
1141 if let Some(growing_arc) = self.growing.read().as_ref() {
1142 let mut growing = growing_arc.write();
1143 if growing.state().is_writable() {
1144 match growing.update_hot_batch_with_metadata(items.iter().copied()) {
1145 Ok(()) => return Ok(()),
1146 Err(SegmentError::NotFound(_)) => {}
1147 Err(other) => return Err(other),
1148 }
1149 }
1150 }
1151
1152 for (entity, modified_columns, metadata) in items {
1153 self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
1154 }
1155 Ok(())
1156 }
1157
1158 pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
1160 if let Some(growing_arc) = self.growing.read().as_ref() {
1163 let mut growing = growing_arc.write();
1164 if growing.contains(id) && growing.state().is_writable() {
1165 let seg_id = growing.id();
1166 let deleted = growing.delete(id)?;
1167 if deleted {
1168 self.entity_segment.write().remove(&id);
1169 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1170 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1171 }
1172 return Ok(deleted);
1173 }
1174 }
1175
1176 let segment_id = self.entity_segment.read().get(&id).copied();
1179
1180 if let Some(seg_id) = segment_id {
1181 if let Some(growing_arc) = self.growing.read().as_ref() {
1182 let mut growing = growing_arc.write();
1183 if growing.id() == seg_id && growing.state().is_writable() {
1184 let deleted = growing.delete(id)?;
1185 if deleted {
1186 self.entity_segment.write().remove(&id);
1187 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1188 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1189 }
1190 return Ok(deleted);
1191 }
1192 }
1193 }
1194
1195 {
1198 let sealed = self.sealed.read();
1199 for segment_arc in sealed.iter() {
1200 let mut seg = segment_arc.write();
1201 let seg_id = seg.id();
1202 if seg.contains(id) {
1203 let deleted = seg.force_delete(id);
1204 drop(seg);
1205 if deleted {
1206 self.entity_segment.write().remove(&id);
1207 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1208 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1209 }
1210 return Ok(deleted);
1211 }
1212 }
1213 }
1214
1215 Ok(false)
1216 }
1217
1218 pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1219 if ids.is_empty() {
1220 return Ok(Vec::new());
1221 }
1222
1223 let mut deleted_ids = Vec::with_capacity(ids.len());
1224
1225 if let Some(growing_arc) = self.growing.read().as_ref() {
1226 let mut growing = growing_arc.write();
1227 if growing.state().is_writable() {
1228 let seg_id = growing.id();
1229 let deleted = growing.delete_batch(ids)?;
1230 if !deleted.is_empty() {
1231 {
1232 let mut entity_segment = self.entity_segment.write();
1233 for id in &deleted {
1234 entity_segment.remove(id);
1235 }
1236 }
1237 self.total_entities_atomic
1238 .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
1239 for id in &deleted {
1240 self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1241 }
1242 deleted_ids.extend(deleted);
1243 }
1244 }
1245 }
1246
1247 if deleted_ids.len() == ids.len() {
1248 return Ok(deleted_ids);
1249 }
1250
1251 let deleted_set: std::collections::HashSet<EntityId> =
1252 deleted_ids.iter().copied().collect();
1253 for &id in ids {
1254 if deleted_set.contains(&id) {
1255 continue;
1256 }
1257 if self.delete(id)? {
1258 deleted_ids.push(id);
1259 }
1260 }
1261
1262 Ok(deleted_ids)
1263 }
1264
1265 pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1267 if let Some(growing_arc) = self.growing.read().as_ref() {
1269 let growing = growing_arc.read();
1270 if growing.contains(id) {
1271 return growing.get_metadata(id);
1272 }
1273 }
1274
1275 let segment_id = self.entity_segment.read().get(&id).copied();
1277
1278 if let Some(seg_id) = segment_id {
1279 if let Some(growing_arc) = self.growing.read().as_ref() {
1280 let growing = growing_arc.read();
1281 if growing.id() == seg_id {
1282 return growing.get_metadata(id);
1283 }
1284 }
1285
1286 let sealed = self.sealed.read();
1287 for segment in sealed.iter() {
1288 if segment.id() == seg_id {
1289 return segment.get_metadata(id);
1290 }
1291 }
1292 }
1293
1294 if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1295 return segment_arc.read().get_metadata(id);
1296 }
1297
1298 None
1299 }
1300
1301 pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1303 if let Some(growing_arc) = self.growing.read().as_ref() {
1306 let mut growing = growing_arc.write();
1307 if growing.contains(id) && growing.state().is_writable() {
1308 return growing.set_metadata(id, metadata);
1309 }
1310 }
1311
1312 let segment_id = self.entity_segment.read().get(&id).copied();
1314
1315 if let Some(seg_id) = segment_id {
1316 if let Some(growing_arc) = self.growing.read().as_ref() {
1317 let mut growing = growing_arc.write();
1318 if growing.id() == seg_id && growing.state().is_writable() {
1319 return growing.set_metadata(id, metadata);
1320 }
1321 }
1322 }
1323
1324 if let Some(entity) = self.get(id) {
1325 return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1326 }
1327
1328 Err(SegmentError::NotFound(id))
1329 }
1330
1331 fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1333 let should_seal = {
1334 let growing_opt = self.growing.read();
1335 if let Some(growing_arc) = growing_opt.as_ref() {
1336 let growing = growing_arc.read();
1337 growing.should_seal(&self.config.segment_config)
1338 || growing.idle_secs() >= self.config.idle_seal_secs
1339 } else {
1340 false
1341 }
1342 };
1343
1344 if should_seal {
1345 self.seal_current()?;
1346 }
1347
1348 Ok(())
1349 }
1350
1351 pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1353 let mut growing_slot = self.growing.write();
1356 if let Some(growing_arc) = growing_slot.take() {
1357 let mut growing = growing_arc.write();
1358 let seg_id = growing.id();
1359 let entity_count = growing.stats().entity_count as u64;
1360
1361 growing.seal()?;
1363
1364 drop(growing); self.sealed.write().push(growing_arc);
1370 drop(growing_slot);
1371
1372 self.mark_sealed_pages_visible(entity_count);
1374
1375 {
1377 let mut stats = self.stats.write();
1378 stats.growing_count = stats.growing_count.saturating_sub(1);
1379 stats.sealed_count += 1;
1380 stats.seal_ops += 1;
1381 }
1382
1383 self.emit(LifecycleEvent::SegmentSealed(seg_id));
1384
1385 return Ok(seg_id);
1386 }
1387
1388 Err(SegmentError::InvalidState(SegmentState::Sealed))
1389 }
1390
1391 pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1393 let has_growing = self.growing.read().is_some();
1394 if has_growing {
1395 Ok(Some(self.seal_current()?))
1396 } else {
1397 Ok(None)
1398 }
1399 }
1400
1401 pub fn all_visible_fraction(&self) -> f64 {
1411 const ROWS_PER_PAGE: u32 = 256;
1412 let sealed = self.sealed.read();
1413 if sealed.is_empty() {
1414 return 0.0;
1415 }
1416 let mut total_pages: u64 = 0;
1417 for seg_arc in sealed.iter() {
1418 let seg = seg_arc.read();
1419 let entity_count = seg.stats().entity_count as u64;
1420 let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1421 total_pages += pages;
1422 }
1423 if total_pages == 0 {
1424 return 0.0;
1425 }
1426 let visible = self.visibility_map.all_visible_count();
1427 (visible as f64 / total_pages as f64).min(1.0)
1428 }
1429
1430 fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1433 const ROWS_PER_PAGE: u32 = 256;
1434 let existing_visible = self.visibility_map.all_visible_count();
1435 let start_page = existing_visible as u32;
1437 let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1438 let end_page = start_page + new_pages as u32;
1439 self.visibility_map.mark_range_visible(start_page, end_page);
1440 }
1441
1442 pub(crate) fn find_unique_key_conflict(
1446 &self,
1447 columns: &[Vec<String>],
1448 key_index: usize,
1449 signatures: &[String],
1450 exclude_id: Option<EntityId>,
1451 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1452 ) -> Option<EntityId> {
1453 if let Some(growing) = self.growing.read().as_ref() {
1454 if let Some(id) = growing.read().find_unique_key_conflict(
1455 columns,
1456 key_index,
1457 signatures,
1458 exclude_id,
1459 reserves_key,
1460 ) {
1461 return Some(id);
1462 }
1463 }
1464 for segment in self.sealed.read().iter() {
1465 if let Some(id) = segment.read().find_unique_key_conflict(
1466 columns,
1467 key_index,
1468 signatures,
1469 exclude_id,
1470 reserves_key,
1471 ) {
1472 return Some(id);
1473 }
1474 }
1475 None
1476 }
1477
1478 pub fn for_each_entity<F>(&self, callback: F)
1484 where
1485 F: FnMut(&UnifiedEntity) -> bool,
1486 {
1487 self.for_each_entity_in_segments(|_| true, || true, callback);
1488 }
1489
1490 pub(crate) fn checkpoint_revisions(&self) -> Vec<(u64, u64)> {
1492 let growing = self.growing.read();
1493 let sealed = self.sealed.read();
1494 growing
1495 .iter()
1496 .chain(sealed.iter())
1497 .map(|segment| {
1498 let segment = segment.read();
1499 (segment.id(), segment.checkpoint_revision())
1500 })
1501 .collect()
1502 }
1503
1504 pub(crate) fn cached_physical_root(&self, seed: u64, compute: impl FnOnce() -> u64) -> u64 {
1506 let segments = self.checkpoint_revisions();
1507 let mut cache = self.physical_root_cache.lock();
1508 if let Some(previous) = cache.as_ref() {
1509 if previous.seed == seed && previous.segments == segments {
1510 return previous.root;
1511 }
1512 }
1513 let root = compute();
1514 *cache = Some(PhysicalRootCacheEntry {
1515 segments,
1516 seed,
1517 root,
1518 });
1519 root
1520 }
1521
1522 pub(crate) fn for_each_entity_in_id_order(
1527 &self,
1528 mut visit: impl FnMut(&UnifiedEntity) -> bool,
1529 ) {
1530 self.for_each_entity_source_in_id_order(|_, entity| visit(entity));
1531 }
1532
1533 pub(crate) fn page_in_id_order(
1535 &self,
1536 offset: usize,
1537 limit: usize,
1538 ) -> (usize, Vec<UnifiedEntity>) {
1539 let end = offset.saturating_add(limit);
1540 let mut total = 0;
1541 let mut items = Vec::with_capacity(limit.min(self.count()));
1542 self.for_each_entity_in_id_order(|entity| {
1543 if total >= offset && total < end {
1544 items.push(entity.clone());
1545 }
1546 total += 1;
1547 true
1548 });
1549 self.restore_owned_row_schemas(items.iter_mut());
1550 (total, items)
1551 }
1552
1553 pub(crate) fn for_each_entity_with_metadata_in_id_order(
1554 &self,
1555 mut visit: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1556 ) {
1557 self.for_each_entity_source_in_id_order(|segment, entity| {
1558 visit(entity, segment.get_metadata(entity.id))
1559 });
1560 }
1561
1562 fn for_each_entity_source_in_id_order(
1563 &self,
1564 mut visit: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
1565 ) {
1566 let segments: Vec<_> = {
1567 let growing = self.growing.read();
1568 let sealed = self.sealed.read();
1569 growing.iter().chain(sealed.iter()).cloned().collect()
1570 };
1571 let mut order = Vec::new();
1572 for (source, segment) in segments.iter().enumerate() {
1573 segment.read().for_each_fast(|entity| {
1574 order.push((entity.id, source));
1575 true
1576 });
1577 }
1578 order.sort_by_key(|(id, _)| id.raw());
1579 for (id, source) in order {
1580 let segment = segments[source].read();
1581 if let Some(entity) = segment.get(id) {
1582 if !visit(&segment, entity) {
1583 break;
1584 }
1585 }
1586 }
1587 }
1588
1589 pub(crate) fn for_each_entity_with_metadata(
1592 &self,
1593 mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1594 ) {
1595 if let Some(growing_arc) = self.growing.read().as_ref() {
1596 let growing = growing_arc.read();
1597 if !growing.for_each_fast(|entity| callback(entity, growing.get_metadata(entity.id))) {
1598 return;
1599 }
1600 }
1601 let sealed = self.sealed.read();
1602 for segment_arc in sealed.iter() {
1603 let segment = segment_arc.read();
1604 if !segment.for_each_fast(|entity| callback(entity, segment.get_metadata(entity.id))) {
1605 return;
1606 }
1607 }
1608 }
1609
1610 fn for_each_entity_in_segments<F>(
1611 &self,
1612 segment_filter: impl Fn(&GrowingSegment) -> bool,
1613 mut before_segment: impl FnMut() -> bool,
1614 mut callback: F,
1615 ) where
1616 F: FnMut(&UnifiedEntity) -> bool,
1617 {
1618 if let Some(growing_arc) = self.growing.read().as_ref() {
1622 if !before_segment() {
1623 return;
1624 }
1625 let growing = if let Some(g) = growing_arc.try_read() {
1626 g
1627 } else {
1628 growing_arc.read()
1629 };
1630 if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
1631 return;
1632 }
1633 }
1634
1635 let sealed = self.sealed.read();
1637 for segment_arc in sealed.iter() {
1638 if !before_segment() {
1639 return;
1640 }
1641 let segment = segment_arc.read();
1642 if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
1643 return;
1644 }
1645 }
1646 }
1647
1648 pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1664 &self,
1665 init: FInit,
1666 fold: FFold,
1667 reduce: FReduce,
1668 ) -> T
1669 where
1670 T: Send,
1671 FInit: Fn() -> T + Send + Sync,
1672 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1673 FReduce: Fn(T, T) -> T + Send + Sync,
1674 {
1675 use rayon::prelude::*;
1676
1677 let mut acc = init();
1680 if let Some(growing_arc) = self.growing.read().as_ref() {
1681 let growing = if let Some(g) = growing_arc.try_read() {
1682 g
1683 } else {
1684 growing_arc.read()
1685 };
1686 growing.for_each_fast(|entity| {
1687 acc = fold(std::mem::replace(&mut acc, init()), entity);
1688 true
1689 });
1690 }
1691
1692 let segments: Vec<_> = {
1696 let sealed = self.sealed.read();
1697 sealed.iter().cloned().collect()
1698 };
1699
1700 if segments.len() <= 1 {
1701 for seg_arc in &segments {
1702 let seg = seg_arc.read();
1703 seg.for_each_fast(|entity| {
1704 acc = fold(std::mem::replace(&mut acc, init()), entity);
1705 true
1706 });
1707 }
1708 return acc;
1709 }
1710
1711 let sealed_acc = segments
1712 .into_par_iter()
1713 .map(|seg_arc| {
1714 let mut local = init();
1715 let seg = seg_arc.read();
1716 seg.for_each_fast(|entity| {
1717 local = fold(std::mem::replace(&mut local, init()), entity);
1718 true
1719 });
1720 local
1721 })
1722 .reduce(&init, &reduce);
1723
1724 reduce(acc, sealed_acc)
1725 }
1726
1727 pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
1736 where
1737 F: FnMut(&UnifiedEntity) -> bool,
1738 {
1739 let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
1740 }
1741
1742 pub fn for_each_entity_zoned_with_stats<F>(
1743 &self,
1744 zone_preds: &[(&str, ZoneColPred<'_>)],
1745 mut callback: F,
1746 ) -> SegmentScanStats
1747 where
1748 F: FnMut(&UnifiedEntity) -> bool,
1749 {
1750 let mut scan_stats = SegmentScanStats::default();
1751 if let Some(growing_arc) = self.growing.read().as_ref() {
1757 scan_stats.segments_total += 1;
1758 scan_stats.segments_scanned += 1;
1759 let growing = if let Some(g) = growing_arc.try_read() {
1760 g
1761 } else {
1762 growing_arc.read()
1763 };
1764 if !growing.for_each_fast(&mut callback) {
1765 return scan_stats;
1766 }
1767 }
1768
1769 let sealed = self.sealed.read();
1771 for segment_arc in sealed.iter() {
1772 scan_stats.segments_total += 1;
1773 let segment = segment_arc.read();
1774 if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
1775 scan_stats.segments_pruned += 1;
1776 continue; }
1778 scan_stats.segments_scanned += 1;
1779 if !segment.for_each_fast(&mut callback) {
1780 return scan_stats;
1781 }
1782 }
1783 scan_stats
1784 }
1785
1786 pub fn query_all_zoned<F>(
1796 &self,
1797 zone_preds: &[(&str, ZoneColPred<'_>)],
1798 filter: F,
1799 ) -> Vec<UnifiedEntity>
1800 where
1801 F: Fn(&UnifiedEntity) -> bool + Sync,
1802 {
1803 self.query_all_zoned_with_stats(zone_preds, filter).0
1804 }
1805
1806 pub fn query_all_zoned_with_stats<F>(
1807 &self,
1808 zone_preds: &[(&str, ZoneColPred<'_>)],
1809 filter: F,
1810 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
1811 where
1812 F: Fn(&UnifiedEntity) -> bool + Sync,
1813 {
1814 let mut results = Vec::new();
1815 let mut scan_stats = SegmentScanStats::default();
1816
1817 if let Some(growing_arc) = self.growing.read().as_ref() {
1820 scan_stats.segments_total += 1;
1821 scan_stats.segments_scanned += 1;
1822 let growing = if let Some(g) = growing_arc.try_read() {
1823 g
1824 } else {
1825 growing_arc.read()
1826 };
1827 results.extend(growing.iter().filter(|e| filter(e)).cloned());
1828 }
1829
1830 let sealed = self.sealed.read();
1832 let surviving: Vec<_> = sealed
1834 .iter()
1835 .filter(|seg_arc| {
1836 scan_stats.segments_total += 1;
1837 if zone_preds.is_empty() {
1838 scan_stats.segments_scanned += 1;
1839 return true;
1840 }
1841 let seg = seg_arc.read();
1842 if seg.can_skip_zone_preds(zone_preds) {
1843 scan_stats.segments_pruned += 1;
1844 false
1845 } else {
1846 scan_stats.segments_scanned += 1;
1847 true
1848 }
1849 })
1850 .collect();
1851
1852 let use_parallel = surviving.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1853
1854 if use_parallel {
1855 let filter_ref = &filter;
1856 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1857 surviving
1858 .iter()
1859 .map(|segment| {
1860 s.spawn(move || {
1861 segment
1862 .read()
1863 .iter()
1864 .filter(|e| filter_ref(e))
1865 .cloned()
1866 .collect::<Vec<_>>()
1867 })
1868 })
1869 .collect::<Vec<_>>()
1870 .into_iter()
1871 .map(|handle| handle.join().unwrap_or_default())
1872 .collect()
1873 });
1874 for batch in segment_results {
1875 results.extend(batch);
1876 }
1877 } else {
1878 for segment_arc in surviving {
1879 let seg = segment_arc.read();
1880 results.extend(seg.iter().filter(|e| filter(e)).cloned());
1881 }
1882 }
1883
1884 (results, scan_stats)
1885 }
1886
1887 fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
1893 match snapshot {
1894 Some(ctx) => entity_visible_with_context(Some(ctx), entity),
1895 None => {
1896 !crate::runtime::ai::moderation::entity_moderation_hidden(entity)
1897 && entity.xmax == 0
1898 }
1899 }
1900 }
1901
1902 pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
1909 where
1910 F: Fn(&UnifiedEntity) -> bool + Sync,
1911 {
1912 self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
1913 }
1914
1915 pub(crate) fn project_key_entity<T>(
1920 &self,
1921 key: &str,
1922 snapshot: Option<&SnapshotContext>,
1923 versioned: bool,
1924 mut project: impl FnMut(&UnifiedEntity) -> T,
1925 ) -> Option<T> {
1926 let mut probe = |segment: &GrowingSegment, best: &mut Option<(u64, T)>| {
1927 if let Some(entity) = segment.find_key_entity(key, versioned, &mut |entity| {
1928 !versioned || Self::scan_entity_visible(snapshot, entity)
1929 }) {
1930 if best.as_ref().is_none_or(|(xmin, _)| entity.xmin >= *xmin) {
1931 *best = Some((entity.xmin, project(entity)));
1932 }
1933 }
1934 };
1935 let mut best = None;
1936 if let Some(growing) = self.growing.read().as_ref() {
1937 probe(&growing.read(), &mut best);
1938 }
1939 for segment in self.sealed.read().iter() {
1940 if !versioned && best.is_some() {
1941 break;
1942 }
1943 probe(&segment.read(), &mut best);
1944 }
1945 best.map(|(_, projected)| projected)
1946 }
1947
1948 pub(crate) fn scan_graph_kind(
1950 &self,
1951 snapshot: Option<&SnapshotContext>,
1952 kind: GraphEntityKind,
1953 ) -> Vec<UnifiedEntity> {
1954 self.query_all_in_segments(
1955 |segment| segment.may_contain_graph_kind(kind),
1956 |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
1957 )
1958 }
1959
1960 pub(crate) fn scan_graph_candidates_for_each<F>(
1965 &self,
1966 snapshot: Option<&SnapshotContext>,
1967 kind: GraphEntityKind,
1968 before_segment: impl FnMut() -> bool,
1969 mut callback: F,
1970 ) where
1971 F: FnMut(&UnifiedEntity) -> bool,
1972 {
1973 self.for_each_entity_in_segments(
1974 |segment| segment.may_contain_graph_kind(kind),
1975 before_segment,
1976 |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
1977 );
1978 }
1979
1980 pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
1985 where
1986 F: FnMut(&UnifiedEntity) -> bool,
1987 {
1988 self.for_each_entity(|entity| {
1989 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
1990 });
1991 }
1992
1993 pub(crate) fn physical_scan_cursor(&self) -> PhysicalScanCursor {
1994 let growing = self.growing.read();
1995 let sealed = self.sealed.read();
1996 PhysicalScanCursor {
1997 sources: growing
1998 .iter()
1999 .chain(sealed.iter())
2000 .map(|segment| (Arc::downgrade(segment), segment.read().scan_cursor()))
2001 .collect(),
2002 source: 0,
2003 }
2004 }
2005
2006 pub(crate) fn physical_scan_batch(
2009 &self,
2010 cursor: &mut PhysicalScanCursor,
2011 work_left: &mut usize,
2012 ) -> Vec<EntityId> {
2013 let mut ids = Vec::with_capacity(super::segment::SCAN_BATCH_SIZE.min(*work_left));
2014 while !cursor.finished() && *work_left > 0 {
2015 let (source, position) = &mut cursor.sources[cursor.source];
2016 if position.finished() {
2017 *work_left -= 1;
2018 cursor.source += 1;
2019 continue;
2020 }
2021 let Some(segment) = source.upgrade() else {
2022 *work_left -= 1;
2023 cursor.source += 1;
2024 continue;
2025 };
2026 segment.read().scan_batch(
2027 position,
2028 &mut || {
2029 if *work_left == 0 {
2030 return false;
2031 }
2032 *work_left -= 1;
2033 true
2034 },
2035 &mut |id| ids.push(id),
2036 );
2037 if position.finished() {
2038 cursor.source += 1;
2039 }
2040 if !ids.is_empty() {
2041 break;
2042 }
2043 }
2044 ids
2045 }
2046
2047 pub(crate) fn scan_batches(
2053 &self,
2054 snapshot: Option<&SnapshotContext>,
2055 mut before_work: impl FnMut() -> bool,
2056 mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
2057 ) -> bool {
2058 let segments = {
2059 let growing = self.growing.read();
2060 let sealed = self.sealed.read();
2061 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
2062 for segment in growing.iter().chain(sealed.iter()) {
2063 if !before_work() {
2064 return false;
2065 }
2066 let cursor = segment.read().scan_cursor();
2067 segments.push((Arc::clone(segment), cursor));
2068 }
2069 segments
2070 };
2071 let mut ids = Vec::new();
2072 let mut batch = Vec::new();
2073 for (segment, mut cursor) in segments {
2074 while !cursor.finished() {
2075 ids.clear();
2076 batch.clear();
2077 let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
2078 ids.reserve_exact(batch_size);
2079 batch.reserve_exact(batch_size);
2080 {
2081 let segment = segment.read();
2082 if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
2083 return false;
2084 }
2085 }
2086 if !before_work() {
2089 return false;
2090 }
2091 self.for_each_id(&ids, |_, entity| {
2092 if Self::scan_entity_visible(snapshot, entity) {
2093 batch.push(entity.clone());
2094 }
2095 });
2096 if !before_work() {
2097 return false;
2098 }
2099 if !consume(&batch) {
2100 return false;
2101 }
2102 }
2103 }
2104 true
2105 }
2106
2107 pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
2109 &self,
2110 snapshot: Option<&SnapshotContext>,
2111 init: FInit,
2112 fold: FFold,
2113 reduce: FReduce,
2114 ) -> T
2115 where
2116 T: Send,
2117 FInit: Fn() -> T + Send + Sync,
2118 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
2119 FReduce: Fn(T, T) -> T + Send + Sync,
2120 {
2121 self.fold_entities_parallel(
2122 init,
2123 |acc, entity| {
2124 if Self::scan_entity_visible(snapshot, entity) {
2125 fold(acc, entity)
2126 } else {
2127 acc
2128 }
2129 },
2130 reduce,
2131 )
2132 }
2133
2134 pub fn scan_for_each_zoned_with_stats<F>(
2136 &self,
2137 snapshot: Option<&SnapshotContext>,
2138 zone_preds: &[(&str, ZoneColPred<'_>)],
2139 mut callback: F,
2140 ) -> SegmentScanStats
2141 where
2142 F: FnMut(&UnifiedEntity) -> bool,
2143 {
2144 self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
2145 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2146 })
2147 }
2148
2149 pub fn scan_zoned_with_stats<F>(
2151 &self,
2152 snapshot: Option<&SnapshotContext>,
2153 zone_preds: &[(&str, ZoneColPred<'_>)],
2154 filter: F,
2155 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2156 where
2157 F: Fn(&UnifiedEntity) -> bool + Sync,
2158 {
2159 self.query_all_zoned_with_stats(zone_preds, |entity| {
2160 Self::scan_entity_visible(snapshot, entity) && filter(entity)
2161 })
2162 }
2163
2164 pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
2167 where
2168 F: Fn(&UnifiedEntity) -> bool + Sync,
2169 {
2170 let mut entities = self.query_all_in_segments(|_| true, filter);
2171 self.restore_owned_row_schemas(entities.iter_mut());
2172 entities
2173 }
2174
2175 fn query_all_in_segments<F>(
2176 &self,
2177 segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
2178 filter: F,
2179 ) -> Vec<UnifiedEntity>
2180 where
2181 F: Fn(&UnifiedEntity) -> bool + Sync,
2182 {
2183 let mut results = Vec::new();
2184
2185 if let Some(growing_arc) = self.growing.read().as_ref() {
2188 let growing = if let Some(g) = growing_arc.try_read() {
2189 g
2190 } else {
2191 growing_arc.read()
2192 };
2193 if segment_filter(&growing) {
2194 results.extend(growing.iter().filter(|e| filter(e)).cloned());
2195 }
2196 }
2197
2198 let sealed = self.sealed.read();
2200 let use_parallel = sealed.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
2201 if use_parallel {
2202 let filter_ref = &filter;
2203 let segment_filter_ref = &segment_filter;
2204 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
2205 sealed
2206 .iter()
2207 .map(|segment| {
2208 s.spawn(move || {
2209 let segment = segment.read();
2210 if segment_filter_ref(&segment) {
2211 segment
2212 .iter()
2213 .filter(|e| filter_ref(e))
2214 .cloned()
2215 .collect::<Vec<_>>()
2216 } else {
2217 Vec::new()
2218 }
2219 })
2220 })
2221 .collect::<Vec<_>>()
2222 .into_iter()
2223 .map(|handle| handle.join().unwrap_or_default())
2224 .collect()
2225 });
2226 for batch in segment_results {
2227 results.extend(batch);
2228 }
2229 } else {
2230 for segment in sealed.iter() {
2231 let seg = segment.read();
2232 if segment_filter(&seg) {
2233 results.extend(seg.iter().filter(|e| filter(e)).cloned());
2234 }
2235 }
2236 }
2237
2238 results
2239 }
2240
2241 pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
2246 if let Some(growing_arc) = self.growing.read().as_ref() {
2247 let growing = growing_arc.read();
2248 if growing.may_contain_exact_key(key) {
2249 return true;
2250 }
2251 } else {
2252 return true;
2253 }
2254
2255 let sealed = self.sealed.read();
2256 for segment_arc in sealed.iter() {
2257 let segment = segment_arc.read();
2258 if segment.may_contain_exact_key(key) {
2259 return true;
2260 }
2261 }
2262
2263 false
2264 }
2265
2266 pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2268 let mut results = Vec::new();
2269
2270 if let Some(growing_arc) = self.growing.read().as_ref() {
2272 let growing = growing_arc.read();
2273 results.extend(growing.filter_metadata(filters));
2274 }
2275
2276 let sealed = self.sealed.read();
2278 for segment in sealed.iter() {
2279 results.extend(segment.filter_metadata(filters));
2280 }
2281
2282 results
2283 }
2284
2285 pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2287 let mut results = Vec::new();
2288
2289 if let Some(growing_arc) = self.growing.read().as_ref() {
2291 let growing = growing_arc.read();
2292 for entity in growing.iter_kind(kind) {
2293 results.push(entity.clone());
2294 }
2295 }
2296
2297 let sealed = self.sealed.read();
2302 for segment_arc in sealed.iter() {
2303 let segment = segment_arc.read();
2304 for entity in segment.iter_kind(kind) {
2305 results.push(entity.clone());
2306 }
2307 }
2308
2309 results
2310 }
2311
2312 pub fn count(&self) -> usize {
2314 self.total_entities_atomic.load(Ordering::Relaxed) as usize
2315 }
2316
2317 pub fn segment_ids(&self) -> Vec<SegmentId> {
2319 let mut ids = Vec::new();
2320
2321 if let Some(growing_arc) = self.growing.read().as_ref() {
2322 ids.push(growing_arc.read().id());
2323 }
2324
2325 let sealed = self.sealed.read();
2326 for segment in sealed.iter() {
2327 ids.push(segment.id());
2328 }
2329
2330 ids.extend(self.archived.read().iter().copied());
2331
2332 ids
2333 }
2334
2335 #[inline]
2348 #[allow(clippy::unused_self)]
2349 fn emit(&self, _event: LifecycleEvent) {}
2350
2351 pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2354 std::mem::take(&mut *self.events.write())
2355 }
2356
2357 pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2359 self.maybe_seal_growing()?;
2361
2362 if self.config.enable_consolidation {
2363 self.consolidation_tick();
2364 }
2365
2366 Ok(())
2367 }
2368
2369 fn consolidation_tick(&self) {
2380 let running = self.consolidation.read().is_some();
2383 if !running && !self.start_consolidation() {
2384 return;
2385 }
2386
2387 let mut guard = self.consolidation.write();
2390 let Some(run) = guard.as_mut() else {
2391 return;
2392 };
2393 self.copy_bounded(run);
2394 if run.cursor >= run.sources.len() {
2395 if let Some(run) = guard.take() {
2396 self.finish_consolidation(run);
2397 }
2398 }
2399 }
2400
2401 fn start_consolidation(&self) -> bool {
2405 let sources = self.select_candidates();
2406 self.start_consolidation_from_sources(sources)
2407 }
2408
2409 fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
2410 if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
2411 return false;
2412 }
2413
2414 let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
2415 let merged = GrowingSegment::new(merged_id, &self.collection);
2416
2417 *self.consolidation.write() = Some(Consolidation {
2418 sources: sources
2419 .into_iter()
2420 .map(|id| SourceProgress {
2421 id,
2422 copied: Vec::new(),
2423 tombstones_at_open: None,
2424 })
2425 .collect(),
2426 merged_id,
2427 merged,
2428 cursor: 0,
2429 pending_ids: Vec::new(),
2430 pending_cursor: 0,
2431 entities_copied: 0,
2432 });
2433
2434 self.stats.write().consolidation.runs_started += 1;
2435 true
2436 }
2437
2438 pub fn pressure_consolidation_tick(&self) -> bool {
2443 if !self.config.enable_consolidation {
2444 return false;
2445 }
2446
2447 let running = self.consolidation.read().is_some();
2448 if !running {
2449 let sources = self.dirty_sealed_segments();
2450 if !self.start_consolidation_from_sources(sources) {
2451 return false;
2452 }
2453 }
2454
2455 self.consolidation_tick();
2456 true
2457 }
2458
2459 fn select_candidates(&self) -> Vec<SegmentId> {
2474 let sealed = self.sealed.read();
2475 if sealed.is_empty() {
2476 return Vec::new();
2477 }
2478
2479 if sealed.len() > self.config.max_sealed_segments {
2480 return sealed.iter().map(|segment| segment.read().id()).collect();
2481 }
2482
2483 let mut live = 0usize;
2484 let mut tombstones = 0usize;
2485 let mut held = 0u64;
2486 let mut reclaimable = 0u64;
2487 let mut dirty = Vec::new();
2488
2489 for segment_arc in sealed.iter() {
2490 let segment = segment_arc.read();
2491 live += segment.live_entity_count();
2492 tombstones += segment.tombstone_count();
2493 held += segment.resident_bytes();
2494 reclaimable += segment.reclaimable_bytes();
2495 if segment.tombstone_count() > 0 {
2496 dirty.push(segment.id());
2497 }
2498 }
2499
2500 let total = live + tombstones;
2501 let tombstone_ratio = if total == 0 {
2502 0.0
2503 } else {
2504 tombstones as f64 / total as f64
2505 };
2506 let fragmentation = if held == 0 {
2507 0.0
2508 } else {
2509 reclaimable as f64 / held as f64
2510 };
2511
2512 if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
2513 || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
2514 {
2515 return dirty;
2516 }
2517
2518 Vec::new()
2519 }
2520
2521 fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
2522 self.sealed
2523 .read()
2524 .iter()
2525 .filter_map(|segment_arc| {
2526 let segment = segment_arc.read();
2527 (segment.reclaimable_bytes() > 0).then(|| segment.id())
2528 })
2529 .collect()
2530 }
2531
2532 fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
2535 if sources.is_empty() {
2536 return false;
2537 }
2538 let sealed = self.sealed.read();
2539 sealed.iter().any(|segment_arc| {
2540 let segment = segment_arc.read();
2541 sources.contains(&segment.id()) && segment.tombstone_count() > 0
2542 })
2543 }
2544
2545 fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
2546 let sealed = self.sealed.read();
2547 sealed
2548 .iter()
2549 .find(|segment| segment.read().id() == id)
2550 .map(Arc::clone)
2551 }
2552
2553 fn copy_bounded(&self, run: &mut Consolidation) {
2558 let mut budget = self.config.consolidation_entities_per_tick;
2559
2560 while budget > 0 && run.cursor < run.sources.len() {
2561 let source_id = run.sources[run.cursor].id;
2562 let Some(source_arc) = self.sealed_segment(source_id) else {
2563 run.cursor += 1;
2567 run.pending_ids.clear();
2568 run.pending_cursor = 0;
2569 continue;
2570 };
2571
2572 if run.sources[run.cursor].tombstones_at_open.is_none() {
2573 let source = source_arc.read();
2574 run.pending_ids = source.live_entity_ids();
2575 run.pending_cursor = 0;
2576 run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
2577 }
2578
2579 {
2580 let source = source_arc.read();
2581 while budget > 0 && run.pending_cursor < run.pending_ids.len() {
2582 let id = run.pending_ids[run.pending_cursor];
2583 run.pending_cursor += 1;
2584 budget -= 1;
2585
2586 let Some(entity) = source.get(id) else {
2589 continue;
2590 };
2591 let metadata = source.get_metadata(id);
2592 run.merged.adopt_entity(entity.clone(), metadata);
2593 run.sources[run.cursor].copied.push(id);
2594 run.entities_copied += 1;
2595 }
2596 }
2597
2598 if run.pending_cursor >= run.pending_ids.len() {
2599 run.cursor += 1;
2600 run.pending_ids.clear();
2601 run.pending_cursor = 0;
2602 }
2603 }
2604 }
2605
2606 fn finish_consolidation(&self, mut run: Consolidation) {
2612 let mut sealed = self.sealed.write();
2613
2614 let mut positions = Vec::with_capacity(run.sources.len());
2615 for source in &run.sources {
2616 if let Some(index) = sealed
2617 .iter()
2618 .position(|segment| segment.read().id() == source.id)
2619 {
2620 positions.push(index);
2621 }
2622 }
2623 if positions.is_empty() {
2624 return;
2625 }
2626 positions.sort_unstable();
2627
2628 let mut stale = Vec::new();
2632 for source in &run.sources {
2633 let Some(source_arc) = sealed
2634 .iter()
2635 .find(|segment| segment.read().id() == source.id)
2636 else {
2637 continue;
2638 };
2639 let segment = source_arc.read();
2640 if source.tombstones_at_open == Some(segment.tombstone_count()) {
2641 continue;
2642 }
2643 stale.extend(
2644 source
2645 .copied
2646 .iter()
2647 .copied()
2648 .filter(|id| !segment.contains(*id)),
2649 );
2650 }
2651 for id in stale {
2652 run.merged.evict_entity(id);
2653 }
2654
2655 let reclaimed_bytes: u64 = positions
2656 .iter()
2657 .map(|&index| sealed[index].read().resident_bytes())
2658 .sum();
2659 let reclaimed_tombstones: u64 = positions
2660 .iter()
2661 .map(|&index| sealed[index].read().tombstone_count() as u64)
2662 .sum();
2663
2664 if run.merged.seal().is_err() {
2665 return;
2666 }
2667 let merged_bytes = run.merged.resident_bytes();
2668 let merged_arc = Arc::new(RwLock::new(run.merged));
2669
2670 {
2671 let mut retired = self.retired.write();
2672 retired.retain(|source| source.strong_count() > 0);
2673 for &index in positions.iter().rev() {
2674 let source = sealed.remove(index);
2675 if Arc::strong_count(&source) > 1 {
2676 retired.push(Arc::downgrade(&source));
2677 }
2678 }
2679 }
2680 let insert_at = positions[0].min(sealed.len());
2681 sealed.insert(insert_at, Arc::clone(&merged_arc));
2682 drop(sealed);
2683
2684 let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
2686 {
2687 let mut entity_segment = self.entity_segment.write();
2688 for segment_id in entity_segment.values_mut() {
2689 if source_ids.contains(segment_id) {
2690 *segment_id = run.merged_id;
2691 }
2692 }
2693 }
2694
2695 {
2696 let mut stats = self.stats.write();
2697 stats.sealed_count = stats
2698 .sealed_count
2699 .saturating_sub(positions.len())
2700 .saturating_add(1);
2701 let consolidation = &mut stats.consolidation;
2702 consolidation.runs_completed += 1;
2703 consolidation.segments_merged += positions.len() as u64;
2704 consolidation.tombstones_reclaimed += reclaimed_tombstones;
2705 consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
2706 }
2707
2708 self.emit(LifecycleEvent::SegmentConsolidated {
2709 source: source_ids,
2710 target: run.merged_id,
2711 });
2712 }
2713
2714 #[cfg(test)]
2717 fn consolidation_progress(&self) -> Option<u64> {
2718 self.consolidation
2719 .read()
2720 .as_ref()
2721 .map(|run| run.entities_copied)
2722 }
2723}
2724
2725impl UnifiedSegment for Arc<RwLock<GrowingSegment>> {
2728 fn id(&self) -> SegmentId {
2729 self.read().id()
2730 }
2731
2732 fn state(&self) -> SegmentState {
2733 self.read().state()
2734 }
2735
2736 fn collection(&self) -> &str {
2737 "unknown"
2739 }
2740
2741 fn stats(&self) -> SegmentStats {
2742 self.read().stats()
2743 }
2744
2745 fn entity_count(&self) -> usize {
2746 self.read().entity_count()
2747 }
2748
2749 fn contains(&self, id: EntityId) -> bool {
2750 self.read().contains(id)
2751 }
2752
2753 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2754 None
2757 }
2758
2759 fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2760 None
2761 }
2762
2763 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2764 self.write().insert(entity)
2765 }
2766
2767 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2768 self.write().update(entity)
2769 }
2770
2771 fn update_hot(
2772 &mut self,
2773 entity: UnifiedEntity,
2774 modified_columns: &[String],
2775 ) -> Result<(), SegmentError> {
2776 self.write().update_hot(entity, modified_columns)
2777 }
2778
2779 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2780 self.write().delete(id)
2781 }
2782
2783 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2784 self.read().get_metadata(id)
2785 }
2786
2787 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2788 self.write().set_metadata(id, metadata)
2789 }
2790
2791 fn seal(&mut self) -> Result<(), SegmentError> {
2792 self.write().seal()
2793 }
2794
2795 fn should_seal(&self, config: &SegmentConfig) -> bool {
2796 self.read().should_seal(config)
2797 }
2798
2799 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2800 Box::new(std::iter::empty())
2802 }
2803
2804 fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2805 Box::new(std::iter::empty())
2806 }
2807
2808 fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2809 self.read().filter_metadata(filters)
2810 }
2811}
2812
2813#[cfg(test)]
2814mod tests {
2815 use std::collections::HashSet;
2816
2817 use super::*;
2818 use crate::runtime::mvcc::SnapshotContext;
2819 use crate::storage::transaction::snapshot::{Snapshot, SnapshotManager};
2820 use crate::storage::unified::entity::{EntityData, EntityKind, RowData};
2821 use reddb_types::Value;
2822
2823 #[test]
2824 fn physical_pages_keep_global_order_across_segments_and_deleted_slots() {
2825 let manager = SegmentManager::new("pages");
2826 for group in [[50, 10], [30, 20]] {
2827 for id in group {
2828 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "pages", id, vec![]);
2829 entity.data = EntityData::Row(RowData::with_names(
2830 vec![Value::Integer(id as i64)],
2831 vec!["id".into()],
2832 ));
2833 manager.insert(entity).unwrap();
2834 }
2835 manager.force_seal().unwrap();
2836 }
2837 manager
2838 .insert(UnifiedEntity::table_row(
2839 EntityId::new(40),
2840 "pages",
2841 40,
2842 vec![],
2843 ))
2844 .unwrap();
2845 let (total, page) = manager.page_in_id_order(1, 2);
2846 assert_eq!(total, 5);
2847 assert_eq!(
2848 page.iter()
2849 .map(|entity| entity.id.raw())
2850 .collect::<Vec<_>>(),
2851 [20, 30]
2852 );
2853 let EntityData::Row(row) = &page[0].data else {
2854 panic!("row")
2855 };
2856 assert_eq!(row.get_field("id"), Some(&Value::Integer(20)));
2857 manager.delete(EntityId::new(30)).unwrap();
2858 let (total, page) = manager.page_in_id_order(1, 2);
2859 assert_eq!(total, 4);
2860 assert_eq!(
2861 page.iter()
2862 .map(|entity| entity.id.raw())
2863 .collect::<Vec<_>>(),
2864 [20, 40]
2865 );
2866 for offset in [4, usize::MAX] {
2867 let (total, page) = manager.page_in_id_order(offset, 2);
2868 assert_eq!(total, 4);
2869 assert!(page.is_empty());
2870 }
2871 }
2872
2873 fn lookup_version(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2874 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2875 entity.data = EntityData::Row(RowData::with_names(
2876 vec![Value::text(key)],
2877 vec!["key".into()],
2878 ));
2879 entity.set_xmin(xmin);
2880 entity.set_xmax(xmax);
2881 entity
2882 }
2883
2884 #[test]
2885 fn key_lookup_matches_snapshot_scan_across_segments() {
2886 let manager = Arc::new(SegmentManager::new("keys"));
2887 for entity in [
2888 lookup_version(1, "changed", 3, 12),
2889 lookup_version(2, "deleted", 3, 7),
2890 ] {
2891 manager.insert(entity).expect("old version");
2892 }
2893 manager.force_seal().expect("seal old versions");
2894 for entity in [
2895 lookup_version(3, "changed", 12, 0),
2896 lookup_version(4, "pending", 4, 0),
2897 lookup_version(5, "own", 15, 0),
2898 lookup_version(6, "live", 3, 0),
2899 ] {
2900 manager.insert(entity).expect("new version");
2901 }
2902 std::thread::spawn(move || {
2903 let snapshot = SnapshotContext {
2904 snapshot: Snapshot {
2905 xid: 10,
2906 in_progress: HashSet::from([4]),
2907 },
2908 manager: Arc::new(SnapshotManager::new()),
2909 own_xids: HashSet::from([15]),
2910 requires_index_fallback: false,
2911 serializable_reader: None,
2912 };
2913 for snapshot in [Some(&snapshot), None] {
2914 for key in ["changed", "deleted", "pending", "own", "live", "absent"] {
2915 let expected = manager
2916 .scan(snapshot, |entity| {
2917 super::super::segment_key_lookup::SegmentKeyLookup::key(entity)
2918 == Some(key)
2919 })
2920 .into_iter()
2921 .max_by_key(|entity| entity.xmin)
2922 .map(|entity| entity.id);
2923 assert_eq!(
2924 manager.project_key_entity(key, snapshot, true, |entity| entity.id),
2925 expected,
2926 "key={key}"
2927 );
2928 }
2929 }
2930 assert_eq!(
2931 manager.project_key_entity("changed", Some(&snapshot), true, |entity| entity.id),
2932 Some(EntityId::new(1))
2933 );
2934 assert_eq!(
2935 manager.project_key_entity("changed", None, true, |entity| entity.id),
2936 Some(EntityId::new(3))
2937 );
2938 assert!(manager
2939 .project_key_entity("deleted", None, true, |entity| entity.id)
2940 .is_none());
2941 })
2942 .join()
2943 .expect("snapshot worker");
2944 }
2945
2946 #[test]
2947 fn key_lookup_survives_paced_consolidation() {
2948 let manager = consolidating_manager(1);
2949 manager
2950 .bulk_insert(
2951 (1..=10)
2952 .map(|id| lookup_version(id, &format!("key-{id}"), 1, 0))
2953 .collect(),
2954 )
2955 .expect("rows");
2956 manager.force_seal().expect("seal");
2957 for id in 1..=4 {
2958 manager.delete(EntityId::new(id)).expect("delete");
2959 }
2960 assert!(manager.pressure_consolidation_tick());
2961 drain_maintenance(&manager);
2962 for id in 1..=10 {
2963 assert_eq!(
2964 manager.project_key_entity(&format!("key-{id}"), None, false, |entity| entity.id),
2965 (id > 4).then_some(EntityId::new(id))
2966 );
2967 }
2968 }
2969
2970 #[test]
2971 #[ignore = "diagnostic of segment fragmentation and retained history; no timing assertions"]
2972 fn key_lookup_fragmentation_history_diagnostic() {
2973 use std::{hint::black_box, time::Instant};
2974 for segments in [1, 16, 64] {
2975 for versions in [1, 100, 1000] {
2976 let manager = SegmentManager::new("keys");
2977 let mut id = 1;
2978 for segment in 0..segments {
2979 let mut rows: Vec<_> = (0..128)
2980 .map(|_| {
2981 let entity = lookup_version(id, &format!("other-{id}"), 1, 0);
2982 id += 1;
2983 entity
2984 })
2985 .collect();
2986 for version in (segment..versions).step_by(segments) {
2987 rows.push(lookup_version(
2988 id,
2989 "history",
2990 version as u64 + 1,
2991 if version + 1 == versions {
2992 0
2993 } else {
2994 version as u64 + 2
2995 },
2996 ));
2997 id += 1;
2998 }
2999 manager.bulk_insert(rows).expect("rows");
3000 manager.force_seal().expect("seal");
3001 }
3002 let started = Instant::now();
3003 for _ in 0..2000 {
3004 let xmin =
3005 manager.project_key_entity(black_box("history"), None, true, |entity| {
3006 entity.xmin
3007 });
3008 assert_eq!(xmin, Some(versions as u64));
3009 }
3010 eprintln!("segments={segments} versions={versions} reads=2000 elapsed_us={} resident_bytes={}", started.elapsed().as_micros(), manager.resident_bytes());
3011 }
3012 }
3013 }
3014
3015 #[test]
3016 fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
3017 let manager = SegmentManager::new("fingerprint");
3018 let row = |id, label: &str| {
3019 UnifiedEntity::table_row(
3020 EntityId::new(id),
3021 "fingerprint",
3022 id,
3023 vec![Value::text(label)],
3024 )
3025 };
3026 manager
3027 .bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
3028 .expect("sealed rows");
3029 manager.force_seal().expect("seal");
3030 manager
3031 .bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
3032 .expect("growing rows");
3033
3034 let mut expected = manager.query_all(|_| true);
3035 expected.sort_by_key(|entity| entity.id.raw());
3036 let expected: Vec<_> = expected
3037 .iter()
3038 .map(|entity| format!("{entity:?}"))
3039 .collect();
3040 let mut stored_addresses = HashSet::new();
3041 manager.for_each_entity(|entity| {
3042 stored_addresses.insert(entity as *const UnifiedEntity as usize);
3043 true
3044 });
3045
3046 let mut actual = Vec::new();
3047 manager.for_each_entity_in_id_order(|entity| {
3048 assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
3049 assert!(
3050 manager.growing.try_write().is_some(),
3051 "topology guard retained"
3052 );
3053 assert!(
3054 manager.sealed.try_write().is_some(),
3055 "topology guard retained"
3056 );
3057 actual.push(format!("{entity:?}"));
3058 true
3059 });
3060 assert_eq!(
3061 actual, expected,
3062 "equal IDs retain their original segment order"
3063 );
3064 }
3065
3066 #[test]
3067 fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
3068 let manager = SegmentManager::new("fingerprint");
3069 for id in [3, 1, 2] {
3070 manager
3071 .insert(UnifiedEntity::table_row(
3072 EntityId::new(id),
3073 "fingerprint",
3074 id,
3075 vec![Value::Integer(id as i64)],
3076 ))
3077 .expect("row");
3078 }
3079 let source = manager.growing.read().as_ref().expect("source").clone();
3080 let handles = Arc::strong_count(&source);
3081 let mut seen = Vec::new();
3082 manager.for_each_entity_in_id_order(|entity| {
3083 seen.push(entity.id.raw());
3084 false
3085 });
3086 assert_eq!(seen, vec![1]);
3087 assert_eq!(Arc::strong_count(&source), handles);
3088 assert!(
3089 source.try_write().is_some(),
3090 "segment guard retained after early stop"
3091 );
3092 }
3093
3094 #[test]
3095 fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
3096 let manager = SegmentManager::with_config(
3097 "vectors",
3098 ManagerConfig {
3099 max_sealed_segments: 1,
3100 consolidation_entities_per_tick: 10_000,
3101 ..Default::default()
3102 },
3103 );
3104 for part in 0..3 {
3105 manager
3106 .bulk_insert(
3107 (1..=300)
3108 .map(|offset| {
3109 UnifiedEntity::vector(
3110 EntityId::new(part * 300 + offset),
3111 "vectors",
3112 vec![1.0, 0.0],
3113 )
3114 })
3115 .collect(),
3116 )
3117 .expect("batch");
3118 if part < 2 {
3119 manager.force_seal().expect("seal");
3120 }
3121 }
3122 let original_growing = manager.growing.read().as_ref().expect("growing").clone();
3123 let mut changed = false;
3124 let mut seen = Vec::new();
3125 assert!(manager.scan_batches(
3126 None,
3127 || true,
3128 |batch| {
3129 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3130 assert!(
3131 manager.growing.try_write().is_some(),
3132 "consumer holds no topology guard"
3133 );
3134 assert!(
3135 original_growing.try_write().is_some(),
3136 "consumer holds no segment guard"
3137 );
3138 for entity in batch {
3139 assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
3140 seen.push(entity.id.raw());
3141 }
3142 if !changed {
3143 changed = true;
3144 manager.force_seal().expect("seal during scan");
3145 manager
3146 .insert(UnifiedEntity::vector(
3147 EntityId::new(5000),
3148 "vectors",
3149 vec![1.0, 0.0],
3150 ))
3151 .expect("new write");
3152 drain_maintenance(&manager);
3153 assert!(manager.stats().consolidation.runs_completed > 0);
3154 assert!(manager
3155 .delete(EntityId::new(1))
3156 .expect("delete after source retirement"));
3157 assert!(manager.get(EntityId::new(1)).is_none());
3158 }
3159 true
3160 }
3161 ));
3162 seen.sort_unstable();
3163 assert_eq!(
3164 seen,
3165 (2..=900).collect::<Vec<_>>(),
3166 "no duplicate, missing or late candidates"
3167 );
3168 let mut after = Vec::new();
3169 assert!(manager.scan_batches(
3170 None,
3171 || true,
3172 |batch| {
3173 after.extend(batch.iter().map(|entity| entity.id.raw()));
3174 true
3175 }
3176 ));
3177 after.sort_unstable();
3178 let mut expected: Vec<_> = (2..=900).collect();
3179 expected.push(5000);
3180 assert_eq!(after, expected, "adopted entities have scan positions");
3181 }
3182
3183 #[test]
3184 fn bounded_scan_applies_snapshot_to_relocated_versions() {
3185 for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
3186 let manager = SegmentManager::with_config(
3187 "vectors",
3188 ManagerConfig {
3189 max_sealed_segments: 1,
3190 ..Default::default()
3191 },
3192 );
3193 for ids in [[1, 2], [3, 4]] {
3194 let entities = ids.map(|id| {
3195 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3196 entity.set_xmin(3);
3197 entity
3198 });
3199 if bulk {
3200 manager
3201 .bulk_insert(entities.into())
3202 .expect("initial versions");
3203 } else {
3204 for entity in entities {
3205 manager.insert(entity).expect("initial version");
3206 }
3207 }
3208 manager.force_seal().expect("seal source");
3209 }
3210 let snapshot = SnapshotContext {
3211 snapshot: Snapshot {
3212 xid: 10,
3213 in_progress: HashSet::new(),
3214 },
3215 manager: Arc::new(SnapshotManager::new()),
3216 own_xids: HashSet::new(),
3217 requires_index_fallback: false,
3218 serializable_reader: None,
3219 };
3220 let mut changed = false;
3221 let mut seen = Vec::new();
3222 assert!(manager.scan_batches(
3223 Some(&snapshot),
3224 || true,
3225 |batch| {
3226 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3227 if !changed {
3228 changed = true;
3229 if consolidate {
3230 drain_maintenance(&manager);
3231 assert!(manager.stats().consolidation.runs_completed > 0);
3232 }
3233 for (id, xmax) in [(3, 8), (4, 12)] {
3234 let mut entity =
3235 manager.get(EntityId::new(id)).expect("retained version");
3236 entity.set_xmax(xmax);
3237 manager
3238 .update(entity)
3239 .expect("update after source retirement");
3240 }
3241 }
3242 true
3243 }
3244 ));
3245 seen.sort_unstable();
3246 assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
3247 let mut again = Vec::new();
3248 assert!(manager.scan_batches(
3249 Some(&snapshot),
3250 || true,
3251 |batch| {
3252 again.extend(batch.iter().map(|entity| entity.id.raw()));
3253 true
3254 }
3255 ));
3256 again.sort_unstable();
3257 assert_eq!(
3258 again,
3259 vec![1, 2, 4],
3260 "old slots do not duplicate relocated versions"
3261 );
3262 let mut current = Vec::new();
3263 assert!(manager.scan_batches(
3264 None,
3265 || true,
3266 |batch| {
3267 current.extend(batch.iter().map(|entity| entity.id.raw()));
3268 true
3269 }
3270 ));
3271 current.sort_unstable();
3272 assert_eq!(
3273 current,
3274 vec![1, 2],
3275 "current view excludes both deleted versions"
3276 );
3277 }
3278 }
3279
3280 #[test]
3281 fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
3282 let manager = SegmentManager::new("mixed");
3283 let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
3284 manager
3285 .bulk_insert((1..=300).map(vector).collect())
3286 .expect("flat");
3287 manager.insert(vector(500)).expect("hash insert");
3288 manager
3289 .bulk_insert(vec![vector(502), vector(505)])
3290 .expect("gapped bulk");
3291 manager.delete(EntityId::new(2)).expect("flat delete");
3292 manager.delete(EntityId::new(505)).expect("hash delete");
3293 let mut changed = false;
3294 let mut seen = Vec::new();
3295 assert!(manager.scan_batches(
3296 None,
3297 || true,
3298 |batch| {
3299 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3300 if !changed {
3301 changed = true;
3302 manager
3303 .insert(vector(700))
3304 .expect("append after cursor boundary");
3305 manager
3306 .update(UnifiedEntity::graph_node(
3307 EntityId::new(500),
3308 "mixed",
3309 "Node",
3310 HashMap::new(),
3311 ))
3312 .expect("structural update");
3313 }
3314 true
3315 }
3316 ));
3317 seen.sort_unstable();
3318 let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
3319 expected.extend([500, 502]);
3320 assert_eq!(seen, expected);
3321 }
3322
3323 #[test]
3324 fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
3325 let manager = SegmentManager::new("vectors");
3326 let mut entities = Vec::new();
3327 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
3328 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3329 entity.set_xmin(xmin);
3330 entity.set_xmax(xmax);
3331 entities.push(entity);
3332 }
3333 manager.bulk_insert(entities).expect("versions");
3334 manager.force_seal().expect("seal");
3335 let snapshot = SnapshotContext {
3336 snapshot: Snapshot {
3337 xid: 10,
3338 in_progress: HashSet::from([4]),
3339 },
3340 manager: Arc::new(SnapshotManager::new()),
3341 own_xids: HashSet::new(),
3342 requires_index_fallback: false,
3343 serializable_reader: None,
3344 };
3345 let mut work = 0;
3346 let mut callbacks = 0;
3347 assert!(!manager.scan_batches(
3348 Some(&snapshot),
3349 || {
3350 work += 1;
3351 work <= 2
3352 },
3353 |_| {
3354 callbacks += 1;
3355 true
3356 }
3357 ));
3358 assert_eq!(callbacks, 0, "a stopped batch is not consumed");
3359 for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
3360 let mut seen = Vec::new();
3361 assert!(manager.scan_batches(
3362 view,
3363 || true,
3364 |batch| {
3365 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3366 true
3367 }
3368 ));
3369 seen.sort_unstable();
3370 assert_eq!(seen, expected);
3371 }
3372 }
3373
3374 #[test]
3375 fn test_manager_basic() {
3376 let manager = SegmentManager::new("test_collection");
3377
3378 let entity = UnifiedEntity::table_row(
3379 manager.next_entity_id(),
3380 "users",
3381 1,
3382 vec![Value::text("Alice".to_string())],
3383 );
3384
3385 let id = manager.insert(entity).unwrap();
3386 assert!(manager.get(id).is_some());
3387 assert_eq!(manager.count(), 1);
3388 }
3389
3390 #[test]
3391 fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
3392 let manager = Arc::new(SegmentManager::new("mixed"));
3393 manager
3394 .bulk_insert(
3395 (1..=128)
3396 .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
3397 .collect(),
3398 )
3399 .expect("vector-only segment");
3400 manager.force_seal().expect("seal vectors");
3401 let mut old =
3402 UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
3403 old.set_xmin(3);
3404 old.set_xmax(12);
3405 let mut new =
3406 UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
3407 new.set_xmin(12);
3408 let edge = UnifiedEntity::graph_edge(
3409 EntityId::new(202),
3410 "mixed",
3411 "a",
3412 "b",
3413 1.0,
3414 Default::default(),
3415 );
3416 manager
3417 .bulk_insert(vec![old, new, edge])
3418 .expect("physical graph versions");
3419 manager.force_seal().expect("seal graph");
3420 manager
3421 .insert(UnifiedEntity::graph_node(
3422 EntityId::new(203),
3423 "mixed",
3424 "person",
3425 Default::default(),
3426 ))
3427 .expect("growing node");
3428 let snapshot = SnapshotContext {
3429 snapshot: Snapshot {
3430 xid: 10,
3431 in_progress: HashSet::new(),
3432 },
3433 manager: Arc::new(SnapshotManager::new()),
3434 own_xids: HashSet::new(),
3435 requires_index_fallback: false,
3436 serializable_reader: None,
3437 };
3438 std::thread::spawn(move || {
3439 for view in [None, Some(&snapshot)] {
3440 for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
3441 let expected: HashSet<_> = manager
3442 .scan(view, |item| kind.matches(&item.kind))
3443 .into_iter()
3444 .map(|item| item.id)
3445 .collect();
3446 let actual: HashSet<_> = manager
3447 .scan_graph_kind(view, kind)
3448 .into_iter()
3449 .map(|item| item.id)
3450 .collect();
3451 assert_eq!(actual, expected);
3452 let mut visited = 0;
3453 let mut candidates = HashSet::new();
3454 manager.scan_graph_candidates_for_each(
3455 view,
3456 kind,
3457 || true,
3458 |item| {
3459 visited += 1;
3460 if kind.matches(&item.kind) {
3461 candidates.insert(item.id);
3462 }
3463 true
3464 },
3465 );
3466 assert_eq!(candidates, expected);
3467 assert_eq!(
3468 visited,
3469 if matches!(kind, GraphEntityKind::Node) {
3470 3
3471 } else {
3472 2
3473 }
3474 );
3475 }
3476 }
3477 let ids: HashSet<_> = manager
3478 .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
3479 .into_iter()
3480 .map(|item| item.id.raw())
3481 .collect();
3482 assert_eq!(ids, HashSet::from([200, 203]));
3483 let ids: HashSet<_> = manager
3484 .scan_graph_kind(None, GraphEntityKind::Node)
3485 .into_iter()
3486 .map(|item| item.id.raw())
3487 .collect();
3488 assert_eq!(ids, HashSet::from([201, 203]));
3489 let mut visits = 0;
3490 manager.scan_graph_candidates_for_each(
3491 None,
3492 GraphEntityKind::Node,
3493 || true,
3494 |_| {
3495 visits += 1;
3496 false
3497 },
3498 );
3499 assert_eq!(
3500 visits, 1,
3501 "early termination must also stop subsequent segments"
3502 );
3503 })
3504 .join()
3505 .expect("snapshot worker");
3506 }
3507
3508 #[test]
3509 fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
3510 let manager = Arc::new(SegmentManager::new("accounts"));
3511 let cases = [
3512 ("legacy row", 0, 0, true),
3513 ("committed creator", 3, 0, true),
3514 ("future creator", 11, 0, false),
3515 ("in-progress creator", 4, 0, false),
3516 ("committed deleter", 3, 7, false),
3517 ("future deleter", 3, 11, true),
3518 ("in-progress deleter", 3, 4, true),
3519 ];
3520
3521 for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
3522 let row_id = index as u64 + 1;
3523 let mut entity = UnifiedEntity::table_row(
3524 EntityId::new(row_id),
3525 "accounts",
3526 row_id,
3527 vec![Value::Integer(row_id as i64)],
3528 );
3529 entity.set_xmin(*xmin);
3530 entity.set_xmax(*xmax);
3531 manager.insert(entity).unwrap();
3532 }
3533
3534 let snapshot = SnapshotContext {
3535 snapshot: Snapshot {
3536 xid: 10,
3537 in_progress: HashSet::from([4]),
3538 },
3539 manager: Arc::new(SnapshotManager::new()),
3540 own_xids: HashSet::new(),
3541 requires_index_fallback: false,
3542 serializable_reader: None,
3543 };
3544
3545 let raw_visible_row_ids: HashSet<u64> = manager
3546 .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
3547 .into_iter()
3548 .filter_map(|entity| match entity.kind {
3549 EntityKind::TableRow { row_id, .. } => Some(row_id),
3550 _ => None,
3551 })
3552 .collect();
3553
3554 let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
3555 .join()
3556 .unwrap();
3557 let visible_row_ids: HashSet<u64> = scanned
3558 .into_iter()
3559 .filter_map(|entity| match entity.kind {
3560 EntityKind::TableRow { row_id, .. } => Some(row_id),
3561 _ => None,
3562 })
3563 .collect();
3564 assert_eq!(visible_row_ids, raw_visible_row_ids);
3565
3566 for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
3567 let row_id = index as u64 + 1;
3568 assert_eq!(
3569 visible_row_ids.contains(&row_id),
3570 *expected_visible,
3571 "{name} visibility mismatch"
3572 );
3573 }
3574 }
3575
3576 #[test]
3577 fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
3578 let manager = SegmentManager::new("accounts");
3582 let mut live =
3583 UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
3584 live.set_xmin(2);
3585 live.set_xmax(0);
3586 manager.insert(live).unwrap();
3587
3588 let mut superseded =
3589 UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
3590 superseded.set_xmin(2);
3591 superseded.set_xmax(9);
3592 manager.insert(superseded).unwrap();
3593
3594 let rows: Vec<u64> = manager
3595 .scan(None, |_| true)
3596 .into_iter()
3597 .filter_map(|entity| match entity.kind {
3598 EntityKind::TableRow { row_id, .. } => Some(row_id),
3599 _ => None,
3600 })
3601 .collect();
3602 assert_eq!(
3603 rows,
3604 vec![1],
3605 "frameless scan must hide superseded versions"
3606 );
3607 }
3608
3609 #[test]
3610 fn query_scan_callers_use_snapshot_api_for_visibility() {
3611 let table = include_str!("../../runtime/query_exec/table.rs");
3612 let aggregate = include_str!("../../runtime/query_exec/aggregate.rs");
3613
3614 assert_eq!(
3615 table.matches(".for_each_entity(|entity|").count(),
3616 1,
3617 "only the visibility-independent schema inspection may use the raw table iterator"
3618 );
3619 assert!(!table.contains(".query_all_zoned_with_stats("));
3620 assert!(!table.contains(".for_each_entity_zoned_with_stats("));
3621 assert!(!table.contains("TableRowMvccReadResolver::captured"));
3622
3623 assert!(!aggregate.contains("TableRowMvccReadResolver"));
3624 assert!(!aggregate.contains(".for_each_entity("));
3625 assert!(!aggregate.contains(".fold_entities_parallel("));
3626 }
3627
3628 #[test]
3629 fn remaining_visibility_scan_callers_use_snapshot_api() {
3630 let aggregate_pushdown =
3631 include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs");
3632 let ask = include_str!("../../runtime/ask_pipeline.rs");
3633 let dml = include_str!("../../runtime/impl_dml.rs");
3634 let dml_targets = include_str!("../../runtime/dml_target_scan.rs");
3635 let events = include_str!("../../runtime/impl_events.rs");
3636 let graph_commands = include_str!("../../runtime/impl_graph_commands.rs");
3637 let graph_dsl = include_str!("../../runtime/graph_dsl.rs");
3638 let graph_tvf = include_str!("../../runtime/graph_tvf.rs");
3639 let kv = include_str!("../../runtime/impl_kv.rs");
3640 let queue = include_str!("../../runtime/impl_queue.rs");
3641 let primary_queue = include_str!("../../runtime/primary_queue_store.rs");
3642 let record_search = include_str!("../../runtime/record_search.rs");
3643 let search = include_str!("../../runtime/impl_search.rs");
3644 let vector = include_str!("../../runtime/query_exec/vector.rs");
3645
3646 assert!(!aggregate_pushdown.contains("TableRowMvccReadResolver"));
3647 assert!(!aggregate_pushdown.contains(".for_each_entity("));
3648 assert!(!ask.contains("for entity in manager.query_all(|_| true)"));
3649 assert!(!dml.contains("let entities = manager.query_all(|_| true);"));
3650 assert_eq!(
3655 dml_targets
3656 .matches("self.candidate_visible(snapshot.as_ref(), entity)")
3657 .count(),
3658 2,
3659 "both raw DML iterations must route through candidate_visible"
3660 );
3661 assert_eq!(
3664 events.matches("manager.query_all(").count(),
3665 1,
3666 "only the visibility-free backfill dedup may use query_all"
3667 );
3668 assert!(!graph_commands.contains("TableRowMvccReadResolver"));
3669 assert!(!graph_dsl.contains("entity_visible_with_context(snap_ctx.as_ref(), e)"));
3670 assert!(!graph_tvf.contains("entity_visible_with_context(snap_ctx.as_ref(), &entity)"));
3671 assert!(!kv.contains("TableRowMvccReadResolver"));
3672 assert!(!queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3673 assert!(!primary_queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3674 assert!(!record_search.contains("TableRowMvccReadResolver"));
3675 assert!(!search.contains("for entity in manager.query_all(|_| true)"));
3676 assert!(!vector.contains("manager.query_all("));
3677 }
3678
3679 fn production_source(src: &str) -> String {
3684 let mut kept: Vec<&str> = Vec::new();
3685 let lines: Vec<&str> = src.lines().collect();
3686 let mut idx = 0;
3687 while idx < lines.len() {
3688 let starts_test_mod = lines[idx].starts_with("#[cfg(test)]")
3689 && lines
3690 .get(idx + 1)
3691 .is_some_and(|next| next.starts_with("mod "));
3692 if !starts_test_mod {
3693 kept.push(lines[idx]);
3694 idx += 1;
3695 continue;
3696 }
3697 idx += 1;
3698 let mut depth: isize = 0;
3699 while idx < lines.len() {
3700 depth += lines[idx].matches('{').count() as isize;
3701 depth -= lines[idx].matches('}').count() as isize;
3702 idx += 1;
3703 if depth <= 0 {
3704 break;
3705 }
3706 }
3707 }
3708 kept.join("\n")
3709 }
3710
3711 fn raw_iteration_sites(src: &str) -> usize {
3713 src.matches("query_all(").count()
3714 + src.matches(".for_each_entity").count()
3715 + src.matches("fold_entities_parallel(").count()
3716 }
3717
3718 #[test]
3719 fn raw_iteration_confined_to_maintenance_allowlist() {
3720 let allowlist: &[(&str, &str, usize, &str)] = &[
3728 (
3729 "runtime/query_exec/table.rs",
3730 include_str!("../../runtime/query_exec/table.rs"),
3731 1,
3732 "documented exception: schema inspection (does the collection ever \
3733 observe these fields?) is visibility-independent",
3734 ),
3735 (
3736 "runtime/query_exec/aggregate.rs",
3737 include_str!("../../runtime/query_exec/aggregate.rs"),
3738 0,
3739 "fully migrated by batch 1",
3740 ),
3741 (
3742 "runtime/query_exec/aggregate_pushdown_dispatch.rs",
3743 include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs"),
3744 0,
3745 "fully migrated by batch 1",
3746 ),
3747 (
3748 "runtime/query_exec/vector.rs",
3749 include_str!("../../runtime/query_exec/vector.rs"),
3750 0,
3751 "fully migrated by batch 2",
3752 ),
3753 (
3754 "runtime/ask_pipeline.rs",
3755 include_str!("../../runtime/ask_pipeline.rs"),
3756 0,
3757 "fully migrated by batch 2",
3758 ),
3759 (
3760 "runtime/impl_dml.rs",
3761 include_str!("../../runtime/impl_dml.rs"),
3762 1,
3763 "documented exception: ON CONFLICT resolves a unique key against \
3764 physical versions on purpose. A snapshot scan cannot see a row an \
3765 uncommitted concurrent writer already holds the key for, which is \
3766 exactly the case that must raise a serialization conflict rather \
3767 than insert a duplicate; the same pass reaps index entries left by \
3768 aborted writers. \
3769 TODO(#2263 follow-up): narrow the candidates through the unique \
3770 index instead of walking the collection",
3771 ),
3772 (
3773 "runtime/dml_target_scan.rs",
3774 include_str!("../../runtime/dml_target_scan.rs"),
3775 2,
3776 "documented exception: kind-conditional visibility the uniform scan \
3777 API cannot express; both iterations gate on candidate_visible",
3778 ),
3779 (
3780 "runtime/impl_events.rs",
3781 include_str!("../../runtime/impl_events.rs"),
3782 1,
3783 "documented exception: backfill dedup must include consumed messages, \
3784 so it is deliberately visibility-free",
3785 ),
3786 (
3787 "runtime/impl_graph_commands.rs",
3788 include_str!("../../runtime/impl_graph_commands.rs"),
3789 0,
3790 "native node properties resolve one snapshot-visible, RLS-admitted version",
3791 ),
3792 (
3793 "runtime/graph_dsl.rs",
3794 include_str!("../../runtime/graph_dsl.rs"),
3795 1,
3796 "remaining legacy debt: unused edge-property helper; search expansion \
3797 now uses snapshot visitors and RLS-gated hydration",
3798 ),
3799 (
3800 "runtime/graph_tvf.rs",
3801 include_str!("../../runtime/graph_tvf.rs"),
3802 0,
3803 "fully migrated by batch 2",
3804 ),
3805 (
3806 "runtime/impl_kv.rs",
3807 include_str!("../../runtime/impl_kv.rs"),
3808 3,
3809 "one non-versioned listing fallback (the versioned arm scans under \
3810 the captured snapshot). Point reads use physical key candidates with \
3811 versioned visibility rechecks. Vault reads use snapshot scans. Offline vault \
3812 export preserves all \
3813 physical versions across tenants; restore's emptiness check must reject \
3814 every physical row, including versions invisible to a snapshot",
3815 ),
3816 (
3817 "runtime/impl_queue.rs",
3818 include_str!("../../runtime/impl_queue.rs"),
3819 10,
3820 "maintenance: red_queue_meta bookkeeping (config, groups, acks, \
3821 pending, dedup TTL reaper, gauge counts) — queue message bodies are \
3822 read through scan. \
3823 TODO(#2138 follow-up): QUEUE PENDING projects meta rows to clients",
3824 ),
3825 (
3826 "runtime/primary_queue_store.rs",
3827 include_str!("../../runtime/primary_queue_store.rs"),
3828 3,
3829 "maintenance: generic red_queue_meta read/delete helpers and the \
3830 next-position allocator, which must not collide with rows invisible \
3831 to the writer's snapshot",
3832 ),
3833 (
3834 "runtime/record_search.rs",
3835 include_str!("../../runtime/record_search.rs"),
3836 0,
3837 "fully migrated by batch 2",
3838 ),
3839 (
3840 "runtime/search_graph.rs",
3841 include_str!("../../runtime/search_graph.rs"),
3842 0,
3843 "search graph preparation uses snapshot visitors and gated hydration",
3844 ),
3845 (
3846 "runtime/impl_search.rs",
3847 include_str!("../../runtime/impl_search.rs"),
3848 1,
3849 "maintenance: ask-audit retention sweep must see every physical row",
3850 ),
3851 ];
3852
3853 for (path, src, allowed, reason) in allowlist {
3854 let found = raw_iteration_sites(&production_source(src));
3855 assert_eq!(
3856 found, *allowed,
3857 "{path}: {found} raw-iteration sites, allowlist permits {allowed} \
3858 ({reason}). Raw iteration bypasses MVCC visibility — use \
3859 SegmentManager::scan / scan_for_each / scan_fold_parallel / \
3860 scan_zoned_with_stats, or extend the allowlist with a maintenance \
3861 justification."
3862 );
3863 }
3864 }
3865
3866 #[test]
3867 fn explicit_snapshot_scan_ignores_ambient_thread_local() {
3868 let manager = Arc::new(SegmentManager::new("accounts"));
3869 let versions = [(0, 0), (3, 0), (11, 0), (4, 0), (3, 7), (3, 11), (3, 4)];
3871 for (index, (xmin, xmax)) in versions.iter().enumerate() {
3872 let row_id = index as u64 + 1;
3873 let mut entity = UnifiedEntity::table_row(
3874 EntityId::new(row_id),
3875 "accounts",
3876 row_id,
3877 vec![Value::Integer(row_id as i64)],
3878 );
3879 entity.set_xmin(*xmin);
3880 entity.set_xmax(*xmax);
3881 manager.insert(entity).unwrap();
3882 }
3883
3884 let explicit = SnapshotContext {
3885 snapshot: Snapshot {
3886 xid: 10,
3887 in_progress: HashSet::from([4]),
3888 },
3889 manager: Arc::new(SnapshotManager::new()),
3890 own_xids: HashSet::new(),
3891 requires_index_fallback: false,
3892 serializable_reader: None,
3893 };
3894 let decoy = SnapshotContext {
3897 snapshot: Snapshot {
3898 xid: 1,
3899 in_progress: HashSet::new(),
3900 },
3901 manager: Arc::new(SnapshotManager::new()),
3902 own_xids: HashSet::new(),
3903 requires_index_fallback: false,
3904 serializable_reader: None,
3905 };
3906
3907 let row_ids = |entities: Vec<UnifiedEntity>| -> HashSet<u64> {
3908 entities
3909 .into_iter()
3910 .filter_map(|entity| match entity.kind {
3911 EntityKind::TableRow { row_id, .. } => Some(row_id),
3912 _ => None,
3913 })
3914 .collect()
3915 };
3916
3917 let guard = crate::runtime::execution_context::CurrentSnapshotGuard::install(decoy);
3918 let on_installing_thread = row_ids(manager.scan(Some(&explicit), |_| true));
3919 drop(guard);
3920
3921 let spawned_manager = Arc::clone(&manager);
3922 let spawned_snapshot = explicit.clone();
3923 let on_fresh_thread = std::thread::spawn(move || {
3924 row_ids(spawned_manager.scan(Some(&spawned_snapshot), |_| true))
3925 })
3926 .join()
3927 .unwrap();
3928
3929 assert_eq!(
3930 on_installing_thread, on_fresh_thread,
3931 "an explicit snapshot must produce the same visible set on any thread"
3932 );
3933 assert_eq!(on_installing_thread, HashSet::from([1, 2, 6, 7]));
3934
3935 let raw_manager = Arc::clone(&manager);
3938 let raw_count = std::thread::spawn(move || raw_manager.query_all(|_| true).len())
3939 .join()
3940 .unwrap();
3941 assert_eq!(
3942 raw_count,
3943 versions.len(),
3944 "raw iteration must expose every physical version"
3945 );
3946 }
3947
3948 #[test]
3949 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
3950 let manager = SegmentManager::new("test_collection");
3951
3952 let entity = UnifiedEntity::table_row(
3953 manager.next_entity_id(),
3954 "users",
3955 1,
3956 vec![Value::text("Alice".to_string())],
3957 );
3958
3959 let id = manager.insert(entity).unwrap();
3960
3961 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3962 assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
3963 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
3964
3965 assert!(manager.delete(id).unwrap());
3966 assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3967 }
3968
3969 #[test]
3970 fn zoned_scan_reports_pruned_segments_and_preserves_results() {
3971 let config = ManagerConfig {
3972 segment_config: SegmentConfig {
3973 max_entities: 2,
3974 ..Default::default()
3975 },
3976 ..Default::default()
3977 };
3978 let manager = SegmentManager::with_config("events", config);
3979
3980 for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
3981 let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
3982 let entity = UnifiedEntity::new(
3983 manager.next_entity_id(),
3984 EntityKind::TableRow {
3985 table: "events".into(),
3986 row_id: (idx + 1) as u64,
3987 },
3988 EntityData::Row(row),
3989 );
3990 manager.insert(entity).unwrap();
3991 }
3992
3993 let full = manager.query_all(|entity| match &entity.data {
3994 EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
3995 _ => false,
3996 });
3997 let probe = Value::Integer(100);
3998 let (pruned, stats) =
3999 manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
4000 match &entity.data {
4001 EntityData::Row(row) => {
4002 matches!(row.get_field("ts"), Some(Value::Integer(100)))
4003 }
4004 _ => false,
4005 }
4006 });
4007
4008 assert_eq!(pruned.len(), full.len());
4009 assert_eq!(pruned[0].id, full[0].id);
4010 assert_eq!(stats.segments_total, 3);
4011 assert_eq!(stats.segments_scanned, 2);
4012 assert_eq!(stats.segments_pruned, 1);
4013 }
4014
4015 #[test]
4016 fn test_manager_auto_seal() {
4017 let config = ManagerConfig {
4018 segment_config: SegmentConfig {
4019 max_entities: 2,
4020 ..Default::default()
4021 },
4022 ..Default::default()
4023 };
4024
4025 let manager = SegmentManager::with_config("test", config);
4026
4027 manager
4029 .insert(UnifiedEntity::vector(
4030 manager.next_entity_id(),
4031 "v",
4032 vec![0.1],
4033 ))
4034 .unwrap();
4035
4036 manager
4038 .insert(UnifiedEntity::vector(
4039 manager.next_entity_id(),
4040 "v",
4041 vec![0.2],
4042 ))
4043 .unwrap();
4044
4045 manager
4047 .insert(UnifiedEntity::vector(
4048 manager.next_entity_id(),
4049 "v",
4050 vec![0.3],
4051 ))
4052 .unwrap();
4053
4054 let stats = manager.stats();
4055 assert_eq!(stats.total_entities, 3);
4056 }
4057
4058 #[test]
4059 fn test_manager_delete() {
4060 let manager = SegmentManager::new("test");
4061
4062 let id = manager
4063 .insert(UnifiedEntity::vector(
4064 manager.next_entity_id(),
4065 "v",
4066 vec![0.1],
4067 ))
4068 .unwrap();
4069
4070 assert!(manager.get(id).is_some());
4071 assert!(manager.delete(id).unwrap());
4072 assert!(manager.get(id).is_none());
4073 }
4074
4075 #[test]
4076 fn test_manager_metadata() {
4077 let manager = SegmentManager::new("test");
4078
4079 let id = manager
4080 .insert(UnifiedEntity::table_row(
4081 manager.next_entity_id(),
4082 "hosts",
4083 1,
4084 vec![Value::text("192.168.1.1".to_string())],
4085 ))
4086 .unwrap();
4087
4088 let mut meta = Metadata::new();
4089 meta.set(
4090 "os",
4091 super::super::metadata::MetadataValue::String("linux".to_string()),
4092 );
4093
4094 manager.set_metadata(id, meta).unwrap();
4095
4096 let retrieved = manager.get_metadata(id).unwrap();
4097 assert!(retrieved.has("os"));
4098 }
4099
4100 #[test]
4101 fn test_manager_query_by_kind() {
4102 let manager = SegmentManager::new("test");
4103
4104 manager
4105 .insert(UnifiedEntity::table_row(
4106 manager.next_entity_id(),
4107 "hosts",
4108 1,
4109 vec![],
4110 ))
4111 .unwrap();
4112
4113 manager
4114 .insert(UnifiedEntity::vector(
4115 manager.next_entity_id(),
4116 "embeddings",
4117 vec![0.1],
4118 ))
4119 .unwrap();
4120
4121 manager
4122 .insert(UnifiedEntity::table_row(
4123 manager.next_entity_id(),
4124 "hosts",
4125 2,
4126 vec![],
4127 ))
4128 .unwrap();
4129
4130 let rows = manager.get_by_kind("table");
4131 assert_eq!(rows.len(), 2);
4132
4133 let vectors = manager.get_by_kind("vector");
4134 assert_eq!(vectors.len(), 1);
4135 }
4136
4137 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
4144 SegmentManager::with_config(
4145 "test",
4146 ManagerConfig {
4147 max_sealed_segments: 1_000,
4148 consolidation_entities_per_tick: entities_per_tick,
4149 ..Default::default()
4150 },
4151 )
4152 }
4153
4154 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
4158 let entities: Vec<UnifiedEntity> = (0..count)
4159 .map(|i| {
4160 UnifiedEntity::table_row(
4161 manager.next_entity_id(),
4162 "users",
4163 i as u64,
4164 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
4165 )
4166 })
4167 .collect();
4168 let ids = manager.bulk_insert(entities).expect("bulk insert");
4169 manager.force_seal().expect("seal");
4170 ids
4171 }
4172
4173 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
4174 let mut ids: Vec<u64> = manager
4175 .query_all(|_| true)
4176 .into_iter()
4177 .map(|entity| entity.id.raw())
4178 .collect();
4179 ids.sort_unstable();
4180 ids
4181 }
4182
4183 fn drain_maintenance(manager: &SegmentManager) -> usize {
4184 let mut ticks = 0;
4185 while manager.consolidation.read().is_some() || ticks == 0 {
4186 manager.run_maintenance().expect("maintenance");
4187 ticks += 1;
4188 assert!(ticks < 1_000, "consolidation failed to converge");
4189 }
4190 ticks
4191 }
4192
4193 #[test]
4194 fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
4195 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4196 let ids = manager
4197 .bulk_insert(
4198 (1..=1024)
4199 .map(|id| {
4200 UnifiedEntity::graph_edge(
4201 EntityId::new(id),
4202 "step",
4203 "2000",
4204 "3000",
4205 1.0,
4206 HashMap::new(),
4207 )
4208 })
4209 .collect(),
4210 )
4211 .expect("edges");
4212 manager.force_seal().expect("seal");
4213 for id in ids.iter().take(300) {
4214 manager.delete(*id).expect("fragment source");
4215 }
4216 let source = Arc::downgrade(&manager.sealed.read()[0]);
4217 let mut seen = Vec::new();
4218 let mut changed = false;
4219 let mut metadata_bytes = 0;
4220 assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
4221 || true,
4222 |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
4223 |batch| {
4224 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4225 assert!(manager.growing.try_write().is_some(), "no growing topology guard");
4226 assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
4227 assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
4228 if !changed {
4229 changed = true;
4230 drain_maintenance(&manager);
4231 assert!(manager.stats().consolidation.runs_completed > 0);
4232 assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
4233 let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
4234 if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
4235 manager.update(updated).expect("update authoritative merged copy");
4236 }
4237 for id in batch {
4238 let entity = manager.get(*id).expect("hydrate current manager");
4239 if id.raw() == 1024 {
4240 assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
4241 }
4242 seen.push(id.raw());
4243 }
4244 Ok(true)
4245 }).expect("cursor"));
4246 assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
4247 assert!(
4248 source.upgrade().is_none(),
4249 "finished cursor releases retired source"
4250 );
4251 assert!(
4252 metadata_bytes < 4096,
4253 "fixed ID buffer plus source handles: {metadata_bytes}"
4254 );
4255 }
4256
4257 #[test]
4258 fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
4259 let manager = SegmentManager::new("graph");
4260 manager
4261 .bulk_insert(
4262 (1..=1024)
4263 .map(|id| {
4264 UnifiedEntity::graph_edge(
4265 EntityId::new(id),
4266 "step",
4267 "2000",
4268 "3000",
4269 1.0,
4270 HashMap::new(),
4271 )
4272 })
4273 .collect(),
4274 )
4275 .expect("edges");
4276 let mut delivered = 0;
4277 let denied = manager.visit_graph_index_batches(
4278 GraphEntityKind::Edge,
4279 &[EntityId::new(2000)],
4280 || true,
4281 |_| Err("out of memory"),
4282 |ids| {
4283 delivered += ids.len();
4284 Ok(true)
4285 },
4286 );
4287 assert_eq!(denied, Err("out of memory"));
4288 assert_eq!(delivered, 0);
4289 let mut work = 0;
4290 assert!(!manager
4291 .visit_graph_index_batches(
4292 GraphEntityKind::Edge,
4293 &[EntityId::new(2000)],
4294 || {
4295 work += 1;
4296 work < 100
4297 },
4298 |_| Ok::<_, ()>(()),
4299 |ids| {
4300 delivered += ids.len();
4301 Ok(true)
4302 }
4303 )
4304 .expect("cancelled cursor"));
4305 assert_eq!(
4306 delivered, 0,
4307 "cancelled partial batch cannot reach consumer"
4308 );
4309 assert!(!manager
4310 .visit_graph_index_batches(
4311 GraphEntityKind::Edge,
4312 &[EntityId::new(2000)],
4313 || true,
4314 |_| Ok::<_, ()>(()),
4315 |ids| {
4316 delivered += ids.len();
4317 Ok(false)
4318 }
4319 )
4320 .expect("consumer stop"));
4321 assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
4322 }
4323
4324 #[test]
4325 fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
4326 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4327 let ids = seal_bulk_segment(&manager, 100);
4328 for id in ids.iter().take(40) {
4329 manager.delete(*id).expect("delete before scan");
4330 }
4331 let source_bytes = manager.resident_bytes();
4332 let source = Arc::downgrade(&manager.sealed.read()[0]);
4333 std::thread::scope(|scope| {
4334 let (ready, waiting) = std::sync::mpsc::channel();
4335 let (release_first, first_wait) = std::sync::mpsc::channel();
4336 let (release_last, last_wait) = std::sync::mpsc::channel();
4337 let reader = |resume: std::sync::mpsc::Receiver<()>,
4338 ready: std::sync::mpsc::Sender<()>| {
4339 assert!(manager.scan_batches(
4340 None,
4341 || true,
4342 |_| {
4343 ready.send(()).expect("signal captured cursor");
4344 let _ = resume.recv();
4345 true
4346 }
4347 ));
4348 };
4349 let ready_first = ready.clone();
4350 let first = scope.spawn(move || reader(first_wait, ready_first));
4351 let last = scope.spawn(move || reader(last_wait, ready));
4352 for _ in 0..2 {
4353 waiting
4354 .recv_timeout(std::time::Duration::from_secs(10))
4355 .expect("reader ready");
4356 }
4357 drain_maintenance(&manager);
4358 assert!(
4359 source.upgrade().is_some(),
4360 "both cursors pin the retired source"
4361 );
4362 let active_bytes = manager.sealed.read()[0].read().resident_bytes();
4363 let held = manager.resident_bytes();
4364 assert!(
4365 held >= active_bytes + source_bytes,
4366 "retired source remains resident"
4367 );
4368 release_first.send(()).expect("release first reader");
4369 first.join().expect("first reader");
4370 assert_eq!(
4371 manager.resident_bytes(),
4372 held,
4373 "one reader still pins the source"
4374 );
4375 release_last.send(()).expect("release last reader");
4376 last.join().expect("last reader");
4377 assert!(
4378 source.upgrade().is_none(),
4379 "accounting must not retain payloads"
4380 );
4381 assert_eq!(held - manager.resident_bytes(), source_bytes);
4382 assert_eq!(
4383 manager.stats().total_memory_bytes as u64,
4384 manager.resident_bytes()
4385 );
4386 });
4387 }
4388
4389 #[test]
4390 fn retained_memory_gates_runtime_growth_until_the_cursor_is_cancelled() {
4391 use crate::storage::memory_pools::MemoryPool;
4392 use crate::{RedDBOptions, RedDBRuntime};
4393
4394 let budget = 128 * 1024;
4395 let runtime =
4396 RedDBRuntime::with_options(RedDBOptions::in_memory().with_memory_budget(budget))
4397 .expect("runtime");
4398 runtime
4399 .db()
4400 .store()
4401 .create_collection("test")
4402 .expect("collection");
4403 let manager = runtime
4404 .db()
4405 .store()
4406 .get_collection("test")
4407 .expect("manager");
4408 let ids = seal_bulk_segment(&manager, 100);
4409 for id in ids.iter().take(40) {
4410 manager.delete(*id).expect("delete before scan");
4411 }
4412 let source_bytes = manager.resident_bytes();
4413 let mut growth = 0;
4414 let mut held_used = 0;
4415 assert!(!manager.scan_batches(
4416 None,
4417 || true,
4418 |_| {
4419 drain_maintenance(&manager);
4420 runtime.refresh_memory_accounting();
4421 held_used = runtime.memory_accounting().total_used_bytes();
4424 growth = budget
4425 .checked_sub(held_used)
4426 .expect("fixture fits while the cursor is open")
4427 + source_bytes / 2;
4428 assert!(growth > 0);
4429 assert!(
4430 runtime
4431 .admit_non_evictable_growth(
4432 MemoryPool::SegmentArena,
4433 "growth with pinned source",
4434 growth,
4435 )
4436 .is_err(),
4437 "retained source must participate in runtime admission"
4438 );
4439 false
4440 }
4441 ));
4442 runtime.refresh_memory_accounting();
4443 assert_eq!(
4444 held_used - runtime.memory_accounting().total_used_bytes(),
4445 source_bytes
4446 );
4447 let _reservation = runtime
4448 .admit_non_evictable_growth(
4449 MemoryPool::SegmentArena,
4450 "growth after cursor cancellation",
4451 growth,
4452 )
4453 .expect("the same growth fits once the cursor releases its source");
4454 }
4455
4456 #[test]
4457 fn retained_memory_includes_an_inflight_consolidation_copy() {
4458 let manager = consolidating_manager(1);
4459 let ids = seal_bulk_segment(&manager, 100);
4460 for id in ids.iter().take(40) {
4461 manager.delete(*id).expect("delete");
4462 }
4463 let before = manager.resident_bytes();
4464 assert!(manager.pressure_consolidation_tick());
4465 let copied = manager
4466 .consolidation
4467 .read()
4468 .as_ref()
4469 .expect("paced merge")
4470 .merged
4471 .resident_bytes();
4472 assert!(copied > 0);
4473 assert!(
4474 manager.resident_bytes() >= before + copied,
4475 "unpublished copy consumes memory"
4476 );
4477 drain_maintenance(&manager);
4478 assert!(
4479 manager.resident_bytes() < before,
4480 "completed merge releases the old storage"
4481 );
4482 }
4483
4484 #[test]
4485 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
4486 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4487 let ids = seal_bulk_segment(&manager, 100);
4488
4489 let bytes_before_delete = manager.resident_bytes();
4490
4491 for id in ids.iter().take(40) {
4493 assert!(manager.delete(*id).expect("delete"));
4494 }
4495
4496 let bytes_with_tombstones = manager.resident_bytes();
4499 assert!(bytes_with_tombstones > bytes_before_delete);
4500
4501 let survivors = live_ids(&manager);
4502 assert_eq!(survivors.len(), 60);
4503
4504 drain_maintenance(&manager);
4505
4506 let stats = manager.stats();
4507 assert_eq!(stats.consolidation.runs_started, 1);
4508 assert_eq!(stats.consolidation.runs_completed, 1);
4509 assert_eq!(stats.consolidation.segments_merged, 1);
4510 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
4511 assert!(stats.consolidation.bytes_reclaimed > 0);
4512
4513 let bytes_after = manager.resident_bytes();
4515 assert!(
4516 bytes_after <= bytes_before_delete,
4517 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
4518 );
4519 assert_eq!(
4520 bytes_with_tombstones - bytes_after,
4521 stats.consolidation.bytes_reclaimed
4522 );
4523 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
4524
4525 assert_eq!(live_ids(&manager), survivors);
4526 assert_eq!(manager.sealed.read().len(), 1);
4527 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
4528 }
4529
4530 #[test]
4531 fn query_results_are_identical_before_during_and_after_consolidation() {
4532 let manager = consolidating_manager(3);
4533 let ids = seal_bulk_segment(&manager, 20);
4534 seal_bulk_segment(&manager, 20);
4535 for id in ids.iter().take(10) {
4536 assert!(manager.delete(*id).expect("delete"));
4537 }
4538
4539 let before = live_ids(&manager);
4540 let survivor = *ids.last().expect("at least one row");
4541 let tombstoned = ids[0];
4542
4543 let mut mid_tick_observations = 0;
4546 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
4547 manager.run_maintenance().expect("maintenance");
4548 assert_eq!(live_ids(&manager), before);
4549 assert!(manager.get(survivor).is_some());
4550 assert!(manager.get(tombstoned).is_none());
4551 assert_eq!(manager.get_by_kind("table").len(), before.len());
4552 mid_tick_observations += 1;
4553 assert!(mid_tick_observations < 1_000);
4554 }
4555
4556 assert!(
4557 mid_tick_observations > 2,
4558 "expected the run to span several ticks, saw {mid_tick_observations}"
4559 );
4560 assert_eq!(live_ids(&manager), before);
4561 assert!(manager.get(survivor).is_some());
4562 assert!(manager.get(tombstoned).is_none());
4563 assert_eq!(manager.get_by_kind("table").len(), before.len());
4564 }
4565
4566 #[test]
4567 fn get_by_kind_reaches_sealed_segments() {
4568 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4569 let ids = seal_bulk_segment(&manager, 5);
4570 assert!(manager.delete(ids[0]).expect("delete"));
4571
4572 assert_eq!(manager.get_by_kind("table").len(), 4);
4576 assert_eq!(manager.get_by_kind("vector").len(), 0);
4577 }
4578
4579 #[test]
4580 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
4581 const PER_TICK: usize = 4;
4582 let manager = consolidating_manager(PER_TICK);
4583 let ids = seal_bulk_segment(&manager, 60);
4584 for id in ids.iter().take(30) {
4585 assert!(manager.delete(*id).expect("delete"));
4586 }
4587
4588 manager.run_maintenance().expect("maintenance");
4590 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
4591
4592 let mut previous = PER_TICK as u64;
4595 let mut ticks = 1;
4596 while let Some(copied) = manager.consolidation_progress() {
4597 assert!(
4598 copied - previous <= PER_TICK as u64,
4599 "tick copied {} entities, bound is {PER_TICK}",
4600 copied - previous
4601 );
4602 previous = copied;
4603 manager.run_maintenance().expect("maintenance");
4604 ticks += 1;
4605 assert!(ticks < 100);
4606 }
4607
4608 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
4611 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4612 assert_eq!(live_ids(&manager).len(), 30);
4613 }
4614
4615 #[test]
4616 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
4617 let manager = SegmentManager::with_config(
4618 "test",
4619 ManagerConfig {
4620 max_sealed_segments: 3,
4621 ..Default::default()
4622 },
4623 );
4624 for _ in 0..4 {
4625 seal_bulk_segment(&manager, 5);
4626 }
4627 assert_eq!(manager.sealed.read().len(), 4);
4628
4629 let before = live_ids(&manager);
4630 drain_maintenance(&manager);
4631
4632 assert_eq!(manager.sealed.read().len(), 1);
4635 assert_eq!(manager.stats().consolidation.segments_merged, 4);
4636 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
4637 assert_eq!(live_ids(&manager), before);
4638 }
4639
4640 #[test]
4641 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
4642 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4643 seal_bulk_segment(&manager, 10);
4644 seal_bulk_segment(&manager, 10);
4645
4646 for _ in 0..5 {
4647 manager.run_maintenance().expect("maintenance");
4648 }
4649
4650 assert_eq!(manager.stats().consolidation.runs_started, 0);
4651 assert_eq!(manager.sealed.read().len(), 2);
4652 }
4653
4654 #[test]
4655 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
4656 use std::sync::atomic::AtomicBool;
4657
4658 let manager = Arc::new(consolidating_manager(2));
4659 let ids = seal_bulk_segment(&manager, 40);
4660 seal_bulk_segment(&manager, 40);
4661 for id in ids.iter().take(20) {
4662 assert!(manager.delete(*id).expect("delete"));
4663 }
4664
4665 let expected = live_ids(&manager);
4666 let stop = Arc::new(AtomicBool::new(false));
4667
4668 let readers: Vec<_> = (0..4)
4669 .map(|_| {
4670 let manager = Arc::clone(&manager);
4671 let stop = Arc::clone(&stop);
4672 let expected = expected.clone();
4673 std::thread::spawn(move || {
4674 let mut scans = 0u64;
4675 loop {
4678 let mut seen: Vec<u64> = manager
4679 .query_all(|_| true)
4680 .into_iter()
4681 .map(|entity| entity.id.raw())
4682 .collect();
4683 seen.sort_unstable();
4684 assert_eq!(seen, expected, "reader saw a torn sealed set");
4687 scans += 1;
4688 if stop.load(Ordering::Relaxed) {
4689 break;
4690 }
4691 }
4692 scans
4693 })
4694 })
4695 .collect();
4696
4697 drain_maintenance(&manager);
4698 stop.store(true, Ordering::Relaxed);
4699
4700 for reader in readers {
4701 let scans = reader.join().expect("reader thread");
4702 assert!(scans > 0, "reader never got to scan");
4703 }
4704
4705 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4706 assert_eq!(live_ids(&manager), expected);
4707 }
4708
4709 #[test]
4710 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
4711 let manager = consolidating_manager(4);
4712 let ids = seal_bulk_segment(&manager, 20);
4713 for id in ids.iter().take(6) {
4714 assert!(manager.delete(*id).expect("delete"));
4715 }
4716
4717 manager.run_maintenance().expect("maintenance");
4719 assert_eq!(manager.consolidation_progress(), Some(4));
4720 let doomed = ids[6];
4721 assert!(manager.delete(doomed).expect("delete"));
4722
4723 let expected = live_ids(&manager);
4724 drain_maintenance(&manager);
4725
4726 assert!(manager.get(doomed).is_none(), "deleted row came back");
4727 assert_eq!(live_ids(&manager), expected);
4728 }
4729
4730 #[test]
4731 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
4732 let manager = consolidating_manager(3);
4733 let ids = seal_bulk_segment(&manager, 30);
4734 for id in ids.iter().take(12) {
4735 assert!(manager.delete(*id).expect("delete"));
4736 }
4737
4738 let before = live_ids(&manager);
4739 let sealed_before: Vec<SegmentId> = manager
4740 .sealed
4741 .read()
4742 .iter()
4743 .map(|segment| segment.read().id())
4744 .collect();
4745
4746 manager.run_maintenance().expect("maintenance");
4747 manager.run_maintenance().expect("maintenance");
4748 assert!(manager.consolidation.read().is_some());
4749
4750 *manager.consolidation.write() = None;
4754
4755 assert_eq!(live_ids(&manager), before);
4756 let sealed_after: Vec<SegmentId> = manager
4757 .sealed
4758 .read()
4759 .iter()
4760 .map(|segment| segment.read().id())
4761 .collect();
4762 assert_eq!(sealed_after, sealed_before);
4763 assert_eq!(manager.stats().consolidation.runs_completed, 0);
4764
4765 drain_maintenance(&manager);
4766 assert_eq!(live_ids(&manager), before);
4767 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4768 }
4769
4770 #[test]
4771 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
4772 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4773
4774 let entities: Vec<UnifiedEntity> = (0..10)
4778 .map(|i| {
4779 UnifiedEntity::table_row(
4780 manager.next_entity_id(),
4781 "users",
4782 i as u64,
4783 vec![Value::Integer(i as i64)],
4784 )
4785 })
4786 .collect();
4787 let ids = manager.bulk_insert(entities).expect("bulk insert");
4788
4789 let mut meta = Metadata::new();
4790 meta.set(
4791 "os",
4792 super::super::metadata::MetadataValue::String("linux".to_string()),
4793 );
4794 manager.set_metadata(ids[9], meta).expect("set metadata");
4795 manager.force_seal().expect("seal");
4796
4797 assert!(manager.delete(ids[0]).expect("delete"));
4798 assert!(manager.delete(ids[1]).expect("delete"));
4799 assert!(manager.delete(ids[2]).expect("delete"));
4800
4801 let before: Vec<UnifiedEntity> = {
4802 let mut entities = manager.query_all(|_| true);
4803 entities.sort_by_key(|entity| entity.id.raw());
4804 entities
4805 };
4806
4807 drain_maintenance(&manager);
4808
4809 let after: Vec<UnifiedEntity> = {
4810 let mut entities = manager.query_all(|_| true);
4811 entities.sort_by_key(|entity| entity.id.raw());
4812 entities
4813 };
4814
4815 let columns = |entity: &UnifiedEntity| match &entity.data {
4816 super::super::entity::EntityData::Row(row) => row.columns.clone(),
4817 _ => Vec::new(),
4818 };
4819
4820 assert_eq!(before.len(), after.len());
4821 for (old, new) in before.iter().zip(after.iter()) {
4822 assert_eq!(old.id, new.id);
4823 assert_eq!(columns(old), columns(new));
4824 assert_eq!(
4825 old.sequence_id, new.sequence_id,
4826 "merge must not re-issue sequence ids"
4827 );
4828 }
4829 assert!(manager
4830 .get_metadata(ids[9])
4831 .expect("metadata survives the merge")
4832 .has("os"));
4833 }
4834
4835 #[test]
4836 #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
4837 fn test_lifecycle_events() {
4838 let manager = SegmentManager::new("test");
4839
4840 manager
4841 .insert(UnifiedEntity::vector(
4842 manager.next_entity_id(),
4843 "v",
4844 vec![0.1],
4845 ))
4846 .unwrap();
4847
4848 let events = manager.drain_events();
4849
4850 assert!(events
4852 .iter()
4853 .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
4854 assert!(events
4855 .iter()
4856 .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
4857 }
4858}