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