1#![allow(clippy::too_many_arguments)] mod backing;
18pub use backing::PagedEntityCursor;
19mod graph_index_merge;
20
21use parking_lot::RwLock;
22use std::collections::HashMap;
23use std::sync::atomic::{AtomicU64, Ordering};
24use std::sync::{Arc, Weak};
25
26use super::entity::{EntityId, EntityKind, UnifiedEntity};
27use super::metadata::{Metadata, MetadataFilter};
28use super::segment::{
29 GraphEntityKind, GrowingSegment, SegmentConfig, SegmentError, SegmentId, SegmentScanCursor,
30 SegmentState, UnifiedSegment, ZoneColPred,
31};
32use super::visibility_map::VisibilityMap;
33use crate::storage::visibility::entity_visible_with_context;
34use crate::transaction::read::SnapshotContext;
35
36pub const CONSOLIDATION_TOMBSTONE_RATIO: f64 = 0.20;
43
44pub const CONSOLIDATION_FRAGMENTATION_RATIO: f64 = 0.30;
49
50pub const CONSOLIDATION_ENTITIES_PER_TICK: usize = 4_096;
56
57#[derive(Debug, Clone)]
59pub struct ManagerConfig {
60 pub segment_config: SegmentConfig,
62 pub max_sealed_segments: usize,
64 pub idle_seal_secs: u64,
66 pub enable_consolidation: bool,
68 pub consolidation_entities_per_tick: usize,
70 pub enable_archival: bool,
72 pub archive_age_secs: u64,
74}
75
76impl Default for ManagerConfig {
77 fn default() -> Self {
78 Self {
79 segment_config: SegmentConfig::default(),
80 max_sealed_segments: 10,
81 idle_seal_secs: 300, enable_consolidation: true,
83 consolidation_entities_per_tick: CONSOLIDATION_ENTITIES_PER_TICK,
84 enable_archival: true,
85 archive_age_secs: 86400 * 7, }
87 }
88}
89
90#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
95pub struct ConsolidationStats {
96 pub runs_started: u64,
98 pub runs_completed: u64,
100 pub segments_merged: u64,
102 pub tombstones_reclaimed: u64,
104 pub bytes_reclaimed: u64,
106}
107
108#[derive(Debug, Clone, Default)]
110pub struct ManagerStats {
111 pub total_entities: usize,
113 pub growing_count: usize,
115 pub sealed_count: usize,
117 pub archived_count: usize,
119 pub total_memory_bytes: usize,
121 pub seal_ops: u64,
123 pub consolidation: ConsolidationStats,
125}
126
127#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
128pub struct SegmentScanStats {
129 pub segments_total: u64,
130 pub segments_scanned: u64,
131 pub segments_pruned: u64,
132}
133
134#[derive(Debug, Clone)]
136pub enum LifecycleEvent {
137 SegmentCreated(SegmentId),
138 SegmentSealed(SegmentId),
139 SegmentConsolidated {
140 source: Vec<SegmentId>,
141 target: SegmentId,
142 },
143 SegmentArchived(SegmentId),
144 EntityInserted(EntityId, SegmentId),
145 EntityDeleted(EntityId, SegmentId),
146}
147
148struct Consolidation {
155 sources: Vec<SourceProgress>,
157 merged_id: SegmentId,
159 merged: GrowingSegment,
161 cursor: usize,
163 pending_ids: Vec<EntityId>,
165 pending_cursor: usize,
167 entities_copied: u64,
169}
170
171struct SourceProgress {
173 id: SegmentId,
174 copied: Vec<EntityId>,
176 tombstones_at_open: Option<usize>,
185}
186
187struct PhysicalRootCacheEntry {
188 segments: Vec<(u64, u64)>,
189 seed: u64,
190 root: u64,
191}
192
193#[doc(hidden)]
196pub struct PhysicalScanCursor {
197 sources: Vec<(Weak<RwLock<GrowingSegment>>, SegmentScanCursor)>,
198 source: usize,
199}
200
201impl PhysicalScanCursor {
202 #[doc(hidden)]
203 pub fn finished(&self) -> bool {
204 self.source == self.sources.len()
205 }
206}
207
208pub struct SegmentManager {
210 row_constraint_lock: Arc<parking_lot::ReentrantMutex<()>>,
213 collection: String,
215 config: ManagerConfig,
217 next_segment_id: AtomicU64,
219 next_entity_id: AtomicU64,
221 next_sequence_id: AtomicU64,
222 next_row_id: AtomicU64,
224 total_entities_atomic: AtomicU64,
228 growing: RwLock<Option<Arc<RwLock<GrowingSegment>>>>,
230 sealed: RwLock<Vec<Arc<RwLock<GrowingSegment>>>>,
232 retired: RwLock<Vec<Weak<RwLock<GrowingSegment>>>>,
234 archived: RwLock<Vec<SegmentId>>,
236 entity_segment: RwLock<HashMap<EntityId, SegmentId>>,
240 column_schema: RwLock<Option<Arc<Vec<String>>>>,
243 recovery_row_schemas: RwLock<Vec<Arc<Vec<String>>>>,
246 stats: RwLock<ManagerStats>,
248 consolidation: RwLock<Option<Consolidation>>,
251 events: RwLock<Vec<LifecycleEvent>>,
253 visibility_map: VisibilityMap,
257 physical_root_cache: parking_lot::Mutex<Option<PhysicalRootCacheEntry>>,
258 backing: RwLock<Option<backing::BackingRead>>,
259}
260
261impl SegmentManager {
262 pub fn new(collection: impl Into<String>) -> Self {
264 Self::with_config(collection, ManagerConfig::default())
265 }
266
267 pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
269 Self {
270 row_constraint_lock: Arc::new(parking_lot::ReentrantMutex::new(())),
271 physical_root_cache: parking_lot::Mutex::new(None),
272 backing: RwLock::new(None),
273 collection: collection.into(),
274 config,
275 next_segment_id: AtomicU64::new(1),
276 next_entity_id: AtomicU64::new(1),
277 next_sequence_id: AtomicU64::new(0),
278 next_row_id: AtomicU64::new(1),
279 total_entities_atomic: AtomicU64::new(0),
280 growing: RwLock::new(None),
281 sealed: RwLock::new(Vec::new()),
282 retired: RwLock::new(Vec::new()),
283 archived: RwLock::new(Vec::new()),
284 entity_segment: RwLock::new(HashMap::new()),
285 column_schema: RwLock::new(None),
286 recovery_row_schemas: RwLock::new(Vec::new()),
287 stats: RwLock::new(ManagerStats::default()),
288 consolidation: RwLock::new(None),
289 events: RwLock::new(Vec::new()),
290 visibility_map: VisibilityMap::new(),
291 }
292 }
293
294 #[doc(hidden)]
295 pub fn row_constraint_lock(&self) -> Arc<parking_lot::ReentrantMutex<()>> {
296 Arc::clone(&self.row_constraint_lock)
297 }
298
299 pub fn get_or_init_schema(
301 &self,
302 named: &HashMap<String, reddb_types::Value>,
303 ) -> Arc<Vec<String>> {
304 {
305 let schema = self.column_schema.read();
306 if let Some(ref s) = *schema {
307 return Arc::clone(s);
308 }
309 }
310 let cols: Vec<String> = named.keys().cloned().collect();
311 let arc = Arc::new(cols);
312 *self.column_schema.write() = Some(Arc::clone(&arc));
313 arc
314 }
315
316 pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
318 self.column_schema.read().clone()
319 }
320
321 #[doc(hidden)]
325 pub fn restore_owned_row_schemas<'a>(
326 &self,
327 entities: impl Iterator<Item = &'a mut UnifiedEntity>,
328 ) {
329 let schema = self.column_schema();
330 for entity in entities {
331 if let super::entity::EntityData::Row(row) = &mut entity.data {
332 if row.named.is_none() && row.schema.is_none() {
333 row.schema = schema.clone();
334 }
335 }
336 }
337 }
338
339 #[doc(hidden)]
340 pub fn set_column_schema_if_empty(&self, columns: Vec<String>) {
341 if columns.is_empty()
342 || columns
343 .iter()
344 .collect::<std::collections::HashSet<_>>()
345 .len()
346 != columns.len()
347 {
348 return;
349 }
350 let mut schema = self.column_schema.write();
351 if schema.is_none() {
352 *schema = Some(Arc::new(columns));
353 }
354 }
355
356 #[doc(hidden)]
360 pub fn compact_recovered_rows(&self) -> usize {
361 if self.recovery_preserves_named_rows() {
362 return 0;
363 }
364 let Some(schema) = self.column_schema() else {
365 return 0;
366 };
367 if schema
368 .iter()
369 .collect::<std::collections::HashSet<_>>()
370 .len()
371 != schema.len()
372 {
373 return 0;
374 }
375 let mut growing = self.growing.write();
376 let mut sealed = self.sealed.write();
377 let mut compacted = 0;
378 for segment in growing.iter_mut().chain(sealed.iter_mut()) {
379 if let Some(segment) = Arc::get_mut(segment) {
380 compacted += segment.get_mut().compact_recovered_rows(&schema);
381 }
382 }
383 compacted
384 }
385
386 pub fn collection(&self) -> &str {
388 &self.collection
389 }
390
391 pub fn config(&self) -> &ManagerConfig {
393 &self.config
394 }
395
396 pub fn stats(&self) -> ManagerStats {
399 let mut s = self.stats.read().clone();
400 s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
401 s.total_memory_bytes = self.resident_bytes() as usize;
402 s
403 }
404
405 pub fn resident_bytes(&self) -> u64 {
408 self.memory_usage().1
409 }
410
411 pub fn memory_bytes(&self) -> u64 {
414 self.memory_usage().0
415 }
416
417 fn memory_usage(&self) -> (u64, u64) {
421 let consolidation = self.consolidation.read();
422 let growing = self.growing.read();
423 let sealed = self.sealed.read();
424 let mut retired = self.retired.write();
425 let mut payload = 0u64;
426 let mut resident = 0u64;
427 {
428 let mut add = |segment: &GrowingSegment| {
429 payload = payload.saturating_add(segment.memory_bytes());
430 resident = resident.saturating_add(segment.resident_bytes());
431 };
432 for segment in growing.iter().chain(sealed.iter()) {
433 add(&segment.read());
434 }
435 retired.retain(|source| {
436 if let Some(source) = source.upgrade() {
437 add(&source.read());
438 true
439 } else {
440 false
441 }
442 });
443 if let Some(run) = consolidation.as_ref() {
444 add(&run.merged);
445 }
446 }
447 let mut bookkeeping =
448 (retired.capacity() * std::mem::size_of::<Weak<RwLock<GrowingSegment>>>()) as u64;
449 if let Some(run) = consolidation.as_ref() {
450 bookkeeping = bookkeeping
451 .saturating_add(
452 (run.sources.capacity() * std::mem::size_of::<SourceProgress>()) as u64,
453 )
454 .saturating_add(
455 (run.pending_ids.capacity() * std::mem::size_of::<EntityId>()) as u64,
456 );
457 for source in &run.sources {
458 bookkeeping = bookkeeping.saturating_add(
459 (source.copied.capacity() * std::mem::size_of::<EntityId>()) as u64,
460 );
461 }
462 }
463 (
464 payload.saturating_add(bookkeeping),
465 resident.saturating_add(bookkeeping),
466 )
467 }
468
469 pub fn reclaimable_bytes(&self) -> u64 {
471 let growing = self.growing.read();
472 let sealed = self.sealed.read();
473 growing
474 .iter()
475 .chain(sealed.iter())
476 .map(|segment| segment.read().reclaimable_bytes())
477 .fold(0, u64::saturating_add)
478 }
479
480 pub fn next_entity_id(&self) -> EntityId {
482 EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
483 }
484
485 pub fn identity_counters(&self) -> (u64, u64) {
487 (
488 self.next_entity_id.load(Ordering::SeqCst),
489 self.next_row_id.load(Ordering::SeqCst),
490 )
491 }
492
493 pub fn restore_identity_counters(&self, entity_id_next: u64, row_id_next: u64) {
494 assert!(
495 entity_id_next > 0 && row_id_next > 0,
496 "invariant: next identities are positive"
497 );
498 self.next_entity_id
499 .fetch_max(entity_id_next, Ordering::SeqCst);
500 self.next_row_id.fetch_max(row_id_next, Ordering::SeqCst);
501 }
502
503 pub fn sequence_id_next(&self) -> u64 {
504 let mut next = self.next_sequence_id.load(Ordering::SeqCst);
505 let growing = self.growing.read();
506 let sealed = self.sealed.read();
507 for segment in growing.iter().chain(sealed.iter()) {
508 next = next.max(segment.read().sequence_id_next());
509 }
510 self.next_sequence_id.fetch_max(next, Ordering::SeqCst);
511 next
512 }
513
514 pub fn restore_sequence_id_next(&self, next: u64) {
515 self.next_sequence_id.fetch_max(next, Ordering::SeqCst);
516 }
517
518 pub fn next_row_id(&self) -> u64 {
520 self.next_row_id.fetch_add(1, Ordering::SeqCst)
521 }
522
523 pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
527 let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
528 start..start + n
529 }
530
531 pub fn register_row_id(&self, id: u64) {
534 let candidate = id.saturating_add(1);
535 let mut current = self.next_row_id.load(Ordering::SeqCst);
536 while candidate > current {
537 match self.next_row_id.compare_exchange(
538 current,
539 candidate,
540 Ordering::SeqCst,
541 Ordering::SeqCst,
542 ) {
543 Ok(_) => break,
544 Err(updated) => current = updated,
545 }
546 }
547 }
548
549 fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
556 {
558 let growing = self.growing.read();
559 if let Some(segment) = growing.as_ref() {
560 return Arc::clone(segment);
561 }
562 }
563
564 let mut growing = self.growing.write();
566 if let Some(segment) = growing.as_ref() {
568 return Arc::clone(segment);
569 }
570
571 let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
572 let segment = GrowingSegment::new(id, &self.collection);
573 segment.restore_sequence_id_next(self.next_sequence_id.load(Ordering::SeqCst));
574 let segment_arc = Arc::new(RwLock::new(segment));
575 *growing = Some(Arc::clone(&segment_arc));
576
577 self.emit(LifecycleEvent::SegmentCreated(id));
578
579 self.stats.write().growing_count += 1;
582
583 segment_arc
584 }
585
586 pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
588 if entity.id.raw() == 0 {
589 entity.id = self.next_entity_id();
590 }
591 self.reject_backing_duplicate(entity.id)?;
592 if self.has_backing() && self.get_with(entity.id, |_| ()).is_some() {
593 return Err(SegmentError::AlreadyExists(entity.id));
594 }
595 self.next_entity_id
596 .fetch_max(entity.id.raw().saturating_add(1), Ordering::SeqCst);
597 self.maybe_seal_growing()?;
599
600 let segment_arc = self.get_or_create_growing();
601 let mut segment = segment_arc.write();
602
603 let entity_id = entity.id;
604 let segment_id = segment.id();
605
606 segment.insert(entity)?;
607
608 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
610
611 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
620
621 Ok(entity_id)
622 }
623
624 fn with_restoration_segment<T>(
628 &self,
629 id: EntityId,
630 restore: impl FnOnce(&mut GrowingSegment) -> Result<T, SegmentError>,
631 ) -> Result<T, SegmentError> {
632 loop {
633 let segment = self.get_or_create_growing();
634 let mut growing = segment.write();
635 if !growing.state().is_writable() {
636 continue;
637 }
638 if !growing.is_tombstoned(id) {
639 return restore(&mut growing);
640 }
641 drop(growing);
642 self.force_seal()?;
643 }
644 }
645
646 fn recovery_preserves_named_rows(&self) -> bool {
647 self.collection.starts_with("red_") || self.collection.starts_with("red.")
651 }
652
653 fn recovery_schema(&self, row: &super::entity::RowData) -> Option<Arc<Vec<String>>> {
654 if self.recovery_preserves_named_rows() {
655 return None;
656 }
657 let named = row.named.as_ref()?;
658 let matches = |schema: &Vec<String>| {
659 schema.len() == named.len() && schema.iter().all(|name| named.contains_key(name))
660 };
661 if let Some(schema) = self.column_schema().filter(|schema| matches(schema)) {
662 return Some(schema);
663 }
664 {
665 let schemas = self.recovery_row_schemas.read();
666 if let Some(schema) = schemas.iter().find(|schema| matches(schema)) {
667 return Some(Arc::clone(schema));
668 }
669 }
670 let mut schemas = self.recovery_row_schemas.write();
671 if let Some(schema) = schemas.iter().find(|schema| matches(schema)) {
672 return Some(Arc::clone(schema));
673 }
674 const MAX_RECOVERY_ROW_SCHEMAS: usize = 16;
675 if schemas.len() >= MAX_RECOVERY_ROW_SCHEMAS {
676 return None;
677 }
678 let mut names: Vec<_> = named.keys().cloned().collect();
679 names.sort_unstable();
680 let schema = Arc::new(names);
681 schemas.push(Arc::clone(&schema));
682 Some(schema)
683 }
684
685 #[doc(hidden)]
689 pub fn restore_snapshot_row_layout(&self, entity: &mut UnifiedEntity) {
690 if let super::entity::EntityData::Row(row) = &mut entity.data {
691 if let Some(schema) = self.recovery_schema(row) {
692 row.compact_with_schema(&schema);
693 }
694 }
695 }
696
697 #[doc(hidden)]
699 pub fn restore_snapshot_entity(
700 &self,
701 mut entity: UnifiedEntity,
702 ) -> Result<EntityId, SegmentError> {
703 self.restore_snapshot_row_layout(&mut entity);
704 self.maybe_seal_growing()?;
705 if self.get(entity.id).is_some() {
706 return Err(SegmentError::AlreadyExists(entity.id));
707 }
708 let entity_id = entity.id;
709 self.next_entity_id
710 .fetch_max(entity_id.raw().saturating_add(1), Ordering::SeqCst);
711 self.with_restoration_segment(entity_id, |segment| {
712 let segment_id = segment.id();
713 segment.adopt_entity(entity, None);
715 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
716 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
717 Ok(entity_id)
718 })
719 }
720
721 pub fn insert_batch(
723 &self,
724 entities: Vec<UnifiedEntity>,
725 ) -> Result<Vec<EntityId>, SegmentError> {
726 let mut ids = Vec::with_capacity(entities.len());
727 for entity in entities {
728 ids.push(self.insert(entity)?);
729 }
730 Ok(ids)
731 }
732
733 pub fn bulk_insert(
736 &self,
737 mut entities: Vec<UnifiedEntity>,
738 ) -> Result<Vec<EntityId>, SegmentError> {
739 if self.has_backing() {
741 let mut identities = std::collections::HashSet::with_capacity(entities.len());
742 for entity in &entities {
743 if entity.id.raw() != 0 {
744 self.reject_backing_duplicate(entity.id)?;
745 if self.get_with(entity.id, |_| ()).is_some() || !identities.insert(entity.id) {
746 return Err(SegmentError::AlreadyExists(entity.id));
747 }
748 }
749 }
750 }
751 for entity in &mut entities {
753 if entity.id.raw() == 0 {
754 entity.id = self.next_entity_id();
755 }
756 self.next_entity_id
757 .fetch_max(entity.id.raw().saturating_add(1), Ordering::SeqCst);
758 if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
759 if *row_id == 0 {
760 *row_id = self.next_row_id();
761 } else {
762 self.register_row_id(*row_id);
763 }
764 }
765 }
766
767 if let Some(first_row) = entities.first() {
770 if let super::entity::EntityData::Row(ref row) = first_row.data {
771 if let Some(ref named) = row.named {
772 let schema = self.get_or_init_schema(named);
773 for entity in &mut entities {
774 if let super::entity::EntityData::Row(ref mut row) = entity.data {
775 if let Some(named) = row.named.take() {
776 let mut cols = Vec::with_capacity(schema.len());
777 for col_name in schema.iter() {
778 cols.push(
779 named
780 .get(col_name)
781 .cloned()
782 .unwrap_or(reddb_types::Value::Null),
783 );
784 }
785 row.columns = cols;
786 row.schema = Some(Arc::clone(&schema));
787 }
788 }
789 }
790 }
791 }
792 }
793
794 let segment_arc = self.get_or_create_growing();
795 let mut segment = segment_arc.write();
796 let ids = segment.bulk_insert(entities)?;
798
799 self.total_entities_atomic
804 .fetch_add(ids.len() as u64, Ordering::Relaxed);
805
806 Ok(ids)
807 }
808
809 pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
811 let mut entity = self.get_with(id, Clone::clone)?;
812 self.restore_owned_row_schemas(std::iter::once(&mut entity));
813 Some(entity)
814 }
815
816 #[doc(hidden)]
819 pub fn get_with<T>(&self, id: EntityId, visit: impl FnOnce(&UnifiedEntity) -> T) -> Option<T> {
820 if let Some(growing_arc) = self.growing.read().as_ref() {
822 let growing = growing_arc.read();
823 if let Some(entity) = growing.get(id) {
824 return Some(visit(entity));
825 }
826 }
827
828 let sealed = self.sealed.read();
830 for segment in sealed.iter() {
831 let seg = segment.read();
832 if let Some(entity) = seg.get(id) {
833 return Some(visit(entity));
834 }
835 }
836
837 None
838 }
839
840 #[doc(hidden)]
844 pub fn graph_node_logical_id(
845 &self,
846 id: EntityId,
847 mut before_segment: impl FnMut() -> bool,
848 ) -> Option<EntityId> {
849 if let Some(growing_arc) = self.growing.read().as_ref() {
850 if !before_segment() {
851 return None;
852 }
853 let growing = growing_arc.read();
854 if let Some(entity) = growing.get(id) {
855 return matches!(entity.kind, EntityKind::GraphNode(_))
856 .then(|| entity.logical_id());
857 }
858 }
859 let sealed = self.sealed.read();
860 for segment in sealed.iter() {
861 if !before_segment() {
862 return None;
863 }
864 let segment = segment.read();
865 if let Some(entity) = segment.get(id) {
866 return matches!(entity.kind, EntityKind::GraphNode(_))
867 .then(|| entity.logical_id());
868 }
869 }
870 None
871 }
872
873 #[doc(hidden)]
878 pub fn visit_graph_index_batches<E>(
879 &self,
880 kind: GraphEntityKind,
881 keys: &[EntityId],
882 mut before_work: impl FnMut() -> bool,
883 mut admit: impl FnMut(usize) -> Result<(), E>,
884 mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
885 ) -> Result<bool, E> {
886 use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
887 admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
888 let mut admitted = 0;
889 let sources = loop {
890 let count = {
891 let growing = self.growing.read();
892 let sealed = self.sealed.read();
893 sealed.len() + usize::from(growing.is_some())
894 };
895 let bytes = count.saturating_mul(std::mem::size_of::<(
896 Arc<RwLock<GrowingSegment>>,
897 GraphIndexCursor,
898 )>());
899 admit(bytes.saturating_sub(admitted))?;
900 admitted = admitted.max(bytes);
901 let mut sources = Vec::with_capacity(count);
902 let growing = self.growing.read();
903 let sealed = self.sealed.read();
904 if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
905 continue;
906 }
907 for segment in growing.iter().chain(sealed.iter()) {
908 if !before_work() {
909 return Ok(false);
910 }
911 let view = segment.read();
912 if !view.may_contain_graph_kind(kind) {
913 continue;
914 }
915 let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
916 return Ok(false);
917 };
918 sources.push((Arc::clone(segment), cursor));
919 }
920 break sources;
921 };
922 let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
923 for (source, mut cursor) in sources {
924 while !cursor.finished(keys) {
925 let mut count = 0;
926 {
927 let segment = source.read();
928 if !segment.may_contain_graph_kind(kind) {
929 break;
930 }
931 if !segment.graph_index_batch(
932 kind,
933 keys,
934 &mut cursor,
935 &mut before_work,
936 &mut |id| {
937 ids[count] = id;
938 count += 1;
939 },
940 ) {
941 return Ok(false);
942 }
943 }
944 if !before_work() || !consume(&ids[..count])? {
945 return Ok(false);
946 }
947 }
948 }
949 Ok(true)
950 }
951
952 pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
957 let schema = self.column_schema();
958 let clone_row = |entity: &UnifiedEntity| {
959 let mut owned = entity.clone();
960 if let super::entity::EntityData::Row(row) = &mut owned.data {
961 if row.schema.is_none() {
962 row.schema = schema.clone();
963 }
964 }
965 owned
966 };
967 let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
968 let mut remaining: Vec<usize> = (0..ids.len()).collect(); if let Some(growing_arc) = self.growing.read().as_ref() {
973 let growing = if let Some(g) = growing_arc.try_read() {
974 g
975 } else {
976 growing_arc.read()
977 };
978 remaining.retain(|&i| {
979 if let Some(entity) = growing.get(ids[i]) {
980 out[i] = Some(clone_row(entity));
981 false } else {
983 true }
985 });
986 }
987
988 if remaining.is_empty() {
989 self.restore_owned_row_schemas(out.iter_mut().flatten());
990 return out;
991 }
992
993 let sealed = self.sealed.read();
995 for segment in sealed.iter() {
996 if remaining.is_empty() {
997 break;
998 }
999 let seg = segment.read();
1000 remaining.retain(|&i| {
1001 if let Some(entity) = seg.get(ids[i]) {
1002 out[i] = Some(clone_row(entity));
1003 false
1004 } else {
1005 true
1006 }
1007 });
1008 }
1009
1010 drop(sealed);
1011 self.restore_owned_row_schemas(out.iter_mut().flatten());
1012 out
1013 }
1014
1015 pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
1028 where
1029 F: FnMut(usize, &UnifiedEntity),
1030 {
1031 thread_local! {
1040 static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
1041 const { std::cell::Cell::new(Vec::new()) };
1042 }
1043
1044 let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
1045 remaining.clear();
1046
1047 if let Some(growing_arc) = self.growing.read().as_ref() {
1048 let growing = if let Some(g) = growing_arc.try_read() {
1049 g
1050 } else {
1051 growing_arc.read()
1052 };
1053 for (i, id) in ids.iter().enumerate() {
1054 if let Some(entity) = growing.get(*id) {
1055 f(i, entity);
1056 } else {
1057 remaining.push(i);
1058 }
1059 }
1060 } else {
1061 remaining.reserve(ids.len());
1062 remaining.extend(0..ids.len());
1063 }
1064
1065 if !remaining.is_empty() {
1066 let sealed = self.sealed.read();
1067 for segment in sealed.iter() {
1068 if remaining.is_empty() {
1069 break;
1070 }
1071 let seg = segment.read();
1072 remaining.retain(|&i| {
1073 if let Some(entity) = seg.get(ids[i]) {
1074 f(i, entity);
1075 false
1076 } else {
1077 true
1078 }
1079 });
1080 }
1081 }
1082
1083 REMAINING_SCRATCH.with(|cell| cell.set(remaining));
1084 }
1085
1086 #[allow(dead_code)] fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
1089 if let Some(growing_arc) = self.growing.read().as_ref() {
1091 let growing = growing_arc.read();
1092 if let Some(entity) = growing.get(id) {
1093 return Some(entity.clone());
1094 }
1095 }
1096
1097 let sealed = self.sealed.read();
1099 for segment in sealed.iter() {
1100 if let Some(entity) = segment.get(id) {
1101 return Some(entity.clone());
1102 }
1103 }
1104
1105 None
1106 }
1107
1108 fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
1109 let sealed = self.sealed.read();
1110 sealed
1111 .iter()
1112 .find(|segment_arc| segment_arc.read().contains(id))
1113 .map(Arc::clone)
1114 }
1115
1116 fn rewrite_sealed_entity_into_growing(
1117 &self,
1118 entity: UnifiedEntity,
1119 metadata: Option<&Metadata>,
1120 ) -> Result<(), SegmentError> {
1121 let entity_id = entity.id;
1122 let sealed_arc = self
1123 .find_sealed_segment_arc(entity_id)
1124 .ok_or(SegmentError::NotFound(entity_id))?;
1125
1126 let metadata_to_apply = {
1127 let mut sealed = sealed_arc.write();
1128 let existing_metadata = sealed.get_metadata(entity_id);
1129 if !sealed.force_delete(entity_id) {
1130 return Err(SegmentError::NotFound(entity_id));
1131 }
1132 metadata.cloned().or(existing_metadata)
1133 };
1134
1135 let growing_id = self.with_restoration_segment(entity_id, |growing| {
1136 growing.insert_preserving_sequence(entity)?;
1137 if let Some(metadata) = metadata_to_apply {
1138 growing.set_metadata(entity_id, metadata)?;
1139 }
1140 Ok(growing.id())
1141 })?;
1142
1143 self.entity_segment.write().insert(entity_id, growing_id);
1144 Ok(())
1145 }
1146
1147 pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
1149 self.materialize_backing_entity(entity.id)?;
1150 let entity_id = entity.id;
1151 let mut entity = Some(entity);
1152
1153 if let Some(growing_arc) = self.growing.read().as_ref() {
1155 let mut growing = growing_arc.write();
1156 if growing.contains(entity_id) && growing.state().is_writable() {
1157 return growing.update(entity.take().expect("entity already moved"));
1158 }
1159 }
1160
1161 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1163 if let Some(seg_id) = segment_id {
1164 if let Some(growing_arc) = self.growing.read().as_ref() {
1165 let mut growing = growing_arc.write();
1166 if growing.id() == seg_id && growing.state().is_writable() {
1167 return growing.update(entity.take().expect("entity already moved"));
1168 }
1169 }
1170 }
1171
1172 if let Some(entity) = entity.take() {
1173 return self.rewrite_sealed_entity_into_growing(entity, None);
1174 }
1175
1176 Err(SegmentError::NotFound(entity_id))
1177 }
1178
1179 pub fn update_with_metadata(
1182 &self,
1183 entity: UnifiedEntity,
1184 metadata: Option<&Metadata>,
1185 ) -> Result<(), SegmentError> {
1186 self.materialize_backing_entity(entity.id)?;
1187 let entity_id = entity.id;
1188 let mut entity = Some(entity);
1189
1190 if let Some(growing_arc) = self.growing.read().as_ref() {
1192 let mut growing = growing_arc.write();
1193 if growing.contains(entity_id) && growing.state().is_writable() {
1194 growing.update(entity.take().expect("entity already moved"))?;
1195 if let Some(metadata) = metadata {
1196 growing.set_metadata(entity_id, metadata.clone())?;
1197 }
1198 return Ok(());
1199 }
1200 }
1201
1202 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1204 if let Some(seg_id) = segment_id {
1205 if let Some(growing_arc) = self.growing.read().as_ref() {
1206 let mut growing = growing_arc.write();
1207 if growing.id() == seg_id && growing.state().is_writable() {
1208 growing.update(entity.take().expect("entity already moved"))?;
1209 if let Some(metadata) = metadata {
1210 growing.set_metadata(entity_id, metadata.clone())?;
1211 }
1212 return Ok(());
1213 }
1214 }
1215 }
1216
1217 if let Some(entity) = entity.take() {
1218 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1219 }
1220
1221 Err(SegmentError::NotFound(entity_id))
1222 }
1223
1224 pub fn update_hot(
1228 &self,
1229 entity: UnifiedEntity,
1230 modified_columns: &[String],
1231 ) -> Result<(), SegmentError> {
1232 self.materialize_backing_entity(entity.id)?;
1233 let entity_id = entity.id;
1234 let mut entity = Some(entity);
1235
1236 if let Some(growing_arc) = self.growing.read().as_ref() {
1237 let mut growing = growing_arc.write();
1238 if growing.contains(entity_id) && growing.state().is_writable() {
1239 return growing.update_hot(
1240 entity.take().expect("entity already moved"),
1241 modified_columns,
1242 );
1243 }
1244 }
1245
1246 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1247 if let Some(seg_id) = segment_id {
1248 if let Some(growing_arc) = self.growing.read().as_ref() {
1249 let mut growing = growing_arc.write();
1250 if growing.id() == seg_id && growing.state().is_writable() {
1251 return growing.update_hot(
1252 entity.take().expect("entity already moved"),
1253 modified_columns,
1254 );
1255 }
1256 }
1257 }
1258
1259 if let Some(entity) = entity.take() {
1260 return self.rewrite_sealed_entity_into_growing(entity, None);
1261 }
1262
1263 Err(SegmentError::NotFound(entity_id))
1264 }
1265
1266 pub fn update_hot_with_metadata(
1269 &self,
1270 entity: UnifiedEntity,
1271 modified_columns: &[String],
1272 metadata: Option<&Metadata>,
1273 ) -> Result<(), SegmentError> {
1274 self.materialize_backing_entity(entity.id)?;
1275 let entity_id = entity.id;
1276 let mut entity = Some(entity);
1277
1278 if let Some(growing_arc) = self.growing.read().as_ref() {
1279 let mut growing = growing_arc.write();
1280 if growing.contains(entity_id) && growing.state().is_writable() {
1281 growing.update_hot(
1282 entity.take().expect("entity already moved"),
1283 modified_columns,
1284 )?;
1285 if let Some(metadata) = metadata {
1286 growing.set_metadata(entity_id, metadata.clone())?;
1287 }
1288 return Ok(());
1289 }
1290 }
1291
1292 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1293 if let Some(seg_id) = segment_id {
1294 if let Some(growing_arc) = self.growing.read().as_ref() {
1295 let mut growing = growing_arc.write();
1296 if growing.id() == seg_id && growing.state().is_writable() {
1297 growing.update_hot(
1298 entity.take().expect("entity already moved"),
1299 modified_columns,
1300 )?;
1301 if let Some(metadata) = metadata {
1302 growing.set_metadata(entity_id, metadata.clone())?;
1303 }
1304 return Ok(());
1305 }
1306 }
1307 }
1308
1309 if let Some(entity) = entity.take() {
1310 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1311 }
1312
1313 Err(SegmentError::NotFound(entity_id))
1314 }
1315
1316 pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
1319 where
1320 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
1321 {
1322 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
1323 items.into_iter().collect();
1324 if items.is_empty() {
1325 return Ok(());
1326 }
1327
1328 if let Some(growing_arc) = self.growing.read().as_ref() {
1329 let mut growing = growing_arc.write();
1330 if growing.state().is_writable() {
1331 match growing.update_hot_batch_with_metadata(items.iter().copied()) {
1332 Ok(()) => return Ok(()),
1333 Err(SegmentError::NotFound(_)) => {}
1334 Err(other) => return Err(other),
1335 }
1336 }
1337 }
1338
1339 for (entity, modified_columns, metadata) in items {
1340 self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
1341 }
1342 Ok(())
1343 }
1344
1345 pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
1347 self.materialize_backing_entity(id)?;
1348 let deleted = self.delete_resident(id)?;
1349 if deleted {
1350 self.remember_backing_delete(id);
1351 }
1352 Ok(deleted)
1353 }
1354
1355 fn delete_resident(&self, id: EntityId) -> Result<bool, SegmentError> {
1356 if let Some(growing_arc) = self.growing.read().as_ref() {
1359 let mut growing = growing_arc.write();
1360 if growing.contains(id) && growing.state().is_writable() {
1361 let seg_id = growing.id();
1362 let deleted = growing.delete(id)?;
1363 if deleted {
1364 self.entity_segment.write().remove(&id);
1365 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1366 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1367 }
1368 return Ok(deleted);
1369 }
1370 }
1371
1372 let segment_id = self.entity_segment.read().get(&id).copied();
1375
1376 if let Some(seg_id) = segment_id {
1377 if let Some(growing_arc) = self.growing.read().as_ref() {
1378 let mut growing = growing_arc.write();
1379 if growing.id() == seg_id && growing.state().is_writable() {
1380 let deleted = growing.delete(id)?;
1381 if deleted {
1382 self.entity_segment.write().remove(&id);
1383 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1384 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1385 }
1386 return Ok(deleted);
1387 }
1388 }
1389 }
1390
1391 {
1394 let sealed = self.sealed.read();
1395 for segment_arc in sealed.iter() {
1396 let mut seg = segment_arc.write();
1397 let seg_id = seg.id();
1398 if seg.contains(id) {
1399 let deleted = seg.force_delete(id);
1400 drop(seg);
1401 if deleted {
1402 self.entity_segment.write().remove(&id);
1403 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1404 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1405 }
1406 return Ok(deleted);
1407 }
1408 }
1409 }
1410
1411 Ok(false)
1412 }
1413
1414 pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1415 if self.has_backing() {
1416 let mut deleted = Vec::new();
1417 for id in ids {
1418 if self.delete(*id)? {
1419 deleted.push(*id);
1420 }
1421 }
1422 return Ok(deleted);
1423 }
1424 if ids.is_empty() {
1425 return Ok(Vec::new());
1426 }
1427
1428 let mut deleted_ids = Vec::with_capacity(ids.len());
1429
1430 if let Some(growing_arc) = self.growing.read().as_ref() {
1431 let mut growing = growing_arc.write();
1432 if growing.state().is_writable() {
1433 let seg_id = growing.id();
1434 let deleted = growing.delete_batch(ids)?;
1435 if !deleted.is_empty() {
1436 {
1437 let mut entity_segment = self.entity_segment.write();
1438 for id in &deleted {
1439 entity_segment.remove(id);
1440 }
1441 }
1442 self.total_entities_atomic
1443 .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
1444 for id in &deleted {
1445 self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1446 }
1447 deleted_ids.extend(deleted);
1448 }
1449 }
1450 }
1451
1452 if deleted_ids.len() == ids.len() {
1453 return Ok(deleted_ids);
1454 }
1455
1456 let deleted_set: std::collections::HashSet<EntityId> =
1457 deleted_ids.iter().copied().collect();
1458 for &id in ids {
1459 if deleted_set.contains(&id) {
1460 continue;
1461 }
1462 if self.delete(id)? {
1463 deleted_ids.push(id);
1464 }
1465 }
1466
1467 Ok(deleted_ids)
1468 }
1469
1470 pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1472 if let Some(growing_arc) = self.growing.read().as_ref() {
1474 let growing = growing_arc.read();
1475 if growing.contains(id) {
1476 return growing.get_metadata(id);
1477 }
1478 }
1479
1480 let segment_id = self.entity_segment.read().get(&id).copied();
1482
1483 if let Some(seg_id) = segment_id {
1484 if let Some(growing_arc) = self.growing.read().as_ref() {
1485 let growing = growing_arc.read();
1486 if growing.id() == seg_id {
1487 return growing.get_metadata(id);
1488 }
1489 }
1490
1491 let sealed = self.sealed.read();
1492 for segment in sealed.iter() {
1493 if segment.id() == seg_id {
1494 return segment.get_metadata(id);
1495 }
1496 }
1497 }
1498
1499 if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1500 return segment_arc.read().get_metadata(id);
1501 }
1502
1503 None
1504 }
1505
1506 pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1508 self.materialize_backing_entity(id)?;
1509 if let Some(growing_arc) = self.growing.read().as_ref() {
1512 let mut growing = growing_arc.write();
1513 if growing.contains(id) && growing.state().is_writable() {
1514 return growing.set_metadata(id, metadata);
1515 }
1516 }
1517
1518 let segment_id = self.entity_segment.read().get(&id).copied();
1520
1521 if let Some(seg_id) = segment_id {
1522 if let Some(growing_arc) = self.growing.read().as_ref() {
1523 let mut growing = growing_arc.write();
1524 if growing.id() == seg_id && growing.state().is_writable() {
1525 return growing.set_metadata(id, metadata);
1526 }
1527 }
1528 }
1529
1530 if let Some(entity) = self.get(id) {
1531 return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1532 }
1533
1534 Err(SegmentError::NotFound(id))
1535 }
1536
1537 fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1539 let should_seal = {
1540 let growing_opt = self.growing.read();
1541 if let Some(growing_arc) = growing_opt.as_ref() {
1542 let growing = growing_arc.read();
1543 growing.should_seal(&self.config.segment_config)
1544 || growing.idle_secs() >= self.config.idle_seal_secs
1545 } else {
1546 false
1547 }
1548 };
1549
1550 if should_seal {
1551 self.seal_current()?;
1552 }
1553
1554 Ok(())
1555 }
1556
1557 pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1559 self.seal_current_if(|_| true)?
1560 .ok_or(SegmentError::InvalidState(SegmentState::Sealed))
1561 }
1562
1563 fn seal_current_if(
1564 &self,
1565 should_seal: impl FnOnce(&GrowingSegment) -> bool,
1566 ) -> Result<Option<SegmentId>, SegmentError> {
1567 let disk_residency = self.has_backing();
1568 let mut growing_slot = self.growing.write();
1571 if let Some(growing_arc) = growing_slot.as_ref() {
1572 let mut growing = growing_arc.write();
1573 if !should_seal(&growing) {
1574 return Ok(None);
1575 }
1576 if disk_residency {
1577 self.next_sequence_id
1578 .fetch_max(growing.sequence_id_next(), Ordering::SeqCst);
1579 }
1580 let seg_id = growing.id();
1581 let entity_count = growing.stats().entity_count as u64;
1582
1583 growing.seal()?;
1585
1586 drop(growing); let growing_arc = growing_slot
1592 .take()
1593 .expect("sealed source remains published");
1594 self.sealed.write().push(growing_arc);
1595 drop(growing_slot);
1596
1597 self.mark_sealed_pages_visible(entity_count);
1599
1600 {
1602 let mut stats = self.stats.write();
1603 stats.growing_count = stats.growing_count.saturating_sub(1);
1604 stats.sealed_count += 1;
1605 stats.seal_ops += 1;
1606 }
1607
1608 self.emit(LifecycleEvent::SegmentSealed(seg_id));
1609
1610 return Ok(Some(seg_id));
1611 }
1612
1613 Ok(None)
1614 }
1615
1616 pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1618 let has_growing = self.growing.read().is_some();
1619 if has_growing {
1620 Ok(Some(self.seal_current()?))
1621 } else {
1622 Ok(None)
1623 }
1624 }
1625
1626 pub fn all_visible_fraction(&self) -> f64 {
1636 const ROWS_PER_PAGE: u32 = 256;
1637 let sealed = self.sealed.read();
1638 if sealed.is_empty() {
1639 return 0.0;
1640 }
1641 let mut total_pages: u64 = 0;
1642 for seg_arc in sealed.iter() {
1643 let seg = seg_arc.read();
1644 let entity_count = seg.stats().entity_count as u64;
1645 let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1646 total_pages += pages;
1647 }
1648 if total_pages == 0 {
1649 return 0.0;
1650 }
1651 let visible = self.visibility_map.all_visible_count();
1652 (visible as f64 / total_pages as f64).min(1.0)
1653 }
1654
1655 fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1658 const ROWS_PER_PAGE: u32 = 256;
1659 let existing_visible = self.visibility_map.all_visible_count();
1660 let start_page = existing_visible as u32;
1662 let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1663 let end_page = start_page + new_pages as u32;
1664 self.visibility_map.mark_range_visible(start_page, end_page);
1665 }
1666
1667 #[doc(hidden)]
1671 pub fn find_unique_key_conflict(
1672 &self,
1673 columns: &[Vec<String>],
1674 key_index: usize,
1675 signatures: &[String],
1676 exclude_id: Option<EntityId>,
1677 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1678 ) -> Option<EntityId> {
1679 if let Some(growing) = self.growing.read().as_ref() {
1680 if let Some(id) = growing.read().find_unique_key_conflict(
1681 columns,
1682 key_index,
1683 signatures,
1684 exclude_id,
1685 reserves_key,
1686 ) {
1687 return Some(id);
1688 }
1689 }
1690 for segment in self.sealed.read().iter() {
1691 if let Some(id) = segment.read().find_unique_key_conflict(
1692 columns,
1693 key_index,
1694 signatures,
1695 exclude_id,
1696 reserves_key,
1697 ) {
1698 return Some(id);
1699 }
1700 }
1701 None
1702 }
1703
1704 pub fn for_each_entity<F>(&self, callback: F)
1710 where
1711 F: FnMut(&UnifiedEntity) -> bool,
1712 {
1713 self.for_each_entity_in_segments(|_| true, || true, callback);
1714 }
1715
1716 #[doc(hidden)]
1718 pub fn checkpoint_revisions(&self) -> Vec<(u64, u64)> {
1719 let growing = self.growing.read();
1720 let sealed = self.sealed.read();
1721 growing
1722 .iter()
1723 .chain(sealed.iter())
1724 .map(|segment| {
1725 let segment = segment.read();
1726 (segment.id(), segment.checkpoint_revision())
1727 })
1728 .collect()
1729 }
1730
1731 #[doc(hidden)]
1732 pub fn checkpoint_mutations_since(&self, before: &[(u64, u64)]) -> Option<Vec<EntityId>> {
1733 let growing = self.growing.read();
1734 let sealed = self.sealed.read();
1735 let previous: HashMap<_, _> = before.iter().copied().collect();
1740 if previous.len() != before.len() {
1741 return None;
1742 }
1743 let mut retained_sources = 0;
1744 let mut changed = std::collections::BTreeSet::new();
1745 for segment in growing.iter().chain(sealed.iter()) {
1746 let segment = segment.read();
1747 let revision = previous.get(&segment.id()).copied();
1748 retained_sources += usize::from(revision.is_some());
1749 let revision = revision.unwrap_or(0);
1752 changed.extend(segment.checkpoint_mutations_since(revision)?);
1753 if changed.len() > super::segment::CHECKPOINT_MUTATIONS_MAX {
1754 return None;
1755 }
1756 }
1757 (retained_sources == previous.len()).then(|| changed.into_iter().collect())
1760 }
1761
1762 #[doc(hidden)]
1763 pub fn checkpoint_tracking_bytes(&self) -> usize {
1764 let growing = self.growing.read();
1765 let sealed = self.sealed.read();
1766 growing
1767 .iter()
1768 .chain(sealed.iter())
1769 .map(|segment| segment.read().checkpoint_tracking_bytes())
1770 .sum()
1771 }
1772
1773 #[doc(hidden)]
1775 pub fn cached_physical_root(&self, seed: u64, compute: impl FnOnce() -> u64) -> u64 {
1776 let segments = self.checkpoint_revisions();
1777 let mut cache = self.physical_root_cache.lock();
1778 if let Some(previous) = cache.as_ref() {
1779 if previous.seed == seed && previous.segments == segments {
1780 return previous.root;
1781 }
1782 }
1783 let root = compute();
1784 *cache = Some(PhysicalRootCacheEntry {
1785 segments,
1786 seed,
1787 root,
1788 });
1789 root
1790 }
1791
1792 #[doc(hidden)]
1797 pub fn for_each_entity_in_id_order(&self, mut visit: impl FnMut(&UnifiedEntity) -> bool) {
1798 self.for_each_entity_source_in_id_order(|_, entity| visit(entity));
1799 }
1800
1801 #[doc(hidden)]
1803 pub fn page_in_id_order(&self, offset: usize, limit: usize) -> (usize, Vec<UnifiedEntity>) {
1804 let end = offset.saturating_add(limit);
1805 let mut total = 0;
1806 let mut items = Vec::with_capacity(limit.min(self.count()));
1807 self.for_each_entity_in_id_order(|entity| {
1808 if total >= offset && total < end {
1809 items.push(entity.clone());
1810 }
1811 total += 1;
1812 true
1813 });
1814 self.restore_owned_row_schemas(items.iter_mut());
1815 (total, items)
1816 }
1817
1818 #[doc(hidden)]
1819 pub fn for_each_entity_with_metadata_in_id_order(
1820 &self,
1821 mut visit: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1822 ) {
1823 self.for_each_entity_source_in_id_order(|segment, entity| {
1824 visit(entity, segment.get_metadata(entity.id))
1825 });
1826 }
1827
1828 fn for_each_entity_source_in_id_order(
1829 &self,
1830 mut visit: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
1831 ) {
1832 let segments: Vec<_> = {
1833 let growing = self.growing.read();
1834 let sealed = self.sealed.read();
1835 growing.iter().chain(sealed.iter()).cloned().collect()
1836 };
1837 let mut order = Vec::new();
1838 for (source, segment) in segments.iter().enumerate() {
1839 segment.read().for_each_fast(|entity| {
1840 order.push((entity.id, source));
1841 true
1842 });
1843 }
1844 order.sort_by_key(|(id, _)| id.raw());
1845 for (id, source) in order {
1846 let segment = segments[source].read();
1847 if let Some(entity) = segment.get(id) {
1848 if !visit(&segment, entity) {
1849 break;
1850 }
1851 }
1852 }
1853 }
1854
1855 #[doc(hidden)]
1858 pub fn for_each_entity_with_metadata(
1859 &self,
1860 mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1861 ) {
1862 self.for_each_entity_source(|segment, entity| {
1863 callback(entity, segment.get_metadata(entity.id))
1864 });
1865 }
1866
1867 fn for_each_entity_source(
1868 &self,
1869 mut callback: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
1870 ) {
1871 if let Some(growing_arc) = self.growing.read().as_ref() {
1872 let growing = growing_arc.read();
1873 if !growing.for_each_fast(|entity| callback(&growing, entity)) {
1874 return;
1875 }
1876 }
1877 let sealed = self.sealed.read();
1878 for segment_arc in sealed.iter() {
1879 let segment = segment_arc.read();
1880 if !segment.for_each_fast(|entity| callback(&segment, entity)) {
1881 return;
1882 }
1883 }
1884 }
1885
1886 fn for_each_entity_in_segments<F>(
1887 &self,
1888 segment_filter: impl Fn(&GrowingSegment) -> bool,
1889 mut before_segment: impl FnMut() -> bool,
1890 mut callback: F,
1891 ) where
1892 F: FnMut(&UnifiedEntity) -> bool,
1893 {
1894 if let Some(growing_arc) = self.growing.read().as_ref() {
1898 if !before_segment() {
1899 return;
1900 }
1901 let growing = if let Some(g) = growing_arc.try_read() {
1902 g
1903 } else {
1904 growing_arc.read()
1905 };
1906 if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
1907 return;
1908 }
1909 }
1910
1911 let sealed = self.sealed.read();
1913 for segment_arc in sealed.iter() {
1914 if !before_segment() {
1915 return;
1916 }
1917 let segment = segment_arc.read();
1918 if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
1919 return;
1920 }
1921 }
1922 }
1923
1924 pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1940 &self,
1941 init: FInit,
1942 fold: FFold,
1943 reduce: FReduce,
1944 ) -> T
1945 where
1946 T: Send,
1947 FInit: Fn() -> T + Send + Sync,
1948 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1949 FReduce: Fn(T, T) -> T + Send + Sync,
1950 {
1951 #[cfg(not(target_arch = "wasm32"))]
1952 use rayon::prelude::*;
1953
1954 let mut acc = init();
1957 if let Some(growing_arc) = self.growing.read().as_ref() {
1958 let growing = if let Some(g) = growing_arc.try_read() {
1959 g
1960 } else {
1961 growing_arc.read()
1962 };
1963 growing.for_each_fast(|entity| {
1964 acc = fold(std::mem::replace(&mut acc, init()), entity);
1965 true
1966 });
1967 }
1968
1969 let segments: Vec<_> = {
1973 let sealed = self.sealed.read();
1974 sealed.iter().cloned().collect()
1975 };
1976
1977 if segments.len() <= 1 {
1978 for seg_arc in &segments {
1979 let seg = seg_arc.read();
1980 seg.for_each_fast(|entity| {
1981 acc = fold(std::mem::replace(&mut acc, init()), entity);
1982 true
1983 });
1984 }
1985 return acc;
1986 }
1987
1988 #[cfg(not(target_arch = "wasm32"))]
1989 let segments = segments.into_par_iter();
1990 #[cfg(target_arch = "wasm32")]
1991 let segments = segments.into_iter();
1992 let partials = segments.map(|seg_arc| {
1993 let mut local = init();
1994 let seg = seg_arc.read();
1995 seg.for_each_fast(|entity| {
1996 local = fold(std::mem::replace(&mut local, init()), entity);
1997 true
1998 });
1999 local
2000 });
2001 #[cfg(not(target_arch = "wasm32"))]
2002 let sealed_acc = partials.reduce(&init, &reduce);
2003 #[cfg(target_arch = "wasm32")]
2004 let sealed_acc = partials.fold(init(), &reduce);
2005
2006 reduce(acc, sealed_acc)
2007 }
2008
2009 pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
2018 where
2019 F: FnMut(&UnifiedEntity) -> bool,
2020 {
2021 let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
2022 }
2023
2024 pub fn for_each_entity_zoned_with_stats<F>(
2025 &self,
2026 zone_preds: &[(&str, ZoneColPred<'_>)],
2027 mut callback: F,
2028 ) -> SegmentScanStats
2029 where
2030 F: FnMut(&UnifiedEntity) -> bool,
2031 {
2032 let mut scan_stats = SegmentScanStats::default();
2033 if let Some(growing_arc) = self.growing.read().as_ref() {
2039 scan_stats.segments_total += 1;
2040 scan_stats.segments_scanned += 1;
2041 let growing = if let Some(g) = growing_arc.try_read() {
2042 g
2043 } else {
2044 growing_arc.read()
2045 };
2046 if !growing.for_each_fast(&mut callback) {
2047 return scan_stats;
2048 }
2049 }
2050
2051 let sealed = self.sealed.read();
2053 for segment_arc in sealed.iter() {
2054 scan_stats.segments_total += 1;
2055 let segment = segment_arc.read();
2056 if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
2057 scan_stats.segments_pruned += 1;
2058 continue; }
2060 scan_stats.segments_scanned += 1;
2061 if !segment.for_each_fast(&mut callback) {
2062 return scan_stats;
2063 }
2064 }
2065 scan_stats
2066 }
2067
2068 pub fn query_all_zoned<F>(
2078 &self,
2079 zone_preds: &[(&str, ZoneColPred<'_>)],
2080 filter: F,
2081 ) -> Vec<UnifiedEntity>
2082 where
2083 F: Fn(&UnifiedEntity) -> bool + Sync,
2084 {
2085 self.query_all_zoned_with_stats(zone_preds, filter).0
2086 }
2087
2088 pub fn query_all_zoned_with_stats<F>(
2089 &self,
2090 zone_preds: &[(&str, ZoneColPred<'_>)],
2091 filter: F,
2092 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2093 where
2094 F: Fn(&UnifiedEntity) -> bool + Sync,
2095 {
2096 self.project_all_zoned_with_stats(zone_preds, |entity| {
2097 filter(entity).then(|| entity.clone())
2098 })
2099 }
2100
2101 fn project_all_zoned_with_stats<T, F>(
2102 &self,
2103 zone_preds: &[(&str, ZoneColPred<'_>)],
2104 project: F,
2105 ) -> (Vec<T>, SegmentScanStats)
2106 where
2107 T: Send,
2108 F: Fn(&UnifiedEntity) -> Option<T> + Sync,
2109 {
2110 let mut results = Vec::new();
2111 let mut scan_stats = SegmentScanStats::default();
2112
2113 if let Some(growing_arc) = self.growing.read().as_ref() {
2116 scan_stats.segments_total += 1;
2117 scan_stats.segments_scanned += 1;
2118 let growing = if let Some(g) = growing_arc.try_read() {
2119 g
2120 } else {
2121 growing_arc.read()
2122 };
2123 results.extend(growing.iter().filter_map(&project));
2124 }
2125
2126 let sealed = self.sealed.read();
2128 let surviving: Vec<_> = sealed
2130 .iter()
2131 .filter(|seg_arc| {
2132 scan_stats.segments_total += 1;
2133 if zone_preds.is_empty() {
2134 scan_stats.segments_scanned += 1;
2135 return true;
2136 }
2137 let seg = seg_arc.read();
2138 if seg.can_skip_zone_preds(zone_preds) {
2139 scan_stats.segments_pruned += 1;
2140 false
2141 } else {
2142 scan_stats.segments_scanned += 1;
2143 true
2144 }
2145 })
2146 .collect();
2147
2148 let use_parallel = surviving.len() > 1 && crate::execution::should_parallelize();
2149
2150 if use_parallel {
2151 let project_ref = &project;
2152 let segment_results: Vec<Vec<T>> = std::thread::scope(|s| {
2153 surviving
2154 .iter()
2155 .map(|segment| {
2156 s.spawn(move || {
2157 segment
2158 .read()
2159 .iter()
2160 .filter_map(project_ref)
2161 .collect::<Vec<_>>()
2162 })
2163 })
2164 .collect::<Vec<_>>()
2165 .into_iter()
2166 .map(|handle| match handle.join() {
2167 Ok(batch) => batch,
2168 Err(panic) => std::panic::resume_unwind(panic),
2169 })
2170 .collect()
2171 });
2172 for batch in segment_results {
2173 results.extend(batch);
2174 }
2175 } else {
2176 for segment_arc in surviving {
2177 let seg = segment_arc.read();
2178 results.extend(seg.iter().filter_map(&project));
2179 }
2180 }
2181
2182 (results, scan_stats)
2183 }
2184
2185 pub fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
2191 match snapshot {
2192 Some(ctx) => entity_visible_with_context(Some(ctx), entity),
2193 None => {
2194 !crate::storage::visibility::entity_moderation_hidden(entity) && entity.xmax == 0
2195 }
2196 }
2197 }
2198
2199 pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
2206 where
2207 F: Fn(&UnifiedEntity) -> bool + Sync,
2208 {
2209 self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
2210 }
2211
2212 #[doc(hidden)]
2217 pub fn project_key_entity<T>(
2218 &self,
2219 key: &str,
2220 snapshot: Option<&SnapshotContext>,
2221 versioned: bool,
2222 mut project: impl FnMut(&UnifiedEntity) -> T,
2223 ) -> Option<T> {
2224 let mut probe = |segment: &GrowingSegment, best: &mut Option<(u64, T)>| {
2225 if let Some(entity) = segment.find_key_entity(key, versioned, &mut |entity| {
2226 !versioned || Self::scan_entity_visible(snapshot, entity)
2227 }) {
2228 if best.as_ref().is_none_or(|(xmin, _)| entity.xmin >= *xmin) {
2229 *best = Some((entity.xmin, project(entity)));
2230 }
2231 }
2232 };
2233 let mut best = None;
2234 if let Some(growing) = self.growing.read().as_ref() {
2235 probe(&growing.read(), &mut best);
2236 }
2237 for segment in self.sealed.read().iter() {
2238 if !versioned && best.is_some() {
2239 break;
2240 }
2241 probe(&segment.read(), &mut best);
2242 }
2243 best.map(|(_, projected)| projected)
2244 }
2245
2246 #[doc(hidden)]
2248 pub fn scan_graph_kind(
2249 &self,
2250 snapshot: Option<&SnapshotContext>,
2251 kind: GraphEntityKind,
2252 ) -> Vec<UnifiedEntity> {
2253 self.query_all_in_segments(
2254 |segment| segment.may_contain_graph_kind(kind),
2255 |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
2256 )
2257 }
2258
2259 #[doc(hidden)]
2264 pub fn scan_graph_candidates_for_each<F>(
2265 &self,
2266 snapshot: Option<&SnapshotContext>,
2267 kind: GraphEntityKind,
2268 before_segment: impl FnMut() -> bool,
2269 mut callback: F,
2270 ) where
2271 F: FnMut(&UnifiedEntity) -> bool,
2272 {
2273 self.for_each_entity_in_segments(
2274 |segment| segment.may_contain_graph_kind(kind),
2275 before_segment,
2276 |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
2277 );
2278 }
2279
2280 pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
2285 where
2286 F: FnMut(&UnifiedEntity) -> bool,
2287 {
2288 self.for_each_entity(|entity| {
2289 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2290 });
2291 }
2292
2293 #[doc(hidden)]
2294 pub fn physical_scan_cursor(&self) -> PhysicalScanCursor {
2295 let growing = self.growing.read();
2296 let sealed = self.sealed.read();
2297 PhysicalScanCursor {
2298 sources: growing
2299 .iter()
2300 .chain(sealed.iter())
2301 .map(|segment| (Arc::downgrade(segment), segment.read().scan_cursor()))
2302 .collect(),
2303 source: 0,
2304 }
2305 }
2306
2307 #[doc(hidden)]
2310 pub fn physical_scan_batch(
2311 &self,
2312 cursor: &mut PhysicalScanCursor,
2313 work_left: &mut usize,
2314 ) -> Vec<EntityId> {
2315 let mut ids = Vec::with_capacity(super::segment::SCAN_BATCH_SIZE.min(*work_left));
2316 while !cursor.finished() && *work_left > 0 {
2317 let (source, position) = &mut cursor.sources[cursor.source];
2318 if position.finished() {
2319 *work_left -= 1;
2320 cursor.source += 1;
2321 continue;
2322 }
2323 let Some(segment) = source.upgrade() else {
2324 *work_left -= 1;
2325 cursor.source += 1;
2326 continue;
2327 };
2328 segment.read().scan_batch(
2329 position,
2330 &mut || {
2331 if *work_left == 0 {
2332 return false;
2333 }
2334 *work_left -= 1;
2335 true
2336 },
2337 &mut |id| ids.push(id),
2338 );
2339 if position.finished() {
2340 cursor.source += 1;
2341 }
2342 if !ids.is_empty() {
2343 break;
2344 }
2345 }
2346 ids
2347 }
2348
2349 #[doc(hidden)]
2355 pub fn scan_batches(
2356 &self,
2357 snapshot: Option<&SnapshotContext>,
2358 mut before_work: impl FnMut() -> bool,
2359 mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
2360 ) -> bool {
2361 let segments = {
2362 let growing = self.growing.read();
2363 let sealed = self.sealed.read();
2364 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
2365 for segment in growing.iter().chain(sealed.iter()) {
2366 if !before_work() {
2367 return false;
2368 }
2369 let cursor = segment.read().scan_cursor();
2370 segments.push((Arc::clone(segment), cursor));
2371 }
2372 segments
2373 };
2374 let mut ids = Vec::new();
2375 let mut batch = Vec::new();
2376 for (segment, mut cursor) in segments {
2377 while !cursor.finished() {
2378 ids.clear();
2379 batch.clear();
2380 let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
2381 ids.reserve_exact(batch_size);
2382 batch.reserve_exact(batch_size);
2383 {
2384 let segment = segment.read();
2385 if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
2386 return false;
2387 }
2388 }
2389 if !before_work() {
2392 return false;
2393 }
2394 self.for_each_id(&ids, |_, entity| {
2395 if Self::scan_entity_visible(snapshot, entity) {
2396 batch.push(entity.clone());
2397 }
2398 });
2399 if !before_work() {
2400 return false;
2401 }
2402 if !consume(&batch) {
2403 return false;
2404 }
2405 }
2406 }
2407 true
2408 }
2409
2410 pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
2412 &self,
2413 snapshot: Option<&SnapshotContext>,
2414 init: FInit,
2415 fold: FFold,
2416 reduce: FReduce,
2417 ) -> T
2418 where
2419 T: Send,
2420 FInit: Fn() -> T + Send + Sync,
2421 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
2422 FReduce: Fn(T, T) -> T + Send + Sync,
2423 {
2424 self.fold_entities_parallel(
2425 init,
2426 |acc, entity| {
2427 if Self::scan_entity_visible(snapshot, entity) {
2428 fold(acc, entity)
2429 } else {
2430 acc
2431 }
2432 },
2433 reduce,
2434 )
2435 }
2436
2437 pub fn scan_for_each_zoned_with_stats<F>(
2439 &self,
2440 snapshot: Option<&SnapshotContext>,
2441 zone_preds: &[(&str, ZoneColPred<'_>)],
2442 mut callback: F,
2443 ) -> SegmentScanStats
2444 where
2445 F: FnMut(&UnifiedEntity) -> bool,
2446 {
2447 self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
2448 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2449 })
2450 }
2451
2452 pub fn scan_zoned_with_stats<F>(
2454 &self,
2455 snapshot: Option<&SnapshotContext>,
2456 zone_preds: &[(&str, ZoneColPred<'_>)],
2457 filter: F,
2458 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2459 where
2460 F: Fn(&UnifiedEntity) -> bool + Sync,
2461 {
2462 self.query_all_zoned_with_stats(zone_preds, |entity| {
2463 Self::scan_entity_visible(snapshot, entity) && filter(entity)
2464 })
2465 }
2466
2467 pub fn scan_project_zoned_with_stats<T, F>(
2471 &self,
2472 snapshot: Option<&SnapshotContext>,
2473 zone_preds: &[(&str, ZoneColPred<'_>)],
2474 project: F,
2475 ) -> (Vec<T>, SegmentScanStats)
2476 where
2477 T: Send,
2478 F: Fn(&UnifiedEntity) -> Option<T> + Sync,
2479 {
2480 self.project_all_zoned_with_stats(zone_preds, |entity| {
2481 if Self::scan_entity_visible(snapshot, entity) {
2482 project(entity)
2483 } else {
2484 None
2485 }
2486 })
2487 }
2488
2489 pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
2492 where
2493 F: Fn(&UnifiedEntity) -> bool + Sync,
2494 {
2495 let mut entities = self.query_all_in_segments(|_| true, filter);
2496 self.restore_owned_row_schemas(entities.iter_mut());
2497 entities
2498 }
2499
2500 fn query_all_in_segments<F>(
2501 &self,
2502 segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
2503 filter: F,
2504 ) -> Vec<UnifiedEntity>
2505 where
2506 F: Fn(&UnifiedEntity) -> bool + Sync,
2507 {
2508 let mut results = Vec::new();
2509
2510 if let Some(growing_arc) = self.growing.read().as_ref() {
2513 let growing = if let Some(g) = growing_arc.try_read() {
2514 g
2515 } else {
2516 growing_arc.read()
2517 };
2518 if segment_filter(&growing) {
2519 results.extend(growing.iter().filter(|e| filter(e)).cloned());
2520 }
2521 }
2522
2523 let sealed = self.sealed.read();
2525 let use_parallel = sealed.len() > 1 && crate::execution::should_parallelize();
2526 if use_parallel {
2527 let filter_ref = &filter;
2528 let segment_filter_ref = &segment_filter;
2529 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
2530 sealed
2531 .iter()
2532 .map(|segment| {
2533 s.spawn(move || {
2534 let segment = segment.read();
2535 if segment_filter_ref(&segment) {
2536 segment
2537 .iter()
2538 .filter(|e| filter_ref(e))
2539 .cloned()
2540 .collect::<Vec<_>>()
2541 } else {
2542 Vec::new()
2543 }
2544 })
2545 })
2546 .collect::<Vec<_>>()
2547 .into_iter()
2548 .map(|handle| handle.join().unwrap_or_default())
2549 .collect()
2550 });
2551 for batch in segment_results {
2552 results.extend(batch);
2553 }
2554 } else {
2555 for segment in sealed.iter() {
2556 let seg = segment.read();
2557 if segment_filter(&seg) {
2558 results.extend(seg.iter().filter(|e| filter(e)).cloned());
2559 }
2560 }
2561 }
2562
2563 results
2564 }
2565
2566 pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
2571 if let Some(growing_arc) = self.growing.read().as_ref() {
2572 let growing = growing_arc.read();
2573 if growing.may_contain_exact_key(key) {
2574 return true;
2575 }
2576 } else {
2577 return true;
2578 }
2579
2580 let sealed = self.sealed.read();
2581 for segment_arc in sealed.iter() {
2582 let segment = segment_arc.read();
2583 if segment.may_contain_exact_key(key) {
2584 return true;
2585 }
2586 }
2587
2588 false
2589 }
2590
2591 pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2593 let mut results = Vec::new();
2594
2595 if let Some(growing_arc) = self.growing.read().as_ref() {
2597 let growing = growing_arc.read();
2598 results.extend(growing.filter_metadata(filters));
2599 }
2600
2601 let sealed = self.sealed.read();
2603 for segment in sealed.iter() {
2604 results.extend(segment.filter_metadata(filters));
2605 }
2606
2607 results
2608 }
2609
2610 pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2612 let mut results = Vec::new();
2613
2614 if let Some(growing_arc) = self.growing.read().as_ref() {
2616 let growing = growing_arc.read();
2617 for entity in growing.iter_kind(kind) {
2618 results.push(entity.clone());
2619 }
2620 }
2621
2622 let sealed = self.sealed.read();
2627 for segment_arc in sealed.iter() {
2628 let segment = segment_arc.read();
2629 for entity in segment.iter_kind(kind) {
2630 results.push(entity.clone());
2631 }
2632 }
2633
2634 results
2635 }
2636
2637 pub fn count(&self) -> usize {
2639 self.total_entities_atomic.load(Ordering::Relaxed) as usize
2640 }
2641
2642 pub fn segment_ids(&self) -> Vec<SegmentId> {
2644 let mut ids = Vec::new();
2645
2646 if let Some(growing_arc) = self.growing.read().as_ref() {
2647 ids.push(growing_arc.read().id());
2648 }
2649
2650 let sealed = self.sealed.read();
2651 for segment in sealed.iter() {
2652 ids.push(segment.id());
2653 }
2654
2655 ids.extend(self.archived.read().iter().copied());
2656
2657 ids
2658 }
2659
2660 #[inline]
2673 #[allow(clippy::unused_self)]
2674 fn emit(&self, _event: LifecycleEvent) {}
2675
2676 pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2679 std::mem::take(&mut *self.events.write())
2680 }
2681
2682 pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2684 self.maybe_seal_growing()?;
2686
2687 if self.config.enable_consolidation {
2688 self.consolidation_tick();
2689 }
2690
2691 Ok(())
2692 }
2693
2694 fn consolidation_tick(&self) {
2705 let running = self.consolidation.read().is_some();
2708 if !running && !self.start_consolidation() {
2709 return;
2710 }
2711
2712 let mut guard = self.consolidation.write();
2715 let Some(run) = guard.as_mut() else {
2716 return;
2717 };
2718 self.copy_bounded(run);
2719 if run.cursor >= run.sources.len() {
2720 if let Some(run) = guard.take() {
2721 self.finish_consolidation(run);
2722 }
2723 }
2724 }
2725
2726 fn start_consolidation(&self) -> bool {
2730 let sources = self.select_candidates();
2731 self.start_consolidation_from_sources(sources)
2732 }
2733
2734 fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
2735 if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
2736 return false;
2737 }
2738
2739 let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
2740 let merged = GrowingSegment::new(merged_id, &self.collection);
2741
2742 *self.consolidation.write() = Some(Consolidation {
2743 sources: sources
2744 .into_iter()
2745 .map(|id| SourceProgress {
2746 id,
2747 copied: Vec::new(),
2748 tombstones_at_open: None,
2749 })
2750 .collect(),
2751 merged_id,
2752 merged,
2753 cursor: 0,
2754 pending_ids: Vec::new(),
2755 pending_cursor: 0,
2756 entities_copied: 0,
2757 });
2758
2759 self.stats.write().consolidation.runs_started += 1;
2760 true
2761 }
2762
2763 pub fn pressure_consolidation_tick(&self) -> bool {
2768 if !self.config.enable_consolidation {
2769 return false;
2770 }
2771
2772 let running = self.consolidation.read().is_some();
2773 if !running {
2774 if self
2778 .seal_current_if(|segment| segment.reclaimable_bytes() > 0)
2779 .is_err()
2780 {
2781 return false;
2782 }
2783 let sources = self.dirty_sealed_segments();
2784 if !self.start_consolidation_from_sources(sources) {
2785 return false;
2786 }
2787 }
2788
2789 self.consolidation_tick();
2790 true
2791 }
2792
2793 fn select_candidates(&self) -> Vec<SegmentId> {
2808 let sealed = self.sealed.read();
2809 if sealed.is_empty() {
2810 return Vec::new();
2811 }
2812
2813 if sealed.len() > self.config.max_sealed_segments {
2814 return sealed.iter().map(|segment| segment.read().id()).collect();
2815 }
2816
2817 let mut live = 0usize;
2818 let mut tombstones = 0usize;
2819 let mut held = 0u64;
2820 let mut reclaimable = 0u64;
2821 let mut dirty = Vec::new();
2822
2823 for segment_arc in sealed.iter() {
2824 let segment = segment_arc.read();
2825 live += segment.live_entity_count();
2826 tombstones += segment.tombstone_count();
2827 held += segment.resident_bytes();
2828 reclaimable += segment.reclaimable_bytes();
2829 if segment.tombstone_count() > 0 {
2830 dirty.push(segment.id());
2831 }
2832 }
2833
2834 let total = live + tombstones;
2835 let tombstone_ratio = if total == 0 {
2836 0.0
2837 } else {
2838 tombstones as f64 / total as f64
2839 };
2840 let fragmentation = if held == 0 {
2841 0.0
2842 } else {
2843 reclaimable as f64 / held as f64
2844 };
2845
2846 if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
2847 || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
2848 {
2849 return dirty;
2850 }
2851
2852 Vec::new()
2853 }
2854
2855 fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
2856 self.sealed
2857 .read()
2858 .iter()
2859 .filter_map(|segment_arc| {
2860 let segment = segment_arc.read();
2861 (segment.reclaimable_bytes() > 0).then(|| segment.id())
2862 })
2863 .collect()
2864 }
2865
2866 fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
2869 if sources.is_empty() {
2870 return false;
2871 }
2872 let sealed = self.sealed.read();
2873 sealed.iter().any(|segment_arc| {
2874 let segment = segment_arc.read();
2875 sources.contains(&segment.id()) && segment.tombstone_count() > 0
2876 })
2877 }
2878
2879 fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
2880 let sealed = self.sealed.read();
2881 sealed
2882 .iter()
2883 .find(|segment| segment.read().id() == id)
2884 .map(Arc::clone)
2885 }
2886
2887 fn copy_bounded(&self, run: &mut Consolidation) {
2892 let mut budget = self.config.consolidation_entities_per_tick;
2893
2894 while budget > 0 && run.cursor < run.sources.len() {
2895 let source_id = run.sources[run.cursor].id;
2896 let Some(source_arc) = self.sealed_segment(source_id) else {
2897 run.cursor += 1;
2901 run.pending_ids.clear();
2902 run.pending_cursor = 0;
2903 continue;
2904 };
2905
2906 if run.sources[run.cursor].tombstones_at_open.is_none() {
2907 let source = source_arc.read();
2908 run.pending_ids = source.live_entity_ids();
2909 run.pending_cursor = 0;
2910 run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
2911 }
2912
2913 {
2914 let source = source_arc.read();
2915 while budget > 0 && run.pending_cursor < run.pending_ids.len() {
2916 let id = run.pending_ids[run.pending_cursor];
2917 run.pending_cursor += 1;
2918 budget -= 1;
2919
2920 let Some(entity) = source.get(id) else {
2923 continue;
2924 };
2925 let metadata = source.get_metadata(id);
2926 run.merged.adopt_entity(entity.clone(), metadata);
2927 run.sources[run.cursor].copied.push(id);
2928 run.entities_copied += 1;
2929 }
2930 }
2931
2932 if run.pending_cursor >= run.pending_ids.len() {
2933 run.cursor += 1;
2934 run.pending_ids.clear();
2935 run.pending_cursor = 0;
2936 }
2937 }
2938 }
2939
2940 fn finish_consolidation(&self, mut run: Consolidation) {
2946 let mut sealed = self.sealed.write();
2947
2948 let mut positions = Vec::with_capacity(run.sources.len());
2949 for source in &run.sources {
2950 if let Some(index) = sealed
2951 .iter()
2952 .position(|segment| segment.read().id() == source.id)
2953 {
2954 positions.push(index);
2955 }
2956 }
2957 if positions.is_empty() {
2958 return;
2959 }
2960 positions.sort_unstable();
2961
2962 let mut stale = Vec::new();
2966 for source in &run.sources {
2967 let Some(source_arc) = sealed
2968 .iter()
2969 .find(|segment| segment.read().id() == source.id)
2970 else {
2971 continue;
2972 };
2973 let segment = source_arc.read();
2974 if source.tombstones_at_open == Some(segment.tombstone_count()) {
2975 continue;
2976 }
2977 stale.extend(
2978 source
2979 .copied
2980 .iter()
2981 .copied()
2982 .filter(|id| !segment.contains(*id)),
2983 );
2984 }
2985 for id in stale {
2986 run.merged.evict_entity(id);
2987 }
2988
2989 let reclaimed_bytes: u64 = positions
2990 .iter()
2991 .map(|&index| sealed[index].read().resident_bytes())
2992 .sum();
2993 let reclaimed_tombstones: u64 = positions
2994 .iter()
2995 .map(|&index| sealed[index].read().tombstone_count() as u64)
2996 .sum();
2997
2998 if run.merged.seal().is_err() {
2999 return;
3000 }
3001 let merged_bytes = run.merged.resident_bytes();
3002 let merged_arc = Arc::new(RwLock::new(run.merged));
3003
3004 {
3005 let mut retired = self.retired.write();
3006 retired.retain(|source| source.strong_count() > 0);
3007 for &index in positions.iter().rev() {
3008 let source = sealed.remove(index);
3009 if Arc::strong_count(&source) > 1 {
3010 retired.push(Arc::downgrade(&source));
3011 }
3012 }
3013 }
3014 let insert_at = positions[0].min(sealed.len());
3015 sealed.insert(insert_at, Arc::clone(&merged_arc));
3016 drop(sealed);
3017
3018 let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
3020 {
3021 let mut entity_segment = self.entity_segment.write();
3022 for segment_id in entity_segment.values_mut() {
3023 if source_ids.contains(segment_id) {
3024 *segment_id = run.merged_id;
3025 }
3026 }
3027 }
3028
3029 {
3030 let mut stats = self.stats.write();
3031 stats.sealed_count = stats
3032 .sealed_count
3033 .saturating_sub(positions.len())
3034 .saturating_add(1);
3035 let consolidation = &mut stats.consolidation;
3036 consolidation.runs_completed += 1;
3037 consolidation.segments_merged += positions.len() as u64;
3038 consolidation.tombstones_reclaimed += reclaimed_tombstones;
3039 consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
3040 }
3041
3042 self.emit(LifecycleEvent::SegmentConsolidated {
3043 source: source_ids,
3044 target: run.merged_id,
3045 });
3046 }
3047
3048 #[cfg(test)]
3051 fn consolidation_progress(&self) -> Option<u64> {
3052 self.consolidation
3053 .read()
3054 .as_ref()
3055 .map(|run| run.entities_copied)
3056 }
3057}
3058
3059#[cfg(test)]
3060mod tests {
3061 use std::collections::HashSet;
3062
3063 use super::*;
3064 use crate::storage::entity::{EntityData, EntityKind, RowData};
3065 use crate::transaction::snapshot::{Snapshot, SnapshotManager};
3066 use reddb_types::Value;
3067
3068 #[test]
3069 fn recovery_view_projection_keeps_schema_bounds_without_retaining_rows_or_identities() {
3070 for collection in ["rows", "red_internal"] {
3071 let manager = SegmentManager::new(collection);
3072 let identities = manager.identity_counters();
3073 for id in 1..=18 {
3074 let field = format!("field_{id}");
3075 let mut entity =
3076 UnifiedEntity::table_row(EntityId::new(id), collection, id, vec![]);
3077 entity.data = EntityData::Row(RowData {
3078 columns: Vec::new(),
3079 named: Some(HashMap::from([
3080 (field.clone(), Value::Integer(id as i64)),
3081 ("id".into(), Value::Integer(id as i64)),
3082 ])),
3083 schema: None,
3084 });
3085 manager.restore_snapshot_row_layout(&mut entity);
3086 let row = entity.data.as_row().expect("row");
3087 assert_eq!(row.get_field(&field), Some(&Value::Integer(id as i64)));
3088 assert_eq!(row.get_field("id"), Some(&Value::Integer(id as i64)));
3089 assert_eq!(
3090 row.named.is_some(),
3091 collection.starts_with("red_") || id > 16
3092 );
3093 }
3094 assert_eq!(manager.count(), 0, "a projection never adopts payloads");
3095 assert_eq!(manager.identity_counters(), identities);
3096 assert_eq!(
3097 manager.recovery_row_schemas.read().len(),
3098 if collection == "rows" { 16 } else { 0 }
3099 );
3100 }
3101 }
3102
3103 #[test]
3104 fn physical_pages_keep_global_order_across_segments_and_deleted_slots() {
3105 let manager = SegmentManager::new("pages");
3106 for group in [[50, 10], [30, 20]] {
3107 for id in group {
3108 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "pages", id, vec![]);
3109 entity.data = EntityData::Row(RowData::with_names(
3110 vec![Value::Integer(id as i64)],
3111 vec!["id".into()],
3112 ));
3113 manager.insert(entity).unwrap();
3114 }
3115 manager.force_seal().unwrap();
3116 }
3117 manager
3118 .insert(UnifiedEntity::table_row(
3119 EntityId::new(40),
3120 "pages",
3121 40,
3122 vec![],
3123 ))
3124 .unwrap();
3125 let (total, page) = manager.page_in_id_order(1, 2);
3126 assert_eq!(total, 5);
3127 assert_eq!(
3128 page.iter()
3129 .map(|entity| entity.id.raw())
3130 .collect::<Vec<_>>(),
3131 [20, 30]
3132 );
3133 let EntityData::Row(row) = &page[0].data else {
3134 panic!("row")
3135 };
3136 assert_eq!(row.get_field("id"), Some(&Value::Integer(20)));
3137 manager.delete(EntityId::new(30)).unwrap();
3138 let (total, page) = manager.page_in_id_order(1, 2);
3139 assert_eq!(total, 4);
3140 assert_eq!(
3141 page.iter()
3142 .map(|entity| entity.id.raw())
3143 .collect::<Vec<_>>(),
3144 [20, 40]
3145 );
3146 for offset in [4, usize::MAX] {
3147 let (total, page) = manager.page_in_id_order(offset, 2);
3148 assert_eq!(total, 4);
3149 assert!(page.is_empty());
3150 }
3151 }
3152
3153 fn lookup_version(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
3154 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
3155 entity.data = EntityData::Row(RowData::with_names(
3156 vec![Value::text(key)],
3157 vec!["key".into()],
3158 ));
3159 entity.set_xmin(xmin);
3160 entity.set_xmax(xmax);
3161 entity
3162 }
3163
3164 #[test]
3165 fn key_lookup_matches_snapshot_scan_across_segments() {
3166 let manager = Arc::new(SegmentManager::new("keys"));
3167 for entity in [
3168 lookup_version(1, "changed", 3, 12),
3169 lookup_version(2, "deleted", 3, 7),
3170 ] {
3171 manager.insert(entity).expect("old version");
3172 }
3173 manager.force_seal().expect("seal old versions");
3174 for entity in [
3175 lookup_version(3, "changed", 12, 0),
3176 lookup_version(4, "pending", 4, 0),
3177 lookup_version(5, "own", 15, 0),
3178 lookup_version(6, "live", 3, 0),
3179 ] {
3180 manager.insert(entity).expect("new version");
3181 }
3182 std::thread::spawn(move || {
3183 let snapshot = SnapshotContext {
3184 snapshot: Snapshot {
3185 xid: 10,
3186 in_progress: HashSet::from([4]),
3187 },
3188 manager: Arc::new(SnapshotManager::new()),
3189 own_xids: HashSet::from([15]),
3190 requires_index_fallback: false,
3191 serializable_reader: None,
3192 };
3193 for snapshot in [Some(&snapshot), None] {
3194 for key in ["changed", "deleted", "pending", "own", "live", "absent"] {
3195 let expected = manager
3196 .scan(snapshot, |entity| {
3197 crate::storage::segment_key_lookup::SegmentKeyLookup::key(entity)
3198 == Some(key)
3199 })
3200 .into_iter()
3201 .max_by_key(|entity| entity.xmin)
3202 .map(|entity| entity.id);
3203 assert_eq!(
3204 manager.project_key_entity(key, snapshot, true, |entity| entity.id),
3205 expected,
3206 "key={key}"
3207 );
3208 }
3209 }
3210 assert_eq!(
3211 manager.project_key_entity("changed", Some(&snapshot), true, |entity| entity.id),
3212 Some(EntityId::new(1))
3213 );
3214 assert_eq!(
3215 manager.project_key_entity("changed", None, true, |entity| entity.id),
3216 Some(EntityId::new(3))
3217 );
3218 assert!(manager
3219 .project_key_entity("deleted", None, true, |entity| entity.id)
3220 .is_none());
3221 })
3222 .join()
3223 .expect("snapshot worker");
3224 }
3225
3226 #[test]
3227 fn key_lookup_survives_paced_consolidation() {
3228 let manager = consolidating_manager(1);
3229 manager
3230 .bulk_insert(
3231 (1..=10)
3232 .map(|id| lookup_version(id, &format!("key-{id}"), 1, 0))
3233 .collect(),
3234 )
3235 .expect("rows");
3236 manager.force_seal().expect("seal");
3237 for id in 1..=4 {
3238 manager.delete(EntityId::new(id)).expect("delete");
3239 }
3240 assert!(manager.pressure_consolidation_tick());
3241 drain_maintenance(&manager);
3242 for id in 1..=10 {
3243 assert_eq!(
3244 manager.project_key_entity(&format!("key-{id}"), None, false, |entity| entity.id),
3245 (id > 4).then_some(EntityId::new(id))
3246 );
3247 }
3248 }
3249
3250 #[test]
3251 #[ignore = "diagnostic of segment fragmentation and retained history; no timing assertions"]
3252 fn key_lookup_fragmentation_history_diagnostic() {
3253 use std::{hint::black_box, time::Instant};
3254 for segments in [1, 16, 64] {
3255 for versions in [1, 100, 1000] {
3256 let manager = SegmentManager::new("keys");
3257 let mut id = 1;
3258 for segment in 0..segments {
3259 let mut rows: Vec<_> = (0..128)
3260 .map(|_| {
3261 let entity = lookup_version(id, &format!("other-{id}"), 1, 0);
3262 id += 1;
3263 entity
3264 })
3265 .collect();
3266 for version in (segment..versions).step_by(segments) {
3267 rows.push(lookup_version(
3268 id,
3269 "history",
3270 version as u64 + 1,
3271 if version + 1 == versions {
3272 0
3273 } else {
3274 version as u64 + 2
3275 },
3276 ));
3277 id += 1;
3278 }
3279 manager.bulk_insert(rows).expect("rows");
3280 manager.force_seal().expect("seal");
3281 }
3282 let started = Instant::now();
3283 for _ in 0..2000 {
3284 let xmin =
3285 manager.project_key_entity(black_box("history"), None, true, |entity| {
3286 entity.xmin
3287 });
3288 assert_eq!(xmin, Some(versions as u64));
3289 }
3290 eprintln!("segments={segments} versions={versions} reads=2000 elapsed_us={} resident_bytes={}", started.elapsed().as_micros(), manager.resident_bytes());
3291 }
3292 }
3293 }
3294
3295 #[test]
3296 fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
3297 let manager = SegmentManager::new("fingerprint");
3298 let row = |id, label: &str| {
3299 UnifiedEntity::table_row(
3300 EntityId::new(id),
3301 "fingerprint",
3302 id,
3303 vec![Value::text(label)],
3304 )
3305 };
3306 manager
3307 .bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
3308 .expect("sealed rows");
3309 manager.force_seal().expect("seal");
3310 manager
3311 .bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
3312 .expect("growing rows");
3313
3314 let mut expected = manager.query_all(|_| true);
3315 expected.sort_by_key(|entity| entity.id.raw());
3316 let expected: Vec<_> = expected
3317 .iter()
3318 .map(|entity| format!("{entity:?}"))
3319 .collect();
3320 let mut stored_addresses = HashSet::new();
3321 manager.for_each_entity(|entity| {
3322 stored_addresses.insert(entity as *const UnifiedEntity as usize);
3323 true
3324 });
3325
3326 let mut actual = Vec::new();
3327 manager.for_each_entity_in_id_order(|entity| {
3328 assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
3329 assert!(
3330 manager.growing.try_write().is_some(),
3331 "topology guard retained"
3332 );
3333 assert!(
3334 manager.sealed.try_write().is_some(),
3335 "topology guard retained"
3336 );
3337 actual.push(format!("{entity:?}"));
3338 true
3339 });
3340 assert_eq!(
3341 actual, expected,
3342 "equal IDs retain their original segment order"
3343 );
3344 }
3345
3346 #[test]
3347 fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
3348 let manager = SegmentManager::new("fingerprint");
3349 for id in [3, 1, 2] {
3350 manager
3351 .insert(UnifiedEntity::table_row(
3352 EntityId::new(id),
3353 "fingerprint",
3354 id,
3355 vec![Value::Integer(id as i64)],
3356 ))
3357 .expect("row");
3358 }
3359 let source = manager.growing.read().as_ref().expect("source").clone();
3360 let handles = Arc::strong_count(&source);
3361 let mut seen = Vec::new();
3362 manager.for_each_entity_in_id_order(|entity| {
3363 seen.push(entity.id.raw());
3364 false
3365 });
3366 assert_eq!(seen, vec![1]);
3367 assert_eq!(Arc::strong_count(&source), handles);
3368 assert!(
3369 source.try_write().is_some(),
3370 "segment guard retained after early stop"
3371 );
3372 }
3373
3374 #[test]
3375 fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
3376 let manager = SegmentManager::with_config(
3377 "vectors",
3378 ManagerConfig {
3379 max_sealed_segments: 1,
3380 consolidation_entities_per_tick: 10_000,
3381 ..Default::default()
3382 },
3383 );
3384 for part in 0..3 {
3385 manager
3386 .bulk_insert(
3387 (1..=300)
3388 .map(|offset| {
3389 UnifiedEntity::vector(
3390 EntityId::new(part * 300 + offset),
3391 "vectors",
3392 vec![1.0, 0.0],
3393 )
3394 })
3395 .collect(),
3396 )
3397 .expect("batch");
3398 if part < 2 {
3399 manager.force_seal().expect("seal");
3400 }
3401 }
3402 let original_growing = manager.growing.read().as_ref().expect("growing").clone();
3403 let mut changed = false;
3404 let mut seen = Vec::new();
3405 assert!(manager.scan_batches(
3406 None,
3407 || true,
3408 |batch| {
3409 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3410 assert!(
3411 manager.growing.try_write().is_some(),
3412 "consumer holds no topology guard"
3413 );
3414 assert!(
3415 original_growing.try_write().is_some(),
3416 "consumer holds no segment guard"
3417 );
3418 for entity in batch {
3419 assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
3420 seen.push(entity.id.raw());
3421 }
3422 if !changed {
3423 changed = true;
3424 manager.force_seal().expect("seal during scan");
3425 manager
3426 .insert(UnifiedEntity::vector(
3427 EntityId::new(5000),
3428 "vectors",
3429 vec![1.0, 0.0],
3430 ))
3431 .expect("new write");
3432 drain_maintenance(&manager);
3433 assert!(manager.stats().consolidation.runs_completed > 0);
3434 assert!(manager
3435 .delete(EntityId::new(1))
3436 .expect("delete after source retirement"));
3437 assert!(manager.get(EntityId::new(1)).is_none());
3438 }
3439 true
3440 }
3441 ));
3442 seen.sort_unstable();
3443 assert_eq!(
3444 seen,
3445 (2..=900).collect::<Vec<_>>(),
3446 "no duplicate, missing or late candidates"
3447 );
3448 let mut after = Vec::new();
3449 assert!(manager.scan_batches(
3450 None,
3451 || true,
3452 |batch| {
3453 after.extend(batch.iter().map(|entity| entity.id.raw()));
3454 true
3455 }
3456 ));
3457 after.sort_unstable();
3458 let mut expected: Vec<_> = (2..=900).collect();
3459 expected.push(5000);
3460 assert_eq!(after, expected, "adopted entities have scan positions");
3461 }
3462
3463 #[test]
3464 fn bounded_scan_applies_snapshot_to_relocated_versions() {
3465 for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
3466 let manager = SegmentManager::with_config(
3467 "vectors",
3468 ManagerConfig {
3469 max_sealed_segments: 1,
3470 ..Default::default()
3471 },
3472 );
3473 for ids in [[1, 2], [3, 4]] {
3474 let entities = ids.map(|id| {
3475 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3476 entity.set_xmin(3);
3477 entity
3478 });
3479 if bulk {
3480 manager
3481 .bulk_insert(entities.into())
3482 .expect("initial versions");
3483 } else {
3484 for entity in entities {
3485 manager.insert(entity).expect("initial version");
3486 }
3487 }
3488 manager.force_seal().expect("seal source");
3489 }
3490 let snapshot = SnapshotContext {
3491 snapshot: Snapshot {
3492 xid: 10,
3493 in_progress: HashSet::new(),
3494 },
3495 manager: Arc::new(SnapshotManager::new()),
3496 own_xids: HashSet::new(),
3497 requires_index_fallback: false,
3498 serializable_reader: None,
3499 };
3500 let mut changed = false;
3501 let mut seen = Vec::new();
3502 assert!(manager.scan_batches(
3503 Some(&snapshot),
3504 || true,
3505 |batch| {
3506 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3507 if !changed {
3508 changed = true;
3509 if consolidate {
3510 drain_maintenance(&manager);
3511 assert!(manager.stats().consolidation.runs_completed > 0);
3512 }
3513 for (id, xmax) in [(3, 8), (4, 12)] {
3514 let mut entity =
3515 manager.get(EntityId::new(id)).expect("retained version");
3516 entity.set_xmax(xmax);
3517 manager
3518 .update(entity)
3519 .expect("update after source retirement");
3520 }
3521 }
3522 true
3523 }
3524 ));
3525 seen.sort_unstable();
3526 assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
3527 let mut again = Vec::new();
3528 assert!(manager.scan_batches(
3529 Some(&snapshot),
3530 || true,
3531 |batch| {
3532 again.extend(batch.iter().map(|entity| entity.id.raw()));
3533 true
3534 }
3535 ));
3536 again.sort_unstable();
3537 assert_eq!(
3538 again,
3539 vec![1, 2, 4],
3540 "old slots do not duplicate relocated versions"
3541 );
3542 let mut current = Vec::new();
3543 assert!(manager.scan_batches(
3544 None,
3545 || true,
3546 |batch| {
3547 current.extend(batch.iter().map(|entity| entity.id.raw()));
3548 true
3549 }
3550 ));
3551 current.sort_unstable();
3552 assert_eq!(
3553 current,
3554 vec![1, 2],
3555 "current view excludes both deleted versions"
3556 );
3557 }
3558 }
3559
3560 #[test]
3561 fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
3562 let manager = SegmentManager::new("mixed");
3563 let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
3564 manager
3565 .bulk_insert((1..=300).map(vector).collect())
3566 .expect("flat");
3567 manager.insert(vector(500)).expect("hash insert");
3568 manager
3569 .bulk_insert(vec![vector(502), vector(505)])
3570 .expect("gapped bulk");
3571 manager.delete(EntityId::new(2)).expect("flat delete");
3572 manager.delete(EntityId::new(505)).expect("hash delete");
3573 let mut changed = false;
3574 let mut seen = Vec::new();
3575 assert!(manager.scan_batches(
3576 None,
3577 || true,
3578 |batch| {
3579 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3580 if !changed {
3581 changed = true;
3582 manager
3583 .insert(vector(700))
3584 .expect("append after cursor boundary");
3585 manager
3586 .update(UnifiedEntity::graph_node(
3587 EntityId::new(500),
3588 "mixed",
3589 "Node",
3590 HashMap::new(),
3591 ))
3592 .expect("structural update");
3593 }
3594 true
3595 }
3596 ));
3597 seen.sort_unstable();
3598 let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
3599 expected.extend([500, 502]);
3600 assert_eq!(seen, expected);
3601 }
3602
3603 #[test]
3604 fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
3605 let manager = SegmentManager::new("vectors");
3606 let mut entities = Vec::new();
3607 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
3608 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3609 entity.set_xmin(xmin);
3610 entity.set_xmax(xmax);
3611 entities.push(entity);
3612 }
3613 manager.bulk_insert(entities).expect("versions");
3614 manager.force_seal().expect("seal");
3615 let snapshot = SnapshotContext {
3616 snapshot: Snapshot {
3617 xid: 10,
3618 in_progress: HashSet::from([4]),
3619 },
3620 manager: Arc::new(SnapshotManager::new()),
3621 own_xids: HashSet::new(),
3622 requires_index_fallback: false,
3623 serializable_reader: None,
3624 };
3625 let mut work = 0;
3626 let mut callbacks = 0;
3627 assert!(!manager.scan_batches(
3628 Some(&snapshot),
3629 || {
3630 work += 1;
3631 work <= 2
3632 },
3633 |_| {
3634 callbacks += 1;
3635 true
3636 }
3637 ));
3638 assert_eq!(callbacks, 0, "a stopped batch is not consumed");
3639 for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
3640 let mut seen = Vec::new();
3641 assert!(manager.scan_batches(
3642 view,
3643 || true,
3644 |batch| {
3645 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3646 true
3647 }
3648 ));
3649 seen.sort_unstable();
3650 assert_eq!(seen, expected);
3651 }
3652 }
3653
3654 #[test]
3655 fn test_manager_basic() {
3656 let manager = SegmentManager::new("test_collection");
3657
3658 let entity = UnifiedEntity::table_row(
3659 manager.next_entity_id(),
3660 "users",
3661 1,
3662 vec![Value::text("Alice".to_string())],
3663 );
3664
3665 let id = manager.insert(entity).unwrap();
3666 assert!(manager.get(id).is_some());
3667 assert_eq!(manager.count(), 1);
3668 }
3669
3670 #[test]
3671 fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
3672 let manager = Arc::new(SegmentManager::new("mixed"));
3673 manager
3674 .bulk_insert(
3675 (1..=128)
3676 .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
3677 .collect(),
3678 )
3679 .expect("vector-only segment");
3680 manager.force_seal().expect("seal vectors");
3681 let mut old =
3682 UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
3683 old.set_xmin(3);
3684 old.set_xmax(12);
3685 let mut new =
3686 UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
3687 new.set_xmin(12);
3688 let edge = UnifiedEntity::graph_edge(
3689 EntityId::new(202),
3690 "mixed",
3691 "a",
3692 "b",
3693 1.0,
3694 Default::default(),
3695 );
3696 manager
3697 .bulk_insert(vec![old, new, edge])
3698 .expect("physical graph versions");
3699 manager.force_seal().expect("seal graph");
3700 manager
3701 .insert(UnifiedEntity::graph_node(
3702 EntityId::new(203),
3703 "mixed",
3704 "person",
3705 Default::default(),
3706 ))
3707 .expect("growing node");
3708 let snapshot = SnapshotContext {
3709 snapshot: Snapshot {
3710 xid: 10,
3711 in_progress: HashSet::new(),
3712 },
3713 manager: Arc::new(SnapshotManager::new()),
3714 own_xids: HashSet::new(),
3715 requires_index_fallback: false,
3716 serializable_reader: None,
3717 };
3718 std::thread::spawn(move || {
3719 for view in [None, Some(&snapshot)] {
3720 for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
3721 let expected: HashSet<_> = manager
3722 .scan(view, |item| kind.matches(&item.kind))
3723 .into_iter()
3724 .map(|item| item.id)
3725 .collect();
3726 let actual: HashSet<_> = manager
3727 .scan_graph_kind(view, kind)
3728 .into_iter()
3729 .map(|item| item.id)
3730 .collect();
3731 assert_eq!(actual, expected);
3732 let mut visited = 0;
3733 let mut candidates = HashSet::new();
3734 manager.scan_graph_candidates_for_each(
3735 view,
3736 kind,
3737 || true,
3738 |item| {
3739 visited += 1;
3740 if kind.matches(&item.kind) {
3741 candidates.insert(item.id);
3742 }
3743 true
3744 },
3745 );
3746 assert_eq!(candidates, expected);
3747 assert_eq!(
3748 visited,
3749 if matches!(kind, GraphEntityKind::Node) {
3750 3
3751 } else {
3752 2
3753 }
3754 );
3755 }
3756 }
3757 let ids: HashSet<_> = manager
3758 .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
3759 .into_iter()
3760 .map(|item| item.id.raw())
3761 .collect();
3762 assert_eq!(ids, HashSet::from([200, 203]));
3763 let ids: HashSet<_> = manager
3764 .scan_graph_kind(None, GraphEntityKind::Node)
3765 .into_iter()
3766 .map(|item| item.id.raw())
3767 .collect();
3768 assert_eq!(ids, HashSet::from([201, 203]));
3769 let mut visits = 0;
3770 manager.scan_graph_candidates_for_each(
3771 None,
3772 GraphEntityKind::Node,
3773 || true,
3774 |_| {
3775 visits += 1;
3776 false
3777 },
3778 );
3779 assert_eq!(
3780 visits, 1,
3781 "early termination must also stop subsequent segments"
3782 );
3783 })
3784 .join()
3785 .expect("snapshot worker");
3786 }
3787
3788 #[test]
3789 fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
3790 let manager = Arc::new(SegmentManager::new("accounts"));
3791 let cases = [
3792 ("legacy row", 0, 0, true),
3793 ("committed creator", 3, 0, true),
3794 ("future creator", 11, 0, false),
3795 ("in-progress creator", 4, 0, false),
3796 ("committed deleter", 3, 7, false),
3797 ("future deleter", 3, 11, true),
3798 ("in-progress deleter", 3, 4, true),
3799 ];
3800
3801 for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
3802 let row_id = index as u64 + 1;
3803 let mut entity = UnifiedEntity::table_row(
3804 EntityId::new(row_id),
3805 "accounts",
3806 row_id,
3807 vec![Value::Integer(row_id as i64)],
3808 );
3809 entity.set_xmin(*xmin);
3810 entity.set_xmax(*xmax);
3811 manager.insert(entity).unwrap();
3812 }
3813
3814 let snapshot = SnapshotContext {
3815 snapshot: Snapshot {
3816 xid: 10,
3817 in_progress: HashSet::from([4]),
3818 },
3819 manager: Arc::new(SnapshotManager::new()),
3820 own_xids: HashSet::new(),
3821 requires_index_fallback: false,
3822 serializable_reader: None,
3823 };
3824
3825 let raw_visible_row_ids: HashSet<u64> = manager
3826 .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
3827 .into_iter()
3828 .filter_map(|entity| match entity.kind {
3829 EntityKind::TableRow { row_id, .. } => Some(row_id),
3830 _ => None,
3831 })
3832 .collect();
3833
3834 let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
3835 .join()
3836 .unwrap();
3837 let visible_row_ids: HashSet<u64> = scanned
3838 .into_iter()
3839 .filter_map(|entity| match entity.kind {
3840 EntityKind::TableRow { row_id, .. } => Some(row_id),
3841 _ => None,
3842 })
3843 .collect();
3844 assert_eq!(visible_row_ids, raw_visible_row_ids);
3845
3846 for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
3847 let row_id = index as u64 + 1;
3848 assert_eq!(
3849 visible_row_ids.contains(&row_id),
3850 *expected_visible,
3851 "{name} visibility mismatch"
3852 );
3853 }
3854 }
3855
3856 #[test]
3857 fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
3858 let manager = SegmentManager::new("accounts");
3862 let mut live =
3863 UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
3864 live.set_xmin(2);
3865 live.set_xmax(0);
3866 manager.insert(live).unwrap();
3867
3868 let mut superseded =
3869 UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
3870 superseded.set_xmin(2);
3871 superseded.set_xmax(9);
3872 manager.insert(superseded).unwrap();
3873
3874 let rows: Vec<u64> = manager
3875 .scan(None, |_| true)
3876 .into_iter()
3877 .filter_map(|entity| match entity.kind {
3878 EntityKind::TableRow { row_id, .. } => Some(row_id),
3879 _ => None,
3880 })
3881 .collect();
3882 assert_eq!(
3883 rows,
3884 vec![1],
3885 "frameless scan must hide superseded versions"
3886 );
3887 }
3888
3889 #[test]
3890 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
3891 let manager = SegmentManager::new("test_collection");
3892
3893 let entity = UnifiedEntity::table_row(
3894 manager.next_entity_id(),
3895 "users",
3896 1,
3897 vec![Value::text("Alice".to_string())],
3898 );
3899
3900 let id = manager.insert(entity).unwrap();
3901
3902 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3903 assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
3904 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
3905
3906 assert!(manager.delete(id).unwrap());
3907 assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3908 }
3909
3910 #[test]
3911 fn zoned_scan_reports_pruned_segments_and_preserves_results() {
3912 let config = ManagerConfig {
3913 segment_config: SegmentConfig {
3914 max_entities: 2,
3915 ..Default::default()
3916 },
3917 ..Default::default()
3918 };
3919 let manager = SegmentManager::with_config("events", config);
3920
3921 for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
3922 let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
3923 let entity = UnifiedEntity::new(
3924 manager.next_entity_id(),
3925 EntityKind::TableRow {
3926 table: "events".into(),
3927 row_id: (idx + 1) as u64,
3928 },
3929 EntityData::Row(row),
3930 );
3931 manager.insert(entity).unwrap();
3932 }
3933
3934 let full = manager.query_all(|entity| match &entity.data {
3935 EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
3936 _ => false,
3937 });
3938 let probe = Value::Integer(100);
3939 let (pruned, stats) =
3940 manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
3941 match &entity.data {
3942 EntityData::Row(row) => {
3943 matches!(row.get_field("ts"), Some(Value::Integer(100)))
3944 }
3945 _ => false,
3946 }
3947 });
3948
3949 assert_eq!(pruned.len(), full.len());
3950 assert_eq!(pruned[0].id, full[0].id);
3951 assert_eq!(stats.segments_total, 3);
3952 assert_eq!(stats.segments_scanned, 2);
3953 assert_eq!(stats.segments_pruned, 1);
3954 }
3955
3956 #[test]
3957 fn projected_zoned_scan_preserves_segment_order_and_pruning() {
3958 let manager = SegmentManager::with_config(
3959 "events",
3960 ManagerConfig {
3961 segment_config: SegmentConfig {
3962 max_entities: 2,
3963 ..Default::default()
3964 },
3965 ..Default::default()
3966 },
3967 );
3968 for (index, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
3969 let mut entity = UnifiedEntity::table_row(
3970 EntityId::new(index as u64 + 1),
3971 "events",
3972 index as u64 + 1,
3973 vec![],
3974 );
3975 entity.data = EntityData::Row(RowData::with_names(
3976 vec![Value::Integer(value)],
3977 vec!["ts".into()],
3978 ));
3979 manager.insert(entity).expect("test row");
3980 }
3981 assert_eq!(manager.stats().sealed_count, 2);
3982 let project = |entity: &UnifiedEntity| match entity.data.as_row()?.get_field("ts")? {
3983 Value::Integer(value) => Some(*value),
3984 _ => None,
3985 };
3986 let (all, stats) = manager.scan_project_zoned_with_stats(None, &[], project);
3987 let (legacy, _) = manager.scan_zoned_with_stats(None, &[], |_| true);
3988 assert_eq!(all, legacy.iter().filter_map(project).collect::<Vec<_>>());
3989 assert_eq!(all[0], 200, "growing segment first");
3990 assert!(all[1..3].iter().all(|value| [10, 20].contains(value)));
3991 assert!(all[3..5].iter().all(|value| [100, 110].contains(value)));
3992 assert_eq!(stats.segments_total, 3);
3993 assert_eq!(stats.segments_scanned, 3);
3994 assert_eq!(stats.segments_pruned, 0);
3995 let probe = Value::Integer(100);
3996 let (pruned, stats) = manager.scan_project_zoned_with_stats(
3997 None,
3998 &[("ts", ZoneColPred::Gte(&probe))],
3999 |entity| project(entity).filter(|value| *value >= 100),
4000 );
4001 assert_eq!(
4002 pruned,
4003 all.into_iter()
4004 .filter(|value| *value >= 100)
4005 .collect::<Vec<_>>()
4006 );
4007 assert_eq!(stats.segments_total, 3);
4008 assert_eq!(stats.segments_scanned, 2);
4009 assert_eq!(stats.segments_pruned, 1);
4010 let (empty, _) = manager.scan_project_zoned_with_stats(None, &[], |_| None::<u64>);
4011 assert!(empty.is_empty(), "projection may reject every candidate");
4012
4013 let failure = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4014 manager.scan_project_zoned_with_stats(None, &[], |entity| {
4015 let value = project(entity)?;
4016 assert_ne!(value, 10, "injected projection worker failure");
4017 Some(value)
4018 })
4019 }));
4020 assert!(
4021 failure.is_err(),
4022 "a failed worker must not return a successful partial scan"
4023 );
4024 }
4025
4026 #[test]
4027 fn projected_zoned_scan_checks_visibility_before_projection_on_workers() {
4028 let manager = SegmentManager::with_config(
4029 "accounts",
4030 ManagerConfig {
4031 segment_config: SegmentConfig {
4032 max_entities: 2,
4033 ..Default::default()
4034 },
4035 ..Default::default()
4036 },
4037 );
4038 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4), (5, 3, 0)] {
4039 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "accounts", id, vec![]);
4040 entity.set_xmin(xmin);
4041 entity.set_xmax(xmax);
4042 manager.insert(entity).expect("versioned test row");
4043 }
4044 let mut hidden = UnifiedEntity::table_row(EntityId::new(6), "accounts", 6, vec![]);
4045 hidden.data = EntityData::Row(RowData::with_names(
4046 vec![Value::text("pending")],
4047 vec!["__moderation_status".into()],
4048 ));
4049 manager.insert(hidden).expect("moderated row");
4050 assert_eq!(manager.stats().sealed_count, 2);
4051 let snapshot = SnapshotContext {
4052 snapshot: Snapshot {
4053 xid: 10,
4054 in_progress: HashSet::from([4]),
4055 },
4056 manager: Arc::new(SnapshotManager::new()),
4057 own_xids: HashSet::new(),
4058 requires_index_fallback: false,
4059 serializable_reader: None,
4060 };
4061 std::thread::spawn(move || {
4062 for (view, expected) in [(Some(&snapshot), vec![5, 1, 3, 4]), (None, vec![5, 2, 3])] {
4063 let called = std::sync::Mutex::new(Vec::new());
4064 let (rows, stats) = manager.scan_project_zoned_with_stats(view, &[], |entity| {
4065 let id = entity.id.raw();
4066 assert!(
4067 expected.contains(&id),
4068 "invisible entity reached projection: {id}"
4069 );
4070 called.lock().expect("callback log").push(id);
4071 Some(id)
4072 });
4073 let (legacy, _) = manager.scan_zoned_with_stats(view, &[], |_| true);
4074 assert_eq!(
4075 rows,
4076 legacy
4077 .iter()
4078 .map(|entity| entity.id.raw())
4079 .collect::<Vec<_>>()
4080 );
4081 assert_eq!(rows[0], 5, "visible growing row first");
4082 let mut called = called.into_inner().expect("callback log");
4083 called.sort_unstable();
4084 let mut expected = expected;
4085 expected.sort_unstable();
4086 assert_eq!(called, expected, "each visible row projected exactly once");
4087 assert_eq!(stats.segments_total, 3);
4088 assert_eq!(stats.segments_scanned, 3);
4089 }
4090 })
4091 .join()
4092 .expect("projected scan worker");
4093 }
4094
4095 #[test]
4096 fn test_manager_auto_seal() {
4097 let config = ManagerConfig {
4098 segment_config: SegmentConfig {
4099 max_entities: 2,
4100 ..Default::default()
4101 },
4102 ..Default::default()
4103 };
4104
4105 let manager = SegmentManager::with_config("test", config);
4106
4107 manager
4109 .insert(UnifiedEntity::vector(
4110 manager.next_entity_id(),
4111 "v",
4112 vec![0.1],
4113 ))
4114 .unwrap();
4115
4116 manager
4118 .insert(UnifiedEntity::vector(
4119 manager.next_entity_id(),
4120 "v",
4121 vec![0.2],
4122 ))
4123 .unwrap();
4124
4125 manager
4127 .insert(UnifiedEntity::vector(
4128 manager.next_entity_id(),
4129 "v",
4130 vec![0.3],
4131 ))
4132 .unwrap();
4133
4134 let stats = manager.stats();
4135 assert_eq!(stats.total_entities, 3);
4136 }
4137
4138 #[test]
4139 fn test_manager_delete() {
4140 let manager = SegmentManager::new("test");
4141
4142 let id = manager
4143 .insert(UnifiedEntity::vector(
4144 manager.next_entity_id(),
4145 "v",
4146 vec![0.1],
4147 ))
4148 .unwrap();
4149
4150 assert!(manager.get(id).is_some());
4151 assert!(manager.delete(id).unwrap());
4152 assert!(manager.get(id).is_none());
4153 }
4154
4155 #[test]
4156 fn test_manager_metadata() {
4157 let manager = SegmentManager::new("test");
4158
4159 let id = manager
4160 .insert(UnifiedEntity::table_row(
4161 manager.next_entity_id(),
4162 "hosts",
4163 1,
4164 vec![Value::text("192.168.1.1".to_string())],
4165 ))
4166 .unwrap();
4167
4168 let mut meta = Metadata::new();
4169 meta.set(
4170 "os",
4171 super::super::metadata::MetadataValue::String("linux".to_string()),
4172 );
4173
4174 manager.set_metadata(id, meta).unwrap();
4175
4176 let retrieved = manager.get_metadata(id).unwrap();
4177 assert!(retrieved.has("os"));
4178 }
4179
4180 #[test]
4181 fn test_manager_query_by_kind() {
4182 let manager = SegmentManager::new("test");
4183
4184 manager
4185 .insert(UnifiedEntity::table_row(
4186 manager.next_entity_id(),
4187 "hosts",
4188 1,
4189 vec![],
4190 ))
4191 .unwrap();
4192
4193 manager
4194 .insert(UnifiedEntity::vector(
4195 manager.next_entity_id(),
4196 "embeddings",
4197 vec![0.1],
4198 ))
4199 .unwrap();
4200
4201 manager
4202 .insert(UnifiedEntity::table_row(
4203 manager.next_entity_id(),
4204 "hosts",
4205 2,
4206 vec![],
4207 ))
4208 .unwrap();
4209
4210 let rows = manager.get_by_kind("table");
4211 assert_eq!(rows.len(), 2);
4212
4213 let vectors = manager.get_by_kind("vector");
4214 assert_eq!(vectors.len(), 1);
4215 }
4216
4217 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
4224 SegmentManager::with_config(
4225 "test",
4226 ManagerConfig {
4227 max_sealed_segments: 1_000,
4228 consolidation_entities_per_tick: entities_per_tick,
4229 ..Default::default()
4230 },
4231 )
4232 }
4233
4234 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
4238 let entities: Vec<UnifiedEntity> = (0..count)
4239 .map(|i| {
4240 UnifiedEntity::table_row(
4241 manager.next_entity_id(),
4242 "users",
4243 i as u64,
4244 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
4245 )
4246 })
4247 .collect();
4248 let ids = manager.bulk_insert(entities).expect("bulk insert");
4249 manager.force_seal().expect("seal");
4250 ids
4251 }
4252
4253 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
4254 let mut ids: Vec<u64> = manager
4255 .query_all(|_| true)
4256 .into_iter()
4257 .map(|entity| entity.id.raw())
4258 .collect();
4259 ids.sort_unstable();
4260 ids
4261 }
4262
4263 fn drain_maintenance(manager: &SegmentManager) -> usize {
4264 let mut ticks = 0;
4265 while manager.consolidation.read().is_some() || ticks == 0 {
4266 manager.run_maintenance().expect("maintenance");
4267 ticks += 1;
4268 assert!(ticks < 1_000, "consolidation failed to converge");
4269 }
4270 ticks
4271 }
4272
4273 #[test]
4274 fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
4275 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4276 let ids = manager
4277 .bulk_insert(
4278 (1..=1024)
4279 .map(|id| {
4280 UnifiedEntity::graph_edge(
4281 EntityId::new(id),
4282 "step",
4283 "2000",
4284 "3000",
4285 1.0,
4286 HashMap::new(),
4287 )
4288 })
4289 .collect(),
4290 )
4291 .expect("edges");
4292 manager.force_seal().expect("seal");
4293 for id in ids.iter().take(300) {
4294 manager.delete(*id).expect("fragment source");
4295 }
4296 let source = Arc::downgrade(&manager.sealed.read()[0]);
4297 let mut seen = Vec::new();
4298 let mut changed = false;
4299 let mut metadata_bytes = 0;
4300 assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
4301 || true,
4302 |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
4303 |batch| {
4304 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4305 assert!(manager.growing.try_write().is_some(), "no growing topology guard");
4306 assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
4307 assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
4308 if !changed {
4309 changed = true;
4310 drain_maintenance(&manager);
4311 assert!(manager.stats().consolidation.runs_completed > 0);
4312 assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
4313 let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
4314 if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
4315 manager.update(updated).expect("update authoritative merged copy");
4316 }
4317 for id in batch {
4318 let entity = manager.get(*id).expect("hydrate current manager");
4319 if id.raw() == 1024 {
4320 assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
4321 }
4322 seen.push(id.raw());
4323 }
4324 Ok(true)
4325 }).expect("cursor"));
4326 assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
4327 assert!(
4328 source.upgrade().is_none(),
4329 "finished cursor releases retired source"
4330 );
4331 assert!(
4332 metadata_bytes < 4096,
4333 "fixed ID buffer plus source handles: {metadata_bytes}"
4334 );
4335 }
4336
4337 #[test]
4338 fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
4339 let manager = SegmentManager::new("graph");
4340 manager
4341 .bulk_insert(
4342 (1..=1024)
4343 .map(|id| {
4344 UnifiedEntity::graph_edge(
4345 EntityId::new(id),
4346 "step",
4347 "2000",
4348 "3000",
4349 1.0,
4350 HashMap::new(),
4351 )
4352 })
4353 .collect(),
4354 )
4355 .expect("edges");
4356 let mut delivered = 0;
4357 let denied = manager.visit_graph_index_batches(
4358 GraphEntityKind::Edge,
4359 &[EntityId::new(2000)],
4360 || true,
4361 |_| Err("out of memory"),
4362 |ids| {
4363 delivered += ids.len();
4364 Ok(true)
4365 },
4366 );
4367 assert_eq!(denied, Err("out of memory"));
4368 assert_eq!(delivered, 0);
4369 let mut work = 0;
4370 assert!(!manager
4371 .visit_graph_index_batches(
4372 GraphEntityKind::Edge,
4373 &[EntityId::new(2000)],
4374 || {
4375 work += 1;
4376 work < 100
4377 },
4378 |_| Ok::<_, ()>(()),
4379 |ids| {
4380 delivered += ids.len();
4381 Ok(true)
4382 }
4383 )
4384 .expect("cancelled cursor"));
4385 assert_eq!(
4386 delivered, 0,
4387 "cancelled partial batch cannot reach consumer"
4388 );
4389 assert!(!manager
4390 .visit_graph_index_batches(
4391 GraphEntityKind::Edge,
4392 &[EntityId::new(2000)],
4393 || true,
4394 |_| Ok::<_, ()>(()),
4395 |ids| {
4396 delivered += ids.len();
4397 Ok(false)
4398 }
4399 )
4400 .expect("consumer stop"));
4401 assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
4402 }
4403
4404 #[test]
4405 fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
4406 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4407 let ids = seal_bulk_segment(&manager, 100);
4408 for id in ids.iter().take(40) {
4409 manager.delete(*id).expect("delete before scan");
4410 }
4411 let source_bytes = manager.resident_bytes();
4412 let source = Arc::downgrade(&manager.sealed.read()[0]);
4413 std::thread::scope(|scope| {
4414 let (ready, waiting) = std::sync::mpsc::channel();
4415 let (release_first, first_wait) = std::sync::mpsc::channel();
4416 let (release_last, last_wait) = std::sync::mpsc::channel();
4417 let reader = |resume: std::sync::mpsc::Receiver<()>,
4418 ready: std::sync::mpsc::Sender<()>| {
4419 assert!(manager.scan_batches(
4420 None,
4421 || true,
4422 |_| {
4423 ready.send(()).expect("signal captured cursor");
4424 let _ = resume.recv();
4425 true
4426 }
4427 ));
4428 };
4429 let ready_first = ready.clone();
4430 let first = scope.spawn(move || reader(first_wait, ready_first));
4431 let last = scope.spawn(move || reader(last_wait, ready));
4432 for _ in 0..2 {
4433 waiting
4434 .recv_timeout(std::time::Duration::from_secs(10))
4435 .expect("reader ready");
4436 }
4437 drain_maintenance(&manager);
4438 assert!(
4439 source.upgrade().is_some(),
4440 "both cursors pin the retired source"
4441 );
4442 let active_bytes = manager.sealed.read()[0].read().resident_bytes();
4443 let held = manager.resident_bytes();
4444 assert!(
4445 held >= active_bytes + source_bytes,
4446 "retired source remains resident"
4447 );
4448 release_first.send(()).expect("release first reader");
4449 first.join().expect("first reader");
4450 assert_eq!(
4451 manager.resident_bytes(),
4452 held,
4453 "one reader still pins the source"
4454 );
4455 release_last.send(()).expect("release last reader");
4456 last.join().expect("last reader");
4457 assert!(
4458 source.upgrade().is_none(),
4459 "accounting must not retain payloads"
4460 );
4461 assert_eq!(held - manager.resident_bytes(), source_bytes);
4462 assert_eq!(
4463 manager.stats().total_memory_bytes as u64,
4464 manager.resident_bytes()
4465 );
4466 });
4467 }
4468
4469 #[test]
4470 fn physical_cursor_tracks_individual_and_mixed_insertions() {
4471 for mode in 0..3 {
4472 let manager = SegmentManager::new("physical_rows");
4473 let row = |id| {
4474 UnifiedEntity::table_row(
4475 EntityId::new(id),
4476 "physical_rows",
4477 id,
4478 vec![Value::Integer(id as i64)],
4479 )
4480 };
4481 match mode {
4482 0 => {
4483 for id in [10, 30, 20] {
4484 manager.insert(row(id)).expect("individual row");
4485 }
4486 }
4487 1 => {
4488 manager.insert(row(10)).expect("individual prefix");
4489 manager
4490 .bulk_insert(vec![row(30), row(31)])
4491 .expect("batch suffix");
4492 }
4493 _ => {
4494 manager
4495 .bulk_insert(vec![row(30), row(31)])
4496 .expect("batch prefix");
4497 manager.insert(row(10)).expect("individual suffix");
4498 }
4499 }
4500 let mut cursor = manager.physical_scan_cursor();
4501 let mut ids = Vec::new();
4502 for _ in 0..32 {
4503 if cursor.finished() {
4504 break;
4505 }
4506 let mut work = 1;
4507 let batch = manager.physical_scan_batch(&mut cursor, &mut work);
4508 assert!(batch.len() <= 1);
4509 ids.extend(batch.into_iter().map(|id| id.raw()));
4510 }
4511 assert!(cursor.finished());
4512 ids.sort_unstable();
4513 assert_eq!(
4514 ids,
4515 if mode == 0 {
4516 vec![10, 20, 30]
4517 } else {
4518 vec![10, 30, 31]
4519 },
4520 "mode={mode}"
4521 );
4522 }
4523 }
4524
4525 fn growing_tombstones_fixture() -> SegmentManager {
4526 let manager = consolidating_manager(1);
4527 let ids = manager
4528 .bulk_insert(
4529 (0..64)
4530 .map(|id| {
4531 UnifiedEntity::table_row(
4532 manager.next_entity_id(),
4533 "users",
4534 id,
4535 vec![Value::Blob(vec![7; 2048])],
4536 )
4537 })
4538 .collect(),
4539 )
4540 .expect("growing rows");
4541 for id in ids.iter().take(48) {
4542 manager.delete(*id).expect("physical tombstone");
4543 }
4544 assert_eq!(manager.stats().sealed_count, 0);
4545 manager
4546 }
4547
4548 #[test]
4549 fn pressure_consolidation_reclaims_unsealed_tombstones() {
4550 let manager = growing_tombstones_fixture();
4551 let before = manager.resident_bytes();
4552 assert!(manager.reclaimable_bytes() > 0);
4553 assert!(manager.pressure_consolidation_tick());
4554 drain_maintenance(&manager);
4555 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
4556 assert!(manager.resident_bytes() < before);
4557 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 48);
4558 }
4559
4560 #[test]
4561 fn pressure_consolidation_preserves_pinned_growing_source_accounting() {
4562 let manager = growing_tombstones_fixture();
4563 let source = Arc::clone(manager.growing.read().as_ref().expect("growing source"));
4564 let weak = Arc::downgrade(&source);
4565 assert!(manager.pressure_consolidation_tick());
4566 drain_maintenance(&manager);
4567 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
4568 let source_bytes = source.read().resident_bytes();
4569 let held = manager.resident_bytes();
4570 assert!(held >= source_bytes);
4571 drop(source);
4572 assert!(weak.upgrade().is_none());
4573 assert_eq!(held - manager.resident_bytes(), source_bytes);
4574 }
4575
4576 #[test]
4577 fn pressure_consolidation_keeps_clean_and_disabled_growing_sources_open() {
4578 let clean = consolidating_manager(1);
4579 clean
4580 .insert(UnifiedEntity::table_row(
4581 clean.next_entity_id(),
4582 "users",
4583 1,
4584 vec![Value::Integer(1)],
4585 ))
4586 .expect("clean row");
4587 assert!(!clean.pressure_consolidation_tick());
4588 assert_eq!(clean.stats().sealed_count, 0);
4589 let mut disabled = growing_tombstones_fixture();
4590 disabled.config.enable_consolidation = false;
4591 assert!(!disabled.pressure_consolidation_tick());
4592 assert_eq!(disabled.stats().sealed_count, 0);
4593 assert_eq!(live_ids(&disabled), (49..=64).collect::<Vec<_>>());
4594 }
4595
4596 #[test]
4597 fn retained_memory_includes_an_inflight_consolidation_copy() {
4598 let manager = consolidating_manager(1);
4599 let ids = seal_bulk_segment(&manager, 100);
4600 for id in ids.iter().take(40) {
4601 manager.delete(*id).expect("delete");
4602 }
4603 let before = manager.resident_bytes();
4604 assert!(manager.pressure_consolidation_tick());
4605 let copied = manager
4606 .consolidation
4607 .read()
4608 .as_ref()
4609 .expect("paced merge")
4610 .merged
4611 .resident_bytes();
4612 assert!(copied > 0);
4613 assert!(
4614 manager.resident_bytes() >= before + copied,
4615 "unpublished copy consumes memory"
4616 );
4617 drain_maintenance(&manager);
4618 assert!(
4619 manager.resident_bytes() < before,
4620 "completed merge releases the old storage"
4621 );
4622 }
4623
4624 #[test]
4625 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
4626 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4627 let ids = seal_bulk_segment(&manager, 100);
4628
4629 let bytes_before_delete = manager.resident_bytes();
4630
4631 for id in ids.iter().take(40) {
4633 assert!(manager.delete(*id).expect("delete"));
4634 }
4635
4636 let bytes_with_tombstones = manager.resident_bytes();
4639 assert!(bytes_with_tombstones > bytes_before_delete);
4640
4641 let survivors = live_ids(&manager);
4642 assert_eq!(survivors.len(), 60);
4643
4644 drain_maintenance(&manager);
4645
4646 let stats = manager.stats();
4647 assert_eq!(stats.consolidation.runs_started, 1);
4648 assert_eq!(stats.consolidation.runs_completed, 1);
4649 assert_eq!(stats.consolidation.segments_merged, 1);
4650 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
4651 assert!(stats.consolidation.bytes_reclaimed > 0);
4652
4653 let bytes_after = manager.resident_bytes();
4655 assert!(
4656 bytes_after <= bytes_before_delete,
4657 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
4658 );
4659 assert_eq!(
4660 bytes_with_tombstones - bytes_after,
4661 stats.consolidation.bytes_reclaimed
4662 );
4663 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
4664
4665 assert_eq!(live_ids(&manager), survivors);
4666 assert_eq!(manager.sealed.read().len(), 1);
4667 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
4668 }
4669
4670 #[test]
4671 fn query_results_are_identical_before_during_and_after_consolidation() {
4672 let manager = consolidating_manager(3);
4673 let ids = seal_bulk_segment(&manager, 20);
4674 seal_bulk_segment(&manager, 20);
4675 for id in ids.iter().take(10) {
4676 assert!(manager.delete(*id).expect("delete"));
4677 }
4678
4679 let before = live_ids(&manager);
4680 let survivor = *ids.last().expect("at least one row");
4681 let tombstoned = ids[0];
4682
4683 let mut mid_tick_observations = 0;
4686 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
4687 manager.run_maintenance().expect("maintenance");
4688 assert_eq!(live_ids(&manager), before);
4689 assert!(manager.get(survivor).is_some());
4690 assert!(manager.get(tombstoned).is_none());
4691 assert_eq!(manager.get_by_kind("table").len(), before.len());
4692 mid_tick_observations += 1;
4693 assert!(mid_tick_observations < 1_000);
4694 }
4695
4696 assert!(
4697 mid_tick_observations > 2,
4698 "expected the run to span several ticks, saw {mid_tick_observations}"
4699 );
4700 assert_eq!(live_ids(&manager), before);
4701 assert!(manager.get(survivor).is_some());
4702 assert!(manager.get(tombstoned).is_none());
4703 assert_eq!(manager.get_by_kind("table").len(), before.len());
4704 }
4705
4706 #[test]
4707 fn get_by_kind_reaches_sealed_segments() {
4708 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4709 let ids = seal_bulk_segment(&manager, 5);
4710 assert!(manager.delete(ids[0]).expect("delete"));
4711
4712 assert_eq!(manager.get_by_kind("table").len(), 4);
4716 assert_eq!(manager.get_by_kind("vector").len(), 0);
4717 }
4718
4719 #[test]
4720 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
4721 const PER_TICK: usize = 4;
4722 let manager = consolidating_manager(PER_TICK);
4723 let ids = seal_bulk_segment(&manager, 60);
4724 for id in ids.iter().take(30) {
4725 assert!(manager.delete(*id).expect("delete"));
4726 }
4727
4728 manager.run_maintenance().expect("maintenance");
4730 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
4731
4732 let mut previous = PER_TICK as u64;
4735 let mut ticks = 1;
4736 while let Some(copied) = manager.consolidation_progress() {
4737 assert!(
4738 copied - previous <= PER_TICK as u64,
4739 "tick copied {} entities, bound is {PER_TICK}",
4740 copied - previous
4741 );
4742 previous = copied;
4743 manager.run_maintenance().expect("maintenance");
4744 ticks += 1;
4745 assert!(ticks < 100);
4746 }
4747
4748 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
4751 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4752 assert_eq!(live_ids(&manager).len(), 30);
4753 }
4754
4755 #[test]
4756 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
4757 let manager = SegmentManager::with_config(
4758 "test",
4759 ManagerConfig {
4760 max_sealed_segments: 3,
4761 ..Default::default()
4762 },
4763 );
4764 for _ in 0..4 {
4765 seal_bulk_segment(&manager, 5);
4766 }
4767 assert_eq!(manager.sealed.read().len(), 4);
4768
4769 let before = live_ids(&manager);
4770 drain_maintenance(&manager);
4771
4772 assert_eq!(manager.sealed.read().len(), 1);
4775 assert_eq!(manager.stats().consolidation.segments_merged, 4);
4776 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
4777 assert_eq!(live_ids(&manager), before);
4778 }
4779
4780 #[test]
4781 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
4782 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4783 seal_bulk_segment(&manager, 10);
4784 seal_bulk_segment(&manager, 10);
4785
4786 for _ in 0..5 {
4787 manager.run_maintenance().expect("maintenance");
4788 }
4789
4790 assert_eq!(manager.stats().consolidation.runs_started, 0);
4791 assert_eq!(manager.sealed.read().len(), 2);
4792 }
4793
4794 #[test]
4795 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
4796 use std::sync::atomic::AtomicBool;
4797
4798 let manager = Arc::new(consolidating_manager(2));
4799 let ids = seal_bulk_segment(&manager, 40);
4800 seal_bulk_segment(&manager, 40);
4801 for id in ids.iter().take(20) {
4802 assert!(manager.delete(*id).expect("delete"));
4803 }
4804
4805 let expected = live_ids(&manager);
4806 let stop = Arc::new(AtomicBool::new(false));
4807
4808 let readers: Vec<_> = (0..4)
4809 .map(|_| {
4810 let manager = Arc::clone(&manager);
4811 let stop = Arc::clone(&stop);
4812 let expected = expected.clone();
4813 std::thread::spawn(move || {
4814 let mut scans = 0u64;
4815 loop {
4818 let mut seen: Vec<u64> = manager
4819 .query_all(|_| true)
4820 .into_iter()
4821 .map(|entity| entity.id.raw())
4822 .collect();
4823 seen.sort_unstable();
4824 assert_eq!(seen, expected, "reader saw a torn sealed set");
4827 scans += 1;
4828 if stop.load(Ordering::Relaxed) {
4829 break;
4830 }
4831 }
4832 scans
4833 })
4834 })
4835 .collect();
4836
4837 drain_maintenance(&manager);
4838 stop.store(true, Ordering::Relaxed);
4839
4840 for reader in readers {
4841 let scans = reader.join().expect("reader thread");
4842 assert!(scans > 0, "reader never got to scan");
4843 }
4844
4845 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4846 assert_eq!(live_ids(&manager), expected);
4847 }
4848
4849 #[test]
4850 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
4851 let manager = consolidating_manager(4);
4852 let ids = seal_bulk_segment(&manager, 20);
4853 for id in ids.iter().take(6) {
4854 assert!(manager.delete(*id).expect("delete"));
4855 }
4856
4857 manager.run_maintenance().expect("maintenance");
4859 assert_eq!(manager.consolidation_progress(), Some(4));
4860 let doomed = ids[6];
4861 assert!(manager.delete(doomed).expect("delete"));
4862
4863 let expected = live_ids(&manager);
4864 drain_maintenance(&manager);
4865
4866 assert!(manager.get(doomed).is_none(), "deleted row came back");
4867 assert_eq!(live_ids(&manager), expected);
4868 }
4869
4870 #[test]
4871 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
4872 let manager = consolidating_manager(3);
4873 let ids = seal_bulk_segment(&manager, 30);
4874 for id in ids.iter().take(12) {
4875 assert!(manager.delete(*id).expect("delete"));
4876 }
4877
4878 let before = live_ids(&manager);
4879 let sealed_before: Vec<SegmentId> = manager
4880 .sealed
4881 .read()
4882 .iter()
4883 .map(|segment| segment.read().id())
4884 .collect();
4885
4886 manager.run_maintenance().expect("maintenance");
4887 manager.run_maintenance().expect("maintenance");
4888 assert!(manager.consolidation.read().is_some());
4889
4890 *manager.consolidation.write() = None;
4894
4895 assert_eq!(live_ids(&manager), before);
4896 let sealed_after: Vec<SegmentId> = manager
4897 .sealed
4898 .read()
4899 .iter()
4900 .map(|segment| segment.read().id())
4901 .collect();
4902 assert_eq!(sealed_after, sealed_before);
4903 assert_eq!(manager.stats().consolidation.runs_completed, 0);
4904
4905 drain_maintenance(&manager);
4906 assert_eq!(live_ids(&manager), before);
4907 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4908 }
4909
4910 #[test]
4911 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
4912 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4913
4914 let entities: Vec<UnifiedEntity> = (0..10)
4918 .map(|i| {
4919 UnifiedEntity::table_row(
4920 manager.next_entity_id(),
4921 "users",
4922 i as u64,
4923 vec![Value::Integer(i as i64)],
4924 )
4925 })
4926 .collect();
4927 let ids = manager.bulk_insert(entities).expect("bulk insert");
4928
4929 let mut meta = Metadata::new();
4930 meta.set(
4931 "os",
4932 super::super::metadata::MetadataValue::String("linux".to_string()),
4933 );
4934 manager.set_metadata(ids[9], meta).expect("set metadata");
4935 manager.force_seal().expect("seal");
4936
4937 assert!(manager.delete(ids[0]).expect("delete"));
4938 assert!(manager.delete(ids[1]).expect("delete"));
4939 assert!(manager.delete(ids[2]).expect("delete"));
4940
4941 let before: Vec<UnifiedEntity> = {
4942 let mut entities = manager.query_all(|_| true);
4943 entities.sort_by_key(|entity| entity.id.raw());
4944 entities
4945 };
4946
4947 drain_maintenance(&manager);
4948
4949 let after: Vec<UnifiedEntity> = {
4950 let mut entities = manager.query_all(|_| true);
4951 entities.sort_by_key(|entity| entity.id.raw());
4952 entities
4953 };
4954
4955 let columns = |entity: &UnifiedEntity| match &entity.data {
4956 super::super::entity::EntityData::Row(row) => row.columns.clone(),
4957 _ => Vec::new(),
4958 };
4959
4960 assert_eq!(before.len(), after.len());
4961 for (old, new) in before.iter().zip(after.iter()) {
4962 assert_eq!(old.id, new.id);
4963 assert_eq!(columns(old), columns(new));
4964 assert_eq!(
4965 old.sequence_id, new.sequence_id,
4966 "merge must not re-issue sequence ids"
4967 );
4968 }
4969 assert!(manager
4970 .get_metadata(ids[9])
4971 .expect("metadata survives the merge")
4972 .has("os"));
4973 }
4974
4975 #[test]
4976 #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
4977 fn test_lifecycle_events() {
4978 let manager = SegmentManager::new("test");
4979
4980 manager
4981 .insert(UnifiedEntity::vector(
4982 manager.next_entity_id(),
4983 "v",
4984 vec![0.1],
4985 ))
4986 .unwrap();
4987
4988 let events = manager.drain_events();
4989
4990 assert!(events
4992 .iter()
4993 .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
4994 assert!(events
4995 .iter()
4996 .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
4997 }
4998}