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(
654 &self,
655 row: &super::entity::RowData,
656 column_schema: Option<Arc<Vec<String>>>,
657 ) -> Option<Arc<Vec<String>>> {
658 if self.recovery_preserves_named_rows() {
659 return None;
660 }
661 let named = row.named.as_ref()?;
662 let matches = |schema: &Vec<String>| {
663 schema.len() == named.len() && schema.iter().all(|name| named.contains_key(name))
664 };
665 if let Some(schema) = column_schema.filter(|schema| matches(schema)) {
666 return Some(schema);
667 }
668 {
669 let schemas = self.recovery_row_schemas.read();
670 if let Some(schema) = schemas.iter().find(|schema| matches(schema)) {
671 return Some(Arc::clone(schema));
672 }
673 }
674 let mut schemas = self.recovery_row_schemas.write();
675 if let Some(schema) = schemas.iter().find(|schema| matches(schema)) {
676 return Some(Arc::clone(schema));
677 }
678 const MAX_RECOVERY_ROW_SCHEMAS: usize = 16;
679 if schemas.len() >= MAX_RECOVERY_ROW_SCHEMAS {
680 return None;
681 }
682 let mut names: Vec<_> = named.keys().cloned().collect();
683 names.sort_unstable();
684 let schema = Arc::new(names);
685 schemas.push(Arc::clone(&schema));
686 Some(schema)
687 }
688
689 #[doc(hidden)]
693 pub fn restore_snapshot_row_layout(&self, entity: &mut UnifiedEntity) {
694 self.restore_recovery_row_layout(entity, self.column_schema());
695 }
696
697 fn restore_recovery_row_layout(
698 &self,
699 entity: &mut UnifiedEntity,
700 column_schema: Option<Arc<Vec<String>>>,
701 ) {
702 if let super::entity::EntityData::Row(row) = &mut entity.data {
703 if let Some(schema) = self.recovery_schema(row, column_schema) {
704 row.compact_with_schema(&schema);
705 }
706 }
707 }
708
709 #[doc(hidden)]
711 pub fn restore_snapshot_entity(
712 &self,
713 mut entity: UnifiedEntity,
714 ) -> Result<EntityId, SegmentError> {
715 self.restore_snapshot_row_layout(&mut entity);
716 self.maybe_seal_growing()?;
717 if self.get(entity.id).is_some() {
718 return Err(SegmentError::AlreadyExists(entity.id));
719 }
720 let entity_id = entity.id;
721 self.next_entity_id
722 .fetch_max(entity_id.raw().saturating_add(1), Ordering::SeqCst);
723 self.with_restoration_segment(entity_id, |segment| {
724 let segment_id = segment.id();
725 segment.adopt_entity(entity, None);
727 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
728 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
729 Ok(entity_id)
730 })
731 }
732
733 pub fn insert_batch(
735 &self,
736 entities: Vec<UnifiedEntity>,
737 ) -> Result<Vec<EntityId>, SegmentError> {
738 let mut ids = Vec::with_capacity(entities.len());
739 for entity in entities {
740 ids.push(self.insert(entity)?);
741 }
742 Ok(ids)
743 }
744
745 pub fn bulk_insert(
748 &self,
749 mut entities: Vec<UnifiedEntity>,
750 ) -> Result<Vec<EntityId>, SegmentError> {
751 if self.has_backing() {
753 let mut identities = std::collections::HashSet::with_capacity(entities.len());
754 for entity in &entities {
755 if entity.id.raw() != 0 {
756 self.reject_backing_duplicate(entity.id)?;
757 if self.get_with(entity.id, |_| ()).is_some() || !identities.insert(entity.id) {
758 return Err(SegmentError::AlreadyExists(entity.id));
759 }
760 }
761 }
762 }
763 for entity in &mut entities {
765 if entity.id.raw() == 0 {
766 entity.id = self.next_entity_id();
767 }
768 self.next_entity_id
769 .fetch_max(entity.id.raw().saturating_add(1), Ordering::SeqCst);
770 if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
771 if *row_id == 0 {
772 *row_id = self.next_row_id();
773 } else {
774 self.register_row_id(*row_id);
775 }
776 }
777 }
778
779 if let Some(first_row) = entities.first() {
782 if let super::entity::EntityData::Row(ref row) = first_row.data {
783 if let Some(ref named) = row.named {
784 let schema = self.get_or_init_schema(named);
785 for entity in &mut entities {
786 if let super::entity::EntityData::Row(ref mut row) = entity.data {
787 if let Some(named) = row.named.take() {
788 let mut cols = Vec::with_capacity(schema.len());
789 for col_name in schema.iter() {
790 cols.push(
791 named
792 .get(col_name)
793 .cloned()
794 .unwrap_or(reddb_types::Value::Null),
795 );
796 }
797 row.columns = cols;
798 row.schema = Some(Arc::clone(&schema));
799 }
800 }
801 }
802 }
803 }
804 }
805
806 let segment_arc = self.get_or_create_growing();
807 let mut segment = segment_arc.write();
808 let ids = segment.bulk_insert(entities)?;
810
811 self.total_entities_atomic
816 .fetch_add(ids.len() as u64, Ordering::Relaxed);
817
818 Ok(ids)
819 }
820
821 pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
823 let mut entity = self.get_with(id, Clone::clone)?;
824 self.restore_owned_row_schemas(std::iter::once(&mut entity));
825 Some(entity)
826 }
827
828 #[doc(hidden)]
831 pub fn get_with<T>(&self, id: EntityId, visit: impl FnOnce(&UnifiedEntity) -> T) -> Option<T> {
832 if let Some(growing_arc) = self.growing.read().as_ref() {
834 let growing = growing_arc.read();
835 if let Some(entity) = growing.get(id) {
836 return Some(visit(entity));
837 }
838 }
839
840 let sealed = self.sealed.read();
842 for segment in sealed.iter() {
843 let seg = segment.read();
844 if let Some(entity) = seg.get(id) {
845 return Some(visit(entity));
846 }
847 }
848
849 None
850 }
851
852 #[doc(hidden)]
856 pub fn graph_node_logical_id(
857 &self,
858 id: EntityId,
859 mut before_segment: impl FnMut() -> bool,
860 ) -> Option<EntityId> {
861 if let Some(growing_arc) = self.growing.read().as_ref() {
862 if !before_segment() {
863 return None;
864 }
865 let growing = growing_arc.read();
866 if let Some(entity) = growing.get(id) {
867 return matches!(entity.kind, EntityKind::GraphNode(_))
868 .then(|| entity.logical_id());
869 }
870 }
871 let sealed = self.sealed.read();
872 for segment in sealed.iter() {
873 if !before_segment() {
874 return None;
875 }
876 let segment = segment.read();
877 if let Some(entity) = segment.get(id) {
878 return matches!(entity.kind, EntityKind::GraphNode(_))
879 .then(|| entity.logical_id());
880 }
881 }
882 None
883 }
884
885 #[doc(hidden)]
890 pub fn visit_graph_index_batches<E>(
891 &self,
892 kind: GraphEntityKind,
893 keys: &[EntityId],
894 mut before_work: impl FnMut() -> bool,
895 mut admit: impl FnMut(usize) -> Result<(), E>,
896 mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
897 ) -> Result<bool, E> {
898 use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
899 admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
900 let mut admitted = 0;
901 let sources = loop {
902 let count = {
903 let growing = self.growing.read();
904 let sealed = self.sealed.read();
905 sealed.len() + usize::from(growing.is_some())
906 };
907 let bytes = count.saturating_mul(std::mem::size_of::<(
908 Arc<RwLock<GrowingSegment>>,
909 GraphIndexCursor,
910 )>());
911 admit(bytes.saturating_sub(admitted))?;
912 admitted = admitted.max(bytes);
913 let mut sources = Vec::with_capacity(count);
914 let growing = self.growing.read();
915 let sealed = self.sealed.read();
916 if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
917 continue;
918 }
919 for segment in growing.iter().chain(sealed.iter()) {
920 if !before_work() {
921 return Ok(false);
922 }
923 let view = segment.read();
924 if !view.may_contain_graph_kind(kind) {
925 continue;
926 }
927 let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
928 return Ok(false);
929 };
930 sources.push((Arc::clone(segment), cursor));
931 }
932 break sources;
933 };
934 let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
935 for (source, mut cursor) in sources {
936 while !cursor.finished(keys) {
937 let mut count = 0;
938 {
939 let segment = source.read();
940 if !segment.may_contain_graph_kind(kind) {
941 break;
942 }
943 if !segment.graph_index_batch(
944 kind,
945 keys,
946 &mut cursor,
947 &mut before_work,
948 &mut |id| {
949 ids[count] = id;
950 count += 1;
951 },
952 ) {
953 return Ok(false);
954 }
955 }
956 if !before_work() || !consume(&ids[..count])? {
957 return Ok(false);
958 }
959 }
960 }
961 Ok(true)
962 }
963
964 pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
969 let schema = self.column_schema();
970 let clone_row = |entity: &UnifiedEntity| {
971 let mut owned = entity.clone();
972 if let super::entity::EntityData::Row(row) = &mut owned.data {
973 if row.schema.is_none() {
974 row.schema = schema.clone();
975 }
976 }
977 owned
978 };
979 let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
980 let mut remaining: Vec<usize> = (0..ids.len()).collect(); if let Some(growing_arc) = self.growing.read().as_ref() {
985 let growing = if let Some(g) = growing_arc.try_read() {
986 g
987 } else {
988 growing_arc.read()
989 };
990 remaining.retain(|&i| {
991 if let Some(entity) = growing.get(ids[i]) {
992 out[i] = Some(clone_row(entity));
993 false } else {
995 true }
997 });
998 }
999
1000 if remaining.is_empty() {
1001 self.restore_owned_row_schemas(out.iter_mut().flatten());
1002 return out;
1003 }
1004
1005 let sealed = self.sealed.read();
1007 for segment in sealed.iter() {
1008 if remaining.is_empty() {
1009 break;
1010 }
1011 let seg = segment.read();
1012 remaining.retain(|&i| {
1013 if let Some(entity) = seg.get(ids[i]) {
1014 out[i] = Some(clone_row(entity));
1015 false
1016 } else {
1017 true
1018 }
1019 });
1020 }
1021
1022 drop(sealed);
1023 self.restore_owned_row_schemas(out.iter_mut().flatten());
1024 out
1025 }
1026
1027 pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
1040 where
1041 F: FnMut(usize, &UnifiedEntity),
1042 {
1043 thread_local! {
1052 static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
1053 const { std::cell::Cell::new(Vec::new()) };
1054 }
1055
1056 let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
1057 remaining.clear();
1058
1059 if let Some(growing_arc) = self.growing.read().as_ref() {
1060 let growing = if let Some(g) = growing_arc.try_read() {
1061 g
1062 } else {
1063 growing_arc.read()
1064 };
1065 for (i, id) in ids.iter().enumerate() {
1066 if let Some(entity) = growing.get(*id) {
1067 f(i, entity);
1068 } else {
1069 remaining.push(i);
1070 }
1071 }
1072 } else {
1073 remaining.reserve(ids.len());
1074 remaining.extend(0..ids.len());
1075 }
1076
1077 if !remaining.is_empty() {
1078 let sealed = self.sealed.read();
1079 for segment in sealed.iter() {
1080 if remaining.is_empty() {
1081 break;
1082 }
1083 let seg = segment.read();
1084 remaining.retain(|&i| {
1085 if let Some(entity) = seg.get(ids[i]) {
1086 f(i, entity);
1087 false
1088 } else {
1089 true
1090 }
1091 });
1092 }
1093 }
1094
1095 REMAINING_SCRATCH.with(|cell| cell.set(remaining));
1096 }
1097
1098 #[allow(dead_code)] fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
1101 if let Some(growing_arc) = self.growing.read().as_ref() {
1103 let growing = growing_arc.read();
1104 if let Some(entity) = growing.get(id) {
1105 return Some(entity.clone());
1106 }
1107 }
1108
1109 let sealed = self.sealed.read();
1111 for segment in sealed.iter() {
1112 if let Some(entity) = segment.get(id) {
1113 return Some(entity.clone());
1114 }
1115 }
1116
1117 None
1118 }
1119
1120 fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
1121 let sealed = self.sealed.read();
1122 sealed
1123 .iter()
1124 .find(|segment_arc| segment_arc.read().contains(id))
1125 .map(Arc::clone)
1126 }
1127
1128 fn rewrite_sealed_entity_into_growing(
1129 &self,
1130 entity: UnifiedEntity,
1131 metadata: Option<&Metadata>,
1132 ) -> Result<(), SegmentError> {
1133 let entity_id = entity.id;
1134 let sealed_arc = self
1135 .find_sealed_segment_arc(entity_id)
1136 .ok_or(SegmentError::NotFound(entity_id))?;
1137
1138 let metadata_to_apply = {
1139 let mut sealed = sealed_arc.write();
1140 let existing_metadata = sealed.get_metadata(entity_id);
1141 if !sealed.force_delete(entity_id) {
1142 return Err(SegmentError::NotFound(entity_id));
1143 }
1144 metadata.cloned().or(existing_metadata)
1145 };
1146
1147 let growing_id = self.with_restoration_segment(entity_id, |growing| {
1148 growing.insert_preserving_sequence(entity)?;
1149 if let Some(metadata) = metadata_to_apply {
1150 growing.set_metadata(entity_id, metadata)?;
1151 }
1152 Ok(growing.id())
1153 })?;
1154
1155 self.entity_segment.write().insert(entity_id, growing_id);
1156 Ok(())
1157 }
1158
1159 pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
1161 self.materialize_backing_entity(entity.id)?;
1162 let entity_id = entity.id;
1163 let mut entity = Some(entity);
1164
1165 if let Some(growing_arc) = self.growing.read().as_ref() {
1167 let mut growing = growing_arc.write();
1168 if growing.contains(entity_id) && growing.state().is_writable() {
1169 return growing.update(entity.take().expect("entity already moved"));
1170 }
1171 }
1172
1173 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1175 if let Some(seg_id) = segment_id {
1176 if let Some(growing_arc) = self.growing.read().as_ref() {
1177 let mut growing = growing_arc.write();
1178 if growing.id() == seg_id && growing.state().is_writable() {
1179 return growing.update(entity.take().expect("entity already moved"));
1180 }
1181 }
1182 }
1183
1184 if let Some(entity) = entity.take() {
1185 return self.rewrite_sealed_entity_into_growing(entity, None);
1186 }
1187
1188 Err(SegmentError::NotFound(entity_id))
1189 }
1190
1191 pub fn update_with_metadata(
1194 &self,
1195 entity: UnifiedEntity,
1196 metadata: Option<&Metadata>,
1197 ) -> Result<(), SegmentError> {
1198 self.materialize_backing_entity(entity.id)?;
1199 let entity_id = entity.id;
1200 let mut entity = Some(entity);
1201
1202 if let Some(growing_arc) = self.growing.read().as_ref() {
1204 let mut growing = growing_arc.write();
1205 if growing.contains(entity_id) && growing.state().is_writable() {
1206 growing.update(entity.take().expect("entity already moved"))?;
1207 if let Some(metadata) = metadata {
1208 growing.set_metadata(entity_id, metadata.clone())?;
1209 }
1210 return Ok(());
1211 }
1212 }
1213
1214 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1216 if let Some(seg_id) = segment_id {
1217 if let Some(growing_arc) = self.growing.read().as_ref() {
1218 let mut growing = growing_arc.write();
1219 if growing.id() == seg_id && growing.state().is_writable() {
1220 growing.update(entity.take().expect("entity already moved"))?;
1221 if let Some(metadata) = metadata {
1222 growing.set_metadata(entity_id, metadata.clone())?;
1223 }
1224 return Ok(());
1225 }
1226 }
1227 }
1228
1229 if let Some(entity) = entity.take() {
1230 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1231 }
1232
1233 Err(SegmentError::NotFound(entity_id))
1234 }
1235
1236 pub fn update_hot(
1240 &self,
1241 entity: UnifiedEntity,
1242 modified_columns: &[String],
1243 ) -> Result<(), SegmentError> {
1244 self.materialize_backing_entity(entity.id)?;
1245 let entity_id = entity.id;
1246 let mut entity = Some(entity);
1247
1248 if let Some(growing_arc) = self.growing.read().as_ref() {
1249 let mut growing = growing_arc.write();
1250 if growing.contains(entity_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 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1259 if let Some(seg_id) = segment_id {
1260 if let Some(growing_arc) = self.growing.read().as_ref() {
1261 let mut growing = growing_arc.write();
1262 if growing.id() == seg_id && growing.state().is_writable() {
1263 return growing.update_hot(
1264 entity.take().expect("entity already moved"),
1265 modified_columns,
1266 );
1267 }
1268 }
1269 }
1270
1271 if let Some(entity) = entity.take() {
1272 return self.rewrite_sealed_entity_into_growing(entity, None);
1273 }
1274
1275 Err(SegmentError::NotFound(entity_id))
1276 }
1277
1278 pub fn update_hot_with_metadata(
1281 &self,
1282 entity: UnifiedEntity,
1283 modified_columns: &[String],
1284 metadata: Option<&Metadata>,
1285 ) -> Result<(), SegmentError> {
1286 self.materialize_backing_entity(entity.id)?;
1287 let entity_id = entity.id;
1288 let mut entity = Some(entity);
1289
1290 if let Some(growing_arc) = self.growing.read().as_ref() {
1291 let mut growing = growing_arc.write();
1292 if growing.contains(entity_id) && growing.state().is_writable() {
1293 growing.update_hot(
1294 entity.take().expect("entity already moved"),
1295 modified_columns,
1296 )?;
1297 if let Some(metadata) = metadata {
1298 growing.set_metadata(entity_id, metadata.clone())?;
1299 }
1300 return Ok(());
1301 }
1302 }
1303
1304 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1305 if let Some(seg_id) = segment_id {
1306 if let Some(growing_arc) = self.growing.read().as_ref() {
1307 let mut growing = growing_arc.write();
1308 if growing.id() == seg_id && growing.state().is_writable() {
1309 growing.update_hot(
1310 entity.take().expect("entity already moved"),
1311 modified_columns,
1312 )?;
1313 if let Some(metadata) = metadata {
1314 growing.set_metadata(entity_id, metadata.clone())?;
1315 }
1316 return Ok(());
1317 }
1318 }
1319 }
1320
1321 if let Some(entity) = entity.take() {
1322 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1323 }
1324
1325 Err(SegmentError::NotFound(entity_id))
1326 }
1327
1328 pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
1331 where
1332 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
1333 {
1334 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
1335 items.into_iter().collect();
1336 if items.is_empty() {
1337 return Ok(());
1338 }
1339
1340 if let Some(growing_arc) = self.growing.read().as_ref() {
1341 let mut growing = growing_arc.write();
1342 if growing.state().is_writable() {
1343 match growing.update_hot_batch_with_metadata(items.iter().copied()) {
1344 Ok(()) => return Ok(()),
1345 Err(SegmentError::NotFound(_)) => {}
1346 Err(other) => return Err(other),
1347 }
1348 }
1349 }
1350
1351 for (entity, modified_columns, metadata) in items {
1352 self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
1353 }
1354 Ok(())
1355 }
1356
1357 pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
1359 self.materialize_backing_entity(id)?;
1360 let deleted = self.delete_resident(id)?;
1361 if deleted {
1362 self.remember_backing_delete(id);
1363 }
1364 Ok(deleted)
1365 }
1366
1367 fn delete_resident(&self, id: EntityId) -> Result<bool, SegmentError> {
1368 if let Some(growing_arc) = self.growing.read().as_ref() {
1371 let mut growing = growing_arc.write();
1372 if growing.contains(id) && growing.state().is_writable() {
1373 let seg_id = growing.id();
1374 let deleted = growing.delete(id)?;
1375 if deleted {
1376 self.entity_segment.write().remove(&id);
1377 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1378 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1379 }
1380 return Ok(deleted);
1381 }
1382 }
1383
1384 let segment_id = self.entity_segment.read().get(&id).copied();
1387
1388 if let Some(seg_id) = segment_id {
1389 if let Some(growing_arc) = self.growing.read().as_ref() {
1390 let mut growing = growing_arc.write();
1391 if growing.id() == seg_id && growing.state().is_writable() {
1392 let deleted = growing.delete(id)?;
1393 if deleted {
1394 self.entity_segment.write().remove(&id);
1395 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1396 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1397 }
1398 return Ok(deleted);
1399 }
1400 }
1401 }
1402
1403 {
1406 let sealed = self.sealed.read();
1407 for segment_arc in sealed.iter() {
1408 let mut seg = segment_arc.write();
1409 let seg_id = seg.id();
1410 if seg.contains(id) {
1411 let deleted = seg.force_delete(id);
1412 drop(seg);
1413 if deleted {
1414 self.entity_segment.write().remove(&id);
1415 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1416 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1417 }
1418 return Ok(deleted);
1419 }
1420 }
1421 }
1422
1423 Ok(false)
1424 }
1425
1426 pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1427 if self.has_backing() {
1428 let mut deleted = Vec::new();
1429 for id in ids {
1430 if self.delete(*id)? {
1431 deleted.push(*id);
1432 }
1433 }
1434 return Ok(deleted);
1435 }
1436 if ids.is_empty() {
1437 return Ok(Vec::new());
1438 }
1439
1440 let mut deleted_ids = Vec::with_capacity(ids.len());
1441
1442 if let Some(growing_arc) = self.growing.read().as_ref() {
1443 let mut growing = growing_arc.write();
1444 if growing.state().is_writable() {
1445 let seg_id = growing.id();
1446 let deleted = growing.delete_batch(ids)?;
1447 if !deleted.is_empty() {
1448 {
1449 let mut entity_segment = self.entity_segment.write();
1450 for id in &deleted {
1451 entity_segment.remove(id);
1452 }
1453 }
1454 self.total_entities_atomic
1455 .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
1456 for id in &deleted {
1457 self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1458 }
1459 deleted_ids.extend(deleted);
1460 }
1461 }
1462 }
1463
1464 if deleted_ids.len() == ids.len() {
1465 return Ok(deleted_ids);
1466 }
1467
1468 let deleted_set: std::collections::HashSet<EntityId> =
1469 deleted_ids.iter().copied().collect();
1470 for &id in ids {
1471 if deleted_set.contains(&id) {
1472 continue;
1473 }
1474 if self.delete(id)? {
1475 deleted_ids.push(id);
1476 }
1477 }
1478
1479 Ok(deleted_ids)
1480 }
1481
1482 pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1484 if let Some(growing_arc) = self.growing.read().as_ref() {
1486 let growing = growing_arc.read();
1487 if growing.contains(id) {
1488 return growing.get_metadata(id);
1489 }
1490 }
1491
1492 let segment_id = self.entity_segment.read().get(&id).copied();
1494
1495 if let Some(seg_id) = segment_id {
1496 if let Some(growing_arc) = self.growing.read().as_ref() {
1497 let growing = growing_arc.read();
1498 if growing.id() == seg_id {
1499 return growing.get_metadata(id);
1500 }
1501 }
1502
1503 let sealed = self.sealed.read();
1504 for segment in sealed.iter() {
1505 if segment.id() == seg_id {
1506 return segment.get_metadata(id);
1507 }
1508 }
1509 }
1510
1511 if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1512 return segment_arc.read().get_metadata(id);
1513 }
1514
1515 None
1516 }
1517
1518 pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1520 self.materialize_backing_entity(id)?;
1521 if let Some(growing_arc) = self.growing.read().as_ref() {
1524 let mut growing = growing_arc.write();
1525 if growing.contains(id) && growing.state().is_writable() {
1526 return growing.set_metadata(id, metadata);
1527 }
1528 }
1529
1530 let segment_id = self.entity_segment.read().get(&id).copied();
1532
1533 if let Some(seg_id) = segment_id {
1534 if let Some(growing_arc) = self.growing.read().as_ref() {
1535 let mut growing = growing_arc.write();
1536 if growing.id() == seg_id && growing.state().is_writable() {
1537 return growing.set_metadata(id, metadata);
1538 }
1539 }
1540 }
1541
1542 if let Some(entity) = self.get(id) {
1543 return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1544 }
1545
1546 Err(SegmentError::NotFound(id))
1547 }
1548
1549 fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1551 let should_seal = {
1552 let growing_opt = self.growing.read();
1553 if let Some(growing_arc) = growing_opt.as_ref() {
1554 let growing = growing_arc.read();
1555 growing.should_seal(&self.config.segment_config)
1556 || growing.idle_secs() >= self.config.idle_seal_secs
1557 } else {
1558 false
1559 }
1560 };
1561
1562 if should_seal {
1563 self.seal_current()?;
1564 }
1565
1566 Ok(())
1567 }
1568
1569 pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1571 self.seal_current_if(|_| true)?
1572 .ok_or(SegmentError::InvalidState(SegmentState::Sealed))
1573 }
1574
1575 fn seal_current_if(
1576 &self,
1577 should_seal: impl FnOnce(&GrowingSegment) -> bool,
1578 ) -> Result<Option<SegmentId>, SegmentError> {
1579 let disk_residency = self.has_backing();
1580 let mut growing_slot = self.growing.write();
1583 if let Some(growing_arc) = growing_slot.as_ref() {
1584 let mut growing = growing_arc.write();
1585 if !should_seal(&growing) {
1586 return Ok(None);
1587 }
1588 if disk_residency {
1589 self.next_sequence_id
1590 .fetch_max(growing.sequence_id_next(), Ordering::SeqCst);
1591 }
1592 let seg_id = growing.id();
1593 let entity_count = growing.stats().entity_count as u64;
1594
1595 growing.seal()?;
1597
1598 drop(growing); let growing_arc = growing_slot
1604 .take()
1605 .expect("sealed source remains published");
1606 self.sealed.write().push(growing_arc);
1607 drop(growing_slot);
1608
1609 self.mark_sealed_pages_visible(entity_count);
1611
1612 {
1614 let mut stats = self.stats.write();
1615 stats.growing_count = stats.growing_count.saturating_sub(1);
1616 stats.sealed_count += 1;
1617 stats.seal_ops += 1;
1618 }
1619
1620 self.emit(LifecycleEvent::SegmentSealed(seg_id));
1621
1622 return Ok(Some(seg_id));
1623 }
1624
1625 Ok(None)
1626 }
1627
1628 pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1630 let has_growing = self.growing.read().is_some();
1631 if has_growing {
1632 Ok(Some(self.seal_current()?))
1633 } else {
1634 Ok(None)
1635 }
1636 }
1637
1638 pub fn all_visible_fraction(&self) -> f64 {
1648 const ROWS_PER_PAGE: u32 = 256;
1649 let sealed = self.sealed.read();
1650 if sealed.is_empty() {
1651 return 0.0;
1652 }
1653 let mut total_pages: u64 = 0;
1654 for seg_arc in sealed.iter() {
1655 let seg = seg_arc.read();
1656 let entity_count = seg.stats().entity_count as u64;
1657 let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1658 total_pages += pages;
1659 }
1660 if total_pages == 0 {
1661 return 0.0;
1662 }
1663 let visible = self.visibility_map.all_visible_count();
1664 (visible as f64 / total_pages as f64).min(1.0)
1665 }
1666
1667 fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1670 const ROWS_PER_PAGE: u32 = 256;
1671 let existing_visible = self.visibility_map.all_visible_count();
1672 let start_page = existing_visible as u32;
1674 let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1675 let end_page = start_page + new_pages as u32;
1676 self.visibility_map.mark_range_visible(start_page, end_page);
1677 }
1678
1679 #[doc(hidden)]
1683 pub fn find_unique_key_conflict(
1684 &self,
1685 columns: &[Vec<String>],
1686 key_index: usize,
1687 signatures: &[String],
1688 exclude_id: Option<EntityId>,
1689 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1690 ) -> Option<EntityId> {
1691 if let Some(growing) = self.growing.read().as_ref() {
1692 if let Some(id) = growing.read().find_unique_key_conflict(
1693 columns,
1694 key_index,
1695 signatures,
1696 exclude_id,
1697 reserves_key,
1698 ) {
1699 return Some(id);
1700 }
1701 }
1702 for segment in self.sealed.read().iter() {
1703 if let Some(id) = segment.read().find_unique_key_conflict(
1704 columns,
1705 key_index,
1706 signatures,
1707 exclude_id,
1708 reserves_key,
1709 ) {
1710 return Some(id);
1711 }
1712 }
1713 None
1714 }
1715
1716 pub fn for_each_entity<F>(&self, callback: F)
1722 where
1723 F: FnMut(&UnifiedEntity) -> bool,
1724 {
1725 self.for_each_entity_in_segments(|_| true, || true, callback);
1726 }
1727
1728 #[doc(hidden)]
1730 pub fn checkpoint_revisions(&self) -> Vec<(u64, u64)> {
1731 let growing = self.growing.read();
1732 let sealed = self.sealed.read();
1733 growing
1734 .iter()
1735 .chain(sealed.iter())
1736 .map(|segment| {
1737 let segment = segment.read();
1738 (segment.id(), segment.checkpoint_revision())
1739 })
1740 .collect()
1741 }
1742
1743 #[doc(hidden)]
1744 pub fn checkpoint_mutations_since(&self, before: &[(u64, u64)]) -> Option<Vec<EntityId>> {
1745 let growing = self.growing.read();
1746 let sealed = self.sealed.read();
1747 let previous: HashMap<_, _> = before.iter().copied().collect();
1752 if previous.len() != before.len() {
1753 return None;
1754 }
1755 let mut retained_sources = 0;
1756 let mut changed = std::collections::BTreeSet::new();
1757 for segment in growing.iter().chain(sealed.iter()) {
1758 let segment = segment.read();
1759 let revision = previous.get(&segment.id()).copied();
1760 retained_sources += usize::from(revision.is_some());
1761 let revision = revision.unwrap_or(0);
1764 changed.extend(segment.checkpoint_mutations_since(revision)?);
1765 if changed.len() > super::segment::CHECKPOINT_MUTATIONS_MAX {
1766 return None;
1767 }
1768 }
1769 (retained_sources == previous.len()).then(|| changed.into_iter().collect())
1772 }
1773
1774 #[doc(hidden)]
1775 pub fn checkpoint_tracking_bytes(&self) -> usize {
1776 let growing = self.growing.read();
1777 let sealed = self.sealed.read();
1778 growing
1779 .iter()
1780 .chain(sealed.iter())
1781 .map(|segment| segment.read().checkpoint_tracking_bytes())
1782 .sum()
1783 }
1784
1785 #[doc(hidden)]
1787 pub fn cached_physical_root(&self, seed: u64, compute: impl FnOnce() -> u64) -> u64 {
1788 let segments = self.checkpoint_revisions();
1789 let mut cache = self.physical_root_cache.lock();
1790 if let Some(previous) = cache.as_ref() {
1791 if previous.seed == seed && previous.segments == segments {
1792 return previous.root;
1793 }
1794 }
1795 let root = compute();
1796 *cache = Some(PhysicalRootCacheEntry {
1797 segments,
1798 seed,
1799 root,
1800 });
1801 root
1802 }
1803
1804 #[doc(hidden)]
1809 pub fn for_each_entity_in_id_order(&self, mut visit: impl FnMut(&UnifiedEntity) -> bool) {
1810 self.for_each_entity_source_in_id_order(|_, entity| visit(entity));
1811 }
1812
1813 #[doc(hidden)]
1815 pub fn page_in_id_order(&self, offset: usize, limit: usize) -> (usize, Vec<UnifiedEntity>) {
1816 let end = offset.saturating_add(limit);
1817 let mut total = 0;
1818 let mut items = Vec::with_capacity(limit.min(self.count()));
1819 self.for_each_entity_in_id_order(|entity| {
1820 if total >= offset && total < end {
1821 items.push(entity.clone());
1822 }
1823 total += 1;
1824 true
1825 });
1826 self.restore_owned_row_schemas(items.iter_mut());
1827 (total, items)
1828 }
1829
1830 #[doc(hidden)]
1831 pub fn for_each_entity_with_metadata_in_id_order(
1832 &self,
1833 mut visit: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1834 ) {
1835 self.for_each_entity_source_in_id_order(|segment, entity| {
1836 visit(entity, segment.get_metadata(entity.id))
1837 });
1838 }
1839
1840 fn for_each_entity_source_in_id_order(
1841 &self,
1842 mut visit: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
1843 ) {
1844 let segments: Vec<_> = {
1845 let growing = self.growing.read();
1846 let sealed = self.sealed.read();
1847 growing.iter().chain(sealed.iter()).cloned().collect()
1848 };
1849 let mut order = Vec::new();
1850 for (source, segment) in segments.iter().enumerate() {
1851 segment.read().for_each_fast(|entity| {
1852 order.push((entity.id, source));
1853 true
1854 });
1855 }
1856 order.sort_by_key(|(id, _)| id.raw());
1857 for (id, source) in order {
1858 let segment = segments[source].read();
1859 if let Some(entity) = segment.get(id) {
1860 if !visit(&segment, entity) {
1861 break;
1862 }
1863 }
1864 }
1865 }
1866
1867 #[doc(hidden)]
1870 pub fn for_each_entity_with_metadata(
1871 &self,
1872 mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1873 ) {
1874 self.for_each_entity_source(|segment, entity| {
1875 callback(entity, segment.get_metadata(entity.id))
1876 });
1877 }
1878
1879 fn for_each_entity_source(
1880 &self,
1881 mut callback: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
1882 ) {
1883 if let Some(growing_arc) = self.growing.read().as_ref() {
1884 let growing = growing_arc.read();
1885 if !growing.for_each_fast(|entity| callback(&growing, entity)) {
1886 return;
1887 }
1888 }
1889 let sealed = self.sealed.read();
1890 for segment_arc in sealed.iter() {
1891 let segment = segment_arc.read();
1892 if !segment.for_each_fast(|entity| callback(&segment, entity)) {
1893 return;
1894 }
1895 }
1896 }
1897
1898 fn for_each_entity_in_segments<F>(
1899 &self,
1900 segment_filter: impl Fn(&GrowingSegment) -> bool,
1901 mut before_segment: impl FnMut() -> bool,
1902 mut callback: F,
1903 ) where
1904 F: FnMut(&UnifiedEntity) -> bool,
1905 {
1906 if let Some(growing_arc) = self.growing.read().as_ref() {
1910 if !before_segment() {
1911 return;
1912 }
1913 let growing = if let Some(g) = growing_arc.try_read() {
1914 g
1915 } else {
1916 growing_arc.read()
1917 };
1918 if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
1919 return;
1920 }
1921 }
1922
1923 let sealed = self.sealed.read();
1925 for segment_arc in sealed.iter() {
1926 if !before_segment() {
1927 return;
1928 }
1929 let segment = segment_arc.read();
1930 if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
1931 return;
1932 }
1933 }
1934 }
1935
1936 pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1952 &self,
1953 init: FInit,
1954 fold: FFold,
1955 reduce: FReduce,
1956 ) -> T
1957 where
1958 T: Send,
1959 FInit: Fn() -> T + Send + Sync,
1960 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1961 FReduce: Fn(T, T) -> T + Send + Sync,
1962 {
1963 #[cfg(not(target_arch = "wasm32"))]
1964 use rayon::prelude::*;
1965
1966 let mut acc = init();
1969 if let Some(growing_arc) = self.growing.read().as_ref() {
1970 let growing = if let Some(g) = growing_arc.try_read() {
1971 g
1972 } else {
1973 growing_arc.read()
1974 };
1975 growing.for_each_fast(|entity| {
1976 acc = fold(std::mem::replace(&mut acc, init()), entity);
1977 true
1978 });
1979 }
1980
1981 let segments: Vec<_> = {
1985 let sealed = self.sealed.read();
1986 sealed.iter().cloned().collect()
1987 };
1988
1989 if segments.len() <= 1 {
1990 for seg_arc in &segments {
1991 let seg = seg_arc.read();
1992 seg.for_each_fast(|entity| {
1993 acc = fold(std::mem::replace(&mut acc, init()), entity);
1994 true
1995 });
1996 }
1997 return acc;
1998 }
1999
2000 #[cfg(not(target_arch = "wasm32"))]
2001 let segments = segments.into_par_iter();
2002 #[cfg(target_arch = "wasm32")]
2003 let segments = segments.into_iter();
2004 let partials = segments.map(|seg_arc| {
2005 let mut local = init();
2006 let seg = seg_arc.read();
2007 seg.for_each_fast(|entity| {
2008 local = fold(std::mem::replace(&mut local, init()), entity);
2009 true
2010 });
2011 local
2012 });
2013 #[cfg(not(target_arch = "wasm32"))]
2014 let sealed_acc = partials.reduce(&init, &reduce);
2015 #[cfg(target_arch = "wasm32")]
2016 let sealed_acc = partials.fold(init(), &reduce);
2017
2018 reduce(acc, sealed_acc)
2019 }
2020
2021 pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
2030 where
2031 F: FnMut(&UnifiedEntity) -> bool,
2032 {
2033 let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
2034 }
2035
2036 pub fn for_each_entity_zoned_with_stats<F>(
2037 &self,
2038 zone_preds: &[(&str, ZoneColPred<'_>)],
2039 mut callback: F,
2040 ) -> SegmentScanStats
2041 where
2042 F: FnMut(&UnifiedEntity) -> bool,
2043 {
2044 let mut scan_stats = SegmentScanStats::default();
2045 if let Some(growing_arc) = self.growing.read().as_ref() {
2051 scan_stats.segments_total += 1;
2052 scan_stats.segments_scanned += 1;
2053 let growing = if let Some(g) = growing_arc.try_read() {
2054 g
2055 } else {
2056 growing_arc.read()
2057 };
2058 if !growing.for_each_fast(&mut callback) {
2059 return scan_stats;
2060 }
2061 }
2062
2063 let sealed = self.sealed.read();
2065 for segment_arc in sealed.iter() {
2066 scan_stats.segments_total += 1;
2067 let segment = segment_arc.read();
2068 if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
2069 scan_stats.segments_pruned += 1;
2070 continue; }
2072 scan_stats.segments_scanned += 1;
2073 if !segment.for_each_fast(&mut callback) {
2074 return scan_stats;
2075 }
2076 }
2077 scan_stats
2078 }
2079
2080 pub fn query_all_zoned<F>(
2090 &self,
2091 zone_preds: &[(&str, ZoneColPred<'_>)],
2092 filter: F,
2093 ) -> Vec<UnifiedEntity>
2094 where
2095 F: Fn(&UnifiedEntity) -> bool + Sync,
2096 {
2097 self.query_all_zoned_with_stats(zone_preds, filter).0
2098 }
2099
2100 pub fn query_all_zoned_with_stats<F>(
2101 &self,
2102 zone_preds: &[(&str, ZoneColPred<'_>)],
2103 filter: F,
2104 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2105 where
2106 F: Fn(&UnifiedEntity) -> bool + Sync,
2107 {
2108 self.project_all_zoned_with_stats(zone_preds, |entity| {
2109 filter(entity).then(|| entity.clone())
2110 })
2111 }
2112
2113 fn project_all_zoned_with_stats<T, F>(
2114 &self,
2115 zone_preds: &[(&str, ZoneColPred<'_>)],
2116 project: F,
2117 ) -> (Vec<T>, SegmentScanStats)
2118 where
2119 T: Send,
2120 F: Fn(&UnifiedEntity) -> Option<T> + Sync,
2121 {
2122 let mut results = Vec::new();
2123 let mut scan_stats = SegmentScanStats::default();
2124
2125 if let Some(growing_arc) = self.growing.read().as_ref() {
2128 scan_stats.segments_total += 1;
2129 scan_stats.segments_scanned += 1;
2130 let growing = if let Some(g) = growing_arc.try_read() {
2131 g
2132 } else {
2133 growing_arc.read()
2134 };
2135 results.extend(growing.iter().filter_map(&project));
2136 }
2137
2138 let sealed = self.sealed.read();
2140 let surviving: Vec<_> = sealed
2142 .iter()
2143 .filter(|seg_arc| {
2144 scan_stats.segments_total += 1;
2145 if zone_preds.is_empty() {
2146 scan_stats.segments_scanned += 1;
2147 return true;
2148 }
2149 let seg = seg_arc.read();
2150 if seg.can_skip_zone_preds(zone_preds) {
2151 scan_stats.segments_pruned += 1;
2152 false
2153 } else {
2154 scan_stats.segments_scanned += 1;
2155 true
2156 }
2157 })
2158 .collect();
2159
2160 let use_parallel = surviving.len() > 1 && crate::execution::should_parallelize();
2161
2162 if use_parallel {
2163 let project_ref = &project;
2164 let segment_results: Vec<Vec<T>> = std::thread::scope(|s| {
2165 surviving
2166 .iter()
2167 .map(|segment| {
2168 s.spawn(move || {
2169 segment
2170 .read()
2171 .iter()
2172 .filter_map(project_ref)
2173 .collect::<Vec<_>>()
2174 })
2175 })
2176 .collect::<Vec<_>>()
2177 .into_iter()
2178 .map(|handle| match handle.join() {
2179 Ok(batch) => batch,
2180 Err(panic) => std::panic::resume_unwind(panic),
2181 })
2182 .collect()
2183 });
2184 for batch in segment_results {
2185 results.extend(batch);
2186 }
2187 } else {
2188 for segment_arc in surviving {
2189 let seg = segment_arc.read();
2190 results.extend(seg.iter().filter_map(&project));
2191 }
2192 }
2193
2194 (results, scan_stats)
2195 }
2196
2197 pub fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
2203 match snapshot {
2204 Some(ctx) => entity_visible_with_context(Some(ctx), entity),
2205 None => {
2206 !crate::storage::visibility::entity_moderation_hidden(entity) && entity.xmax == 0
2207 }
2208 }
2209 }
2210
2211 pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
2218 where
2219 F: Fn(&UnifiedEntity) -> bool + Sync,
2220 {
2221 self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
2222 }
2223
2224 #[doc(hidden)]
2229 pub fn project_key_entity<T>(
2230 &self,
2231 key: &str,
2232 snapshot: Option<&SnapshotContext>,
2233 versioned: bool,
2234 mut project: impl FnMut(&UnifiedEntity) -> T,
2235 ) -> Option<T> {
2236 let mut probe = |segment: &GrowingSegment, best: &mut Option<(u64, T)>| {
2237 if let Some(entity) = segment.find_key_entity(key, versioned, &mut |entity| {
2238 !versioned || Self::scan_entity_visible(snapshot, entity)
2239 }) {
2240 if best.as_ref().is_none_or(|(xmin, _)| entity.xmin >= *xmin) {
2241 *best = Some((entity.xmin, project(entity)));
2242 }
2243 }
2244 };
2245 let mut best = None;
2246 if let Some(growing) = self.growing.read().as_ref() {
2247 probe(&growing.read(), &mut best);
2248 }
2249 for segment in self.sealed.read().iter() {
2250 if !versioned && best.is_some() {
2251 break;
2252 }
2253 probe(&segment.read(), &mut best);
2254 }
2255 best.map(|(_, projected)| projected)
2256 }
2257
2258 #[doc(hidden)]
2261 pub fn visit_queue_dedup(
2262 &self,
2263 queue: &str,
2264 key: Option<&str>,
2265 now: u64,
2266 snapshot: Option<&SnapshotContext>,
2267 mut visit: impl FnMut(&UnifiedEntity) -> bool,
2268 ) {
2269 let mut visible = |entity: &UnifiedEntity| {
2270 (key.is_some() && !Self::scan_entity_visible(snapshot, entity)) || visit(entity)
2271 };
2272 if let Some(growing) = self.growing.read().as_ref() {
2273 if !growing
2274 .read()
2275 .visit_queue_dedup(queue, key, now, &mut visible)
2276 {
2277 return;
2278 }
2279 }
2280 for segment in self.sealed.read().iter() {
2281 if !segment
2282 .read()
2283 .visit_queue_dedup(queue, key, now, &mut visible)
2284 {
2285 return;
2286 }
2287 }
2288 }
2289
2290 #[doc(hidden)]
2292 pub fn scan_graph_kind(
2293 &self,
2294 snapshot: Option<&SnapshotContext>,
2295 kind: GraphEntityKind,
2296 ) -> Vec<UnifiedEntity> {
2297 self.query_all_in_segments(
2298 |segment| segment.may_contain_graph_kind(kind),
2299 |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
2300 )
2301 }
2302
2303 #[doc(hidden)]
2307 pub fn scan_graph_batches(
2308 &self,
2309 snapshot: Option<&SnapshotContext>,
2310 kind: GraphEntityKind,
2311 mut before_work: impl FnMut() -> bool,
2312 mut consume: impl FnMut(Vec<UnifiedEntity>) -> bool,
2313 ) -> bool {
2314 let segments = {
2315 let growing = self.growing.read();
2316 let sealed = self.sealed.read();
2317 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
2318 for segment in growing.iter().chain(sealed.iter()) {
2319 if !before_work() {
2320 return false;
2321 }
2322 let view = segment.read();
2323 if view.may_contain_graph_kind(kind) {
2324 segments.push((Arc::clone(segment), view.scan_cursor()));
2325 }
2326 }
2327 segments
2328 };
2329 let mut ids = [EntityId::new(0); super::segment::SCAN_BATCH_SIZE];
2330 for (segment, mut cursor) in segments {
2331 while !cursor.finished() {
2332 let mut count = 0;
2333 {
2334 let view = segment.read();
2335 if !view.scan_batch(&mut cursor, &mut before_work, &mut |id| {
2336 ids[count] = id;
2337 count += 1;
2338 }) {
2339 return false;
2340 }
2341 }
2342 if !before_work() {
2343 return false;
2344 }
2345 let mut batch = vec![None; count];
2348 self.for_each_id(&ids[..count], |position, entity| {
2349 if kind.matches(&entity.kind) && Self::scan_entity_visible(snapshot, entity) {
2350 batch[position] = Some(entity.clone());
2351 }
2352 });
2353 if !before_work() || !consume(batch.into_iter().flatten().collect()) {
2354 return false;
2355 }
2356 }
2357 }
2358 true
2359 }
2360
2361 #[doc(hidden)]
2366 pub fn scan_graph_candidates_for_each<F>(
2367 &self,
2368 snapshot: Option<&SnapshotContext>,
2369 kind: GraphEntityKind,
2370 before_segment: impl FnMut() -> bool,
2371 mut callback: F,
2372 ) where
2373 F: FnMut(&UnifiedEntity) -> bool,
2374 {
2375 self.for_each_entity_in_segments(
2376 |segment| segment.may_contain_graph_kind(kind),
2377 before_segment,
2378 |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
2379 );
2380 }
2381
2382 pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
2387 where
2388 F: FnMut(&UnifiedEntity) -> bool,
2389 {
2390 self.for_each_entity(|entity| {
2391 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2392 });
2393 }
2394
2395 #[doc(hidden)]
2396 pub fn physical_scan_cursor(&self) -> PhysicalScanCursor {
2397 let growing = self.growing.read();
2398 let sealed = self.sealed.read();
2399 PhysicalScanCursor {
2400 sources: growing
2401 .iter()
2402 .chain(sealed.iter())
2403 .map(|segment| (Arc::downgrade(segment), segment.read().scan_cursor()))
2404 .collect(),
2405 source: 0,
2406 }
2407 }
2408
2409 #[doc(hidden)]
2412 pub fn physical_scan_batch(
2413 &self,
2414 cursor: &mut PhysicalScanCursor,
2415 work_left: &mut usize,
2416 ) -> Vec<EntityId> {
2417 let mut ids = Vec::with_capacity(super::segment::SCAN_BATCH_SIZE.min(*work_left));
2418 while !cursor.finished() && *work_left > 0 {
2419 let (source, position) = &mut cursor.sources[cursor.source];
2420 if position.finished() {
2421 *work_left -= 1;
2422 cursor.source += 1;
2423 continue;
2424 }
2425 let Some(segment) = source.upgrade() else {
2426 *work_left -= 1;
2427 cursor.source += 1;
2428 continue;
2429 };
2430 segment.read().scan_batch(
2431 position,
2432 &mut || {
2433 if *work_left == 0 {
2434 return false;
2435 }
2436 *work_left -= 1;
2437 true
2438 },
2439 &mut |id| ids.push(id),
2440 );
2441 if position.finished() {
2442 cursor.source += 1;
2443 }
2444 if !ids.is_empty() {
2445 break;
2446 }
2447 }
2448 ids
2449 }
2450
2451 #[doc(hidden)]
2457 pub fn scan_batches(
2458 &self,
2459 snapshot: Option<&SnapshotContext>,
2460 mut before_work: impl FnMut() -> bool,
2461 mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
2462 ) -> bool {
2463 let segments = {
2464 let growing = self.growing.read();
2465 let sealed = self.sealed.read();
2466 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
2467 for segment in growing.iter().chain(sealed.iter()) {
2468 if !before_work() {
2469 return false;
2470 }
2471 let cursor = segment.read().scan_cursor();
2472 segments.push((Arc::clone(segment), cursor));
2473 }
2474 segments
2475 };
2476 let mut ids = Vec::new();
2477 let mut batch = Vec::new();
2478 for (segment, mut cursor) in segments {
2479 while !cursor.finished() {
2480 ids.clear();
2481 batch.clear();
2482 let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
2483 ids.reserve_exact(batch_size);
2484 batch.reserve_exact(batch_size);
2485 {
2486 let segment = segment.read();
2487 if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
2488 return false;
2489 }
2490 }
2491 if !before_work() {
2494 return false;
2495 }
2496 self.for_each_id(&ids, |_, entity| {
2497 if Self::scan_entity_visible(snapshot, entity) {
2498 batch.push(entity.clone());
2499 }
2500 });
2501 if !before_work() {
2502 return false;
2503 }
2504 if !consume(&batch) {
2505 return false;
2506 }
2507 }
2508 }
2509 true
2510 }
2511
2512 pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
2514 &self,
2515 snapshot: Option<&SnapshotContext>,
2516 init: FInit,
2517 fold: FFold,
2518 reduce: FReduce,
2519 ) -> T
2520 where
2521 T: Send,
2522 FInit: Fn() -> T + Send + Sync,
2523 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
2524 FReduce: Fn(T, T) -> T + Send + Sync,
2525 {
2526 self.fold_entities_parallel(
2527 init,
2528 |acc, entity| {
2529 if Self::scan_entity_visible(snapshot, entity) {
2530 fold(acc, entity)
2531 } else {
2532 acc
2533 }
2534 },
2535 reduce,
2536 )
2537 }
2538
2539 pub fn scan_for_each_zoned_with_stats<F>(
2541 &self,
2542 snapshot: Option<&SnapshotContext>,
2543 zone_preds: &[(&str, ZoneColPred<'_>)],
2544 mut callback: F,
2545 ) -> SegmentScanStats
2546 where
2547 F: FnMut(&UnifiedEntity) -> bool,
2548 {
2549 self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
2550 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2551 })
2552 }
2553
2554 pub fn scan_zoned_with_stats<F>(
2556 &self,
2557 snapshot: Option<&SnapshotContext>,
2558 zone_preds: &[(&str, ZoneColPred<'_>)],
2559 filter: F,
2560 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2561 where
2562 F: Fn(&UnifiedEntity) -> bool + Sync,
2563 {
2564 self.query_all_zoned_with_stats(zone_preds, |entity| {
2565 Self::scan_entity_visible(snapshot, entity) && filter(entity)
2566 })
2567 }
2568
2569 pub fn scan_project_zoned_with_stats<T, F>(
2573 &self,
2574 snapshot: Option<&SnapshotContext>,
2575 zone_preds: &[(&str, ZoneColPred<'_>)],
2576 project: F,
2577 ) -> (Vec<T>, SegmentScanStats)
2578 where
2579 T: Send,
2580 F: Fn(&UnifiedEntity) -> Option<T> + Sync,
2581 {
2582 self.project_all_zoned_with_stats(zone_preds, |entity| {
2583 if Self::scan_entity_visible(snapshot, entity) {
2584 project(entity)
2585 } else {
2586 None
2587 }
2588 })
2589 }
2590
2591 pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
2594 where
2595 F: Fn(&UnifiedEntity) -> bool + Sync,
2596 {
2597 let mut entities = self.query_all_in_segments(|_| true, filter);
2598 self.restore_owned_row_schemas(entities.iter_mut());
2599 entities
2600 }
2601
2602 fn query_all_in_segments<F>(
2603 &self,
2604 segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
2605 filter: F,
2606 ) -> Vec<UnifiedEntity>
2607 where
2608 F: Fn(&UnifiedEntity) -> bool + Sync,
2609 {
2610 let mut results = Vec::new();
2611
2612 if let Some(growing_arc) = self.growing.read().as_ref() {
2615 let growing = if let Some(g) = growing_arc.try_read() {
2616 g
2617 } else {
2618 growing_arc.read()
2619 };
2620 if segment_filter(&growing) {
2621 results.extend(growing.iter().filter(|e| filter(e)).cloned());
2622 }
2623 }
2624
2625 let sealed = self.sealed.read();
2627 let use_parallel = sealed.len() > 1 && crate::execution::should_parallelize();
2628 if use_parallel {
2629 let filter_ref = &filter;
2630 let segment_filter_ref = &segment_filter;
2631 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
2632 sealed
2633 .iter()
2634 .map(|segment| {
2635 s.spawn(move || {
2636 let segment = segment.read();
2637 if segment_filter_ref(&segment) {
2638 segment
2639 .iter()
2640 .filter(|e| filter_ref(e))
2641 .cloned()
2642 .collect::<Vec<_>>()
2643 } else {
2644 Vec::new()
2645 }
2646 })
2647 })
2648 .collect::<Vec<_>>()
2649 .into_iter()
2650 .map(|handle| handle.join().unwrap_or_default())
2651 .collect()
2652 });
2653 for batch in segment_results {
2654 results.extend(batch);
2655 }
2656 } else {
2657 for segment in sealed.iter() {
2658 let seg = segment.read();
2659 if segment_filter(&seg) {
2660 results.extend(seg.iter().filter(|e| filter(e)).cloned());
2661 }
2662 }
2663 }
2664
2665 results
2666 }
2667
2668 pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
2673 if let Some(growing_arc) = self.growing.read().as_ref() {
2674 let growing = growing_arc.read();
2675 if growing.may_contain_exact_key(key) {
2676 return true;
2677 }
2678 } else {
2679 return true;
2680 }
2681
2682 let sealed = self.sealed.read();
2683 for segment_arc in sealed.iter() {
2684 let segment = segment_arc.read();
2685 if segment.may_contain_exact_key(key) {
2686 return true;
2687 }
2688 }
2689
2690 false
2691 }
2692
2693 pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2695 let mut results = Vec::new();
2696
2697 if let Some(growing_arc) = self.growing.read().as_ref() {
2699 let growing = growing_arc.read();
2700 results.extend(growing.filter_metadata(filters));
2701 }
2702
2703 let sealed = self.sealed.read();
2705 for segment in sealed.iter() {
2706 results.extend(segment.filter_metadata(filters));
2707 }
2708
2709 results
2710 }
2711
2712 pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2714 let mut results = Vec::new();
2715
2716 if let Some(growing_arc) = self.growing.read().as_ref() {
2718 let growing = growing_arc.read();
2719 for entity in growing.iter_kind(kind) {
2720 results.push(entity.clone());
2721 }
2722 }
2723
2724 let sealed = self.sealed.read();
2729 for segment_arc in sealed.iter() {
2730 let segment = segment_arc.read();
2731 for entity in segment.iter_kind(kind) {
2732 results.push(entity.clone());
2733 }
2734 }
2735
2736 results
2737 }
2738
2739 pub fn count(&self) -> usize {
2741 self.total_entities_atomic.load(Ordering::Relaxed) as usize
2742 }
2743
2744 pub fn segment_ids(&self) -> Vec<SegmentId> {
2746 let mut ids = Vec::new();
2747
2748 if let Some(growing_arc) = self.growing.read().as_ref() {
2749 ids.push(growing_arc.read().id());
2750 }
2751
2752 let sealed = self.sealed.read();
2753 for segment in sealed.iter() {
2754 ids.push(segment.id());
2755 }
2756
2757 ids.extend(self.archived.read().iter().copied());
2758
2759 ids
2760 }
2761
2762 #[inline]
2775 #[allow(clippy::unused_self)]
2776 fn emit(&self, _event: LifecycleEvent) {}
2777
2778 pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2781 std::mem::take(&mut *self.events.write())
2782 }
2783
2784 pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2786 self.maybe_seal_growing()?;
2788
2789 if self.config.enable_consolidation {
2790 self.consolidation_tick();
2791 }
2792
2793 Ok(())
2794 }
2795
2796 fn consolidation_tick(&self) {
2807 let running = self.consolidation.read().is_some();
2810 if !running && !self.start_consolidation() {
2811 return;
2812 }
2813
2814 let mut guard = self.consolidation.write();
2817 let Some(run) = guard.as_mut() else {
2818 return;
2819 };
2820 self.copy_bounded(run);
2821 if run.cursor >= run.sources.len() {
2822 if let Some(run) = guard.take() {
2823 self.finish_consolidation(run);
2824 }
2825 }
2826 }
2827
2828 fn start_consolidation(&self) -> bool {
2832 let sources = self.select_candidates();
2833 self.start_consolidation_from_sources(sources)
2834 }
2835
2836 fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
2837 if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
2838 return false;
2839 }
2840
2841 let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
2842 let merged = GrowingSegment::new(merged_id, &self.collection);
2843
2844 *self.consolidation.write() = Some(Consolidation {
2845 sources: sources
2846 .into_iter()
2847 .map(|id| SourceProgress {
2848 id,
2849 copied: Vec::new(),
2850 tombstones_at_open: None,
2851 })
2852 .collect(),
2853 merged_id,
2854 merged,
2855 cursor: 0,
2856 pending_ids: Vec::new(),
2857 pending_cursor: 0,
2858 entities_copied: 0,
2859 });
2860
2861 self.stats.write().consolidation.runs_started += 1;
2862 true
2863 }
2864
2865 pub fn pressure_consolidation_tick(&self) -> bool {
2870 if !self.config.enable_consolidation {
2871 return false;
2872 }
2873
2874 let running = self.consolidation.read().is_some();
2875 if !running {
2876 if self
2880 .seal_current_if(|segment| segment.reclaimable_bytes() > 0)
2881 .is_err()
2882 {
2883 return false;
2884 }
2885 let sources = self.dirty_sealed_segments();
2886 if !self.start_consolidation_from_sources(sources) {
2887 return false;
2888 }
2889 }
2890
2891 self.consolidation_tick();
2892 true
2893 }
2894
2895 fn select_candidates(&self) -> Vec<SegmentId> {
2910 let sealed = self.sealed.read();
2911 if sealed.is_empty() {
2912 return Vec::new();
2913 }
2914
2915 if sealed.len() > self.config.max_sealed_segments {
2916 return sealed.iter().map(|segment| segment.read().id()).collect();
2917 }
2918
2919 let mut live = 0usize;
2920 let mut tombstones = 0usize;
2921 let mut held = 0u64;
2922 let mut reclaimable = 0u64;
2923 let mut dirty = Vec::new();
2924
2925 for segment_arc in sealed.iter() {
2926 let segment = segment_arc.read();
2927 live += segment.live_entity_count();
2928 tombstones += segment.tombstone_count();
2929 held += segment.resident_bytes();
2930 reclaimable += segment.reclaimable_bytes();
2931 if segment.tombstone_count() > 0 {
2932 dirty.push(segment.id());
2933 }
2934 }
2935
2936 let total = live + tombstones;
2937 let tombstone_ratio = if total == 0 {
2938 0.0
2939 } else {
2940 tombstones as f64 / total as f64
2941 };
2942 let fragmentation = if held == 0 {
2943 0.0
2944 } else {
2945 reclaimable as f64 / held as f64
2946 };
2947
2948 if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
2949 || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
2950 {
2951 return dirty;
2952 }
2953
2954 Vec::new()
2955 }
2956
2957 fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
2958 self.sealed
2959 .read()
2960 .iter()
2961 .filter_map(|segment_arc| {
2962 let segment = segment_arc.read();
2963 (segment.reclaimable_bytes() > 0).then(|| segment.id())
2964 })
2965 .collect()
2966 }
2967
2968 fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
2971 if sources.is_empty() {
2972 return false;
2973 }
2974 let sealed = self.sealed.read();
2975 sealed.iter().any(|segment_arc| {
2976 let segment = segment_arc.read();
2977 sources.contains(&segment.id()) && segment.tombstone_count() > 0
2978 })
2979 }
2980
2981 fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
2982 let sealed = self.sealed.read();
2983 sealed
2984 .iter()
2985 .find(|segment| segment.read().id() == id)
2986 .map(Arc::clone)
2987 }
2988
2989 fn copy_bounded(&self, run: &mut Consolidation) {
2994 let mut budget = self.config.consolidation_entities_per_tick;
2995
2996 while budget > 0 && run.cursor < run.sources.len() {
2997 let source_id = run.sources[run.cursor].id;
2998 let Some(source_arc) = self.sealed_segment(source_id) else {
2999 run.cursor += 1;
3003 run.pending_ids.clear();
3004 run.pending_cursor = 0;
3005 continue;
3006 };
3007
3008 if run.sources[run.cursor].tombstones_at_open.is_none() {
3009 let source = source_arc.read();
3010 run.pending_ids = source.live_entity_ids();
3011 run.pending_cursor = 0;
3012 run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
3013 }
3014
3015 {
3016 let source = source_arc.read();
3017 while budget > 0 && run.pending_cursor < run.pending_ids.len() {
3018 let id = run.pending_ids[run.pending_cursor];
3019 run.pending_cursor += 1;
3020 budget -= 1;
3021
3022 let Some(entity) = source.get(id) else {
3025 continue;
3026 };
3027 let metadata = source.get_metadata(id);
3028 run.merged.adopt_entity(entity.clone(), metadata);
3029 run.sources[run.cursor].copied.push(id);
3030 run.entities_copied += 1;
3031 }
3032 }
3033
3034 if run.pending_cursor >= run.pending_ids.len() {
3035 run.cursor += 1;
3036 run.pending_ids.clear();
3037 run.pending_cursor = 0;
3038 }
3039 }
3040 }
3041
3042 fn finish_consolidation(&self, mut run: Consolidation) {
3048 let mut sealed = self.sealed.write();
3049
3050 let mut positions = Vec::with_capacity(run.sources.len());
3051 for source in &run.sources {
3052 if let Some(index) = sealed
3053 .iter()
3054 .position(|segment| segment.read().id() == source.id)
3055 {
3056 positions.push(index);
3057 }
3058 }
3059 if positions.is_empty() {
3060 return;
3061 }
3062 positions.sort_unstable();
3063
3064 let mut stale = Vec::new();
3068 for source in &run.sources {
3069 let Some(source_arc) = sealed
3070 .iter()
3071 .find(|segment| segment.read().id() == source.id)
3072 else {
3073 continue;
3074 };
3075 let segment = source_arc.read();
3076 if source.tombstones_at_open == Some(segment.tombstone_count()) {
3077 continue;
3078 }
3079 stale.extend(
3080 source
3081 .copied
3082 .iter()
3083 .copied()
3084 .filter(|id| !segment.contains(*id)),
3085 );
3086 }
3087 for id in stale {
3088 run.merged.evict_entity(id);
3089 }
3090
3091 let reclaimed_bytes: u64 = positions
3092 .iter()
3093 .map(|&index| sealed[index].read().resident_bytes())
3094 .sum();
3095 let reclaimed_tombstones: u64 = positions
3096 .iter()
3097 .map(|&index| sealed[index].read().tombstone_count() as u64)
3098 .sum();
3099
3100 if run.merged.seal().is_err() {
3101 return;
3102 }
3103 let merged_bytes = run.merged.resident_bytes();
3104 let merged_arc = Arc::new(RwLock::new(run.merged));
3105
3106 {
3107 let mut retired = self.retired.write();
3108 retired.retain(|source| source.strong_count() > 0);
3109 for &index in positions.iter().rev() {
3110 let source = sealed.remove(index);
3111 if Arc::strong_count(&source) > 1 {
3112 retired.push(Arc::downgrade(&source));
3113 }
3114 }
3115 }
3116 let insert_at = positions[0].min(sealed.len());
3117 sealed.insert(insert_at, Arc::clone(&merged_arc));
3118 drop(sealed);
3119
3120 let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
3122 {
3123 let mut entity_segment = self.entity_segment.write();
3124 for segment_id in entity_segment.values_mut() {
3125 if source_ids.contains(segment_id) {
3126 *segment_id = run.merged_id;
3127 }
3128 }
3129 }
3130
3131 {
3132 let mut stats = self.stats.write();
3133 stats.sealed_count = stats
3134 .sealed_count
3135 .saturating_sub(positions.len())
3136 .saturating_add(1);
3137 let consolidation = &mut stats.consolidation;
3138 consolidation.runs_completed += 1;
3139 consolidation.segments_merged += positions.len() as u64;
3140 consolidation.tombstones_reclaimed += reclaimed_tombstones;
3141 consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
3142 }
3143
3144 self.emit(LifecycleEvent::SegmentConsolidated {
3145 source: source_ids,
3146 target: run.merged_id,
3147 });
3148 }
3149
3150 #[cfg(test)]
3153 fn consolidation_progress(&self) -> Option<u64> {
3154 self.consolidation
3155 .read()
3156 .as_ref()
3157 .map(|run| run.entities_copied)
3158 }
3159}
3160
3161#[cfg(test)]
3162mod tests {
3163 use std::collections::HashSet;
3164
3165 use super::*;
3166 use crate::storage::entity::{EntityData, EntityKind, RowData};
3167 use crate::transaction::snapshot::{Snapshot, SnapshotManager};
3168 use reddb_types::Value;
3169
3170 #[test]
3171 fn recovery_view_projection_keeps_schema_bounds_without_retaining_rows_or_identities() {
3172 for collection in ["rows", "red_internal"] {
3173 let manager = SegmentManager::new(collection);
3174 let identities = manager.identity_counters();
3175 for id in 1..=18 {
3176 let field = format!("field_{id}");
3177 let mut entity =
3178 UnifiedEntity::table_row(EntityId::new(id), collection, id, vec![]);
3179 entity.data = EntityData::Row(RowData {
3180 columns: Vec::new(),
3181 named: Some(HashMap::from([
3182 (field.clone(), Value::Integer(id as i64)),
3183 ("id".into(), Value::Integer(id as i64)),
3184 ])),
3185 schema: None,
3186 });
3187 manager.restore_snapshot_row_layout(&mut entity);
3188 let row = entity.data.as_row().expect("row");
3189 assert_eq!(row.get_field(&field), Some(&Value::Integer(id as i64)));
3190 assert_eq!(row.get_field("id"), Some(&Value::Integer(id as i64)));
3191 assert_eq!(
3192 row.named.is_some(),
3193 collection.starts_with("red_") || id > 16
3194 );
3195 }
3196 assert_eq!(manager.count(), 0, "a projection never adopts payloads");
3197 assert_eq!(manager.identity_counters(), identities);
3198 assert_eq!(
3199 manager.recovery_row_schemas.read().len(),
3200 if collection == "rows" { 16 } else { 0 }
3201 );
3202 }
3203 }
3204
3205 #[test]
3206 fn physical_pages_keep_global_order_across_segments_and_deleted_slots() {
3207 let manager = SegmentManager::new("pages");
3208 for group in [[50, 10], [30, 20]] {
3209 for id in group {
3210 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "pages", id, vec![]);
3211 entity.data = EntityData::Row(RowData::with_names(
3212 vec![Value::Integer(id as i64)],
3213 vec!["id".into()],
3214 ));
3215 manager.insert(entity).unwrap();
3216 }
3217 manager.force_seal().unwrap();
3218 }
3219 manager
3220 .insert(UnifiedEntity::table_row(
3221 EntityId::new(40),
3222 "pages",
3223 40,
3224 vec![],
3225 ))
3226 .unwrap();
3227 let (total, page) = manager.page_in_id_order(1, 2);
3228 assert_eq!(total, 5);
3229 assert_eq!(
3230 page.iter()
3231 .map(|entity| entity.id.raw())
3232 .collect::<Vec<_>>(),
3233 [20, 30]
3234 );
3235 let EntityData::Row(row) = &page[0].data else {
3236 panic!("row")
3237 };
3238 assert_eq!(row.get_field("id"), Some(&Value::Integer(20)));
3239 manager.delete(EntityId::new(30)).unwrap();
3240 let (total, page) = manager.page_in_id_order(1, 2);
3241 assert_eq!(total, 4);
3242 assert_eq!(
3243 page.iter()
3244 .map(|entity| entity.id.raw())
3245 .collect::<Vec<_>>(),
3246 [20, 40]
3247 );
3248 for offset in [4, usize::MAX] {
3249 let (total, page) = manager.page_in_id_order(offset, 2);
3250 assert_eq!(total, 4);
3251 assert!(page.is_empty());
3252 }
3253 }
3254
3255 fn lookup_version(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
3256 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
3257 entity.data = EntityData::Row(RowData::with_names(
3258 vec![Value::text(key)],
3259 vec!["key".into()],
3260 ));
3261 entity.set_xmin(xmin);
3262 entity.set_xmax(xmax);
3263 entity
3264 }
3265
3266 #[test]
3267 fn key_lookup_matches_snapshot_scan_across_segments() {
3268 let manager = Arc::new(SegmentManager::new("keys"));
3269 for entity in [
3270 lookup_version(1, "changed", 3, 12),
3271 lookup_version(2, "deleted", 3, 7),
3272 ] {
3273 manager.insert(entity).expect("old version");
3274 }
3275 manager.force_seal().expect("seal old versions");
3276 for entity in [
3277 lookup_version(3, "changed", 12, 0),
3278 lookup_version(4, "pending", 4, 0),
3279 lookup_version(5, "own", 15, 0),
3280 lookup_version(6, "live", 3, 0),
3281 ] {
3282 manager.insert(entity).expect("new version");
3283 }
3284 std::thread::spawn(move || {
3285 let snapshot = SnapshotContext {
3286 snapshot: Snapshot {
3287 xid: 10,
3288 in_progress: HashSet::from([4]),
3289 },
3290 manager: Arc::new(SnapshotManager::new()),
3291 own_xids: HashSet::from([15]),
3292 requires_index_fallback: false,
3293 serializable_reader: None,
3294 };
3295 for snapshot in [Some(&snapshot), None] {
3296 for key in ["changed", "deleted", "pending", "own", "live", "absent"] {
3297 let expected = manager
3298 .scan(snapshot, |entity| {
3299 crate::storage::segment_key_lookup::SegmentKeyLookup::key(entity)
3300 == Some(key)
3301 })
3302 .into_iter()
3303 .max_by_key(|entity| entity.xmin)
3304 .map(|entity| entity.id);
3305 assert_eq!(
3306 manager.project_key_entity(key, snapshot, true, |entity| entity.id),
3307 expected,
3308 "key={key}"
3309 );
3310 }
3311 }
3312 assert_eq!(
3313 manager.project_key_entity("changed", Some(&snapshot), true, |entity| entity.id),
3314 Some(EntityId::new(1))
3315 );
3316 assert_eq!(
3317 manager.project_key_entity("changed", None, true, |entity| entity.id),
3318 Some(EntityId::new(3))
3319 );
3320 assert!(manager
3321 .project_key_entity("deleted", None, true, |entity| entity.id)
3322 .is_none());
3323 })
3324 .join()
3325 .expect("snapshot worker");
3326 }
3327
3328 #[test]
3329 fn key_lookup_survives_paced_consolidation() {
3330 let manager = consolidating_manager(1);
3331 manager
3332 .bulk_insert(
3333 (1..=10)
3334 .map(|id| lookup_version(id, &format!("key-{id}"), 1, 0))
3335 .collect(),
3336 )
3337 .expect("rows");
3338 manager.force_seal().expect("seal");
3339 for id in 1..=4 {
3340 manager.delete(EntityId::new(id)).expect("delete");
3341 }
3342 assert!(manager.pressure_consolidation_tick());
3343 drain_maintenance(&manager);
3344 for id in 1..=10 {
3345 assert_eq!(
3346 manager.project_key_entity(&format!("key-{id}"), None, false, |entity| entity.id),
3347 (id > 4).then_some(EntityId::new(id))
3348 );
3349 }
3350 }
3351
3352 #[test]
3353 #[ignore = "diagnostic of segment fragmentation and retained history; no timing assertions"]
3354 fn key_lookup_fragmentation_history_diagnostic() {
3355 use std::{hint::black_box, time::Instant};
3356 for segments in [1, 16, 64] {
3357 for versions in [1, 100, 1000] {
3358 let manager = SegmentManager::new("keys");
3359 let mut id = 1;
3360 for segment in 0..segments {
3361 let mut rows: Vec<_> = (0..128)
3362 .map(|_| {
3363 let entity = lookup_version(id, &format!("other-{id}"), 1, 0);
3364 id += 1;
3365 entity
3366 })
3367 .collect();
3368 for version in (segment..versions).step_by(segments) {
3369 rows.push(lookup_version(
3370 id,
3371 "history",
3372 version as u64 + 1,
3373 if version + 1 == versions {
3374 0
3375 } else {
3376 version as u64 + 2
3377 },
3378 ));
3379 id += 1;
3380 }
3381 manager.bulk_insert(rows).expect("rows");
3382 manager.force_seal().expect("seal");
3383 }
3384 let started = Instant::now();
3385 for _ in 0..2000 {
3386 let xmin =
3387 manager.project_key_entity(black_box("history"), None, true, |entity| {
3388 entity.xmin
3389 });
3390 assert_eq!(xmin, Some(versions as u64));
3391 }
3392 eprintln!("segments={segments} versions={versions} reads=2000 elapsed_us={} resident_bytes={}", started.elapsed().as_micros(), manager.resident_bytes());
3393 }
3394 }
3395 }
3396
3397 #[test]
3398 fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
3399 let manager = SegmentManager::new("fingerprint");
3400 let row = |id, label: &str| {
3401 UnifiedEntity::table_row(
3402 EntityId::new(id),
3403 "fingerprint",
3404 id,
3405 vec![Value::text(label)],
3406 )
3407 };
3408 manager
3409 .bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
3410 .expect("sealed rows");
3411 manager.force_seal().expect("seal");
3412 manager
3413 .bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
3414 .expect("growing rows");
3415
3416 let mut expected = manager.query_all(|_| true);
3417 expected.sort_by_key(|entity| entity.id.raw());
3418 let expected: Vec<_> = expected
3419 .iter()
3420 .map(|entity| format!("{entity:?}"))
3421 .collect();
3422 let mut stored_addresses = HashSet::new();
3423 manager.for_each_entity(|entity| {
3424 stored_addresses.insert(entity as *const UnifiedEntity as usize);
3425 true
3426 });
3427
3428 let mut actual = Vec::new();
3429 manager.for_each_entity_in_id_order(|entity| {
3430 assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
3431 assert!(
3432 manager.growing.try_write().is_some(),
3433 "topology guard retained"
3434 );
3435 assert!(
3436 manager.sealed.try_write().is_some(),
3437 "topology guard retained"
3438 );
3439 actual.push(format!("{entity:?}"));
3440 true
3441 });
3442 assert_eq!(
3443 actual, expected,
3444 "equal IDs retain their original segment order"
3445 );
3446 }
3447
3448 #[test]
3449 fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
3450 let manager = SegmentManager::new("fingerprint");
3451 for id in [3, 1, 2] {
3452 manager
3453 .insert(UnifiedEntity::table_row(
3454 EntityId::new(id),
3455 "fingerprint",
3456 id,
3457 vec![Value::Integer(id as i64)],
3458 ))
3459 .expect("row");
3460 }
3461 let source = manager.growing.read().as_ref().expect("source").clone();
3462 let handles = Arc::strong_count(&source);
3463 let mut seen = Vec::new();
3464 manager.for_each_entity_in_id_order(|entity| {
3465 seen.push(entity.id.raw());
3466 false
3467 });
3468 assert_eq!(seen, vec![1]);
3469 assert_eq!(Arc::strong_count(&source), handles);
3470 assert!(
3471 source.try_write().is_some(),
3472 "segment guard retained after early stop"
3473 );
3474 }
3475
3476 #[test]
3477 fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
3478 let manager = SegmentManager::with_config(
3479 "vectors",
3480 ManagerConfig {
3481 max_sealed_segments: 1,
3482 consolidation_entities_per_tick: 10_000,
3483 ..Default::default()
3484 },
3485 );
3486 for part in 0..3 {
3487 manager
3488 .bulk_insert(
3489 (1..=300)
3490 .map(|offset| {
3491 UnifiedEntity::vector(
3492 EntityId::new(part * 300 + offset),
3493 "vectors",
3494 vec![1.0, 0.0],
3495 )
3496 })
3497 .collect(),
3498 )
3499 .expect("batch");
3500 if part < 2 {
3501 manager.force_seal().expect("seal");
3502 }
3503 }
3504 let original_growing = manager.growing.read().as_ref().expect("growing").clone();
3505 let mut changed = false;
3506 let mut seen = Vec::new();
3507 assert!(manager.scan_batches(
3508 None,
3509 || true,
3510 |batch| {
3511 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3512 assert!(
3513 manager.growing.try_write().is_some(),
3514 "consumer holds no topology guard"
3515 );
3516 assert!(
3517 original_growing.try_write().is_some(),
3518 "consumer holds no segment guard"
3519 );
3520 for entity in batch {
3521 assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
3522 seen.push(entity.id.raw());
3523 }
3524 if !changed {
3525 changed = true;
3526 manager.force_seal().expect("seal during scan");
3527 manager
3528 .insert(UnifiedEntity::vector(
3529 EntityId::new(5000),
3530 "vectors",
3531 vec![1.0, 0.0],
3532 ))
3533 .expect("new write");
3534 drain_maintenance(&manager);
3535 assert!(manager.stats().consolidation.runs_completed > 0);
3536 assert!(manager
3537 .delete(EntityId::new(1))
3538 .expect("delete after source retirement"));
3539 assert!(manager.get(EntityId::new(1)).is_none());
3540 }
3541 true
3542 }
3543 ));
3544 seen.sort_unstable();
3545 assert_eq!(
3546 seen,
3547 (2..=900).collect::<Vec<_>>(),
3548 "no duplicate, missing or late candidates"
3549 );
3550 let mut after = Vec::new();
3551 assert!(manager.scan_batches(
3552 None,
3553 || true,
3554 |batch| {
3555 after.extend(batch.iter().map(|entity| entity.id.raw()));
3556 true
3557 }
3558 ));
3559 after.sort_unstable();
3560 let mut expected: Vec<_> = (2..=900).collect();
3561 expected.push(5000);
3562 assert_eq!(after, expected, "adopted entities have scan positions");
3563 }
3564
3565 #[test]
3566 fn bounded_scan_applies_snapshot_to_relocated_versions() {
3567 for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
3568 let manager = SegmentManager::with_config(
3569 "vectors",
3570 ManagerConfig {
3571 max_sealed_segments: 1,
3572 ..Default::default()
3573 },
3574 );
3575 for ids in [[1, 2], [3, 4]] {
3576 let entities = ids.map(|id| {
3577 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3578 entity.set_xmin(3);
3579 entity
3580 });
3581 if bulk {
3582 manager
3583 .bulk_insert(entities.into())
3584 .expect("initial versions");
3585 } else {
3586 for entity in entities {
3587 manager.insert(entity).expect("initial version");
3588 }
3589 }
3590 manager.force_seal().expect("seal source");
3591 }
3592 let snapshot = SnapshotContext {
3593 snapshot: Snapshot {
3594 xid: 10,
3595 in_progress: HashSet::new(),
3596 },
3597 manager: Arc::new(SnapshotManager::new()),
3598 own_xids: HashSet::new(),
3599 requires_index_fallback: false,
3600 serializable_reader: None,
3601 };
3602 let mut changed = false;
3603 let mut seen = Vec::new();
3604 assert!(manager.scan_batches(
3605 Some(&snapshot),
3606 || true,
3607 |batch| {
3608 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3609 if !changed {
3610 changed = true;
3611 if consolidate {
3612 drain_maintenance(&manager);
3613 assert!(manager.stats().consolidation.runs_completed > 0);
3614 }
3615 for (id, xmax) in [(3, 8), (4, 12)] {
3616 let mut entity =
3617 manager.get(EntityId::new(id)).expect("retained version");
3618 entity.set_xmax(xmax);
3619 manager
3620 .update(entity)
3621 .expect("update after source retirement");
3622 }
3623 }
3624 true
3625 }
3626 ));
3627 seen.sort_unstable();
3628 assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
3629 let mut again = Vec::new();
3630 assert!(manager.scan_batches(
3631 Some(&snapshot),
3632 || true,
3633 |batch| {
3634 again.extend(batch.iter().map(|entity| entity.id.raw()));
3635 true
3636 }
3637 ));
3638 again.sort_unstable();
3639 assert_eq!(
3640 again,
3641 vec![1, 2, 4],
3642 "old slots do not duplicate relocated versions"
3643 );
3644 let mut current = Vec::new();
3645 assert!(manager.scan_batches(
3646 None,
3647 || true,
3648 |batch| {
3649 current.extend(batch.iter().map(|entity| entity.id.raw()));
3650 true
3651 }
3652 ));
3653 current.sort_unstable();
3654 assert_eq!(
3655 current,
3656 vec![1, 2],
3657 "current view excludes both deleted versions"
3658 );
3659 }
3660 }
3661
3662 #[test]
3663 fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
3664 let manager = SegmentManager::new("mixed");
3665 let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
3666 manager
3667 .bulk_insert((1..=300).map(vector).collect())
3668 .expect("flat");
3669 manager.insert(vector(500)).expect("hash insert");
3670 manager
3671 .bulk_insert(vec![vector(502), vector(505)])
3672 .expect("gapped bulk");
3673 manager.delete(EntityId::new(2)).expect("flat delete");
3674 manager.delete(EntityId::new(505)).expect("hash delete");
3675 let mut changed = false;
3676 let mut seen = Vec::new();
3677 assert!(manager.scan_batches(
3678 None,
3679 || true,
3680 |batch| {
3681 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3682 if !changed {
3683 changed = true;
3684 manager
3685 .insert(vector(700))
3686 .expect("append after cursor boundary");
3687 manager
3688 .update(UnifiedEntity::graph_node(
3689 EntityId::new(500),
3690 "mixed",
3691 "Node",
3692 HashMap::new(),
3693 ))
3694 .expect("structural update");
3695 }
3696 true
3697 }
3698 ));
3699 seen.sort_unstable();
3700 let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
3701 expected.extend([500, 502]);
3702 assert_eq!(seen, expected);
3703 }
3704
3705 #[test]
3706 fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
3707 let manager = SegmentManager::new("vectors");
3708 let mut entities = Vec::new();
3709 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
3710 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3711 entity.set_xmin(xmin);
3712 entity.set_xmax(xmax);
3713 entities.push(entity);
3714 }
3715 manager.bulk_insert(entities).expect("versions");
3716 manager.force_seal().expect("seal");
3717 let snapshot = SnapshotContext {
3718 snapshot: Snapshot {
3719 xid: 10,
3720 in_progress: HashSet::from([4]),
3721 },
3722 manager: Arc::new(SnapshotManager::new()),
3723 own_xids: HashSet::new(),
3724 requires_index_fallback: false,
3725 serializable_reader: None,
3726 };
3727 let mut work = 0;
3728 let mut callbacks = 0;
3729 assert!(!manager.scan_batches(
3730 Some(&snapshot),
3731 || {
3732 work += 1;
3733 work <= 2
3734 },
3735 |_| {
3736 callbacks += 1;
3737 true
3738 }
3739 ));
3740 assert_eq!(callbacks, 0, "a stopped batch is not consumed");
3741 for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
3742 let mut seen = Vec::new();
3743 assert!(manager.scan_batches(
3744 view,
3745 || true,
3746 |batch| {
3747 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3748 true
3749 }
3750 ));
3751 seen.sort_unstable();
3752 assert_eq!(seen, expected);
3753 }
3754 }
3755
3756 #[test]
3757 fn test_manager_basic() {
3758 let manager = SegmentManager::new("test_collection");
3759
3760 let entity = UnifiedEntity::table_row(
3761 manager.next_entity_id(),
3762 "users",
3763 1,
3764 vec![Value::text("Alice".to_string())],
3765 );
3766
3767 let id = manager.insert(entity).unwrap();
3768 assert!(manager.get(id).is_some());
3769 assert_eq!(manager.count(), 1);
3770 }
3771
3772 #[test]
3773 fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
3774 let manager = Arc::new(SegmentManager::new("mixed"));
3775 manager
3776 .bulk_insert(
3777 (1..=128)
3778 .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
3779 .collect(),
3780 )
3781 .expect("vector-only segment");
3782 manager.force_seal().expect("seal vectors");
3783 let mut old =
3784 UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
3785 old.set_xmin(3);
3786 old.set_xmax(12);
3787 let mut new =
3788 UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
3789 new.set_xmin(12);
3790 let edge = UnifiedEntity::graph_edge(
3791 EntityId::new(202),
3792 "mixed",
3793 "a",
3794 "b",
3795 1.0,
3796 Default::default(),
3797 );
3798 manager
3799 .bulk_insert(vec![old, new, edge])
3800 .expect("physical graph versions");
3801 manager.force_seal().expect("seal graph");
3802 manager
3803 .insert(UnifiedEntity::graph_node(
3804 EntityId::new(203),
3805 "mixed",
3806 "person",
3807 Default::default(),
3808 ))
3809 .expect("growing node");
3810 let snapshot = SnapshotContext {
3811 snapshot: Snapshot {
3812 xid: 10,
3813 in_progress: HashSet::new(),
3814 },
3815 manager: Arc::new(SnapshotManager::new()),
3816 own_xids: HashSet::new(),
3817 requires_index_fallback: false,
3818 serializable_reader: None,
3819 };
3820 std::thread::spawn(move || {
3821 for view in [None, Some(&snapshot)] {
3822 for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
3823 let expected: HashSet<_> = manager
3824 .scan(view, |item| kind.matches(&item.kind))
3825 .into_iter()
3826 .map(|item| item.id)
3827 .collect();
3828 let actual: HashSet<_> = manager
3829 .scan_graph_kind(view, kind)
3830 .into_iter()
3831 .map(|item| item.id)
3832 .collect();
3833 assert_eq!(actual, expected);
3834 let mut batched = HashSet::new();
3835 assert!(manager.scan_graph_batches(
3836 view,
3837 kind,
3838 || true,
3839 |batch| {
3840 batched.extend(batch.into_iter().map(|item| item.id));
3841 true
3842 }
3843 ));
3844 assert_eq!(batched, expected);
3845 let mut visited = 0;
3846 let mut candidates = HashSet::new();
3847 manager.scan_graph_candidates_for_each(
3848 view,
3849 kind,
3850 || true,
3851 |item| {
3852 visited += 1;
3853 if kind.matches(&item.kind) {
3854 candidates.insert(item.id);
3855 }
3856 true
3857 },
3858 );
3859 assert_eq!(candidates, expected);
3860 assert_eq!(
3861 visited,
3862 if matches!(kind, GraphEntityKind::Node) {
3863 3
3864 } else {
3865 2
3866 }
3867 );
3868 }
3869 }
3870 let ids: HashSet<_> = manager
3871 .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
3872 .into_iter()
3873 .map(|item| item.id.raw())
3874 .collect();
3875 assert_eq!(ids, HashSet::from([200, 203]));
3876 let ids: HashSet<_> = manager
3877 .scan_graph_kind(None, GraphEntityKind::Node)
3878 .into_iter()
3879 .map(|item| item.id.raw())
3880 .collect();
3881 assert_eq!(ids, HashSet::from([201, 203]));
3882 let mut visits = 0;
3883 manager.scan_graph_candidates_for_each(
3884 None,
3885 GraphEntityKind::Node,
3886 || true,
3887 |_| {
3888 visits += 1;
3889 false
3890 },
3891 );
3892 assert_eq!(
3893 visits, 1,
3894 "early termination must also stop subsequent segments"
3895 );
3896 })
3897 .join()
3898 .expect("snapshot worker");
3899 }
3900
3901 #[test]
3902 fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
3903 let manager = Arc::new(SegmentManager::new("accounts"));
3904 let cases = [
3905 ("legacy row", 0, 0, true),
3906 ("committed creator", 3, 0, true),
3907 ("future creator", 11, 0, false),
3908 ("in-progress creator", 4, 0, false),
3909 ("committed deleter", 3, 7, false),
3910 ("future deleter", 3, 11, true),
3911 ("in-progress deleter", 3, 4, true),
3912 ];
3913
3914 for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
3915 let row_id = index as u64 + 1;
3916 let mut entity = UnifiedEntity::table_row(
3917 EntityId::new(row_id),
3918 "accounts",
3919 row_id,
3920 vec![Value::Integer(row_id as i64)],
3921 );
3922 entity.set_xmin(*xmin);
3923 entity.set_xmax(*xmax);
3924 manager.insert(entity).unwrap();
3925 }
3926
3927 let snapshot = SnapshotContext {
3928 snapshot: Snapshot {
3929 xid: 10,
3930 in_progress: HashSet::from([4]),
3931 },
3932 manager: Arc::new(SnapshotManager::new()),
3933 own_xids: HashSet::new(),
3934 requires_index_fallback: false,
3935 serializable_reader: None,
3936 };
3937
3938 let raw_visible_row_ids: HashSet<u64> = manager
3939 .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
3940 .into_iter()
3941 .filter_map(|entity| match entity.kind {
3942 EntityKind::TableRow { row_id, .. } => Some(row_id),
3943 _ => None,
3944 })
3945 .collect();
3946
3947 let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
3948 .join()
3949 .unwrap();
3950 let visible_row_ids: HashSet<u64> = scanned
3951 .into_iter()
3952 .filter_map(|entity| match entity.kind {
3953 EntityKind::TableRow { row_id, .. } => Some(row_id),
3954 _ => None,
3955 })
3956 .collect();
3957 assert_eq!(visible_row_ids, raw_visible_row_ids);
3958
3959 for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
3960 let row_id = index as u64 + 1;
3961 assert_eq!(
3962 visible_row_ids.contains(&row_id),
3963 *expected_visible,
3964 "{name} visibility mismatch"
3965 );
3966 }
3967 }
3968
3969 #[test]
3970 fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
3971 let manager = SegmentManager::new("accounts");
3975 let mut live =
3976 UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
3977 live.set_xmin(2);
3978 live.set_xmax(0);
3979 manager.insert(live).unwrap();
3980
3981 let mut superseded =
3982 UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
3983 superseded.set_xmin(2);
3984 superseded.set_xmax(9);
3985 manager.insert(superseded).unwrap();
3986
3987 let rows: Vec<u64> = manager
3988 .scan(None, |_| true)
3989 .into_iter()
3990 .filter_map(|entity| match entity.kind {
3991 EntityKind::TableRow { row_id, .. } => Some(row_id),
3992 _ => None,
3993 })
3994 .collect();
3995 assert_eq!(
3996 rows,
3997 vec![1],
3998 "frameless scan must hide superseded versions"
3999 );
4000 }
4001
4002 #[test]
4003 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
4004 let manager = SegmentManager::new("test_collection");
4005
4006 let entity = UnifiedEntity::table_row(
4007 manager.next_entity_id(),
4008 "users",
4009 1,
4010 vec![Value::text("Alice".to_string())],
4011 );
4012
4013 let id = manager.insert(entity).unwrap();
4014
4015 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
4016 assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
4017 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
4018
4019 assert!(manager.delete(id).unwrap());
4020 assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
4021 }
4022
4023 #[test]
4024 fn zoned_scan_reports_pruned_segments_and_preserves_results() {
4025 let config = ManagerConfig {
4026 segment_config: SegmentConfig {
4027 max_entities: 2,
4028 ..Default::default()
4029 },
4030 ..Default::default()
4031 };
4032 let manager = SegmentManager::with_config("events", config);
4033
4034 for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
4035 let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
4036 let entity = UnifiedEntity::new(
4037 manager.next_entity_id(),
4038 EntityKind::TableRow {
4039 table: "events".into(),
4040 row_id: (idx + 1) as u64,
4041 },
4042 EntityData::Row(row),
4043 );
4044 manager.insert(entity).unwrap();
4045 }
4046
4047 let full = manager.query_all(|entity| match &entity.data {
4048 EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
4049 _ => false,
4050 });
4051 let probe = Value::Integer(100);
4052 let (pruned, stats) =
4053 manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
4054 match &entity.data {
4055 EntityData::Row(row) => {
4056 matches!(row.get_field("ts"), Some(Value::Integer(100)))
4057 }
4058 _ => false,
4059 }
4060 });
4061
4062 assert_eq!(pruned.len(), full.len());
4063 assert_eq!(pruned[0].id, full[0].id);
4064 assert_eq!(stats.segments_total, 3);
4065 assert_eq!(stats.segments_scanned, 2);
4066 assert_eq!(stats.segments_pruned, 1);
4067 }
4068
4069 #[test]
4070 fn projected_zoned_scan_preserves_segment_order_and_pruning() {
4071 let manager = SegmentManager::with_config(
4072 "events",
4073 ManagerConfig {
4074 segment_config: SegmentConfig {
4075 max_entities: 2,
4076 ..Default::default()
4077 },
4078 ..Default::default()
4079 },
4080 );
4081 for (index, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
4082 let mut entity = UnifiedEntity::table_row(
4083 EntityId::new(index as u64 + 1),
4084 "events",
4085 index as u64 + 1,
4086 vec![],
4087 );
4088 entity.data = EntityData::Row(RowData::with_names(
4089 vec![Value::Integer(value)],
4090 vec!["ts".into()],
4091 ));
4092 manager.insert(entity).expect("test row");
4093 }
4094 assert_eq!(manager.stats().sealed_count, 2);
4095 let project = |entity: &UnifiedEntity| match entity.data.as_row()?.get_field("ts")? {
4096 Value::Integer(value) => Some(*value),
4097 _ => None,
4098 };
4099 let (all, stats) = manager.scan_project_zoned_with_stats(None, &[], project);
4100 let (legacy, _) = manager.scan_zoned_with_stats(None, &[], |_| true);
4101 assert_eq!(all, legacy.iter().filter_map(project).collect::<Vec<_>>());
4102 assert_eq!(all[0], 200, "growing segment first");
4103 assert!(all[1..3].iter().all(|value| [10, 20].contains(value)));
4104 assert!(all[3..5].iter().all(|value| [100, 110].contains(value)));
4105 assert_eq!(stats.segments_total, 3);
4106 assert_eq!(stats.segments_scanned, 3);
4107 assert_eq!(stats.segments_pruned, 0);
4108 let probe = Value::Integer(100);
4109 let (pruned, stats) = manager.scan_project_zoned_with_stats(
4110 None,
4111 &[("ts", ZoneColPred::Gte(&probe))],
4112 |entity| project(entity).filter(|value| *value >= 100),
4113 );
4114 assert_eq!(
4115 pruned,
4116 all.into_iter()
4117 .filter(|value| *value >= 100)
4118 .collect::<Vec<_>>()
4119 );
4120 assert_eq!(stats.segments_total, 3);
4121 assert_eq!(stats.segments_scanned, 2);
4122 assert_eq!(stats.segments_pruned, 1);
4123 let (empty, _) = manager.scan_project_zoned_with_stats(None, &[], |_| None::<u64>);
4124 assert!(empty.is_empty(), "projection may reject every candidate");
4125
4126 let failure = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4127 manager.scan_project_zoned_with_stats(None, &[], |entity| {
4128 let value = project(entity)?;
4129 assert_ne!(value, 10, "injected projection worker failure");
4130 Some(value)
4131 })
4132 }));
4133 assert!(
4134 failure.is_err(),
4135 "a failed worker must not return a successful partial scan"
4136 );
4137 }
4138
4139 #[test]
4140 fn projected_zoned_scan_checks_visibility_before_projection_on_workers() {
4141 let manager = SegmentManager::with_config(
4142 "accounts",
4143 ManagerConfig {
4144 segment_config: SegmentConfig {
4145 max_entities: 2,
4146 ..Default::default()
4147 },
4148 ..Default::default()
4149 },
4150 );
4151 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4), (5, 3, 0)] {
4152 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "accounts", id, vec![]);
4153 entity.set_xmin(xmin);
4154 entity.set_xmax(xmax);
4155 manager.insert(entity).expect("versioned test row");
4156 }
4157 let mut hidden = UnifiedEntity::table_row(EntityId::new(6), "accounts", 6, vec![]);
4158 hidden.data = EntityData::Row(RowData::with_names(
4159 vec![Value::text("pending")],
4160 vec!["__moderation_status".into()],
4161 ));
4162 manager.insert(hidden).expect("moderated row");
4163 assert_eq!(manager.stats().sealed_count, 2);
4164 let snapshot = SnapshotContext {
4165 snapshot: Snapshot {
4166 xid: 10,
4167 in_progress: HashSet::from([4]),
4168 },
4169 manager: Arc::new(SnapshotManager::new()),
4170 own_xids: HashSet::new(),
4171 requires_index_fallback: false,
4172 serializable_reader: None,
4173 };
4174 std::thread::spawn(move || {
4175 for (view, expected) in [(Some(&snapshot), vec![5, 1, 3, 4]), (None, vec![5, 2, 3])] {
4176 let called = std::sync::Mutex::new(Vec::new());
4177 let (rows, stats) = manager.scan_project_zoned_with_stats(view, &[], |entity| {
4178 let id = entity.id.raw();
4179 assert!(
4180 expected.contains(&id),
4181 "invisible entity reached projection: {id}"
4182 );
4183 called.lock().expect("callback log").push(id);
4184 Some(id)
4185 });
4186 let (legacy, _) = manager.scan_zoned_with_stats(view, &[], |_| true);
4187 assert_eq!(
4188 rows,
4189 legacy
4190 .iter()
4191 .map(|entity| entity.id.raw())
4192 .collect::<Vec<_>>()
4193 );
4194 assert_eq!(rows[0], 5, "visible growing row first");
4195 let mut called = called.into_inner().expect("callback log");
4196 called.sort_unstable();
4197 let mut expected = expected;
4198 expected.sort_unstable();
4199 assert_eq!(called, expected, "each visible row projected exactly once");
4200 assert_eq!(stats.segments_total, 3);
4201 assert_eq!(stats.segments_scanned, 3);
4202 }
4203 })
4204 .join()
4205 .expect("projected scan worker");
4206 }
4207
4208 #[test]
4209 fn test_manager_auto_seal() {
4210 let config = ManagerConfig {
4211 segment_config: SegmentConfig {
4212 max_entities: 2,
4213 ..Default::default()
4214 },
4215 ..Default::default()
4216 };
4217
4218 let manager = SegmentManager::with_config("test", config);
4219
4220 manager
4222 .insert(UnifiedEntity::vector(
4223 manager.next_entity_id(),
4224 "v",
4225 vec![0.1],
4226 ))
4227 .unwrap();
4228
4229 manager
4231 .insert(UnifiedEntity::vector(
4232 manager.next_entity_id(),
4233 "v",
4234 vec![0.2],
4235 ))
4236 .unwrap();
4237
4238 manager
4240 .insert(UnifiedEntity::vector(
4241 manager.next_entity_id(),
4242 "v",
4243 vec![0.3],
4244 ))
4245 .unwrap();
4246
4247 let stats = manager.stats();
4248 assert_eq!(stats.total_entities, 3);
4249 }
4250
4251 #[test]
4252 fn test_manager_delete() {
4253 let manager = SegmentManager::new("test");
4254
4255 let id = manager
4256 .insert(UnifiedEntity::vector(
4257 manager.next_entity_id(),
4258 "v",
4259 vec![0.1],
4260 ))
4261 .unwrap();
4262
4263 assert!(manager.get(id).is_some());
4264 assert!(manager.delete(id).unwrap());
4265 assert!(manager.get(id).is_none());
4266 }
4267
4268 #[test]
4269 fn test_manager_metadata() {
4270 let manager = SegmentManager::new("test");
4271
4272 let id = manager
4273 .insert(UnifiedEntity::table_row(
4274 manager.next_entity_id(),
4275 "hosts",
4276 1,
4277 vec![Value::text("192.168.1.1".to_string())],
4278 ))
4279 .unwrap();
4280
4281 let mut meta = Metadata::new();
4282 meta.set(
4283 "os",
4284 super::super::metadata::MetadataValue::String("linux".to_string()),
4285 );
4286
4287 manager.set_metadata(id, meta).unwrap();
4288
4289 let retrieved = manager.get_metadata(id).unwrap();
4290 assert!(retrieved.has("os"));
4291 }
4292
4293 #[test]
4294 fn test_manager_query_by_kind() {
4295 let manager = SegmentManager::new("test");
4296
4297 manager
4298 .insert(UnifiedEntity::table_row(
4299 manager.next_entity_id(),
4300 "hosts",
4301 1,
4302 vec![],
4303 ))
4304 .unwrap();
4305
4306 manager
4307 .insert(UnifiedEntity::vector(
4308 manager.next_entity_id(),
4309 "embeddings",
4310 vec![0.1],
4311 ))
4312 .unwrap();
4313
4314 manager
4315 .insert(UnifiedEntity::table_row(
4316 manager.next_entity_id(),
4317 "hosts",
4318 2,
4319 vec![],
4320 ))
4321 .unwrap();
4322
4323 let rows = manager.get_by_kind("table");
4324 assert_eq!(rows.len(), 2);
4325
4326 let vectors = manager.get_by_kind("vector");
4327 assert_eq!(vectors.len(), 1);
4328 }
4329
4330 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
4337 SegmentManager::with_config(
4338 "test",
4339 ManagerConfig {
4340 max_sealed_segments: 1_000,
4341 consolidation_entities_per_tick: entities_per_tick,
4342 ..Default::default()
4343 },
4344 )
4345 }
4346
4347 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
4351 let entities: Vec<UnifiedEntity> = (0..count)
4352 .map(|i| {
4353 UnifiedEntity::table_row(
4354 manager.next_entity_id(),
4355 "users",
4356 i as u64,
4357 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
4358 )
4359 })
4360 .collect();
4361 let ids = manager.bulk_insert(entities).expect("bulk insert");
4362 manager.force_seal().expect("seal");
4363 ids
4364 }
4365
4366 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
4367 let mut ids: Vec<u64> = manager
4368 .query_all(|_| true)
4369 .into_iter()
4370 .map(|entity| entity.id.raw())
4371 .collect();
4372 ids.sort_unstable();
4373 ids
4374 }
4375
4376 fn drain_maintenance(manager: &SegmentManager) -> usize {
4377 let mut ticks = 0;
4378 while manager.consolidation.read().is_some() || ticks == 0 {
4379 manager.run_maintenance().expect("maintenance");
4380 ticks += 1;
4381 assert!(ticks < 1_000, "consolidation failed to converge");
4382 }
4383 ticks
4384 }
4385
4386 #[test]
4387 fn graph_batches_bound_payloads_and_hydrate_current_entities_after_sealing() {
4388 let manager = SegmentManager::new("graph");
4389 let node =
4390 |id| UnifiedEntity::graph_node(EntityId::new(id), "graph", "node", HashMap::new());
4391 manager
4392 .bulk_insert((1..=1024).map(node).collect())
4393 .expect("nodes");
4394 let mut hidden = node(1025);
4395 hidden.set_xmax(5);
4396 manager.insert(hidden).expect("hidden version");
4397 manager
4398 .insert(UnifiedEntity::vector(EntityId::new(1026), "graph", vec![]))
4399 .expect("mixed kind");
4400 let mut work = 0;
4401 let mut delivered = 0;
4402 assert!(!manager.scan_graph_batches(
4403 None,
4404 GraphEntityKind::Node,
4405 || {
4406 work += 1;
4407 work < 100
4408 },
4409 |batch| {
4410 delivered += batch.len();
4411 true
4412 }
4413 ));
4414 assert_eq!(delivered, 0, "cancelled partial batch must not escape");
4415 let mut seen = Vec::new();
4416 assert!(manager.scan_graph_batches(None, GraphEntityKind::Node, || true, |batch| {
4417 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4418 assert!(manager.growing.try_write().is_some(), "no topology guard");
4419 assert!(manager.sealed.try_write().is_some(), "no sealed guard");
4420 if seen.is_empty() {
4421 manager.force_seal().expect("seal during delivery");
4422 manager.insert(node(2000)).expect("append beyond captured boundary");
4423 let mut updated = node(900);
4424 if let EntityKind::GraphNode(node) = &mut updated.kind {
4425 node.label = "updated".into();
4426 }
4427 manager.update(updated).expect("update unread candidate");
4428 }
4429 for entity in batch {
4430 if entity.id.raw() == 900 {
4431 assert!(matches!(&entity.kind, EntityKind::GraphNode(node) if node.label == "updated"));
4432 }
4433 seen.push(entity.id.raw());
4434 }
4435 true
4436 }));
4437 assert_eq!(seen, (1..=1024).collect::<Vec<_>>());
4438 }
4439
4440 #[test]
4441 fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
4442 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4443 let ids = manager
4444 .bulk_insert(
4445 (1..=1024)
4446 .map(|id| {
4447 UnifiedEntity::graph_edge(
4448 EntityId::new(id),
4449 "step",
4450 "2000",
4451 "3000",
4452 1.0,
4453 HashMap::new(),
4454 )
4455 })
4456 .collect(),
4457 )
4458 .expect("edges");
4459 manager.force_seal().expect("seal");
4460 for id in ids.iter().take(300) {
4461 manager.delete(*id).expect("fragment source");
4462 }
4463 let source = Arc::downgrade(&manager.sealed.read()[0]);
4464 let mut seen = Vec::new();
4465 let mut changed = false;
4466 let mut metadata_bytes = 0;
4467 assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
4468 || true,
4469 |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
4470 |batch| {
4471 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4472 assert!(manager.growing.try_write().is_some(), "no growing topology guard");
4473 assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
4474 assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
4475 if !changed {
4476 changed = true;
4477 drain_maintenance(&manager);
4478 assert!(manager.stats().consolidation.runs_completed > 0);
4479 assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
4480 let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
4481 if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
4482 manager.update(updated).expect("update authoritative merged copy");
4483 }
4484 for id in batch {
4485 let entity = manager.get(*id).expect("hydrate current manager");
4486 if id.raw() == 1024 {
4487 assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
4488 }
4489 seen.push(id.raw());
4490 }
4491 Ok(true)
4492 }).expect("cursor"));
4493 assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
4494 assert!(
4495 source.upgrade().is_none(),
4496 "finished cursor releases retired source"
4497 );
4498 assert!(
4499 metadata_bytes < 4096,
4500 "fixed ID buffer plus source handles: {metadata_bytes}"
4501 );
4502 }
4503
4504 #[test]
4505 fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
4506 let manager = SegmentManager::new("graph");
4507 manager
4508 .bulk_insert(
4509 (1..=1024)
4510 .map(|id| {
4511 UnifiedEntity::graph_edge(
4512 EntityId::new(id),
4513 "step",
4514 "2000",
4515 "3000",
4516 1.0,
4517 HashMap::new(),
4518 )
4519 })
4520 .collect(),
4521 )
4522 .expect("edges");
4523 let mut delivered = 0;
4524 let denied = manager.visit_graph_index_batches(
4525 GraphEntityKind::Edge,
4526 &[EntityId::new(2000)],
4527 || true,
4528 |_| Err("out of memory"),
4529 |ids| {
4530 delivered += ids.len();
4531 Ok(true)
4532 },
4533 );
4534 assert_eq!(denied, Err("out of memory"));
4535 assert_eq!(delivered, 0);
4536 let mut work = 0;
4537 assert!(!manager
4538 .visit_graph_index_batches(
4539 GraphEntityKind::Edge,
4540 &[EntityId::new(2000)],
4541 || {
4542 work += 1;
4543 work < 100
4544 },
4545 |_| Ok::<_, ()>(()),
4546 |ids| {
4547 delivered += ids.len();
4548 Ok(true)
4549 }
4550 )
4551 .expect("cancelled cursor"));
4552 assert_eq!(
4553 delivered, 0,
4554 "cancelled partial batch cannot reach consumer"
4555 );
4556 assert!(!manager
4557 .visit_graph_index_batches(
4558 GraphEntityKind::Edge,
4559 &[EntityId::new(2000)],
4560 || true,
4561 |_| Ok::<_, ()>(()),
4562 |ids| {
4563 delivered += ids.len();
4564 Ok(false)
4565 }
4566 )
4567 .expect("consumer stop"));
4568 assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
4569 }
4570
4571 #[test]
4572 fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
4573 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4574 let ids = seal_bulk_segment(&manager, 100);
4575 for id in ids.iter().take(40) {
4576 manager.delete(*id).expect("delete before scan");
4577 }
4578 let source_bytes = manager.resident_bytes();
4579 let source = Arc::downgrade(&manager.sealed.read()[0]);
4580 std::thread::scope(|scope| {
4581 let (ready, waiting) = std::sync::mpsc::channel();
4582 let (release_first, first_wait) = std::sync::mpsc::channel();
4583 let (release_last, last_wait) = std::sync::mpsc::channel();
4584 let reader = |resume: std::sync::mpsc::Receiver<()>,
4585 ready: std::sync::mpsc::Sender<()>| {
4586 assert!(manager.scan_batches(
4587 None,
4588 || true,
4589 |_| {
4590 ready.send(()).expect("signal captured cursor");
4591 let _ = resume.recv();
4592 true
4593 }
4594 ));
4595 };
4596 let ready_first = ready.clone();
4597 let first = scope.spawn(move || reader(first_wait, ready_first));
4598 let last = scope.spawn(move || reader(last_wait, ready));
4599 for _ in 0..2 {
4600 waiting
4601 .recv_timeout(std::time::Duration::from_secs(10))
4602 .expect("reader ready");
4603 }
4604 drain_maintenance(&manager);
4605 assert!(
4606 source.upgrade().is_some(),
4607 "both cursors pin the retired source"
4608 );
4609 let active_bytes = manager.sealed.read()[0].read().resident_bytes();
4610 let held = manager.resident_bytes();
4611 assert!(
4612 held >= active_bytes + source_bytes,
4613 "retired source remains resident"
4614 );
4615 release_first.send(()).expect("release first reader");
4616 first.join().expect("first reader");
4617 assert_eq!(
4618 manager.resident_bytes(),
4619 held,
4620 "one reader still pins the source"
4621 );
4622 release_last.send(()).expect("release last reader");
4623 last.join().expect("last reader");
4624 assert!(
4625 source.upgrade().is_none(),
4626 "accounting must not retain payloads"
4627 );
4628 assert_eq!(held - manager.resident_bytes(), source_bytes);
4629 assert_eq!(
4630 manager.stats().total_memory_bytes as u64,
4631 manager.resident_bytes()
4632 );
4633 });
4634 }
4635
4636 #[test]
4637 fn physical_cursor_tracks_individual_and_mixed_insertions() {
4638 for mode in 0..3 {
4639 let manager = SegmentManager::new("physical_rows");
4640 let row = |id| {
4641 UnifiedEntity::table_row(
4642 EntityId::new(id),
4643 "physical_rows",
4644 id,
4645 vec![Value::Integer(id as i64)],
4646 )
4647 };
4648 match mode {
4649 0 => {
4650 for id in [10, 30, 20] {
4651 manager.insert(row(id)).expect("individual row");
4652 }
4653 }
4654 1 => {
4655 manager.insert(row(10)).expect("individual prefix");
4656 manager
4657 .bulk_insert(vec![row(30), row(31)])
4658 .expect("batch suffix");
4659 }
4660 _ => {
4661 manager
4662 .bulk_insert(vec![row(30), row(31)])
4663 .expect("batch prefix");
4664 manager.insert(row(10)).expect("individual suffix");
4665 }
4666 }
4667 let mut cursor = manager.physical_scan_cursor();
4668 let mut ids = Vec::new();
4669 for _ in 0..32 {
4670 if cursor.finished() {
4671 break;
4672 }
4673 let mut work = 1;
4674 let batch = manager.physical_scan_batch(&mut cursor, &mut work);
4675 assert!(batch.len() <= 1);
4676 ids.extend(batch.into_iter().map(|id| id.raw()));
4677 }
4678 assert!(cursor.finished());
4679 ids.sort_unstable();
4680 assert_eq!(
4681 ids,
4682 if mode == 0 {
4683 vec![10, 20, 30]
4684 } else {
4685 vec![10, 30, 31]
4686 },
4687 "mode={mode}"
4688 );
4689 }
4690 }
4691
4692 fn growing_tombstones_fixture() -> SegmentManager {
4693 let manager = consolidating_manager(1);
4694 let ids = manager
4695 .bulk_insert(
4696 (0..64)
4697 .map(|id| {
4698 UnifiedEntity::table_row(
4699 manager.next_entity_id(),
4700 "users",
4701 id,
4702 vec![Value::Blob(vec![7; 2048])],
4703 )
4704 })
4705 .collect(),
4706 )
4707 .expect("growing rows");
4708 for id in ids.iter().take(48) {
4709 manager.delete(*id).expect("physical tombstone");
4710 }
4711 assert_eq!(manager.stats().sealed_count, 0);
4712 manager
4713 }
4714
4715 #[test]
4716 fn pressure_consolidation_reclaims_unsealed_tombstones() {
4717 let manager = growing_tombstones_fixture();
4718 let before = manager.resident_bytes();
4719 assert!(manager.reclaimable_bytes() > 0);
4720 assert!(manager.pressure_consolidation_tick());
4721 drain_maintenance(&manager);
4722 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
4723 assert!(manager.resident_bytes() < before);
4724 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 48);
4725 }
4726
4727 #[test]
4728 fn pressure_consolidation_preserves_pinned_growing_source_accounting() {
4729 let manager = growing_tombstones_fixture();
4730 let source = Arc::clone(manager.growing.read().as_ref().expect("growing source"));
4731 let weak = Arc::downgrade(&source);
4732 assert!(manager.pressure_consolidation_tick());
4733 drain_maintenance(&manager);
4734 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
4735 let source_bytes = source.read().resident_bytes();
4736 let held = manager.resident_bytes();
4737 assert!(held >= source_bytes);
4738 drop(source);
4739 assert!(weak.upgrade().is_none());
4740 assert_eq!(held - manager.resident_bytes(), source_bytes);
4741 }
4742
4743 #[test]
4744 fn pressure_consolidation_keeps_clean_and_disabled_growing_sources_open() {
4745 let clean = consolidating_manager(1);
4746 clean
4747 .insert(UnifiedEntity::table_row(
4748 clean.next_entity_id(),
4749 "users",
4750 1,
4751 vec![Value::Integer(1)],
4752 ))
4753 .expect("clean row");
4754 assert!(!clean.pressure_consolidation_tick());
4755 assert_eq!(clean.stats().sealed_count, 0);
4756 let mut disabled = growing_tombstones_fixture();
4757 disabled.config.enable_consolidation = false;
4758 assert!(!disabled.pressure_consolidation_tick());
4759 assert_eq!(disabled.stats().sealed_count, 0);
4760 assert_eq!(live_ids(&disabled), (49..=64).collect::<Vec<_>>());
4761 }
4762
4763 #[test]
4764 fn retained_memory_includes_an_inflight_consolidation_copy() {
4765 let manager = consolidating_manager(1);
4766 let ids = seal_bulk_segment(&manager, 100);
4767 for id in ids.iter().take(40) {
4768 manager.delete(*id).expect("delete");
4769 }
4770 let before = manager.resident_bytes();
4771 assert!(manager.pressure_consolidation_tick());
4772 let copied = manager
4773 .consolidation
4774 .read()
4775 .as_ref()
4776 .expect("paced merge")
4777 .merged
4778 .resident_bytes();
4779 assert!(copied > 0);
4780 assert!(
4781 manager.resident_bytes() >= before + copied,
4782 "unpublished copy consumes memory"
4783 );
4784 drain_maintenance(&manager);
4785 assert!(
4786 manager.resident_bytes() < before,
4787 "completed merge releases the old storage"
4788 );
4789 }
4790
4791 #[test]
4792 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
4793 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4794 let ids = seal_bulk_segment(&manager, 100);
4795
4796 let bytes_before_delete = manager.resident_bytes();
4797
4798 for id in ids.iter().take(40) {
4800 assert!(manager.delete(*id).expect("delete"));
4801 }
4802
4803 let bytes_with_tombstones = manager.resident_bytes();
4806 assert!(bytes_with_tombstones > bytes_before_delete);
4807
4808 let survivors = live_ids(&manager);
4809 assert_eq!(survivors.len(), 60);
4810
4811 drain_maintenance(&manager);
4812
4813 let stats = manager.stats();
4814 assert_eq!(stats.consolidation.runs_started, 1);
4815 assert_eq!(stats.consolidation.runs_completed, 1);
4816 assert_eq!(stats.consolidation.segments_merged, 1);
4817 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
4818 assert!(stats.consolidation.bytes_reclaimed > 0);
4819
4820 let bytes_after = manager.resident_bytes();
4822 assert!(
4823 bytes_after <= bytes_before_delete,
4824 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
4825 );
4826 assert_eq!(
4827 bytes_with_tombstones - bytes_after,
4828 stats.consolidation.bytes_reclaimed
4829 );
4830 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
4831
4832 assert_eq!(live_ids(&manager), survivors);
4833 assert_eq!(manager.sealed.read().len(), 1);
4834 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
4835 }
4836
4837 #[test]
4838 fn query_results_are_identical_before_during_and_after_consolidation() {
4839 let manager = consolidating_manager(3);
4840 let ids = seal_bulk_segment(&manager, 20);
4841 seal_bulk_segment(&manager, 20);
4842 for id in ids.iter().take(10) {
4843 assert!(manager.delete(*id).expect("delete"));
4844 }
4845
4846 let before = live_ids(&manager);
4847 let survivor = *ids.last().expect("at least one row");
4848 let tombstoned = ids[0];
4849
4850 let mut mid_tick_observations = 0;
4853 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
4854 manager.run_maintenance().expect("maintenance");
4855 assert_eq!(live_ids(&manager), before);
4856 assert!(manager.get(survivor).is_some());
4857 assert!(manager.get(tombstoned).is_none());
4858 assert_eq!(manager.get_by_kind("table").len(), before.len());
4859 mid_tick_observations += 1;
4860 assert!(mid_tick_observations < 1_000);
4861 }
4862
4863 assert!(
4864 mid_tick_observations > 2,
4865 "expected the run to span several ticks, saw {mid_tick_observations}"
4866 );
4867 assert_eq!(live_ids(&manager), before);
4868 assert!(manager.get(survivor).is_some());
4869 assert!(manager.get(tombstoned).is_none());
4870 assert_eq!(manager.get_by_kind("table").len(), before.len());
4871 }
4872
4873 #[test]
4874 fn get_by_kind_reaches_sealed_segments() {
4875 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4876 let ids = seal_bulk_segment(&manager, 5);
4877 assert!(manager.delete(ids[0]).expect("delete"));
4878
4879 assert_eq!(manager.get_by_kind("table").len(), 4);
4883 assert_eq!(manager.get_by_kind("vector").len(), 0);
4884 }
4885
4886 #[test]
4887 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
4888 const PER_TICK: usize = 4;
4889 let manager = consolidating_manager(PER_TICK);
4890 let ids = seal_bulk_segment(&manager, 60);
4891 for id in ids.iter().take(30) {
4892 assert!(manager.delete(*id).expect("delete"));
4893 }
4894
4895 manager.run_maintenance().expect("maintenance");
4897 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
4898
4899 let mut previous = PER_TICK as u64;
4902 let mut ticks = 1;
4903 while let Some(copied) = manager.consolidation_progress() {
4904 assert!(
4905 copied - previous <= PER_TICK as u64,
4906 "tick copied {} entities, bound is {PER_TICK}",
4907 copied - previous
4908 );
4909 previous = copied;
4910 manager.run_maintenance().expect("maintenance");
4911 ticks += 1;
4912 assert!(ticks < 100);
4913 }
4914
4915 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
4918 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4919 assert_eq!(live_ids(&manager).len(), 30);
4920 }
4921
4922 #[test]
4923 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
4924 let manager = SegmentManager::with_config(
4925 "test",
4926 ManagerConfig {
4927 max_sealed_segments: 3,
4928 ..Default::default()
4929 },
4930 );
4931 for _ in 0..4 {
4932 seal_bulk_segment(&manager, 5);
4933 }
4934 assert_eq!(manager.sealed.read().len(), 4);
4935
4936 let before = live_ids(&manager);
4937 drain_maintenance(&manager);
4938
4939 assert_eq!(manager.sealed.read().len(), 1);
4942 assert_eq!(manager.stats().consolidation.segments_merged, 4);
4943 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
4944 assert_eq!(live_ids(&manager), before);
4945 }
4946
4947 #[test]
4948 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
4949 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4950 seal_bulk_segment(&manager, 10);
4951 seal_bulk_segment(&manager, 10);
4952
4953 for _ in 0..5 {
4954 manager.run_maintenance().expect("maintenance");
4955 }
4956
4957 assert_eq!(manager.stats().consolidation.runs_started, 0);
4958 assert_eq!(manager.sealed.read().len(), 2);
4959 }
4960
4961 #[test]
4962 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
4963 use std::sync::atomic::AtomicBool;
4964
4965 let manager = Arc::new(consolidating_manager(2));
4966 let ids = seal_bulk_segment(&manager, 40);
4967 seal_bulk_segment(&manager, 40);
4968 for id in ids.iter().take(20) {
4969 assert!(manager.delete(*id).expect("delete"));
4970 }
4971
4972 let expected = live_ids(&manager);
4973 let stop = Arc::new(AtomicBool::new(false));
4974
4975 let readers: Vec<_> = (0..4)
4976 .map(|_| {
4977 let manager = Arc::clone(&manager);
4978 let stop = Arc::clone(&stop);
4979 let expected = expected.clone();
4980 std::thread::spawn(move || {
4981 let mut scans = 0u64;
4982 loop {
4985 let mut seen: Vec<u64> = manager
4986 .query_all(|_| true)
4987 .into_iter()
4988 .map(|entity| entity.id.raw())
4989 .collect();
4990 seen.sort_unstable();
4991 assert_eq!(seen, expected, "reader saw a torn sealed set");
4994 scans += 1;
4995 if stop.load(Ordering::Relaxed) {
4996 break;
4997 }
4998 }
4999 scans
5000 })
5001 })
5002 .collect();
5003
5004 drain_maintenance(&manager);
5005 stop.store(true, Ordering::Relaxed);
5006
5007 for reader in readers {
5008 let scans = reader.join().expect("reader thread");
5009 assert!(scans > 0, "reader never got to scan");
5010 }
5011
5012 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5013 assert_eq!(live_ids(&manager), expected);
5014 }
5015
5016 #[test]
5017 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
5018 let manager = consolidating_manager(4);
5019 let ids = seal_bulk_segment(&manager, 20);
5020 for id in ids.iter().take(6) {
5021 assert!(manager.delete(*id).expect("delete"));
5022 }
5023
5024 manager.run_maintenance().expect("maintenance");
5026 assert_eq!(manager.consolidation_progress(), Some(4));
5027 let doomed = ids[6];
5028 assert!(manager.delete(doomed).expect("delete"));
5029
5030 let expected = live_ids(&manager);
5031 drain_maintenance(&manager);
5032
5033 assert!(manager.get(doomed).is_none(), "deleted row came back");
5034 assert_eq!(live_ids(&manager), expected);
5035 }
5036
5037 #[test]
5038 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
5039 let manager = consolidating_manager(3);
5040 let ids = seal_bulk_segment(&manager, 30);
5041 for id in ids.iter().take(12) {
5042 assert!(manager.delete(*id).expect("delete"));
5043 }
5044
5045 let before = live_ids(&manager);
5046 let sealed_before: Vec<SegmentId> = manager
5047 .sealed
5048 .read()
5049 .iter()
5050 .map(|segment| segment.read().id())
5051 .collect();
5052
5053 manager.run_maintenance().expect("maintenance");
5054 manager.run_maintenance().expect("maintenance");
5055 assert!(manager.consolidation.read().is_some());
5056
5057 *manager.consolidation.write() = None;
5061
5062 assert_eq!(live_ids(&manager), before);
5063 let sealed_after: Vec<SegmentId> = manager
5064 .sealed
5065 .read()
5066 .iter()
5067 .map(|segment| segment.read().id())
5068 .collect();
5069 assert_eq!(sealed_after, sealed_before);
5070 assert_eq!(manager.stats().consolidation.runs_completed, 0);
5071
5072 drain_maintenance(&manager);
5073 assert_eq!(live_ids(&manager), before);
5074 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5075 }
5076
5077 #[test]
5078 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
5079 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5080
5081 let entities: Vec<UnifiedEntity> = (0..10)
5085 .map(|i| {
5086 UnifiedEntity::table_row(
5087 manager.next_entity_id(),
5088 "users",
5089 i as u64,
5090 vec![Value::Integer(i as i64)],
5091 )
5092 })
5093 .collect();
5094 let ids = manager.bulk_insert(entities).expect("bulk insert");
5095
5096 let mut meta = Metadata::new();
5097 meta.set(
5098 "os",
5099 super::super::metadata::MetadataValue::String("linux".to_string()),
5100 );
5101 manager.set_metadata(ids[9], meta).expect("set metadata");
5102 manager.force_seal().expect("seal");
5103
5104 assert!(manager.delete(ids[0]).expect("delete"));
5105 assert!(manager.delete(ids[1]).expect("delete"));
5106 assert!(manager.delete(ids[2]).expect("delete"));
5107
5108 let before: Vec<UnifiedEntity> = {
5109 let mut entities = manager.query_all(|_| true);
5110 entities.sort_by_key(|entity| entity.id.raw());
5111 entities
5112 };
5113
5114 drain_maintenance(&manager);
5115
5116 let after: Vec<UnifiedEntity> = {
5117 let mut entities = manager.query_all(|_| true);
5118 entities.sort_by_key(|entity| entity.id.raw());
5119 entities
5120 };
5121
5122 let columns = |entity: &UnifiedEntity| match &entity.data {
5123 super::super::entity::EntityData::Row(row) => row.columns.clone(),
5124 _ => Vec::new(),
5125 };
5126
5127 assert_eq!(before.len(), after.len());
5128 for (old, new) in before.iter().zip(after.iter()) {
5129 assert_eq!(old.id, new.id);
5130 assert_eq!(columns(old), columns(new));
5131 assert_eq!(
5132 old.sequence_id, new.sequence_id,
5133 "merge must not re-issue sequence ids"
5134 );
5135 }
5136 assert!(manager
5137 .get_metadata(ids[9])
5138 .expect("metadata survives the merge")
5139 .has("os"));
5140 }
5141
5142 #[test]
5143 #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
5144 fn test_lifecycle_events() {
5145 let manager = SegmentManager::new("test");
5146
5147 manager
5148 .insert(UnifiedEntity::vector(
5149 manager.next_entity_id(),
5150 "v",
5151 vec![0.1],
5152 ))
5153 .unwrap();
5154
5155 let events = manager.drain_events();
5156
5157 assert!(events
5159 .iter()
5160 .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
5161 assert!(events
5162 .iter()
5163 .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
5164 }
5165}