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