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, SegmentState,
26 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 struct SegmentManager {
190 row_constraint_lock: Arc<parking_lot::ReentrantMutex<()>>,
193 collection: String,
195 config: ManagerConfig,
197 next_segment_id: AtomicU64,
199 next_entity_id: AtomicU64,
201 next_row_id: AtomicU64,
203 total_entities_atomic: AtomicU64,
207 growing: RwLock<Option<Arc<RwLock<GrowingSegment>>>>,
209 sealed: RwLock<Vec<Arc<RwLock<GrowingSegment>>>>,
211 retired: RwLock<Vec<Weak<RwLock<GrowingSegment>>>>,
213 archived: RwLock<Vec<SegmentId>>,
215 entity_segment: RwLock<HashMap<EntityId, SegmentId>>,
219 column_schema: RwLock<Option<Arc<Vec<String>>>>,
222 stats: RwLock<ManagerStats>,
224 consolidation: RwLock<Option<Consolidation>>,
227 events: RwLock<Vec<LifecycleEvent>>,
229 visibility_map: VisibilityMap,
233 physical_root_cache: parking_lot::Mutex<Option<PhysicalRootCacheEntry>>,
234}
235
236impl SegmentManager {
237 pub fn new(collection: impl Into<String>) -> Self {
239 Self::with_config(collection, ManagerConfig::default())
240 }
241
242 pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
244 Self {
245 row_constraint_lock: Arc::new(parking_lot::ReentrantMutex::new(())),
246 physical_root_cache: parking_lot::Mutex::new(None),
247 collection: collection.into(),
248 config,
249 next_segment_id: AtomicU64::new(1),
250 next_entity_id: AtomicU64::new(1),
251 next_row_id: AtomicU64::new(1),
252 total_entities_atomic: AtomicU64::new(0),
253 growing: RwLock::new(None),
254 sealed: RwLock::new(Vec::new()),
255 retired: RwLock::new(Vec::new()),
256 archived: RwLock::new(Vec::new()),
257 entity_segment: RwLock::new(HashMap::new()),
258 column_schema: RwLock::new(None),
259 stats: RwLock::new(ManagerStats::default()),
260 consolidation: RwLock::new(None),
261 events: RwLock::new(Vec::new()),
262 visibility_map: VisibilityMap::new(),
263 }
264 }
265
266 pub(crate) fn row_constraint_lock(&self) -> Arc<parking_lot::ReentrantMutex<()>> {
267 Arc::clone(&self.row_constraint_lock)
268 }
269
270 pub fn get_or_init_schema(
272 &self,
273 named: &HashMap<String, reddb_types::Value>,
274 ) -> Arc<Vec<String>> {
275 {
276 let schema = self.column_schema.read();
277 if let Some(ref s) = *schema {
278 return Arc::clone(s);
279 }
280 }
281 let cols: Vec<String> = named.keys().cloned().collect();
282 let arc = Arc::new(cols);
283 *self.column_schema.write() = Some(Arc::clone(&arc));
284 arc
285 }
286
287 pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
289 self.column_schema.read().clone()
290 }
291
292 pub(crate) fn set_column_schema_if_empty(&self, columns: Vec<String>) {
293 if columns.is_empty() {
294 return;
295 }
296 let mut schema = self.column_schema.write();
297 if schema.is_none() {
298 *schema = Some(Arc::new(columns));
299 }
300 }
301
302 pub fn collection(&self) -> &str {
304 &self.collection
305 }
306
307 pub fn config(&self) -> &ManagerConfig {
309 &self.config
310 }
311
312 pub fn stats(&self) -> ManagerStats {
315 let mut s = self.stats.read().clone();
316 s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
317 s.total_memory_bytes = self.resident_bytes() as usize;
318 s
319 }
320
321 pub fn resident_bytes(&self) -> u64 {
324 self.memory_usage().1
325 }
326
327 pub fn memory_bytes(&self) -> u64 {
330 self.memory_usage().0
331 }
332
333 fn memory_usage(&self) -> (u64, u64) {
337 let consolidation = self.consolidation.read();
338 let growing = self.growing.read();
339 let sealed = self.sealed.read();
340 let mut retired = self.retired.write();
341 let mut payload = 0u64;
342 let mut resident = 0u64;
343 {
344 let mut add = |segment: &GrowingSegment| {
345 payload = payload.saturating_add(segment.memory_bytes());
346 resident = resident.saturating_add(segment.resident_bytes());
347 };
348 for segment in growing.iter().chain(sealed.iter()) {
349 add(&segment.read());
350 }
351 retired.retain(|source| {
352 if let Some(source) = source.upgrade() {
353 add(&source.read());
354 true
355 } else {
356 false
357 }
358 });
359 if let Some(run) = consolidation.as_ref() {
360 add(&run.merged);
361 }
362 }
363 let mut bookkeeping =
364 (retired.capacity() * std::mem::size_of::<Weak<RwLock<GrowingSegment>>>()) as u64;
365 if let Some(run) = consolidation.as_ref() {
366 bookkeeping = bookkeeping
367 .saturating_add(
368 (run.sources.capacity() * std::mem::size_of::<SourceProgress>()) as u64,
369 )
370 .saturating_add(
371 (run.pending_ids.capacity() * std::mem::size_of::<EntityId>()) as u64,
372 );
373 for source in &run.sources {
374 bookkeeping = bookkeeping.saturating_add(
375 (source.copied.capacity() * std::mem::size_of::<EntityId>()) as u64,
376 );
377 }
378 }
379 (
380 payload.saturating_add(bookkeeping),
381 resident.saturating_add(bookkeeping),
382 )
383 }
384
385 pub fn reclaimable_bytes(&self) -> u64 {
387 self.sealed
388 .read()
389 .iter()
390 .map(|segment| segment.read().reclaimable_bytes())
391 .fold(0, u64::saturating_add)
392 }
393
394 pub fn next_entity_id(&self) -> EntityId {
396 EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
397 }
398
399 pub fn next_row_id(&self) -> u64 {
401 self.next_row_id.fetch_add(1, Ordering::SeqCst)
402 }
403
404 pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
408 let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
409 start..start + n
410 }
411
412 pub fn register_row_id(&self, id: u64) {
415 let candidate = id.saturating_add(1);
416 let mut current = self.next_row_id.load(Ordering::SeqCst);
417 while candidate > current {
418 match self.next_row_id.compare_exchange(
419 current,
420 candidate,
421 Ordering::SeqCst,
422 Ordering::SeqCst,
423 ) {
424 Ok(_) => break,
425 Err(updated) => current = updated,
426 }
427 }
428 }
429
430 fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
437 {
439 let growing = self.growing.read();
440 if let Some(segment) = growing.as_ref() {
441 return Arc::clone(segment);
442 }
443 }
444
445 let mut growing = self.growing.write();
447 if let Some(segment) = growing.as_ref() {
449 return Arc::clone(segment);
450 }
451
452 let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
453 let segment = GrowingSegment::new(id, &self.collection);
454 let segment_arc = Arc::new(RwLock::new(segment));
455 *growing = Some(Arc::clone(&segment_arc));
456
457 self.emit(LifecycleEvent::SegmentCreated(id));
458
459 self.stats.write().growing_count += 1;
462
463 segment_arc
464 }
465
466 pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
468 self.maybe_seal_growing()?;
470
471 let segment_arc = self.get_or_create_growing();
472 let mut segment = segment_arc.write();
473
474 if entity.id.raw() == 0 {
476 entity.id = self.next_entity_id();
477 }
478
479 let entity_id = entity.id;
480 let segment_id = segment.id();
481
482 segment.insert(entity)?;
483
484 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
486
487 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
496
497 Ok(entity_id)
498 }
499
500 pub(crate) fn restore_snapshot_entity(
502 &self,
503 entity: UnifiedEntity,
504 ) -> Result<EntityId, SegmentError> {
505 self.maybe_seal_growing()?;
506 if self.get(entity.id).is_some() {
507 return Err(SegmentError::AlreadyExists(entity.id));
508 }
509 let segment_arc = self.get_or_create_growing();
510 let mut segment = segment_arc.write();
511 let entity_id = entity.id;
512 let segment_id = segment.id();
513 segment.adopt_entity(entity, None);
516 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
517 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
518 Ok(entity_id)
519 }
520
521 pub fn insert_batch(
523 &self,
524 entities: Vec<UnifiedEntity>,
525 ) -> Result<Vec<EntityId>, SegmentError> {
526 let mut ids = Vec::with_capacity(entities.len());
527 for entity in entities {
528 ids.push(self.insert(entity)?);
529 }
530 Ok(ids)
531 }
532
533 pub fn bulk_insert(
536 &self,
537 mut entities: Vec<UnifiedEntity>,
538 ) -> Result<Vec<EntityId>, SegmentError> {
539 for entity in &mut entities {
541 if entity.id.raw() == 0 {
542 entity.id = self.next_entity_id();
543 }
544 if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
545 if *row_id == 0 {
546 *row_id = self.next_row_id();
547 } else {
548 self.register_row_id(*row_id);
549 }
550 }
551 }
552
553 if let Some(first_row) = entities.first() {
556 if let super::entity::EntityData::Row(ref row) = first_row.data {
557 if let Some(ref named) = row.named {
558 let schema = self.get_or_init_schema(named);
559 for entity in &mut entities {
560 if let super::entity::EntityData::Row(ref mut row) = entity.data {
561 if let Some(named) = row.named.take() {
562 let mut cols = Vec::with_capacity(schema.len());
563 for col_name in schema.iter() {
564 cols.push(
565 named
566 .get(col_name)
567 .cloned()
568 .unwrap_or(reddb_types::Value::Null),
569 );
570 }
571 row.columns = cols;
572 row.schema = Some(Arc::clone(&schema));
573 }
574 }
575 }
576 }
577 }
578 }
579
580 let segment_arc = self.get_or_create_growing();
581 let mut segment = segment_arc.write();
582 let segment_id = segment.id();
583
584 let ids = segment.bulk_insert(entities)?;
586
587 self.total_entities_atomic
592 .fetch_add(ids.len() as u64, Ordering::Relaxed);
593
594 Ok(ids)
595 }
596
597 pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
599 self.get_with(id, Clone::clone)
600 }
601
602 pub(crate) fn get_with<T>(
605 &self,
606 id: EntityId,
607 visit: impl FnOnce(&UnifiedEntity) -> T,
608 ) -> Option<T> {
609 if let Some(growing_arc) = self.growing.read().as_ref() {
611 let growing = growing_arc.read();
612 if let Some(entity) = growing.get(id) {
613 return Some(visit(entity));
614 }
615 }
616
617 let sealed = self.sealed.read();
619 for segment in sealed.iter() {
620 let seg = segment.read();
621 if let Some(entity) = seg.get(id) {
622 return Some(visit(entity));
623 }
624 }
625
626 None
627 }
628
629 pub(crate) fn graph_node_logical_id(
633 &self,
634 id: EntityId,
635 mut before_segment: impl FnMut() -> bool,
636 ) -> Option<EntityId> {
637 if let Some(growing_arc) = self.growing.read().as_ref() {
638 if !before_segment() {
639 return None;
640 }
641 let growing = growing_arc.read();
642 if let Some(entity) = growing.get(id) {
643 return matches!(entity.kind, EntityKind::GraphNode(_))
644 .then(|| entity.logical_id());
645 }
646 }
647 let sealed = self.sealed.read();
648 for segment in sealed.iter() {
649 if !before_segment() {
650 return None;
651 }
652 let segment = segment.read();
653 if let Some(entity) = segment.get(id) {
654 return matches!(entity.kind, EntityKind::GraphNode(_))
655 .then(|| entity.logical_id());
656 }
657 }
658 None
659 }
660
661 pub(crate) fn visit_graph_index_batches<E>(
666 &self,
667 kind: GraphEntityKind,
668 keys: &[EntityId],
669 mut before_work: impl FnMut() -> bool,
670 mut admit: impl FnMut(usize) -> Result<(), E>,
671 mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
672 ) -> Result<bool, E> {
673 use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
674 admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
675 let mut admitted = 0;
676 let sources = loop {
677 let count = {
678 let growing = self.growing.read();
679 let sealed = self.sealed.read();
680 sealed.len() + usize::from(growing.is_some())
681 };
682 let bytes = count.saturating_mul(std::mem::size_of::<(
683 Arc<RwLock<GrowingSegment>>,
684 GraphIndexCursor,
685 )>());
686 admit(bytes.saturating_sub(admitted))?;
687 admitted = admitted.max(bytes);
688 let mut sources = Vec::with_capacity(count);
689 let growing = self.growing.read();
690 let sealed = self.sealed.read();
691 if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
692 continue;
693 }
694 for segment in growing.iter().chain(sealed.iter()) {
695 if !before_work() {
696 return Ok(false);
697 }
698 let view = segment.read();
699 if !view.may_contain_graph_kind(kind) {
700 continue;
701 }
702 let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
703 return Ok(false);
704 };
705 sources.push((Arc::clone(segment), cursor));
706 }
707 break sources;
708 };
709 let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
710 for (source, mut cursor) in sources {
711 while !cursor.finished(keys) {
712 let mut count = 0;
713 {
714 let segment = source.read();
715 if !segment.may_contain_graph_kind(kind) {
716 break;
717 }
718 if !segment.graph_index_batch(
719 kind,
720 keys,
721 &mut cursor,
722 &mut before_work,
723 &mut |id| {
724 ids[count] = id;
725 count += 1;
726 },
727 ) {
728 return Ok(false);
729 }
730 }
731 if !before_work() || !consume(&ids[..count])? {
732 return Ok(false);
733 }
734 }
735 }
736 Ok(true)
737 }
738
739 pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
744 let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
745 let mut remaining: Vec<usize> = (0..ids.len()).collect(); if let Some(growing_arc) = self.growing.read().as_ref() {
750 let growing = if let Some(g) = growing_arc.try_read() {
751 g
752 } else {
753 growing_arc.read()
754 };
755 remaining.retain(|&i| {
756 if let Some(entity) = growing.get(ids[i]) {
757 out[i] = Some(entity.clone());
758 false } else {
760 true }
762 });
763 }
764
765 if remaining.is_empty() {
766 return out;
767 }
768
769 let sealed = self.sealed.read();
771 for segment in sealed.iter() {
772 if remaining.is_empty() {
773 break;
774 }
775 let seg = segment.read();
776 remaining.retain(|&i| {
777 if let Some(entity) = seg.get(ids[i]) {
778 out[i] = Some(entity.clone());
779 false
780 } else {
781 true
782 }
783 });
784 }
785
786 out
787 }
788
789 pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
802 where
803 F: FnMut(usize, &UnifiedEntity),
804 {
805 thread_local! {
814 static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
815 const { std::cell::Cell::new(Vec::new()) };
816 }
817
818 let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
819 remaining.clear();
820
821 if let Some(growing_arc) = self.growing.read().as_ref() {
822 let growing = if let Some(g) = growing_arc.try_read() {
823 g
824 } else {
825 growing_arc.read()
826 };
827 for (i, id) in ids.iter().enumerate() {
828 if let Some(entity) = growing.get(*id) {
829 f(i, entity);
830 } else {
831 remaining.push(i);
832 }
833 }
834 } else {
835 remaining.reserve(ids.len());
836 remaining.extend(0..ids.len());
837 }
838
839 if !remaining.is_empty() {
840 let sealed = self.sealed.read();
841 for segment in sealed.iter() {
842 if remaining.is_empty() {
843 break;
844 }
845 let seg = segment.read();
846 remaining.retain(|&i| {
847 if let Some(entity) = seg.get(ids[i]) {
848 f(i, entity);
849 false
850 } else {
851 true
852 }
853 });
854 }
855 }
856
857 REMAINING_SCRATCH.with(|cell| cell.set(remaining));
858 }
859
860 fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
862 if let Some(growing_arc) = self.growing.read().as_ref() {
864 let growing = growing_arc.read();
865 if let Some(entity) = growing.get(id) {
866 return Some(entity.clone());
867 }
868 }
869
870 let sealed = self.sealed.read();
872 for segment in sealed.iter() {
873 if let Some(entity) = segment.get(id) {
874 return Some(entity.clone());
875 }
876 }
877
878 None
879 }
880
881 fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
882 let sealed = self.sealed.read();
883 sealed
884 .iter()
885 .find(|segment_arc| segment_arc.read().contains(id))
886 .map(Arc::clone)
887 }
888
889 fn rewrite_sealed_entity_into_growing(
890 &self,
891 entity: UnifiedEntity,
892 metadata: Option<&Metadata>,
893 ) -> Result<(), SegmentError> {
894 let entity_id = entity.id;
895 let sealed_arc = self
896 .find_sealed_segment_arc(entity_id)
897 .ok_or(SegmentError::NotFound(entity_id))?;
898
899 let metadata_to_apply = {
900 let mut sealed = sealed_arc.write();
901 let existing_metadata = sealed.get_metadata(entity_id);
902 if !sealed.force_delete(entity_id) {
903 return Err(SegmentError::NotFound(entity_id));
904 }
905 metadata.cloned().or(existing_metadata)
906 };
907
908 let growing_arc = self.get_or_create_growing();
909 let growing_id = {
910 let mut growing = growing_arc.write();
911 growing.insert(entity)?;
912 if let Some(metadata) = metadata_to_apply {
913 growing.set_metadata(entity_id, metadata)?;
914 }
915 growing.id()
916 };
917
918 self.entity_segment.write().insert(entity_id, growing_id);
919 Ok(())
920 }
921
922 pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
924 let entity_id = entity.id;
925 let mut entity = Some(entity);
926
927 if let Some(growing_arc) = self.growing.read().as_ref() {
929 let mut growing = growing_arc.write();
930 if growing.contains(entity_id) && growing.state().is_writable() {
931 return growing.update(entity.take().expect("entity already moved"));
932 }
933 }
934
935 let segment_id = self.entity_segment.read().get(&entity_id).copied();
937 if let Some(seg_id) = segment_id {
938 if let Some(growing_arc) = self.growing.read().as_ref() {
939 let mut growing = growing_arc.write();
940 if growing.id() == seg_id && growing.state().is_writable() {
941 return growing.update(entity.take().expect("entity already moved"));
942 }
943 }
944 }
945
946 if let Some(entity) = entity.take() {
947 return self.rewrite_sealed_entity_into_growing(entity, None);
948 }
949
950 Err(SegmentError::NotFound(entity_id))
951 }
952
953 pub fn update_with_metadata(
956 &self,
957 entity: UnifiedEntity,
958 metadata: Option<&Metadata>,
959 ) -> Result<(), SegmentError> {
960 let entity_id = entity.id;
961 let mut entity = Some(entity);
962
963 if let Some(growing_arc) = self.growing.read().as_ref() {
965 let mut growing = growing_arc.write();
966 if growing.contains(entity_id) && growing.state().is_writable() {
967 growing.update(entity.take().expect("entity already moved"))?;
968 if let Some(metadata) = metadata {
969 growing.set_metadata(entity_id, metadata.clone())?;
970 }
971 return Ok(());
972 }
973 }
974
975 let segment_id = self.entity_segment.read().get(&entity_id).copied();
977 if let Some(seg_id) = segment_id {
978 if let Some(growing_arc) = self.growing.read().as_ref() {
979 let mut growing = growing_arc.write();
980 if growing.id() == seg_id && growing.state().is_writable() {
981 growing.update(entity.take().expect("entity already moved"))?;
982 if let Some(metadata) = metadata {
983 growing.set_metadata(entity_id, metadata.clone())?;
984 }
985 return Ok(());
986 }
987 }
988 }
989
990 if let Some(entity) = entity.take() {
991 return self.rewrite_sealed_entity_into_growing(entity, metadata);
992 }
993
994 Err(SegmentError::NotFound(entity_id))
995 }
996
997 pub fn update_hot(
1001 &self,
1002 entity: UnifiedEntity,
1003 modified_columns: &[String],
1004 ) -> Result<(), SegmentError> {
1005 let entity_id = entity.id;
1006 let mut entity = Some(entity);
1007
1008 if let Some(growing_arc) = self.growing.read().as_ref() {
1009 let mut growing = growing_arc.write();
1010 if growing.contains(entity_id) && growing.state().is_writable() {
1011 return growing.update_hot(
1012 entity.take().expect("entity already moved"),
1013 modified_columns,
1014 );
1015 }
1016 }
1017
1018 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 return growing.update_hot(
1024 entity.take().expect("entity already moved"),
1025 modified_columns,
1026 );
1027 }
1028 }
1029 }
1030
1031 if let Some(entity) = entity.take() {
1032 return self.rewrite_sealed_entity_into_growing(entity, None);
1033 }
1034
1035 Err(SegmentError::NotFound(entity_id))
1036 }
1037
1038 pub fn update_hot_with_metadata(
1041 &self,
1042 entity: UnifiedEntity,
1043 modified_columns: &[String],
1044 metadata: Option<&Metadata>,
1045 ) -> Result<(), SegmentError> {
1046 let entity_id = entity.id;
1047 let mut entity = Some(entity);
1048
1049 if let Some(growing_arc) = self.growing.read().as_ref() {
1050 let mut growing = growing_arc.write();
1051 if growing.contains(entity_id) && growing.state().is_writable() {
1052 growing.update_hot(
1053 entity.take().expect("entity already moved"),
1054 modified_columns,
1055 )?;
1056 if let Some(metadata) = metadata {
1057 growing.set_metadata(entity_id, metadata.clone())?;
1058 }
1059 return Ok(());
1060 }
1061 }
1062
1063 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1064 if let Some(seg_id) = segment_id {
1065 if let Some(growing_arc) = self.growing.read().as_ref() {
1066 let mut growing = growing_arc.write();
1067 if growing.id() == seg_id && growing.state().is_writable() {
1068 growing.update_hot(
1069 entity.take().expect("entity already moved"),
1070 modified_columns,
1071 )?;
1072 if let Some(metadata) = metadata {
1073 growing.set_metadata(entity_id, metadata.clone())?;
1074 }
1075 return Ok(());
1076 }
1077 }
1078 }
1079
1080 if let Some(entity) = entity.take() {
1081 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1082 }
1083
1084 Err(SegmentError::NotFound(entity_id))
1085 }
1086
1087 pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
1090 where
1091 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
1092 {
1093 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
1094 items.into_iter().collect();
1095 if items.is_empty() {
1096 return Ok(());
1097 }
1098
1099 if let Some(growing_arc) = self.growing.read().as_ref() {
1100 let mut growing = growing_arc.write();
1101 if growing.state().is_writable() {
1102 match growing.update_hot_batch_with_metadata(items.iter().copied()) {
1103 Ok(()) => return Ok(()),
1104 Err(SegmentError::NotFound(_)) => {}
1105 Err(other) => return Err(other),
1106 }
1107 }
1108 }
1109
1110 for (entity, modified_columns, metadata) in items {
1111 self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
1112 }
1113 Ok(())
1114 }
1115
1116 pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
1118 if let Some(growing_arc) = self.growing.read().as_ref() {
1121 let mut growing = growing_arc.write();
1122 if growing.contains(id) && growing.state().is_writable() {
1123 let seg_id = growing.id();
1124 let deleted = growing.delete(id)?;
1125 if deleted {
1126 self.entity_segment.write().remove(&id);
1127 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1128 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1129 }
1130 return Ok(deleted);
1131 }
1132 }
1133
1134 let segment_id = self.entity_segment.read().get(&id).copied();
1137
1138 if let Some(seg_id) = segment_id {
1139 if let Some(growing_arc) = self.growing.read().as_ref() {
1140 let mut growing = growing_arc.write();
1141 if growing.id() == seg_id && growing.state().is_writable() {
1142 let deleted = growing.delete(id)?;
1143 if deleted {
1144 self.entity_segment.write().remove(&id);
1145 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1146 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1147 }
1148 return Ok(deleted);
1149 }
1150 }
1151 }
1152
1153 {
1156 let sealed = self.sealed.read();
1157 for segment_arc in sealed.iter() {
1158 let mut seg = segment_arc.write();
1159 let seg_id = seg.id();
1160 if seg.contains(id) {
1161 let deleted = seg.force_delete(id);
1162 drop(seg);
1163 if deleted {
1164 self.entity_segment.write().remove(&id);
1165 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1166 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1167 }
1168 return Ok(deleted);
1169 }
1170 }
1171 }
1172
1173 Ok(false)
1174 }
1175
1176 pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1177 if ids.is_empty() {
1178 return Ok(Vec::new());
1179 }
1180
1181 let mut deleted_ids = Vec::with_capacity(ids.len());
1182
1183 if let Some(growing_arc) = self.growing.read().as_ref() {
1184 let mut growing = growing_arc.write();
1185 if growing.state().is_writable() {
1186 let seg_id = growing.id();
1187 let deleted = growing.delete_batch(ids)?;
1188 if !deleted.is_empty() {
1189 {
1190 let mut entity_segment = self.entity_segment.write();
1191 for id in &deleted {
1192 entity_segment.remove(id);
1193 }
1194 }
1195 self.total_entities_atomic
1196 .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
1197 for id in &deleted {
1198 self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1199 }
1200 deleted_ids.extend(deleted);
1201 }
1202 }
1203 }
1204
1205 if deleted_ids.len() == ids.len() {
1206 return Ok(deleted_ids);
1207 }
1208
1209 let deleted_set: std::collections::HashSet<EntityId> =
1210 deleted_ids.iter().copied().collect();
1211 for &id in ids {
1212 if deleted_set.contains(&id) {
1213 continue;
1214 }
1215 if self.delete(id)? {
1216 deleted_ids.push(id);
1217 }
1218 }
1219
1220 Ok(deleted_ids)
1221 }
1222
1223 pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1225 if let Some(growing_arc) = self.growing.read().as_ref() {
1227 let growing = growing_arc.read();
1228 if growing.contains(id) {
1229 return growing.get_metadata(id);
1230 }
1231 }
1232
1233 let segment_id = self.entity_segment.read().get(&id).copied();
1235
1236 if let Some(seg_id) = segment_id {
1237 if let Some(growing_arc) = self.growing.read().as_ref() {
1238 let growing = growing_arc.read();
1239 if growing.id() == seg_id {
1240 return growing.get_metadata(id);
1241 }
1242 }
1243
1244 let sealed = self.sealed.read();
1245 for segment in sealed.iter() {
1246 if segment.id() == seg_id {
1247 return segment.get_metadata(id);
1248 }
1249 }
1250 }
1251
1252 if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1253 return segment_arc.read().get_metadata(id);
1254 }
1255
1256 None
1257 }
1258
1259 pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1261 if let Some(growing_arc) = self.growing.read().as_ref() {
1264 let mut growing = growing_arc.write();
1265 if growing.contains(id) && growing.state().is_writable() {
1266 return growing.set_metadata(id, metadata);
1267 }
1268 }
1269
1270 let segment_id = self.entity_segment.read().get(&id).copied();
1272
1273 if let Some(seg_id) = segment_id {
1274 if let Some(growing_arc) = self.growing.read().as_ref() {
1275 let mut growing = growing_arc.write();
1276 if growing.id() == seg_id && growing.state().is_writable() {
1277 return growing.set_metadata(id, metadata);
1278 }
1279 }
1280 }
1281
1282 if let Some(entity) = self.get(id) {
1283 return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1284 }
1285
1286 Err(SegmentError::NotFound(id))
1287 }
1288
1289 fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1291 let should_seal = {
1292 let growing_opt = self.growing.read();
1293 if let Some(growing_arc) = growing_opt.as_ref() {
1294 let growing = growing_arc.read();
1295 growing.should_seal(&self.config.segment_config)
1296 || growing.idle_secs() >= self.config.idle_seal_secs
1297 } else {
1298 false
1299 }
1300 };
1301
1302 if should_seal {
1303 self.seal_current()?;
1304 }
1305
1306 Ok(())
1307 }
1308
1309 pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1311 let mut growing_slot = self.growing.write();
1314 if let Some(growing_arc) = growing_slot.take() {
1315 let mut growing = growing_arc.write();
1316 let seg_id = growing.id();
1317 let entity_count = growing.stats().entity_count as u64;
1318
1319 growing.seal()?;
1321
1322 drop(growing); self.sealed.write().push(growing_arc);
1328 drop(growing_slot);
1329
1330 self.mark_sealed_pages_visible(entity_count);
1332
1333 {
1335 let mut stats = self.stats.write();
1336 stats.growing_count = stats.growing_count.saturating_sub(1);
1337 stats.sealed_count += 1;
1338 stats.seal_ops += 1;
1339 }
1340
1341 self.emit(LifecycleEvent::SegmentSealed(seg_id));
1342
1343 return Ok(seg_id);
1344 }
1345
1346 Err(SegmentError::InvalidState(SegmentState::Sealed))
1347 }
1348
1349 pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1351 let has_growing = self.growing.read().is_some();
1352 if has_growing {
1353 Ok(Some(self.seal_current()?))
1354 } else {
1355 Ok(None)
1356 }
1357 }
1358
1359 pub fn all_visible_fraction(&self) -> f64 {
1369 const ROWS_PER_PAGE: u32 = 256;
1370 let sealed = self.sealed.read();
1371 if sealed.is_empty() {
1372 return 0.0;
1373 }
1374 let mut total_pages: u64 = 0;
1375 for seg_arc in sealed.iter() {
1376 let seg = seg_arc.read();
1377 let entity_count = seg.stats().entity_count as u64;
1378 let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1379 total_pages += pages;
1380 }
1381 if total_pages == 0 {
1382 return 0.0;
1383 }
1384 let visible = self.visibility_map.all_visible_count();
1385 (visible as f64 / total_pages as f64).min(1.0)
1386 }
1387
1388 fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1391 const ROWS_PER_PAGE: u32 = 256;
1392 let existing_visible = self.visibility_map.all_visible_count();
1393 let start_page = existing_visible as u32;
1395 let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1396 let end_page = start_page + new_pages as u32;
1397 self.visibility_map.mark_range_visible(start_page, end_page);
1398 }
1399
1400 pub(crate) fn find_unique_key_conflict(
1404 &self,
1405 columns: &[Vec<String>],
1406 key_index: usize,
1407 signatures: &[String],
1408 exclude_id: Option<EntityId>,
1409 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1410 ) -> Option<EntityId> {
1411 if let Some(growing) = self.growing.read().as_ref() {
1412 if let Some(id) = growing.read().find_unique_key_conflict(
1413 columns,
1414 key_index,
1415 signatures,
1416 exclude_id,
1417 reserves_key,
1418 ) {
1419 return Some(id);
1420 }
1421 }
1422 for segment in self.sealed.read().iter() {
1423 if let Some(id) = segment.read().find_unique_key_conflict(
1424 columns,
1425 key_index,
1426 signatures,
1427 exclude_id,
1428 reserves_key,
1429 ) {
1430 return Some(id);
1431 }
1432 }
1433 None
1434 }
1435
1436 pub fn for_each_entity<F>(&self, callback: F)
1442 where
1443 F: FnMut(&UnifiedEntity) -> bool,
1444 {
1445 self.for_each_entity_in_segments(|_| true, || true, callback);
1446 }
1447
1448 pub(crate) fn checkpoint_revisions(&self) -> Vec<(u64, u64)> {
1450 let growing = self.growing.read();
1451 let sealed = self.sealed.read();
1452 growing
1453 .iter()
1454 .chain(sealed.iter())
1455 .map(|segment| {
1456 let segment = segment.read();
1457 (segment.id(), segment.checkpoint_revision())
1458 })
1459 .collect()
1460 }
1461
1462 pub(crate) fn cached_physical_root(&self, seed: u64, compute: impl FnOnce() -> u64) -> u64 {
1464 let segments = self.checkpoint_revisions();
1465 let mut cache = self.physical_root_cache.lock();
1466 if let Some(previous) = cache.as_ref() {
1467 if previous.seed == seed && previous.segments == segments {
1468 return previous.root;
1469 }
1470 }
1471 let root = compute();
1472 *cache = Some(PhysicalRootCacheEntry {
1473 segments,
1474 seed,
1475 root,
1476 });
1477 root
1478 }
1479
1480 pub(crate) fn for_each_entity_in_id_order(
1485 &self,
1486 mut visit: impl FnMut(&UnifiedEntity) -> bool,
1487 ) {
1488 self.for_each_entity_source_in_id_order(|_, entity| visit(entity));
1489 }
1490
1491 pub(crate) fn for_each_entity_with_metadata_in_id_order(
1492 &self,
1493 mut visit: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1494 ) {
1495 self.for_each_entity_source_in_id_order(|segment, entity| {
1496 visit(entity, segment.get_metadata(entity.id))
1497 });
1498 }
1499
1500 fn for_each_entity_source_in_id_order(
1501 &self,
1502 mut visit: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
1503 ) {
1504 let segments: Vec<_> = {
1505 let growing = self.growing.read();
1506 let sealed = self.sealed.read();
1507 growing.iter().chain(sealed.iter()).cloned().collect()
1508 };
1509 let mut order = Vec::new();
1510 for (source, segment) in segments.iter().enumerate() {
1511 segment.read().for_each_fast(|entity| {
1512 order.push((entity.id, source));
1513 true
1514 });
1515 }
1516 order.sort_by_key(|(id, _)| id.raw());
1517 for (id, source) in order {
1518 let segment = segments[source].read();
1519 if let Some(entity) = segment.get(id) {
1520 if !visit(&segment, entity) {
1521 break;
1522 }
1523 }
1524 }
1525 }
1526
1527 pub(crate) fn for_each_entity_with_metadata(
1530 &self,
1531 mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1532 ) {
1533 if let Some(growing_arc) = self.growing.read().as_ref() {
1534 let growing = growing_arc.read();
1535 if !growing.for_each_fast(|entity| callback(entity, growing.get_metadata(entity.id))) {
1536 return;
1537 }
1538 }
1539 let sealed = self.sealed.read();
1540 for segment_arc in sealed.iter() {
1541 let segment = segment_arc.read();
1542 if !segment.for_each_fast(|entity| callback(entity, segment.get_metadata(entity.id))) {
1543 return;
1544 }
1545 }
1546 }
1547
1548 fn for_each_entity_in_segments<F>(
1549 &self,
1550 segment_filter: impl Fn(&GrowingSegment) -> bool,
1551 mut before_segment: impl FnMut() -> bool,
1552 mut callback: F,
1553 ) where
1554 F: FnMut(&UnifiedEntity) -> bool,
1555 {
1556 if let Some(growing_arc) = self.growing.read().as_ref() {
1560 if !before_segment() {
1561 return;
1562 }
1563 let growing = if let Some(g) = growing_arc.try_read() {
1564 g
1565 } else {
1566 growing_arc.read()
1567 };
1568 if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
1569 return;
1570 }
1571 }
1572
1573 let sealed = self.sealed.read();
1575 for segment_arc in sealed.iter() {
1576 if !before_segment() {
1577 return;
1578 }
1579 let segment = segment_arc.read();
1580 if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
1581 return;
1582 }
1583 }
1584 }
1585
1586 pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1602 &self,
1603 init: FInit,
1604 fold: FFold,
1605 reduce: FReduce,
1606 ) -> T
1607 where
1608 T: Send,
1609 FInit: Fn() -> T + Send + Sync,
1610 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1611 FReduce: Fn(T, T) -> T + Send + Sync,
1612 {
1613 use rayon::prelude::*;
1614
1615 let mut acc = init();
1618 if let Some(growing_arc) = self.growing.read().as_ref() {
1619 let growing = if let Some(g) = growing_arc.try_read() {
1620 g
1621 } else {
1622 growing_arc.read()
1623 };
1624 growing.for_each_fast(|entity| {
1625 acc = fold(std::mem::replace(&mut acc, init()), entity);
1626 true
1627 });
1628 }
1629
1630 let segments: Vec<_> = {
1634 let sealed = self.sealed.read();
1635 sealed.iter().cloned().collect()
1636 };
1637
1638 if segments.len() <= 1 {
1639 for seg_arc in &segments {
1640 let seg = seg_arc.read();
1641 seg.for_each_fast(|entity| {
1642 acc = fold(std::mem::replace(&mut acc, init()), entity);
1643 true
1644 });
1645 }
1646 return acc;
1647 }
1648
1649 let sealed_acc = segments
1650 .into_par_iter()
1651 .map(|seg_arc| {
1652 let mut local = init();
1653 let seg = seg_arc.read();
1654 seg.for_each_fast(|entity| {
1655 local = fold(std::mem::replace(&mut local, init()), entity);
1656 true
1657 });
1658 local
1659 })
1660 .reduce(&init, &reduce);
1661
1662 reduce(acc, sealed_acc)
1663 }
1664
1665 pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
1674 where
1675 F: FnMut(&UnifiedEntity) -> bool,
1676 {
1677 let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
1678 }
1679
1680 pub fn for_each_entity_zoned_with_stats<F>(
1681 &self,
1682 zone_preds: &[(&str, ZoneColPred<'_>)],
1683 mut callback: F,
1684 ) -> SegmentScanStats
1685 where
1686 F: FnMut(&UnifiedEntity) -> bool,
1687 {
1688 let mut scan_stats = SegmentScanStats::default();
1689 if let Some(growing_arc) = self.growing.read().as_ref() {
1695 scan_stats.segments_total += 1;
1696 scan_stats.segments_scanned += 1;
1697 let growing = if let Some(g) = growing_arc.try_read() {
1698 g
1699 } else {
1700 growing_arc.read()
1701 };
1702 if !growing.for_each_fast(&mut callback) {
1703 return scan_stats;
1704 }
1705 }
1706
1707 let sealed = self.sealed.read();
1709 for segment_arc in sealed.iter() {
1710 scan_stats.segments_total += 1;
1711 let segment = segment_arc.read();
1712 if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
1713 scan_stats.segments_pruned += 1;
1714 continue; }
1716 scan_stats.segments_scanned += 1;
1717 if !segment.for_each_fast(&mut callback) {
1718 return scan_stats;
1719 }
1720 }
1721 scan_stats
1722 }
1723
1724 pub fn query_all_zoned<F>(
1734 &self,
1735 zone_preds: &[(&str, ZoneColPred<'_>)],
1736 filter: F,
1737 ) -> Vec<UnifiedEntity>
1738 where
1739 F: Fn(&UnifiedEntity) -> bool + Sync,
1740 {
1741 self.query_all_zoned_with_stats(zone_preds, filter).0
1742 }
1743
1744 pub fn query_all_zoned_with_stats<F>(
1745 &self,
1746 zone_preds: &[(&str, ZoneColPred<'_>)],
1747 filter: F,
1748 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
1749 where
1750 F: Fn(&UnifiedEntity) -> bool + Sync,
1751 {
1752 let mut results = Vec::new();
1753 let mut scan_stats = SegmentScanStats::default();
1754
1755 if let Some(growing_arc) = self.growing.read().as_ref() {
1758 scan_stats.segments_total += 1;
1759 scan_stats.segments_scanned += 1;
1760 let growing = if let Some(g) = growing_arc.try_read() {
1761 g
1762 } else {
1763 growing_arc.read()
1764 };
1765 results.extend(growing.iter().filter(|e| filter(e)).cloned());
1766 }
1767
1768 let sealed = self.sealed.read();
1770 let surviving: Vec<_> = sealed
1772 .iter()
1773 .filter(|seg_arc| {
1774 scan_stats.segments_total += 1;
1775 if zone_preds.is_empty() {
1776 scan_stats.segments_scanned += 1;
1777 return true;
1778 }
1779 let seg = seg_arc.read();
1780 if seg.can_skip_zone_preds(zone_preds) {
1781 scan_stats.segments_pruned += 1;
1782 false
1783 } else {
1784 scan_stats.segments_scanned += 1;
1785 true
1786 }
1787 })
1788 .collect();
1789
1790 let use_parallel = surviving.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1791
1792 if use_parallel {
1793 let filter_ref = &filter;
1794 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1795 surviving
1796 .iter()
1797 .map(|segment| {
1798 s.spawn(move || {
1799 segment
1800 .read()
1801 .iter()
1802 .filter(|e| filter_ref(e))
1803 .cloned()
1804 .collect::<Vec<_>>()
1805 })
1806 })
1807 .collect::<Vec<_>>()
1808 .into_iter()
1809 .map(|handle| handle.join().unwrap_or_default())
1810 .collect()
1811 });
1812 for batch in segment_results {
1813 results.extend(batch);
1814 }
1815 } else {
1816 for segment_arc in surviving {
1817 let seg = segment_arc.read();
1818 results.extend(seg.iter().filter(|e| filter(e)).cloned());
1819 }
1820 }
1821
1822 (results, scan_stats)
1823 }
1824
1825 fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
1831 match snapshot {
1832 Some(ctx) => entity_visible_with_context(Some(ctx), entity),
1833 None => {
1834 !crate::runtime::ai::moderation::entity_moderation_hidden(entity)
1835 && entity.xmax == 0
1836 }
1837 }
1838 }
1839
1840 pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
1847 where
1848 F: Fn(&UnifiedEntity) -> bool + Sync,
1849 {
1850 self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
1851 }
1852
1853 pub(crate) fn scan_graph_kind(
1855 &self,
1856 snapshot: Option<&SnapshotContext>,
1857 kind: GraphEntityKind,
1858 ) -> Vec<UnifiedEntity> {
1859 self.query_all_in_segments(
1860 |segment| segment.may_contain_graph_kind(kind),
1861 |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
1862 )
1863 }
1864
1865 pub(crate) fn scan_graph_candidates_for_each<F>(
1870 &self,
1871 snapshot: Option<&SnapshotContext>,
1872 kind: GraphEntityKind,
1873 before_segment: impl FnMut() -> bool,
1874 mut callback: F,
1875 ) where
1876 F: FnMut(&UnifiedEntity) -> bool,
1877 {
1878 self.for_each_entity_in_segments(
1879 |segment| segment.may_contain_graph_kind(kind),
1880 before_segment,
1881 |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
1882 );
1883 }
1884
1885 pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
1890 where
1891 F: FnMut(&UnifiedEntity) -> bool,
1892 {
1893 self.for_each_entity(|entity| {
1894 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
1895 });
1896 }
1897
1898 pub(crate) fn scan_batches(
1904 &self,
1905 snapshot: Option<&SnapshotContext>,
1906 mut before_work: impl FnMut() -> bool,
1907 mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
1908 ) -> bool {
1909 let segments = {
1910 let growing = self.growing.read();
1911 let sealed = self.sealed.read();
1912 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
1913 for segment in growing.iter().chain(sealed.iter()) {
1914 if !before_work() {
1915 return false;
1916 }
1917 let cursor = segment.read().scan_cursor();
1918 segments.push((Arc::clone(segment), cursor));
1919 }
1920 segments
1921 };
1922 let mut ids = Vec::new();
1923 let mut batch = Vec::new();
1924 for (segment, mut cursor) in segments {
1925 while !cursor.finished() {
1926 ids.clear();
1927 batch.clear();
1928 let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
1929 ids.reserve_exact(batch_size);
1930 batch.reserve_exact(batch_size);
1931 {
1932 let segment = segment.read();
1933 if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
1934 return false;
1935 }
1936 }
1937 if !before_work() {
1940 return false;
1941 }
1942 self.for_each_id(&ids, |_, entity| {
1943 if Self::scan_entity_visible(snapshot, entity) {
1944 batch.push(entity.clone());
1945 }
1946 });
1947 if !before_work() {
1948 return false;
1949 }
1950 if !consume(&batch) {
1951 return false;
1952 }
1953 }
1954 }
1955 true
1956 }
1957
1958 pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
1960 &self,
1961 snapshot: Option<&SnapshotContext>,
1962 init: FInit,
1963 fold: FFold,
1964 reduce: FReduce,
1965 ) -> T
1966 where
1967 T: Send,
1968 FInit: Fn() -> T + Send + Sync,
1969 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1970 FReduce: Fn(T, T) -> T + Send + Sync,
1971 {
1972 self.fold_entities_parallel(
1973 init,
1974 |acc, entity| {
1975 if Self::scan_entity_visible(snapshot, entity) {
1976 fold(acc, entity)
1977 } else {
1978 acc
1979 }
1980 },
1981 reduce,
1982 )
1983 }
1984
1985 pub fn scan_for_each_zoned_with_stats<F>(
1987 &self,
1988 snapshot: Option<&SnapshotContext>,
1989 zone_preds: &[(&str, ZoneColPred<'_>)],
1990 mut callback: F,
1991 ) -> SegmentScanStats
1992 where
1993 F: FnMut(&UnifiedEntity) -> bool,
1994 {
1995 self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
1996 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
1997 })
1998 }
1999
2000 pub fn scan_zoned_with_stats<F>(
2002 &self,
2003 snapshot: Option<&SnapshotContext>,
2004 zone_preds: &[(&str, ZoneColPred<'_>)],
2005 filter: F,
2006 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2007 where
2008 F: Fn(&UnifiedEntity) -> bool + Sync,
2009 {
2010 self.query_all_zoned_with_stats(zone_preds, |entity| {
2011 Self::scan_entity_visible(snapshot, entity) && filter(entity)
2012 })
2013 }
2014
2015 pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
2018 where
2019 F: Fn(&UnifiedEntity) -> bool + Sync,
2020 {
2021 self.query_all_in_segments(|_| true, filter)
2022 }
2023
2024 fn query_all_in_segments<F>(
2025 &self,
2026 segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
2027 filter: F,
2028 ) -> Vec<UnifiedEntity>
2029 where
2030 F: Fn(&UnifiedEntity) -> bool + Sync,
2031 {
2032 let mut results = Vec::new();
2033
2034 if let Some(growing_arc) = self.growing.read().as_ref() {
2037 let growing = if let Some(g) = growing_arc.try_read() {
2038 g
2039 } else {
2040 growing_arc.read()
2041 };
2042 if segment_filter(&growing) {
2043 results.extend(growing.iter().filter(|e| filter(e)).cloned());
2044 }
2045 }
2046
2047 let sealed = self.sealed.read();
2049 let use_parallel = sealed.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
2050 if use_parallel {
2051 let filter_ref = &filter;
2052 let segment_filter_ref = &segment_filter;
2053 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
2054 sealed
2055 .iter()
2056 .map(|segment| {
2057 s.spawn(move || {
2058 let segment = segment.read();
2059 if segment_filter_ref(&segment) {
2060 segment
2061 .iter()
2062 .filter(|e| filter_ref(e))
2063 .cloned()
2064 .collect::<Vec<_>>()
2065 } else {
2066 Vec::new()
2067 }
2068 })
2069 })
2070 .collect::<Vec<_>>()
2071 .into_iter()
2072 .map(|handle| handle.join().unwrap_or_default())
2073 .collect()
2074 });
2075 for batch in segment_results {
2076 results.extend(batch);
2077 }
2078 } else {
2079 for segment in sealed.iter() {
2080 let seg = segment.read();
2081 if segment_filter(&seg) {
2082 results.extend(seg.iter().filter(|e| filter(e)).cloned());
2083 }
2084 }
2085 }
2086
2087 results
2088 }
2089
2090 pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
2095 if let Some(growing_arc) = self.growing.read().as_ref() {
2096 let growing = growing_arc.read();
2097 if growing.may_contain_exact_key(key) {
2098 return true;
2099 }
2100 } else {
2101 return true;
2102 }
2103
2104 let sealed = self.sealed.read();
2105 for segment_arc in sealed.iter() {
2106 let segment = segment_arc.read();
2107 if segment.may_contain_exact_key(key) {
2108 return true;
2109 }
2110 }
2111
2112 false
2113 }
2114
2115 pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2117 let mut results = Vec::new();
2118
2119 if let Some(growing_arc) = self.growing.read().as_ref() {
2121 let growing = growing_arc.read();
2122 results.extend(growing.filter_metadata(filters));
2123 }
2124
2125 let sealed = self.sealed.read();
2127 for segment in sealed.iter() {
2128 results.extend(segment.filter_metadata(filters));
2129 }
2130
2131 results
2132 }
2133
2134 pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2136 let mut results = Vec::new();
2137
2138 if let Some(growing_arc) = self.growing.read().as_ref() {
2140 let growing = growing_arc.read();
2141 for entity in growing.iter_kind(kind) {
2142 results.push(entity.clone());
2143 }
2144 }
2145
2146 let sealed = self.sealed.read();
2151 for segment_arc in sealed.iter() {
2152 let segment = segment_arc.read();
2153 for entity in segment.iter_kind(kind) {
2154 results.push(entity.clone());
2155 }
2156 }
2157
2158 results
2159 }
2160
2161 pub fn count(&self) -> usize {
2163 self.total_entities_atomic.load(Ordering::Relaxed) as usize
2164 }
2165
2166 pub fn segment_ids(&self) -> Vec<SegmentId> {
2168 let mut ids = Vec::new();
2169
2170 if let Some(growing_arc) = self.growing.read().as_ref() {
2171 ids.push(growing_arc.read().id());
2172 }
2173
2174 let sealed = self.sealed.read();
2175 for segment in sealed.iter() {
2176 ids.push(segment.id());
2177 }
2178
2179 ids.extend(self.archived.read().iter().copied());
2180
2181 ids
2182 }
2183
2184 #[inline]
2197 #[allow(clippy::unused_self)]
2198 fn emit(&self, _event: LifecycleEvent) {}
2199
2200 pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2203 std::mem::take(&mut *self.events.write())
2204 }
2205
2206 pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2208 self.maybe_seal_growing()?;
2210
2211 if self.config.enable_consolidation {
2212 self.consolidation_tick();
2213 }
2214
2215 Ok(())
2216 }
2217
2218 fn consolidation_tick(&self) {
2229 let running = self.consolidation.read().is_some();
2232 if !running && !self.start_consolidation() {
2233 return;
2234 }
2235
2236 let mut guard = self.consolidation.write();
2239 let Some(run) = guard.as_mut() else {
2240 return;
2241 };
2242 self.copy_bounded(run);
2243 if run.cursor >= run.sources.len() {
2244 if let Some(run) = guard.take() {
2245 self.finish_consolidation(run);
2246 }
2247 }
2248 }
2249
2250 fn start_consolidation(&self) -> bool {
2254 let sources = self.select_candidates();
2255 self.start_consolidation_from_sources(sources)
2256 }
2257
2258 fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
2259 if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
2260 return false;
2261 }
2262
2263 let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
2264 let merged = GrowingSegment::new(merged_id, &self.collection);
2265
2266 *self.consolidation.write() = Some(Consolidation {
2267 sources: sources
2268 .into_iter()
2269 .map(|id| SourceProgress {
2270 id,
2271 copied: Vec::new(),
2272 tombstones_at_open: None,
2273 })
2274 .collect(),
2275 merged_id,
2276 merged,
2277 cursor: 0,
2278 pending_ids: Vec::new(),
2279 pending_cursor: 0,
2280 entities_copied: 0,
2281 });
2282
2283 self.stats.write().consolidation.runs_started += 1;
2284 true
2285 }
2286
2287 pub fn pressure_consolidation_tick(&self) -> bool {
2292 if !self.config.enable_consolidation {
2293 return false;
2294 }
2295
2296 let running = self.consolidation.read().is_some();
2297 if !running {
2298 let sources = self.dirty_sealed_segments();
2299 if !self.start_consolidation_from_sources(sources) {
2300 return false;
2301 }
2302 }
2303
2304 self.consolidation_tick();
2305 true
2306 }
2307
2308 fn select_candidates(&self) -> Vec<SegmentId> {
2323 let sealed = self.sealed.read();
2324 if sealed.is_empty() {
2325 return Vec::new();
2326 }
2327
2328 if sealed.len() > self.config.max_sealed_segments {
2329 return sealed.iter().map(|segment| segment.read().id()).collect();
2330 }
2331
2332 let mut live = 0usize;
2333 let mut tombstones = 0usize;
2334 let mut held = 0u64;
2335 let mut reclaimable = 0u64;
2336 let mut dirty = Vec::new();
2337
2338 for segment_arc in sealed.iter() {
2339 let segment = segment_arc.read();
2340 live += segment.live_entity_count();
2341 tombstones += segment.tombstone_count();
2342 held += segment.resident_bytes();
2343 reclaimable += segment.reclaimable_bytes();
2344 if segment.tombstone_count() > 0 {
2345 dirty.push(segment.id());
2346 }
2347 }
2348
2349 let total = live + tombstones;
2350 let tombstone_ratio = if total == 0 {
2351 0.0
2352 } else {
2353 tombstones as f64 / total as f64
2354 };
2355 let fragmentation = if held == 0 {
2356 0.0
2357 } else {
2358 reclaimable as f64 / held as f64
2359 };
2360
2361 if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
2362 || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
2363 {
2364 return dirty;
2365 }
2366
2367 Vec::new()
2368 }
2369
2370 fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
2371 self.sealed
2372 .read()
2373 .iter()
2374 .filter_map(|segment_arc| {
2375 let segment = segment_arc.read();
2376 (segment.reclaimable_bytes() > 0).then(|| segment.id())
2377 })
2378 .collect()
2379 }
2380
2381 fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
2384 if sources.is_empty() {
2385 return false;
2386 }
2387 let sealed = self.sealed.read();
2388 sealed.iter().any(|segment_arc| {
2389 let segment = segment_arc.read();
2390 sources.contains(&segment.id()) && segment.tombstone_count() > 0
2391 })
2392 }
2393
2394 fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
2395 let sealed = self.sealed.read();
2396 sealed
2397 .iter()
2398 .find(|segment| segment.read().id() == id)
2399 .map(Arc::clone)
2400 }
2401
2402 fn copy_bounded(&self, run: &mut Consolidation) {
2407 let mut budget = self.config.consolidation_entities_per_tick;
2408
2409 while budget > 0 && run.cursor < run.sources.len() {
2410 let source_id = run.sources[run.cursor].id;
2411 let Some(source_arc) = self.sealed_segment(source_id) else {
2412 run.cursor += 1;
2416 run.pending_ids.clear();
2417 run.pending_cursor = 0;
2418 continue;
2419 };
2420
2421 if run.sources[run.cursor].tombstones_at_open.is_none() {
2422 let source = source_arc.read();
2423 run.pending_ids = source.live_entity_ids();
2424 run.pending_cursor = 0;
2425 run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
2426 }
2427
2428 {
2429 let source = source_arc.read();
2430 while budget > 0 && run.pending_cursor < run.pending_ids.len() {
2431 let id = run.pending_ids[run.pending_cursor];
2432 run.pending_cursor += 1;
2433 budget -= 1;
2434
2435 let Some(entity) = source.get(id) else {
2438 continue;
2439 };
2440 let metadata = source.get_metadata(id);
2441 run.merged.adopt_entity(entity.clone(), metadata);
2442 run.sources[run.cursor].copied.push(id);
2443 run.entities_copied += 1;
2444 }
2445 }
2446
2447 if run.pending_cursor >= run.pending_ids.len() {
2448 run.cursor += 1;
2449 run.pending_ids.clear();
2450 run.pending_cursor = 0;
2451 }
2452 }
2453 }
2454
2455 fn finish_consolidation(&self, mut run: Consolidation) {
2461 let mut sealed = self.sealed.write();
2462
2463 let mut positions = Vec::with_capacity(run.sources.len());
2464 for source in &run.sources {
2465 if let Some(index) = sealed
2466 .iter()
2467 .position(|segment| segment.read().id() == source.id)
2468 {
2469 positions.push(index);
2470 }
2471 }
2472 if positions.is_empty() {
2473 return;
2474 }
2475 positions.sort_unstable();
2476
2477 let mut stale = Vec::new();
2481 for source in &run.sources {
2482 let Some(source_arc) = sealed
2483 .iter()
2484 .find(|segment| segment.read().id() == source.id)
2485 else {
2486 continue;
2487 };
2488 let segment = source_arc.read();
2489 if source.tombstones_at_open == Some(segment.tombstone_count()) {
2490 continue;
2491 }
2492 stale.extend(
2493 source
2494 .copied
2495 .iter()
2496 .copied()
2497 .filter(|id| !segment.contains(*id)),
2498 );
2499 }
2500 for id in stale {
2501 run.merged.evict_entity(id);
2502 }
2503
2504 let reclaimed_bytes: u64 = positions
2505 .iter()
2506 .map(|&index| sealed[index].read().resident_bytes())
2507 .sum();
2508 let reclaimed_tombstones: u64 = positions
2509 .iter()
2510 .map(|&index| sealed[index].read().tombstone_count() as u64)
2511 .sum();
2512
2513 if run.merged.seal().is_err() {
2514 return;
2515 }
2516 let merged_bytes = run.merged.resident_bytes();
2517 let merged_arc = Arc::new(RwLock::new(run.merged));
2518
2519 {
2520 let mut retired = self.retired.write();
2521 retired.retain(|source| source.strong_count() > 0);
2522 for &index in positions.iter().rev() {
2523 let source = sealed.remove(index);
2524 if Arc::strong_count(&source) > 1 {
2525 retired.push(Arc::downgrade(&source));
2526 }
2527 }
2528 }
2529 let insert_at = positions[0].min(sealed.len());
2530 sealed.insert(insert_at, Arc::clone(&merged_arc));
2531 drop(sealed);
2532
2533 let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
2535 {
2536 let mut entity_segment = self.entity_segment.write();
2537 for segment_id in entity_segment.values_mut() {
2538 if source_ids.contains(segment_id) {
2539 *segment_id = run.merged_id;
2540 }
2541 }
2542 }
2543
2544 {
2545 let mut stats = self.stats.write();
2546 stats.sealed_count = stats
2547 .sealed_count
2548 .saturating_sub(positions.len())
2549 .saturating_add(1);
2550 let consolidation = &mut stats.consolidation;
2551 consolidation.runs_completed += 1;
2552 consolidation.segments_merged += positions.len() as u64;
2553 consolidation.tombstones_reclaimed += reclaimed_tombstones;
2554 consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
2555 }
2556
2557 self.emit(LifecycleEvent::SegmentConsolidated {
2558 source: source_ids,
2559 target: run.merged_id,
2560 });
2561 }
2562
2563 #[cfg(test)]
2566 fn consolidation_progress(&self) -> Option<u64> {
2567 self.consolidation
2568 .read()
2569 .as_ref()
2570 .map(|run| run.entities_copied)
2571 }
2572}
2573
2574impl UnifiedSegment for Arc<RwLock<GrowingSegment>> {
2577 fn id(&self) -> SegmentId {
2578 self.read().id()
2579 }
2580
2581 fn state(&self) -> SegmentState {
2582 self.read().state()
2583 }
2584
2585 fn collection(&self) -> &str {
2586 "unknown"
2588 }
2589
2590 fn stats(&self) -> SegmentStats {
2591 self.read().stats()
2592 }
2593
2594 fn entity_count(&self) -> usize {
2595 self.read().entity_count()
2596 }
2597
2598 fn contains(&self, id: EntityId) -> bool {
2599 self.read().contains(id)
2600 }
2601
2602 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2603 None
2606 }
2607
2608 fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2609 None
2610 }
2611
2612 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2613 self.write().insert(entity)
2614 }
2615
2616 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2617 self.write().update(entity)
2618 }
2619
2620 fn update_hot(
2621 &mut self,
2622 entity: UnifiedEntity,
2623 modified_columns: &[String],
2624 ) -> Result<(), SegmentError> {
2625 self.write().update_hot(entity, modified_columns)
2626 }
2627
2628 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2629 self.write().delete(id)
2630 }
2631
2632 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2633 self.read().get_metadata(id)
2634 }
2635
2636 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2637 self.write().set_metadata(id, metadata)
2638 }
2639
2640 fn seal(&mut self) -> Result<(), SegmentError> {
2641 self.write().seal()
2642 }
2643
2644 fn should_seal(&self, config: &SegmentConfig) -> bool {
2645 self.read().should_seal(config)
2646 }
2647
2648 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2649 Box::new(std::iter::empty())
2651 }
2652
2653 fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2654 Box::new(std::iter::empty())
2655 }
2656
2657 fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2658 self.read().filter_metadata(filters)
2659 }
2660}
2661
2662#[cfg(test)]
2663mod tests {
2664 use std::collections::HashSet;
2665
2666 use super::*;
2667 use crate::runtime::mvcc::SnapshotContext;
2668 use crate::storage::transaction::snapshot::{Snapshot, SnapshotManager};
2669 use crate::storage::unified::entity::{EntityData, EntityKind, RowData};
2670 use reddb_types::Value;
2671
2672 #[test]
2673 fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
2674 let manager = SegmentManager::new("fingerprint");
2675 let row = |id, label: &str| {
2676 UnifiedEntity::table_row(
2677 EntityId::new(id),
2678 "fingerprint",
2679 id,
2680 vec![Value::text(label)],
2681 )
2682 };
2683 manager
2684 .bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
2685 .expect("sealed rows");
2686 manager.force_seal().expect("seal");
2687 manager
2688 .bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
2689 .expect("growing rows");
2690
2691 let mut expected = manager.query_all(|_| true);
2692 expected.sort_by_key(|entity| entity.id.raw());
2693 let expected: Vec<_> = expected
2694 .iter()
2695 .map(|entity| format!("{entity:?}"))
2696 .collect();
2697 let mut stored_addresses = HashSet::new();
2698 manager.for_each_entity(|entity| {
2699 stored_addresses.insert(entity as *const UnifiedEntity as usize);
2700 true
2701 });
2702
2703 let mut actual = Vec::new();
2704 manager.for_each_entity_in_id_order(|entity| {
2705 assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
2706 assert!(
2707 manager.growing.try_write().is_some(),
2708 "topology guard retained"
2709 );
2710 assert!(
2711 manager.sealed.try_write().is_some(),
2712 "topology guard retained"
2713 );
2714 actual.push(format!("{entity:?}"));
2715 true
2716 });
2717 assert_eq!(
2718 actual, expected,
2719 "equal IDs retain their original segment order"
2720 );
2721 }
2722
2723 #[test]
2724 fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
2725 let manager = SegmentManager::new("fingerprint");
2726 for id in [3, 1, 2] {
2727 manager
2728 .insert(UnifiedEntity::table_row(
2729 EntityId::new(id),
2730 "fingerprint",
2731 id,
2732 vec![Value::Integer(id as i64)],
2733 ))
2734 .expect("row");
2735 }
2736 let source = manager.growing.read().as_ref().expect("source").clone();
2737 let handles = Arc::strong_count(&source);
2738 let mut seen = Vec::new();
2739 manager.for_each_entity_in_id_order(|entity| {
2740 seen.push(entity.id.raw());
2741 false
2742 });
2743 assert_eq!(seen, vec![1]);
2744 assert_eq!(Arc::strong_count(&source), handles);
2745 assert!(
2746 source.try_write().is_some(),
2747 "segment guard retained after early stop"
2748 );
2749 }
2750
2751 #[test]
2752 fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
2753 let manager = SegmentManager::with_config(
2754 "vectors",
2755 ManagerConfig {
2756 max_sealed_segments: 1,
2757 consolidation_entities_per_tick: 10_000,
2758 ..Default::default()
2759 },
2760 );
2761 for part in 0..3 {
2762 manager
2763 .bulk_insert(
2764 (1..=300)
2765 .map(|offset| {
2766 UnifiedEntity::vector(
2767 EntityId::new(part * 300 + offset),
2768 "vectors",
2769 vec![1.0, 0.0],
2770 )
2771 })
2772 .collect(),
2773 )
2774 .expect("batch");
2775 if part < 2 {
2776 manager.force_seal().expect("seal");
2777 }
2778 }
2779 let original_growing = manager.growing.read().as_ref().expect("growing").clone();
2780 let mut changed = false;
2781 let mut seen = Vec::new();
2782 assert!(manager.scan_batches(
2783 None,
2784 || true,
2785 |batch| {
2786 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
2787 assert!(
2788 manager.growing.try_write().is_some(),
2789 "consumer holds no topology guard"
2790 );
2791 assert!(
2792 original_growing.try_write().is_some(),
2793 "consumer holds no segment guard"
2794 );
2795 for entity in batch {
2796 assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
2797 seen.push(entity.id.raw());
2798 }
2799 if !changed {
2800 changed = true;
2801 manager.force_seal().expect("seal during scan");
2802 manager
2803 .insert(UnifiedEntity::vector(
2804 EntityId::new(5000),
2805 "vectors",
2806 vec![1.0, 0.0],
2807 ))
2808 .expect("new write");
2809 drain_maintenance(&manager);
2810 assert!(manager.stats().consolidation.runs_completed > 0);
2811 assert!(manager
2812 .delete(EntityId::new(1))
2813 .expect("delete after source retirement"));
2814 assert!(manager.get(EntityId::new(1)).is_none());
2815 }
2816 true
2817 }
2818 ));
2819 seen.sort_unstable();
2820 assert_eq!(
2821 seen,
2822 (2..=900).collect::<Vec<_>>(),
2823 "no duplicate, missing or late candidates"
2824 );
2825 let mut after = Vec::new();
2826 assert!(manager.scan_batches(
2827 None,
2828 || true,
2829 |batch| {
2830 after.extend(batch.iter().map(|entity| entity.id.raw()));
2831 true
2832 }
2833 ));
2834 after.sort_unstable();
2835 let mut expected: Vec<_> = (2..=900).collect();
2836 expected.push(5000);
2837 assert_eq!(after, expected, "adopted entities have scan positions");
2838 }
2839
2840 #[test]
2841 fn bounded_scan_applies_snapshot_to_relocated_versions() {
2842 for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
2843 let manager = SegmentManager::with_config(
2844 "vectors",
2845 ManagerConfig {
2846 max_sealed_segments: 1,
2847 ..Default::default()
2848 },
2849 );
2850 for ids in [[1, 2], [3, 4]] {
2851 let entities = ids.map(|id| {
2852 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
2853 entity.set_xmin(3);
2854 entity
2855 });
2856 if bulk {
2857 manager
2858 .bulk_insert(entities.into())
2859 .expect("initial versions");
2860 } else {
2861 for entity in entities {
2862 manager.insert(entity).expect("initial version");
2863 }
2864 }
2865 manager.force_seal().expect("seal source");
2866 }
2867 let snapshot = SnapshotContext {
2868 snapshot: Snapshot {
2869 xid: 10,
2870 in_progress: HashSet::new(),
2871 },
2872 manager: Arc::new(SnapshotManager::new()),
2873 own_xids: HashSet::new(),
2874 requires_index_fallback: false,
2875 serializable_reader: None,
2876 };
2877 let mut changed = false;
2878 let mut seen = Vec::new();
2879 assert!(manager.scan_batches(
2880 Some(&snapshot),
2881 || true,
2882 |batch| {
2883 seen.extend(batch.iter().map(|entity| entity.id.raw()));
2884 if !changed {
2885 changed = true;
2886 if consolidate {
2887 drain_maintenance(&manager);
2888 assert!(manager.stats().consolidation.runs_completed > 0);
2889 }
2890 for (id, xmax) in [(3, 8), (4, 12)] {
2891 let mut entity =
2892 manager.get(EntityId::new(id)).expect("retained version");
2893 entity.set_xmax(xmax);
2894 manager
2895 .update(entity)
2896 .expect("update after source retirement");
2897 }
2898 }
2899 true
2900 }
2901 ));
2902 seen.sort_unstable();
2903 assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
2904 let mut again = Vec::new();
2905 assert!(manager.scan_batches(
2906 Some(&snapshot),
2907 || true,
2908 |batch| {
2909 again.extend(batch.iter().map(|entity| entity.id.raw()));
2910 true
2911 }
2912 ));
2913 again.sort_unstable();
2914 assert_eq!(
2915 again,
2916 vec![1, 2, 4],
2917 "old slots do not duplicate relocated versions"
2918 );
2919 let mut current = Vec::new();
2920 assert!(manager.scan_batches(
2921 None,
2922 || true,
2923 |batch| {
2924 current.extend(batch.iter().map(|entity| entity.id.raw()));
2925 true
2926 }
2927 ));
2928 current.sort_unstable();
2929 assert_eq!(
2930 current,
2931 vec![1, 2],
2932 "current view excludes both deleted versions"
2933 );
2934 }
2935 }
2936
2937 #[test]
2938 fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
2939 let manager = SegmentManager::new("mixed");
2940 let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
2941 manager
2942 .bulk_insert((1..=300).map(vector).collect())
2943 .expect("flat");
2944 manager.insert(vector(500)).expect("hash insert");
2945 manager
2946 .bulk_insert(vec![vector(502), vector(505)])
2947 .expect("gapped bulk");
2948 manager.delete(EntityId::new(2)).expect("flat delete");
2949 manager.delete(EntityId::new(505)).expect("hash delete");
2950 let mut changed = false;
2951 let mut seen = Vec::new();
2952 assert!(manager.scan_batches(
2953 None,
2954 || true,
2955 |batch| {
2956 seen.extend(batch.iter().map(|entity| entity.id.raw()));
2957 if !changed {
2958 changed = true;
2959 manager
2960 .insert(vector(700))
2961 .expect("append after cursor boundary");
2962 manager
2963 .update(UnifiedEntity::graph_node(
2964 EntityId::new(500),
2965 "mixed",
2966 "Node",
2967 HashMap::new(),
2968 ))
2969 .expect("structural update");
2970 }
2971 true
2972 }
2973 ));
2974 seen.sort_unstable();
2975 let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
2976 expected.extend([500, 502]);
2977 assert_eq!(seen, expected);
2978 }
2979
2980 #[test]
2981 fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
2982 let manager = SegmentManager::new("vectors");
2983 let mut entities = Vec::new();
2984 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
2985 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
2986 entity.set_xmin(xmin);
2987 entity.set_xmax(xmax);
2988 entities.push(entity);
2989 }
2990 manager.bulk_insert(entities).expect("versions");
2991 manager.force_seal().expect("seal");
2992 let snapshot = SnapshotContext {
2993 snapshot: Snapshot {
2994 xid: 10,
2995 in_progress: HashSet::from([4]),
2996 },
2997 manager: Arc::new(SnapshotManager::new()),
2998 own_xids: HashSet::new(),
2999 requires_index_fallback: false,
3000 serializable_reader: None,
3001 };
3002 let mut work = 0;
3003 let mut callbacks = 0;
3004 assert!(!manager.scan_batches(
3005 Some(&snapshot),
3006 || {
3007 work += 1;
3008 work <= 2
3009 },
3010 |_| {
3011 callbacks += 1;
3012 true
3013 }
3014 ));
3015 assert_eq!(callbacks, 0, "a stopped batch is not consumed");
3016 for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
3017 let mut seen = Vec::new();
3018 assert!(manager.scan_batches(
3019 view,
3020 || true,
3021 |batch| {
3022 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3023 true
3024 }
3025 ));
3026 seen.sort_unstable();
3027 assert_eq!(seen, expected);
3028 }
3029 }
3030
3031 #[test]
3032 fn test_manager_basic() {
3033 let manager = SegmentManager::new("test_collection");
3034
3035 let entity = UnifiedEntity::table_row(
3036 manager.next_entity_id(),
3037 "users",
3038 1,
3039 vec![Value::text("Alice".to_string())],
3040 );
3041
3042 let id = manager.insert(entity).unwrap();
3043 assert!(manager.get(id).is_some());
3044 assert_eq!(manager.count(), 1);
3045 }
3046
3047 #[test]
3048 fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
3049 let manager = Arc::new(SegmentManager::new("mixed"));
3050 manager
3051 .bulk_insert(
3052 (1..=128)
3053 .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
3054 .collect(),
3055 )
3056 .expect("vector-only segment");
3057 manager.force_seal().expect("seal vectors");
3058 let mut old =
3059 UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
3060 old.set_xmin(3);
3061 old.set_xmax(12);
3062 let mut new =
3063 UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
3064 new.set_xmin(12);
3065 let edge = UnifiedEntity::graph_edge(
3066 EntityId::new(202),
3067 "mixed",
3068 "a",
3069 "b",
3070 1.0,
3071 Default::default(),
3072 );
3073 manager
3074 .bulk_insert(vec![old, new, edge])
3075 .expect("physical graph versions");
3076 manager.force_seal().expect("seal graph");
3077 manager
3078 .insert(UnifiedEntity::graph_node(
3079 EntityId::new(203),
3080 "mixed",
3081 "person",
3082 Default::default(),
3083 ))
3084 .expect("growing node");
3085 let snapshot = SnapshotContext {
3086 snapshot: Snapshot {
3087 xid: 10,
3088 in_progress: HashSet::new(),
3089 },
3090 manager: Arc::new(SnapshotManager::new()),
3091 own_xids: HashSet::new(),
3092 requires_index_fallback: false,
3093 serializable_reader: None,
3094 };
3095 std::thread::spawn(move || {
3096 for view in [None, Some(&snapshot)] {
3097 for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
3098 let expected: HashSet<_> = manager
3099 .scan(view, |item| kind.matches(&item.kind))
3100 .into_iter()
3101 .map(|item| item.id)
3102 .collect();
3103 let actual: HashSet<_> = manager
3104 .scan_graph_kind(view, kind)
3105 .into_iter()
3106 .map(|item| item.id)
3107 .collect();
3108 assert_eq!(actual, expected);
3109 let mut visited = 0;
3110 let mut candidates = HashSet::new();
3111 manager.scan_graph_candidates_for_each(
3112 view,
3113 kind,
3114 || true,
3115 |item| {
3116 visited += 1;
3117 if kind.matches(&item.kind) {
3118 candidates.insert(item.id);
3119 }
3120 true
3121 },
3122 );
3123 assert_eq!(candidates, expected);
3124 assert_eq!(
3125 visited,
3126 if matches!(kind, GraphEntityKind::Node) {
3127 3
3128 } else {
3129 2
3130 }
3131 );
3132 }
3133 }
3134 let ids: HashSet<_> = manager
3135 .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
3136 .into_iter()
3137 .map(|item| item.id.raw())
3138 .collect();
3139 assert_eq!(ids, HashSet::from([200, 203]));
3140 let ids: HashSet<_> = manager
3141 .scan_graph_kind(None, GraphEntityKind::Node)
3142 .into_iter()
3143 .map(|item| item.id.raw())
3144 .collect();
3145 assert_eq!(ids, HashSet::from([201, 203]));
3146 let mut visits = 0;
3147 manager.scan_graph_candidates_for_each(
3148 None,
3149 GraphEntityKind::Node,
3150 || true,
3151 |_| {
3152 visits += 1;
3153 false
3154 },
3155 );
3156 assert_eq!(
3157 visits, 1,
3158 "early termination must also stop subsequent segments"
3159 );
3160 })
3161 .join()
3162 .expect("snapshot worker");
3163 }
3164
3165 #[test]
3166 fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
3167 let manager = Arc::new(SegmentManager::new("accounts"));
3168 let cases = [
3169 ("legacy row", 0, 0, true),
3170 ("committed creator", 3, 0, true),
3171 ("future creator", 11, 0, false),
3172 ("in-progress creator", 4, 0, false),
3173 ("committed deleter", 3, 7, false),
3174 ("future deleter", 3, 11, true),
3175 ("in-progress deleter", 3, 4, true),
3176 ];
3177
3178 for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
3179 let row_id = index as u64 + 1;
3180 let mut entity = UnifiedEntity::table_row(
3181 EntityId::new(row_id),
3182 "accounts",
3183 row_id,
3184 vec![Value::Integer(row_id as i64)],
3185 );
3186 entity.set_xmin(*xmin);
3187 entity.set_xmax(*xmax);
3188 manager.insert(entity).unwrap();
3189 }
3190
3191 let snapshot = SnapshotContext {
3192 snapshot: Snapshot {
3193 xid: 10,
3194 in_progress: HashSet::from([4]),
3195 },
3196 manager: Arc::new(SnapshotManager::new()),
3197 own_xids: HashSet::new(),
3198 requires_index_fallback: false,
3199 serializable_reader: None,
3200 };
3201
3202 let raw_visible_row_ids: HashSet<u64> = manager
3203 .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
3204 .into_iter()
3205 .filter_map(|entity| match entity.kind {
3206 EntityKind::TableRow { row_id, .. } => Some(row_id),
3207 _ => None,
3208 })
3209 .collect();
3210
3211 let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
3212 .join()
3213 .unwrap();
3214 let visible_row_ids: HashSet<u64> = scanned
3215 .into_iter()
3216 .filter_map(|entity| match entity.kind {
3217 EntityKind::TableRow { row_id, .. } => Some(row_id),
3218 _ => None,
3219 })
3220 .collect();
3221 assert_eq!(visible_row_ids, raw_visible_row_ids);
3222
3223 for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
3224 let row_id = index as u64 + 1;
3225 assert_eq!(
3226 visible_row_ids.contains(&row_id),
3227 *expected_visible,
3228 "{name} visibility mismatch"
3229 );
3230 }
3231 }
3232
3233 #[test]
3234 fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
3235 let manager = SegmentManager::new("accounts");
3239 let mut live =
3240 UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
3241 live.set_xmin(2);
3242 live.set_xmax(0);
3243 manager.insert(live).unwrap();
3244
3245 let mut superseded =
3246 UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
3247 superseded.set_xmin(2);
3248 superseded.set_xmax(9);
3249 manager.insert(superseded).unwrap();
3250
3251 let rows: Vec<u64> = manager
3252 .scan(None, |_| true)
3253 .into_iter()
3254 .filter_map(|entity| match entity.kind {
3255 EntityKind::TableRow { row_id, .. } => Some(row_id),
3256 _ => None,
3257 })
3258 .collect();
3259 assert_eq!(
3260 rows,
3261 vec![1],
3262 "frameless scan must hide superseded versions"
3263 );
3264 }
3265
3266 #[test]
3267 fn query_scan_callers_use_snapshot_api_for_visibility() {
3268 let table = include_str!("../../runtime/query_exec/table.rs");
3269 let aggregate = include_str!("../../runtime/query_exec/aggregate.rs");
3270
3271 assert_eq!(
3272 table.matches(".for_each_entity(|entity|").count(),
3273 1,
3274 "only the visibility-independent schema inspection may use the raw table iterator"
3275 );
3276 assert!(!table.contains(".query_all_zoned_with_stats("));
3277 assert!(!table.contains(".for_each_entity_zoned_with_stats("));
3278 assert!(!table.contains("TableRowMvccReadResolver::captured"));
3279
3280 assert!(!aggregate.contains("TableRowMvccReadResolver"));
3281 assert!(!aggregate.contains(".for_each_entity("));
3282 assert!(!aggregate.contains(".fold_entities_parallel("));
3283 }
3284
3285 #[test]
3286 fn remaining_visibility_scan_callers_use_snapshot_api() {
3287 let aggregate_pushdown =
3288 include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs");
3289 let ask = include_str!("../../runtime/ask_pipeline.rs");
3290 let dml = include_str!("../../runtime/impl_dml.rs");
3291 let dml_targets = include_str!("../../runtime/dml_target_scan.rs");
3292 let events = include_str!("../../runtime/impl_events.rs");
3293 let graph_commands = include_str!("../../runtime/impl_graph_commands.rs");
3294 let graph_dsl = include_str!("../../runtime/graph_dsl.rs");
3295 let graph_tvf = include_str!("../../runtime/graph_tvf.rs");
3296 let kv = include_str!("../../runtime/impl_kv.rs");
3297 let queue = include_str!("../../runtime/impl_queue.rs");
3298 let primary_queue = include_str!("../../runtime/primary_queue_store.rs");
3299 let record_search = include_str!("../../runtime/record_search.rs");
3300 let search = include_str!("../../runtime/impl_search.rs");
3301 let vector = include_str!("../../runtime/query_exec/vector.rs");
3302
3303 assert!(!aggregate_pushdown.contains("TableRowMvccReadResolver"));
3304 assert!(!aggregate_pushdown.contains(".for_each_entity("));
3305 assert!(!ask.contains("for entity in manager.query_all(|_| true)"));
3306 assert!(!dml.contains("let entities = manager.query_all(|_| true);"));
3307 assert_eq!(
3312 dml_targets
3313 .matches("self.candidate_visible(snapshot.as_ref(), entity)")
3314 .count(),
3315 2,
3316 "both raw DML iterations must route through candidate_visible"
3317 );
3318 assert_eq!(
3321 events.matches("manager.query_all(").count(),
3322 1,
3323 "only the visibility-free backfill dedup may use query_all"
3324 );
3325 assert!(!graph_commands.contains("TableRowMvccReadResolver"));
3326 assert!(!graph_dsl.contains("entity_visible_with_context(snap_ctx.as_ref(), e)"));
3327 assert!(!graph_tvf.contains("entity_visible_with_context(snap_ctx.as_ref(), &entity)"));
3328 assert!(!kv.contains("TableRowMvccReadResolver"));
3329 assert!(!queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3330 assert!(!primary_queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3331 assert!(!record_search.contains("TableRowMvccReadResolver"));
3332 assert!(!search.contains("for entity in manager.query_all(|_| true)"));
3333 assert!(!vector.contains("manager.query_all("));
3334 }
3335
3336 fn production_source(src: &str) -> String {
3341 let mut kept: Vec<&str> = Vec::new();
3342 let lines: Vec<&str> = src.lines().collect();
3343 let mut idx = 0;
3344 while idx < lines.len() {
3345 let starts_test_mod = lines[idx].starts_with("#[cfg(test)]")
3346 && lines
3347 .get(idx + 1)
3348 .is_some_and(|next| next.starts_with("mod "));
3349 if !starts_test_mod {
3350 kept.push(lines[idx]);
3351 idx += 1;
3352 continue;
3353 }
3354 idx += 1;
3355 let mut depth: isize = 0;
3356 while idx < lines.len() {
3357 depth += lines[idx].matches('{').count() as isize;
3358 depth -= lines[idx].matches('}').count() as isize;
3359 idx += 1;
3360 if depth <= 0 {
3361 break;
3362 }
3363 }
3364 }
3365 kept.join("\n")
3366 }
3367
3368 fn raw_iteration_sites(src: &str) -> usize {
3370 src.matches("query_all(").count()
3371 + src.matches(".for_each_entity").count()
3372 + src.matches("fold_entities_parallel(").count()
3373 }
3374
3375 #[test]
3376 fn raw_iteration_confined_to_maintenance_allowlist() {
3377 let allowlist: &[(&str, &str, usize, &str)] = &[
3385 (
3386 "runtime/query_exec/table.rs",
3387 include_str!("../../runtime/query_exec/table.rs"),
3388 1,
3389 "documented exception: schema inspection (does the collection ever \
3390 observe these fields?) is visibility-independent",
3391 ),
3392 (
3393 "runtime/query_exec/aggregate.rs",
3394 include_str!("../../runtime/query_exec/aggregate.rs"),
3395 0,
3396 "fully migrated by batch 1",
3397 ),
3398 (
3399 "runtime/query_exec/aggregate_pushdown_dispatch.rs",
3400 include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs"),
3401 0,
3402 "fully migrated by batch 1",
3403 ),
3404 (
3405 "runtime/query_exec/vector.rs",
3406 include_str!("../../runtime/query_exec/vector.rs"),
3407 0,
3408 "fully migrated by batch 2",
3409 ),
3410 (
3411 "runtime/ask_pipeline.rs",
3412 include_str!("../../runtime/ask_pipeline.rs"),
3413 0,
3414 "fully migrated by batch 2",
3415 ),
3416 (
3417 "runtime/impl_dml.rs",
3418 include_str!("../../runtime/impl_dml.rs"),
3419 1,
3420 "documented exception: ON CONFLICT resolves a unique key against \
3421 physical versions on purpose. A snapshot scan cannot see a row an \
3422 uncommitted concurrent writer already holds the key for, which is \
3423 exactly the case that must raise a serialization conflict rather \
3424 than insert a duplicate; the same pass reaps index entries left by \
3425 aborted writers. \
3426 TODO(#2263 follow-up): narrow the candidates through the unique \
3427 index instead of walking the collection",
3428 ),
3429 (
3430 "runtime/dml_target_scan.rs",
3431 include_str!("../../runtime/dml_target_scan.rs"),
3432 2,
3433 "documented exception: kind-conditional visibility the uniform scan \
3434 API cannot express; both iterations gate on candidate_visible",
3435 ),
3436 (
3437 "runtime/impl_events.rs",
3438 include_str!("../../runtime/impl_events.rs"),
3439 1,
3440 "documented exception: backfill dedup must include consumed messages, \
3441 so it is deliberately visibility-free",
3442 ),
3443 (
3444 "runtime/impl_graph_commands.rs",
3445 include_str!("../../runtime/impl_graph_commands.rs"),
3446 0,
3447 "native node properties resolve one snapshot-visible, RLS-admitted version",
3448 ),
3449 (
3450 "runtime/graph_dsl.rs",
3451 include_str!("../../runtime/graph_dsl.rs"),
3452 1,
3453 "remaining legacy debt: unused edge-property helper; search expansion \
3454 now uses snapshot visitors and RLS-gated hydration",
3455 ),
3456 (
3457 "runtime/graph_tvf.rs",
3458 include_str!("../../runtime/graph_tvf.rs"),
3459 0,
3460 "fully migrated by batch 2",
3461 ),
3462 (
3463 "runtime/impl_kv.rs",
3464 include_str!("../../runtime/impl_kv.rs"),
3465 4,
3466 "two non-versioned-collection fallbacks (the versioned arms scan under \
3467 the captured snapshot). Vault reads use snapshot scans. Offline vault \
3468 export preserves all \
3469 physical versions across tenants; restore's emptiness check must reject \
3470 every physical row, including versions invisible to a snapshot",
3471 ),
3472 (
3473 "runtime/impl_queue.rs",
3474 include_str!("../../runtime/impl_queue.rs"),
3475 10,
3476 "maintenance: red_queue_meta bookkeeping (config, groups, acks, \
3477 pending, dedup TTL reaper, gauge counts) — queue message bodies are \
3478 read through scan. \
3479 TODO(#2138 follow-up): QUEUE PENDING projects meta rows to clients",
3480 ),
3481 (
3482 "runtime/primary_queue_store.rs",
3483 include_str!("../../runtime/primary_queue_store.rs"),
3484 3,
3485 "maintenance: generic red_queue_meta read/delete helpers and the \
3486 next-position allocator, which must not collide with rows invisible \
3487 to the writer's snapshot",
3488 ),
3489 (
3490 "runtime/record_search.rs",
3491 include_str!("../../runtime/record_search.rs"),
3492 0,
3493 "fully migrated by batch 2",
3494 ),
3495 (
3496 "runtime/search_graph.rs",
3497 include_str!("../../runtime/search_graph.rs"),
3498 0,
3499 "search graph preparation uses snapshot visitors and gated hydration",
3500 ),
3501 (
3502 "runtime/impl_search.rs",
3503 include_str!("../../runtime/impl_search.rs"),
3504 1,
3505 "maintenance: ask-audit retention sweep must see every physical row",
3506 ),
3507 ];
3508
3509 for (path, src, allowed, reason) in allowlist {
3510 let found = raw_iteration_sites(&production_source(src));
3511 assert_eq!(
3512 found, *allowed,
3513 "{path}: {found} raw-iteration sites, allowlist permits {allowed} \
3514 ({reason}). Raw iteration bypasses MVCC visibility — use \
3515 SegmentManager::scan / scan_for_each / scan_fold_parallel / \
3516 scan_zoned_with_stats, or extend the allowlist with a maintenance \
3517 justification."
3518 );
3519 }
3520 }
3521
3522 #[test]
3523 fn explicit_snapshot_scan_ignores_ambient_thread_local() {
3524 let manager = Arc::new(SegmentManager::new("accounts"));
3525 let versions = [(0, 0), (3, 0), (11, 0), (4, 0), (3, 7), (3, 11), (3, 4)];
3527 for (index, (xmin, xmax)) in versions.iter().enumerate() {
3528 let row_id = index as u64 + 1;
3529 let mut entity = UnifiedEntity::table_row(
3530 EntityId::new(row_id),
3531 "accounts",
3532 row_id,
3533 vec![Value::Integer(row_id as i64)],
3534 );
3535 entity.set_xmin(*xmin);
3536 entity.set_xmax(*xmax);
3537 manager.insert(entity).unwrap();
3538 }
3539
3540 let explicit = SnapshotContext {
3541 snapshot: Snapshot {
3542 xid: 10,
3543 in_progress: HashSet::from([4]),
3544 },
3545 manager: Arc::new(SnapshotManager::new()),
3546 own_xids: HashSet::new(),
3547 requires_index_fallback: false,
3548 serializable_reader: None,
3549 };
3550 let decoy = SnapshotContext {
3553 snapshot: Snapshot {
3554 xid: 1,
3555 in_progress: HashSet::new(),
3556 },
3557 manager: Arc::new(SnapshotManager::new()),
3558 own_xids: HashSet::new(),
3559 requires_index_fallback: false,
3560 serializable_reader: None,
3561 };
3562
3563 let row_ids = |entities: Vec<UnifiedEntity>| -> HashSet<u64> {
3564 entities
3565 .into_iter()
3566 .filter_map(|entity| match entity.kind {
3567 EntityKind::TableRow { row_id, .. } => Some(row_id),
3568 _ => None,
3569 })
3570 .collect()
3571 };
3572
3573 let guard = crate::runtime::execution_context::CurrentSnapshotGuard::install(decoy);
3574 let on_installing_thread = row_ids(manager.scan(Some(&explicit), |_| true));
3575 drop(guard);
3576
3577 let spawned_manager = Arc::clone(&manager);
3578 let spawned_snapshot = explicit.clone();
3579 let on_fresh_thread = std::thread::spawn(move || {
3580 row_ids(spawned_manager.scan(Some(&spawned_snapshot), |_| true))
3581 })
3582 .join()
3583 .unwrap();
3584
3585 assert_eq!(
3586 on_installing_thread, on_fresh_thread,
3587 "an explicit snapshot must produce the same visible set on any thread"
3588 );
3589 assert_eq!(on_installing_thread, HashSet::from([1, 2, 6, 7]));
3590
3591 let raw_manager = Arc::clone(&manager);
3594 let raw_count = std::thread::spawn(move || raw_manager.query_all(|_| true).len())
3595 .join()
3596 .unwrap();
3597 assert_eq!(
3598 raw_count,
3599 versions.len(),
3600 "raw iteration must expose every physical version"
3601 );
3602 }
3603
3604 #[test]
3605 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
3606 let manager = SegmentManager::new("test_collection");
3607
3608 let entity = UnifiedEntity::table_row(
3609 manager.next_entity_id(),
3610 "users",
3611 1,
3612 vec![Value::text("Alice".to_string())],
3613 );
3614
3615 let id = manager.insert(entity).unwrap();
3616
3617 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3618 assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
3619 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
3620
3621 assert!(manager.delete(id).unwrap());
3622 assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3623 }
3624
3625 #[test]
3626 fn zoned_scan_reports_pruned_segments_and_preserves_results() {
3627 let config = ManagerConfig {
3628 segment_config: SegmentConfig {
3629 max_entities: 2,
3630 ..Default::default()
3631 },
3632 ..Default::default()
3633 };
3634 let manager = SegmentManager::with_config("events", config);
3635
3636 for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
3637 let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
3638 let entity = UnifiedEntity::new(
3639 manager.next_entity_id(),
3640 EntityKind::TableRow {
3641 table: "events".into(),
3642 row_id: (idx + 1) as u64,
3643 },
3644 EntityData::Row(row),
3645 );
3646 manager.insert(entity).unwrap();
3647 }
3648
3649 let full = manager.query_all(|entity| match &entity.data {
3650 EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
3651 _ => false,
3652 });
3653 let probe = Value::Integer(100);
3654 let (pruned, stats) =
3655 manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
3656 match &entity.data {
3657 EntityData::Row(row) => {
3658 matches!(row.get_field("ts"), Some(Value::Integer(100)))
3659 }
3660 _ => false,
3661 }
3662 });
3663
3664 assert_eq!(pruned.len(), full.len());
3665 assert_eq!(pruned[0].id, full[0].id);
3666 assert_eq!(stats.segments_total, 3);
3667 assert_eq!(stats.segments_scanned, 2);
3668 assert_eq!(stats.segments_pruned, 1);
3669 }
3670
3671 #[test]
3672 fn test_manager_auto_seal() {
3673 let config = ManagerConfig {
3674 segment_config: SegmentConfig {
3675 max_entities: 2,
3676 ..Default::default()
3677 },
3678 ..Default::default()
3679 };
3680
3681 let manager = SegmentManager::with_config("test", config);
3682
3683 manager
3685 .insert(UnifiedEntity::vector(
3686 manager.next_entity_id(),
3687 "v",
3688 vec![0.1],
3689 ))
3690 .unwrap();
3691
3692 manager
3694 .insert(UnifiedEntity::vector(
3695 manager.next_entity_id(),
3696 "v",
3697 vec![0.2],
3698 ))
3699 .unwrap();
3700
3701 manager
3703 .insert(UnifiedEntity::vector(
3704 manager.next_entity_id(),
3705 "v",
3706 vec![0.3],
3707 ))
3708 .unwrap();
3709
3710 let stats = manager.stats();
3711 assert_eq!(stats.total_entities, 3);
3712 }
3713
3714 #[test]
3715 fn test_manager_delete() {
3716 let manager = SegmentManager::new("test");
3717
3718 let id = manager
3719 .insert(UnifiedEntity::vector(
3720 manager.next_entity_id(),
3721 "v",
3722 vec![0.1],
3723 ))
3724 .unwrap();
3725
3726 assert!(manager.get(id).is_some());
3727 assert!(manager.delete(id).unwrap());
3728 assert!(manager.get(id).is_none());
3729 }
3730
3731 #[test]
3732 fn test_manager_metadata() {
3733 let manager = SegmentManager::new("test");
3734
3735 let id = manager
3736 .insert(UnifiedEntity::table_row(
3737 manager.next_entity_id(),
3738 "hosts",
3739 1,
3740 vec![Value::text("192.168.1.1".to_string())],
3741 ))
3742 .unwrap();
3743
3744 let mut meta = Metadata::new();
3745 meta.set(
3746 "os",
3747 super::super::metadata::MetadataValue::String("linux".to_string()),
3748 );
3749
3750 manager.set_metadata(id, meta).unwrap();
3751
3752 let retrieved = manager.get_metadata(id).unwrap();
3753 assert!(retrieved.has("os"));
3754 }
3755
3756 #[test]
3757 fn test_manager_query_by_kind() {
3758 let manager = SegmentManager::new("test");
3759
3760 manager
3761 .insert(UnifiedEntity::table_row(
3762 manager.next_entity_id(),
3763 "hosts",
3764 1,
3765 vec![],
3766 ))
3767 .unwrap();
3768
3769 manager
3770 .insert(UnifiedEntity::vector(
3771 manager.next_entity_id(),
3772 "embeddings",
3773 vec![0.1],
3774 ))
3775 .unwrap();
3776
3777 manager
3778 .insert(UnifiedEntity::table_row(
3779 manager.next_entity_id(),
3780 "hosts",
3781 2,
3782 vec![],
3783 ))
3784 .unwrap();
3785
3786 let rows = manager.get_by_kind("table");
3787 assert_eq!(rows.len(), 2);
3788
3789 let vectors = manager.get_by_kind("vector");
3790 assert_eq!(vectors.len(), 1);
3791 }
3792
3793 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
3800 SegmentManager::with_config(
3801 "test",
3802 ManagerConfig {
3803 max_sealed_segments: 1_000,
3804 consolidation_entities_per_tick: entities_per_tick,
3805 ..Default::default()
3806 },
3807 )
3808 }
3809
3810 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
3814 let entities: Vec<UnifiedEntity> = (0..count)
3815 .map(|i| {
3816 UnifiedEntity::table_row(
3817 manager.next_entity_id(),
3818 "users",
3819 i as u64,
3820 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
3821 )
3822 })
3823 .collect();
3824 let ids = manager.bulk_insert(entities).expect("bulk insert");
3825 manager.force_seal().expect("seal");
3826 ids
3827 }
3828
3829 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
3830 let mut ids: Vec<u64> = manager
3831 .query_all(|_| true)
3832 .into_iter()
3833 .map(|entity| entity.id.raw())
3834 .collect();
3835 ids.sort_unstable();
3836 ids
3837 }
3838
3839 fn drain_maintenance(manager: &SegmentManager) -> usize {
3840 let mut ticks = 0;
3841 while manager.consolidation.read().is_some() || ticks == 0 {
3842 manager.run_maintenance().expect("maintenance");
3843 ticks += 1;
3844 assert!(ticks < 1_000, "consolidation failed to converge");
3845 }
3846 ticks
3847 }
3848
3849 #[test]
3850 fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
3851 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3852 let ids = manager
3853 .bulk_insert(
3854 (1..=1024)
3855 .map(|id| {
3856 UnifiedEntity::graph_edge(
3857 EntityId::new(id),
3858 "step",
3859 "2000",
3860 "3000",
3861 1.0,
3862 HashMap::new(),
3863 )
3864 })
3865 .collect(),
3866 )
3867 .expect("edges");
3868 manager.force_seal().expect("seal");
3869 for id in ids.iter().take(300) {
3870 manager.delete(*id).expect("fragment source");
3871 }
3872 let source = Arc::downgrade(&manager.sealed.read()[0]);
3873 let mut seen = Vec::new();
3874 let mut changed = false;
3875 let mut metadata_bytes = 0;
3876 assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
3877 || true,
3878 |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
3879 |batch| {
3880 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3881 assert!(manager.growing.try_write().is_some(), "no growing topology guard");
3882 assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
3883 assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
3884 if !changed {
3885 changed = true;
3886 drain_maintenance(&manager);
3887 assert!(manager.stats().consolidation.runs_completed > 0);
3888 assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
3889 let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
3890 if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
3891 manager.update(updated).expect("update authoritative merged copy");
3892 }
3893 for id in batch {
3894 let entity = manager.get(*id).expect("hydrate current manager");
3895 if id.raw() == 1024 {
3896 assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
3897 }
3898 seen.push(id.raw());
3899 }
3900 Ok(true)
3901 }).expect("cursor"));
3902 assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
3903 assert!(
3904 source.upgrade().is_none(),
3905 "finished cursor releases retired source"
3906 );
3907 assert!(
3908 metadata_bytes < 4096,
3909 "fixed ID buffer plus source handles: {metadata_bytes}"
3910 );
3911 }
3912
3913 #[test]
3914 fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
3915 let manager = SegmentManager::new("graph");
3916 manager
3917 .bulk_insert(
3918 (1..=1024)
3919 .map(|id| {
3920 UnifiedEntity::graph_edge(
3921 EntityId::new(id),
3922 "step",
3923 "2000",
3924 "3000",
3925 1.0,
3926 HashMap::new(),
3927 )
3928 })
3929 .collect(),
3930 )
3931 .expect("edges");
3932 let mut delivered = 0;
3933 let denied = manager.visit_graph_index_batches(
3934 GraphEntityKind::Edge,
3935 &[EntityId::new(2000)],
3936 || true,
3937 |_| Err("out of memory"),
3938 |ids| {
3939 delivered += ids.len();
3940 Ok(true)
3941 },
3942 );
3943 assert_eq!(denied, Err("out of memory"));
3944 assert_eq!(delivered, 0);
3945 let mut work = 0;
3946 assert!(!manager
3947 .visit_graph_index_batches(
3948 GraphEntityKind::Edge,
3949 &[EntityId::new(2000)],
3950 || {
3951 work += 1;
3952 work < 100
3953 },
3954 |_| Ok::<_, ()>(()),
3955 |ids| {
3956 delivered += ids.len();
3957 Ok(true)
3958 }
3959 )
3960 .expect("cancelled cursor"));
3961 assert_eq!(
3962 delivered, 0,
3963 "cancelled partial batch cannot reach consumer"
3964 );
3965 assert!(!manager
3966 .visit_graph_index_batches(
3967 GraphEntityKind::Edge,
3968 &[EntityId::new(2000)],
3969 || true,
3970 |_| Ok::<_, ()>(()),
3971 |ids| {
3972 delivered += ids.len();
3973 Ok(false)
3974 }
3975 )
3976 .expect("consumer stop"));
3977 assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
3978 }
3979
3980 #[test]
3981 fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
3982 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3983 let ids = seal_bulk_segment(&manager, 100);
3984 for id in ids.iter().take(40) {
3985 manager.delete(*id).expect("delete before scan");
3986 }
3987 let source_bytes = manager.resident_bytes();
3988 let source = Arc::downgrade(&manager.sealed.read()[0]);
3989 std::thread::scope(|scope| {
3990 let (ready, waiting) = std::sync::mpsc::channel();
3991 let (release_first, first_wait) = std::sync::mpsc::channel();
3992 let (release_last, last_wait) = std::sync::mpsc::channel();
3993 let reader = |resume: std::sync::mpsc::Receiver<()>,
3994 ready: std::sync::mpsc::Sender<()>| {
3995 assert!(manager.scan_batches(
3996 None,
3997 || true,
3998 |_| {
3999 ready.send(()).expect("signal captured cursor");
4000 let _ = resume.recv();
4001 true
4002 }
4003 ));
4004 };
4005 let ready_first = ready.clone();
4006 let first = scope.spawn(move || reader(first_wait, ready_first));
4007 let last = scope.spawn(move || reader(last_wait, ready));
4008 for _ in 0..2 {
4009 waiting
4010 .recv_timeout(std::time::Duration::from_secs(10))
4011 .expect("reader ready");
4012 }
4013 drain_maintenance(&manager);
4014 assert!(
4015 source.upgrade().is_some(),
4016 "both cursors pin the retired source"
4017 );
4018 let active_bytes = manager.sealed.read()[0].read().resident_bytes();
4019 let held = manager.resident_bytes();
4020 assert!(
4021 held >= active_bytes + source_bytes,
4022 "retired source remains resident"
4023 );
4024 release_first.send(()).expect("release first reader");
4025 first.join().expect("first reader");
4026 assert_eq!(
4027 manager.resident_bytes(),
4028 held,
4029 "one reader still pins the source"
4030 );
4031 release_last.send(()).expect("release last reader");
4032 last.join().expect("last reader");
4033 assert!(
4034 source.upgrade().is_none(),
4035 "accounting must not retain payloads"
4036 );
4037 assert_eq!(held - manager.resident_bytes(), source_bytes);
4038 assert_eq!(
4039 manager.stats().total_memory_bytes as u64,
4040 manager.resident_bytes()
4041 );
4042 });
4043 }
4044
4045 #[test]
4046 fn retained_memory_gates_runtime_growth_until_the_cursor_is_cancelled() {
4047 use crate::storage::memory_pools::MemoryPool;
4048 use crate::{RedDBOptions, RedDBRuntime};
4049
4050 let budget = 128 * 1024;
4051 let runtime =
4052 RedDBRuntime::with_options(RedDBOptions::in_memory().with_memory_budget(budget))
4053 .expect("runtime");
4054 runtime
4055 .db()
4056 .store()
4057 .create_collection("test")
4058 .expect("collection");
4059 let manager = runtime
4060 .db()
4061 .store()
4062 .get_collection("test")
4063 .expect("manager");
4064 let ids = seal_bulk_segment(&manager, 100);
4065 for id in ids.iter().take(40) {
4066 manager.delete(*id).expect("delete before scan");
4067 }
4068 let source_bytes = manager.resident_bytes();
4069 let mut growth = 0;
4070 let mut held_used = 0;
4071 assert!(!manager.scan_batches(
4072 None,
4073 || true,
4074 |_| {
4075 drain_maintenance(&manager);
4076 runtime.refresh_memory_accounting();
4077 held_used = runtime.memory_accounting().total_used_bytes();
4080 growth = budget
4081 .checked_sub(held_used)
4082 .expect("fixture fits while the cursor is open")
4083 + source_bytes / 2;
4084 assert!(growth > 0);
4085 assert!(
4086 runtime
4087 .admit_non_evictable_growth(
4088 MemoryPool::SegmentArena,
4089 "growth with pinned source",
4090 growth,
4091 )
4092 .is_err(),
4093 "retained source must participate in runtime admission"
4094 );
4095 false
4096 }
4097 ));
4098 runtime.refresh_memory_accounting();
4099 assert_eq!(
4100 held_used - runtime.memory_accounting().total_used_bytes(),
4101 source_bytes
4102 );
4103 let _reservation = runtime
4104 .admit_non_evictable_growth(
4105 MemoryPool::SegmentArena,
4106 "growth after cursor cancellation",
4107 growth,
4108 )
4109 .expect("the same growth fits once the cursor releases its source");
4110 }
4111
4112 #[test]
4113 fn retained_memory_includes_an_inflight_consolidation_copy() {
4114 let manager = consolidating_manager(1);
4115 let ids = seal_bulk_segment(&manager, 100);
4116 for id in ids.iter().take(40) {
4117 manager.delete(*id).expect("delete");
4118 }
4119 let before = manager.resident_bytes();
4120 assert!(manager.pressure_consolidation_tick());
4121 let copied = manager
4122 .consolidation
4123 .read()
4124 .as_ref()
4125 .expect("paced merge")
4126 .merged
4127 .resident_bytes();
4128 assert!(copied > 0);
4129 assert!(
4130 manager.resident_bytes() >= before + copied,
4131 "unpublished copy consumes memory"
4132 );
4133 drain_maintenance(&manager);
4134 assert!(
4135 manager.resident_bytes() < before,
4136 "completed merge releases the old storage"
4137 );
4138 }
4139
4140 #[test]
4141 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
4142 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4143 let ids = seal_bulk_segment(&manager, 100);
4144
4145 let bytes_before_delete = manager.resident_bytes();
4146
4147 for id in ids.iter().take(40) {
4149 assert!(manager.delete(*id).expect("delete"));
4150 }
4151
4152 let bytes_with_tombstones = manager.resident_bytes();
4155 assert!(bytes_with_tombstones > bytes_before_delete);
4156
4157 let survivors = live_ids(&manager);
4158 assert_eq!(survivors.len(), 60);
4159
4160 drain_maintenance(&manager);
4161
4162 let stats = manager.stats();
4163 assert_eq!(stats.consolidation.runs_started, 1);
4164 assert_eq!(stats.consolidation.runs_completed, 1);
4165 assert_eq!(stats.consolidation.segments_merged, 1);
4166 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
4167 assert!(stats.consolidation.bytes_reclaimed > 0);
4168
4169 let bytes_after = manager.resident_bytes();
4171 assert!(
4172 bytes_after <= bytes_before_delete,
4173 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
4174 );
4175 assert_eq!(
4176 bytes_with_tombstones - bytes_after,
4177 stats.consolidation.bytes_reclaimed
4178 );
4179 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
4180
4181 assert_eq!(live_ids(&manager), survivors);
4182 assert_eq!(manager.sealed.read().len(), 1);
4183 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
4184 }
4185
4186 #[test]
4187 fn query_results_are_identical_before_during_and_after_consolidation() {
4188 let manager = consolidating_manager(3);
4189 let ids = seal_bulk_segment(&manager, 20);
4190 seal_bulk_segment(&manager, 20);
4191 for id in ids.iter().take(10) {
4192 assert!(manager.delete(*id).expect("delete"));
4193 }
4194
4195 let before = live_ids(&manager);
4196 let survivor = *ids.last().expect("at least one row");
4197 let tombstoned = ids[0];
4198
4199 let mut mid_tick_observations = 0;
4202 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
4203 manager.run_maintenance().expect("maintenance");
4204 assert_eq!(live_ids(&manager), before);
4205 assert!(manager.get(survivor).is_some());
4206 assert!(manager.get(tombstoned).is_none());
4207 assert_eq!(manager.get_by_kind("table").len(), before.len());
4208 mid_tick_observations += 1;
4209 assert!(mid_tick_observations < 1_000);
4210 }
4211
4212 assert!(
4213 mid_tick_observations > 2,
4214 "expected the run to span several ticks, saw {mid_tick_observations}"
4215 );
4216 assert_eq!(live_ids(&manager), before);
4217 assert!(manager.get(survivor).is_some());
4218 assert!(manager.get(tombstoned).is_none());
4219 assert_eq!(manager.get_by_kind("table").len(), before.len());
4220 }
4221
4222 #[test]
4223 fn get_by_kind_reaches_sealed_segments() {
4224 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4225 let ids = seal_bulk_segment(&manager, 5);
4226 assert!(manager.delete(ids[0]).expect("delete"));
4227
4228 assert_eq!(manager.get_by_kind("table").len(), 4);
4232 assert_eq!(manager.get_by_kind("vector").len(), 0);
4233 }
4234
4235 #[test]
4236 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
4237 const PER_TICK: usize = 4;
4238 let manager = consolidating_manager(PER_TICK);
4239 let ids = seal_bulk_segment(&manager, 60);
4240 for id in ids.iter().take(30) {
4241 assert!(manager.delete(*id).expect("delete"));
4242 }
4243
4244 manager.run_maintenance().expect("maintenance");
4246 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
4247
4248 let mut previous = PER_TICK as u64;
4251 let mut ticks = 1;
4252 while let Some(copied) = manager.consolidation_progress() {
4253 assert!(
4254 copied - previous <= PER_TICK as u64,
4255 "tick copied {} entities, bound is {PER_TICK}",
4256 copied - previous
4257 );
4258 previous = copied;
4259 manager.run_maintenance().expect("maintenance");
4260 ticks += 1;
4261 assert!(ticks < 100);
4262 }
4263
4264 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
4267 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4268 assert_eq!(live_ids(&manager).len(), 30);
4269 }
4270
4271 #[test]
4272 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
4273 let manager = SegmentManager::with_config(
4274 "test",
4275 ManagerConfig {
4276 max_sealed_segments: 3,
4277 ..Default::default()
4278 },
4279 );
4280 for _ in 0..4 {
4281 seal_bulk_segment(&manager, 5);
4282 }
4283 assert_eq!(manager.sealed.read().len(), 4);
4284
4285 let before = live_ids(&manager);
4286 drain_maintenance(&manager);
4287
4288 assert_eq!(manager.sealed.read().len(), 1);
4291 assert_eq!(manager.stats().consolidation.segments_merged, 4);
4292 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
4293 assert_eq!(live_ids(&manager), before);
4294 }
4295
4296 #[test]
4297 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
4298 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4299 seal_bulk_segment(&manager, 10);
4300 seal_bulk_segment(&manager, 10);
4301
4302 for _ in 0..5 {
4303 manager.run_maintenance().expect("maintenance");
4304 }
4305
4306 assert_eq!(manager.stats().consolidation.runs_started, 0);
4307 assert_eq!(manager.sealed.read().len(), 2);
4308 }
4309
4310 #[test]
4311 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
4312 use std::sync::atomic::AtomicBool;
4313
4314 let manager = Arc::new(consolidating_manager(2));
4315 let ids = seal_bulk_segment(&manager, 40);
4316 seal_bulk_segment(&manager, 40);
4317 for id in ids.iter().take(20) {
4318 assert!(manager.delete(*id).expect("delete"));
4319 }
4320
4321 let expected = live_ids(&manager);
4322 let stop = Arc::new(AtomicBool::new(false));
4323
4324 let readers: Vec<_> = (0..4)
4325 .map(|_| {
4326 let manager = Arc::clone(&manager);
4327 let stop = Arc::clone(&stop);
4328 let expected = expected.clone();
4329 std::thread::spawn(move || {
4330 let mut scans = 0u64;
4331 loop {
4334 let mut seen: Vec<u64> = manager
4335 .query_all(|_| true)
4336 .into_iter()
4337 .map(|entity| entity.id.raw())
4338 .collect();
4339 seen.sort_unstable();
4340 assert_eq!(seen, expected, "reader saw a torn sealed set");
4343 scans += 1;
4344 if stop.load(Ordering::Relaxed) {
4345 break;
4346 }
4347 }
4348 scans
4349 })
4350 })
4351 .collect();
4352
4353 drain_maintenance(&manager);
4354 stop.store(true, Ordering::Relaxed);
4355
4356 for reader in readers {
4357 let scans = reader.join().expect("reader thread");
4358 assert!(scans > 0, "reader never got to scan");
4359 }
4360
4361 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4362 assert_eq!(live_ids(&manager), expected);
4363 }
4364
4365 #[test]
4366 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
4367 let manager = consolidating_manager(4);
4368 let ids = seal_bulk_segment(&manager, 20);
4369 for id in ids.iter().take(6) {
4370 assert!(manager.delete(*id).expect("delete"));
4371 }
4372
4373 manager.run_maintenance().expect("maintenance");
4375 assert_eq!(manager.consolidation_progress(), Some(4));
4376 let doomed = ids[6];
4377 assert!(manager.delete(doomed).expect("delete"));
4378
4379 let expected = live_ids(&manager);
4380 drain_maintenance(&manager);
4381
4382 assert!(manager.get(doomed).is_none(), "deleted row came back");
4383 assert_eq!(live_ids(&manager), expected);
4384 }
4385
4386 #[test]
4387 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
4388 let manager = consolidating_manager(3);
4389 let ids = seal_bulk_segment(&manager, 30);
4390 for id in ids.iter().take(12) {
4391 assert!(manager.delete(*id).expect("delete"));
4392 }
4393
4394 let before = live_ids(&manager);
4395 let sealed_before: Vec<SegmentId> = manager
4396 .sealed
4397 .read()
4398 .iter()
4399 .map(|segment| segment.read().id())
4400 .collect();
4401
4402 manager.run_maintenance().expect("maintenance");
4403 manager.run_maintenance().expect("maintenance");
4404 assert!(manager.consolidation.read().is_some());
4405
4406 *manager.consolidation.write() = None;
4410
4411 assert_eq!(live_ids(&manager), before);
4412 let sealed_after: Vec<SegmentId> = manager
4413 .sealed
4414 .read()
4415 .iter()
4416 .map(|segment| segment.read().id())
4417 .collect();
4418 assert_eq!(sealed_after, sealed_before);
4419 assert_eq!(manager.stats().consolidation.runs_completed, 0);
4420
4421 drain_maintenance(&manager);
4422 assert_eq!(live_ids(&manager), before);
4423 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4424 }
4425
4426 #[test]
4427 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
4428 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4429
4430 let entities: Vec<UnifiedEntity> = (0..10)
4434 .map(|i| {
4435 UnifiedEntity::table_row(
4436 manager.next_entity_id(),
4437 "users",
4438 i as u64,
4439 vec![Value::Integer(i as i64)],
4440 )
4441 })
4442 .collect();
4443 let ids = manager.bulk_insert(entities).expect("bulk insert");
4444
4445 let mut meta = Metadata::new();
4446 meta.set(
4447 "os",
4448 super::super::metadata::MetadataValue::String("linux".to_string()),
4449 );
4450 manager.set_metadata(ids[9], meta).expect("set metadata");
4451 manager.force_seal().expect("seal");
4452
4453 assert!(manager.delete(ids[0]).expect("delete"));
4454 assert!(manager.delete(ids[1]).expect("delete"));
4455 assert!(manager.delete(ids[2]).expect("delete"));
4456
4457 let before: Vec<UnifiedEntity> = {
4458 let mut entities = manager.query_all(|_| true);
4459 entities.sort_by_key(|entity| entity.id.raw());
4460 entities
4461 };
4462
4463 drain_maintenance(&manager);
4464
4465 let after: Vec<UnifiedEntity> = {
4466 let mut entities = manager.query_all(|_| true);
4467 entities.sort_by_key(|entity| entity.id.raw());
4468 entities
4469 };
4470
4471 let columns = |entity: &UnifiedEntity| match &entity.data {
4472 super::super::entity::EntityData::Row(row) => row.columns.clone(),
4473 _ => Vec::new(),
4474 };
4475
4476 assert_eq!(before.len(), after.len());
4477 for (old, new) in before.iter().zip(after.iter()) {
4478 assert_eq!(old.id, new.id);
4479 assert_eq!(columns(old), columns(new));
4480 assert_eq!(
4481 old.sequence_id, new.sequence_id,
4482 "merge must not re-issue sequence ids"
4483 );
4484 }
4485 assert!(manager
4486 .get_metadata(ids[9])
4487 .expect("metadata survives the merge")
4488 .has("os"));
4489 }
4490
4491 #[test]
4492 #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
4493 fn test_lifecycle_events() {
4494 let manager = SegmentManager::new("test");
4495
4496 manager
4497 .insert(UnifiedEntity::vector(
4498 manager.next_entity_id(),
4499 "v",
4500 vec![0.1],
4501 ))
4502 .unwrap();
4503
4504 let events = manager.drain_events();
4505
4506 assert!(events
4508 .iter()
4509 .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
4510 assert!(events
4511 .iter()
4512 .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
4513 }
4514}