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