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