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