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 stats: RwLock<ManagerStats>,
242 consolidation: RwLock<Option<Consolidation>>,
245 events: RwLock<Vec<LifecycleEvent>>,
247 visibility_map: VisibilityMap,
251 physical_root_cache: parking_lot::Mutex<Option<PhysicalRootCacheEntry>>,
252}
253
254impl SegmentManager {
255 pub fn new(collection: impl Into<String>) -> Self {
257 Self::with_config(collection, ManagerConfig::default())
258 }
259
260 pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
262 Self {
263 row_constraint_lock: Arc::new(parking_lot::ReentrantMutex::new(())),
264 physical_root_cache: parking_lot::Mutex::new(None),
265 collection: collection.into(),
266 config,
267 next_segment_id: AtomicU64::new(1),
268 next_entity_id: AtomicU64::new(1),
269 next_row_id: AtomicU64::new(1),
270 total_entities_atomic: AtomicU64::new(0),
271 growing: RwLock::new(None),
272 sealed: RwLock::new(Vec::new()),
273 retired: RwLock::new(Vec::new()),
274 archived: RwLock::new(Vec::new()),
275 entity_segment: RwLock::new(HashMap::new()),
276 column_schema: RwLock::new(None),
277 stats: RwLock::new(ManagerStats::default()),
278 consolidation: RwLock::new(None),
279 events: RwLock::new(Vec::new()),
280 visibility_map: VisibilityMap::new(),
281 }
282 }
283
284 #[doc(hidden)]
285 pub fn row_constraint_lock(&self) -> Arc<parking_lot::ReentrantMutex<()>> {
286 Arc::clone(&self.row_constraint_lock)
287 }
288
289 pub fn get_or_init_schema(
291 &self,
292 named: &HashMap<String, reddb_types::Value>,
293 ) -> Arc<Vec<String>> {
294 {
295 let schema = self.column_schema.read();
296 if let Some(ref s) = *schema {
297 return Arc::clone(s);
298 }
299 }
300 let cols: Vec<String> = named.keys().cloned().collect();
301 let arc = Arc::new(cols);
302 *self.column_schema.write() = Some(Arc::clone(&arc));
303 arc
304 }
305
306 pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
308 self.column_schema.read().clone()
309 }
310
311 fn restore_owned_row_schemas<'a>(&self, entities: impl Iterator<Item = &'a mut UnifiedEntity>) {
315 let schema = self.column_schema();
316 for entity in entities {
317 if let super::entity::EntityData::Row(row) = &mut entity.data {
318 if row.named.is_none() && row.schema.is_none() {
319 row.schema = schema.clone();
320 }
321 }
322 }
323 }
324
325 #[doc(hidden)]
326 pub fn set_column_schema_if_empty(&self, columns: Vec<String>) {
327 if columns.is_empty() {
328 return;
329 }
330 let mut schema = self.column_schema.write();
331 if schema.is_none() {
332 *schema = Some(Arc::new(columns));
333 }
334 }
335
336 pub fn collection(&self) -> &str {
338 &self.collection
339 }
340
341 pub fn config(&self) -> &ManagerConfig {
343 &self.config
344 }
345
346 pub fn stats(&self) -> ManagerStats {
349 let mut s = self.stats.read().clone();
350 s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
351 s.total_memory_bytes = self.resident_bytes() as usize;
352 s
353 }
354
355 pub fn resident_bytes(&self) -> u64 {
358 self.memory_usage().1
359 }
360
361 pub fn memory_bytes(&self) -> u64 {
364 self.memory_usage().0
365 }
366
367 fn memory_usage(&self) -> (u64, u64) {
371 let consolidation = self.consolidation.read();
372 let growing = self.growing.read();
373 let sealed = self.sealed.read();
374 let mut retired = self.retired.write();
375 let mut payload = 0u64;
376 let mut resident = 0u64;
377 {
378 let mut add = |segment: &GrowingSegment| {
379 payload = payload.saturating_add(segment.memory_bytes());
380 resident = resident.saturating_add(segment.resident_bytes());
381 };
382 for segment in growing.iter().chain(sealed.iter()) {
383 add(&segment.read());
384 }
385 retired.retain(|source| {
386 if let Some(source) = source.upgrade() {
387 add(&source.read());
388 true
389 } else {
390 false
391 }
392 });
393 if let Some(run) = consolidation.as_ref() {
394 add(&run.merged);
395 }
396 }
397 let mut bookkeeping =
398 (retired.capacity() * std::mem::size_of::<Weak<RwLock<GrowingSegment>>>()) as u64;
399 if let Some(run) = consolidation.as_ref() {
400 bookkeeping = bookkeeping
401 .saturating_add(
402 (run.sources.capacity() * std::mem::size_of::<SourceProgress>()) as u64,
403 )
404 .saturating_add(
405 (run.pending_ids.capacity() * std::mem::size_of::<EntityId>()) as u64,
406 );
407 for source in &run.sources {
408 bookkeeping = bookkeeping.saturating_add(
409 (source.copied.capacity() * std::mem::size_of::<EntityId>()) as u64,
410 );
411 }
412 }
413 (
414 payload.saturating_add(bookkeeping),
415 resident.saturating_add(bookkeeping),
416 )
417 }
418
419 pub fn reclaimable_bytes(&self) -> u64 {
421 self.sealed
422 .read()
423 .iter()
424 .map(|segment| segment.read().reclaimable_bytes())
425 .fold(0, u64::saturating_add)
426 }
427
428 pub fn next_entity_id(&self) -> EntityId {
430 EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
431 }
432
433 pub fn next_row_id(&self) -> u64 {
435 self.next_row_id.fetch_add(1, Ordering::SeqCst)
436 }
437
438 pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
442 let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
443 start..start + n
444 }
445
446 pub fn register_row_id(&self, id: u64) {
449 let candidate = id.saturating_add(1);
450 let mut current = self.next_row_id.load(Ordering::SeqCst);
451 while candidate > current {
452 match self.next_row_id.compare_exchange(
453 current,
454 candidate,
455 Ordering::SeqCst,
456 Ordering::SeqCst,
457 ) {
458 Ok(_) => break,
459 Err(updated) => current = updated,
460 }
461 }
462 }
463
464 fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
471 {
473 let growing = self.growing.read();
474 if let Some(segment) = growing.as_ref() {
475 return Arc::clone(segment);
476 }
477 }
478
479 let mut growing = self.growing.write();
481 if let Some(segment) = growing.as_ref() {
483 return Arc::clone(segment);
484 }
485
486 let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
487 let segment = GrowingSegment::new(id, &self.collection);
488 let segment_arc = Arc::new(RwLock::new(segment));
489 *growing = Some(Arc::clone(&segment_arc));
490
491 self.emit(LifecycleEvent::SegmentCreated(id));
492
493 self.stats.write().growing_count += 1;
496
497 segment_arc
498 }
499
500 pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
502 self.maybe_seal_growing()?;
504
505 let segment_arc = self.get_or_create_growing();
506 let mut segment = segment_arc.write();
507
508 if entity.id.raw() == 0 {
510 entity.id = self.next_entity_id();
511 }
512
513 let entity_id = entity.id;
514 let segment_id = segment.id();
515
516 segment.insert(entity)?;
517
518 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
520
521 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
530
531 Ok(entity_id)
532 }
533
534 fn with_restoration_segment<T>(
538 &self,
539 id: EntityId,
540 restore: impl FnOnce(&mut GrowingSegment) -> Result<T, SegmentError>,
541 ) -> Result<T, SegmentError> {
542 loop {
543 let segment = self.get_or_create_growing();
544 let mut growing = segment.write();
545 if !growing.state().is_writable() {
546 continue;
547 }
548 if !growing.is_tombstoned(id) {
549 return restore(&mut growing);
550 }
551 drop(growing);
552 self.force_seal()?;
553 }
554 }
555
556 #[doc(hidden)]
558 pub fn restore_snapshot_entity(&self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
559 self.maybe_seal_growing()?;
560 if self.get(entity.id).is_some() {
561 return Err(SegmentError::AlreadyExists(entity.id));
562 }
563 let entity_id = entity.id;
564 self.with_restoration_segment(entity_id, |segment| {
565 let segment_id = segment.id();
566 segment.adopt_entity(entity, None);
568 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
569 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
570 Ok(entity_id)
571 })
572 }
573
574 pub fn insert_batch(
576 &self,
577 entities: Vec<UnifiedEntity>,
578 ) -> Result<Vec<EntityId>, SegmentError> {
579 let mut ids = Vec::with_capacity(entities.len());
580 for entity in entities {
581 ids.push(self.insert(entity)?);
582 }
583 Ok(ids)
584 }
585
586 pub fn bulk_insert(
589 &self,
590 mut entities: Vec<UnifiedEntity>,
591 ) -> Result<Vec<EntityId>, SegmentError> {
592 for entity in &mut entities {
594 if entity.id.raw() == 0 {
595 entity.id = self.next_entity_id();
596 }
597 if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
598 if *row_id == 0 {
599 *row_id = self.next_row_id();
600 } else {
601 self.register_row_id(*row_id);
602 }
603 }
604 }
605
606 if let Some(first_row) = entities.first() {
609 if let super::entity::EntityData::Row(ref row) = first_row.data {
610 if let Some(ref named) = row.named {
611 let schema = self.get_or_init_schema(named);
612 for entity in &mut entities {
613 if let super::entity::EntityData::Row(ref mut row) = entity.data {
614 if let Some(named) = row.named.take() {
615 let mut cols = Vec::with_capacity(schema.len());
616 for col_name in schema.iter() {
617 cols.push(
618 named
619 .get(col_name)
620 .cloned()
621 .unwrap_or(reddb_types::Value::Null),
622 );
623 }
624 row.columns = cols;
625 row.schema = Some(Arc::clone(&schema));
626 }
627 }
628 }
629 }
630 }
631 }
632
633 let segment_arc = self.get_or_create_growing();
634 let mut segment = segment_arc.write();
635 let ids = segment.bulk_insert(entities)?;
637
638 self.total_entities_atomic
643 .fetch_add(ids.len() as u64, Ordering::Relaxed);
644
645 Ok(ids)
646 }
647
648 pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
650 let mut entity = self.get_with(id, Clone::clone)?;
651 self.restore_owned_row_schemas(std::iter::once(&mut entity));
652 Some(entity)
653 }
654
655 #[doc(hidden)]
658 pub fn get_with<T>(&self, id: EntityId, visit: impl FnOnce(&UnifiedEntity) -> T) -> Option<T> {
659 if let Some(growing_arc) = self.growing.read().as_ref() {
661 let growing = growing_arc.read();
662 if let Some(entity) = growing.get(id) {
663 return Some(visit(entity));
664 }
665 }
666
667 let sealed = self.sealed.read();
669 for segment in sealed.iter() {
670 let seg = segment.read();
671 if let Some(entity) = seg.get(id) {
672 return Some(visit(entity));
673 }
674 }
675
676 None
677 }
678
679 #[doc(hidden)]
683 pub fn graph_node_logical_id(
684 &self,
685 id: EntityId,
686 mut before_segment: impl FnMut() -> bool,
687 ) -> Option<EntityId> {
688 if let Some(growing_arc) = self.growing.read().as_ref() {
689 if !before_segment() {
690 return None;
691 }
692 let growing = growing_arc.read();
693 if let Some(entity) = growing.get(id) {
694 return matches!(entity.kind, EntityKind::GraphNode(_))
695 .then(|| entity.logical_id());
696 }
697 }
698 let sealed = self.sealed.read();
699 for segment in sealed.iter() {
700 if !before_segment() {
701 return None;
702 }
703 let segment = segment.read();
704 if let Some(entity) = segment.get(id) {
705 return matches!(entity.kind, EntityKind::GraphNode(_))
706 .then(|| entity.logical_id());
707 }
708 }
709 None
710 }
711
712 #[doc(hidden)]
717 pub fn visit_graph_index_batches<E>(
718 &self,
719 kind: GraphEntityKind,
720 keys: &[EntityId],
721 mut before_work: impl FnMut() -> bool,
722 mut admit: impl FnMut(usize) -> Result<(), E>,
723 mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
724 ) -> Result<bool, E> {
725 use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
726 admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
727 let mut admitted = 0;
728 let sources = loop {
729 let count = {
730 let growing = self.growing.read();
731 let sealed = self.sealed.read();
732 sealed.len() + usize::from(growing.is_some())
733 };
734 let bytes = count.saturating_mul(std::mem::size_of::<(
735 Arc<RwLock<GrowingSegment>>,
736 GraphIndexCursor,
737 )>());
738 admit(bytes.saturating_sub(admitted))?;
739 admitted = admitted.max(bytes);
740 let mut sources = Vec::with_capacity(count);
741 let growing = self.growing.read();
742 let sealed = self.sealed.read();
743 if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
744 continue;
745 }
746 for segment in growing.iter().chain(sealed.iter()) {
747 if !before_work() {
748 return Ok(false);
749 }
750 let view = segment.read();
751 if !view.may_contain_graph_kind(kind) {
752 continue;
753 }
754 let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
755 return Ok(false);
756 };
757 sources.push((Arc::clone(segment), cursor));
758 }
759 break sources;
760 };
761 let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
762 for (source, mut cursor) in sources {
763 while !cursor.finished(keys) {
764 let mut count = 0;
765 {
766 let segment = source.read();
767 if !segment.may_contain_graph_kind(kind) {
768 break;
769 }
770 if !segment.graph_index_batch(
771 kind,
772 keys,
773 &mut cursor,
774 &mut before_work,
775 &mut |id| {
776 ids[count] = id;
777 count += 1;
778 },
779 ) {
780 return Ok(false);
781 }
782 }
783 if !before_work() || !consume(&ids[..count])? {
784 return Ok(false);
785 }
786 }
787 }
788 Ok(true)
789 }
790
791 pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
796 let schema = self.column_schema();
797 let clone_row = |entity: &UnifiedEntity| {
798 let mut owned = entity.clone();
799 if let super::entity::EntityData::Row(row) = &mut owned.data {
800 if row.schema.is_none() {
801 row.schema = schema.clone();
802 }
803 }
804 owned
805 };
806 let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
807 let mut remaining: Vec<usize> = (0..ids.len()).collect(); if let Some(growing_arc) = self.growing.read().as_ref() {
812 let growing = if let Some(g) = growing_arc.try_read() {
813 g
814 } else {
815 growing_arc.read()
816 };
817 remaining.retain(|&i| {
818 if let Some(entity) = growing.get(ids[i]) {
819 out[i] = Some(clone_row(entity));
820 false } else {
822 true }
824 });
825 }
826
827 if remaining.is_empty() {
828 self.restore_owned_row_schemas(out.iter_mut().flatten());
829 return out;
830 }
831
832 let sealed = self.sealed.read();
834 for segment in sealed.iter() {
835 if remaining.is_empty() {
836 break;
837 }
838 let seg = segment.read();
839 remaining.retain(|&i| {
840 if let Some(entity) = seg.get(ids[i]) {
841 out[i] = Some(clone_row(entity));
842 false
843 } else {
844 true
845 }
846 });
847 }
848
849 drop(sealed);
850 self.restore_owned_row_schemas(out.iter_mut().flatten());
851 out
852 }
853
854 pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
867 where
868 F: FnMut(usize, &UnifiedEntity),
869 {
870 thread_local! {
879 static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
880 const { std::cell::Cell::new(Vec::new()) };
881 }
882
883 let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
884 remaining.clear();
885
886 if let Some(growing_arc) = self.growing.read().as_ref() {
887 let growing = if let Some(g) = growing_arc.try_read() {
888 g
889 } else {
890 growing_arc.read()
891 };
892 for (i, id) in ids.iter().enumerate() {
893 if let Some(entity) = growing.get(*id) {
894 f(i, entity);
895 } else {
896 remaining.push(i);
897 }
898 }
899 } else {
900 remaining.reserve(ids.len());
901 remaining.extend(0..ids.len());
902 }
903
904 if !remaining.is_empty() {
905 let sealed = self.sealed.read();
906 for segment in sealed.iter() {
907 if remaining.is_empty() {
908 break;
909 }
910 let seg = segment.read();
911 remaining.retain(|&i| {
912 if let Some(entity) = seg.get(ids[i]) {
913 f(i, entity);
914 false
915 } else {
916 true
917 }
918 });
919 }
920 }
921
922 REMAINING_SCRATCH.with(|cell| cell.set(remaining));
923 }
924
925 #[allow(dead_code)] fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
928 if let Some(growing_arc) = self.growing.read().as_ref() {
930 let growing = growing_arc.read();
931 if let Some(entity) = growing.get(id) {
932 return Some(entity.clone());
933 }
934 }
935
936 let sealed = self.sealed.read();
938 for segment in sealed.iter() {
939 if let Some(entity) = segment.get(id) {
940 return Some(entity.clone());
941 }
942 }
943
944 None
945 }
946
947 fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
948 let sealed = self.sealed.read();
949 sealed
950 .iter()
951 .find(|segment_arc| segment_arc.read().contains(id))
952 .map(Arc::clone)
953 }
954
955 fn rewrite_sealed_entity_into_growing(
956 &self,
957 entity: UnifiedEntity,
958 metadata: Option<&Metadata>,
959 ) -> Result<(), SegmentError> {
960 let entity_id = entity.id;
961 let sealed_arc = self
962 .find_sealed_segment_arc(entity_id)
963 .ok_or(SegmentError::NotFound(entity_id))?;
964
965 let metadata_to_apply = {
966 let mut sealed = sealed_arc.write();
967 let existing_metadata = sealed.get_metadata(entity_id);
968 if !sealed.force_delete(entity_id) {
969 return Err(SegmentError::NotFound(entity_id));
970 }
971 metadata.cloned().or(existing_metadata)
972 };
973
974 let growing_id = self.with_restoration_segment(entity_id, |growing| {
975 growing.insert_preserving_sequence(entity)?;
976 if let Some(metadata) = metadata_to_apply {
977 growing.set_metadata(entity_id, metadata)?;
978 }
979 Ok(growing.id())
980 })?;
981
982 self.entity_segment.write().insert(entity_id, growing_id);
983 Ok(())
984 }
985
986 pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
988 let entity_id = entity.id;
989 let mut entity = Some(entity);
990
991 if let Some(growing_arc) = self.growing.read().as_ref() {
993 let mut growing = growing_arc.write();
994 if growing.contains(entity_id) && growing.state().is_writable() {
995 return growing.update(entity.take().expect("entity already moved"));
996 }
997 }
998
999 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1001 if let Some(seg_id) = segment_id {
1002 if let Some(growing_arc) = self.growing.read().as_ref() {
1003 let mut growing = growing_arc.write();
1004 if growing.id() == seg_id && growing.state().is_writable() {
1005 return growing.update(entity.take().expect("entity already moved"));
1006 }
1007 }
1008 }
1009
1010 if let Some(entity) = entity.take() {
1011 return self.rewrite_sealed_entity_into_growing(entity, None);
1012 }
1013
1014 Err(SegmentError::NotFound(entity_id))
1015 }
1016
1017 pub fn update_with_metadata(
1020 &self,
1021 entity: UnifiedEntity,
1022 metadata: Option<&Metadata>,
1023 ) -> Result<(), SegmentError> {
1024 let entity_id = entity.id;
1025 let mut entity = Some(entity);
1026
1027 if let Some(growing_arc) = self.growing.read().as_ref() {
1029 let mut growing = growing_arc.write();
1030 if growing.contains(entity_id) && growing.state().is_writable() {
1031 growing.update(entity.take().expect("entity already moved"))?;
1032 if let Some(metadata) = metadata {
1033 growing.set_metadata(entity_id, metadata.clone())?;
1034 }
1035 return Ok(());
1036 }
1037 }
1038
1039 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1041 if let Some(seg_id) = segment_id {
1042 if let Some(growing_arc) = self.growing.read().as_ref() {
1043 let mut growing = growing_arc.write();
1044 if growing.id() == seg_id && growing.state().is_writable() {
1045 growing.update(entity.take().expect("entity already moved"))?;
1046 if let Some(metadata) = metadata {
1047 growing.set_metadata(entity_id, metadata.clone())?;
1048 }
1049 return Ok(());
1050 }
1051 }
1052 }
1053
1054 if let Some(entity) = entity.take() {
1055 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1056 }
1057
1058 Err(SegmentError::NotFound(entity_id))
1059 }
1060
1061 pub fn update_hot(
1065 &self,
1066 entity: UnifiedEntity,
1067 modified_columns: &[String],
1068 ) -> Result<(), SegmentError> {
1069 let entity_id = entity.id;
1070 let mut entity = Some(entity);
1071
1072 if let Some(growing_arc) = self.growing.read().as_ref() {
1073 let mut growing = growing_arc.write();
1074 if growing.contains(entity_id) && growing.state().is_writable() {
1075 return growing.update_hot(
1076 entity.take().expect("entity already moved"),
1077 modified_columns,
1078 );
1079 }
1080 }
1081
1082 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1083 if let Some(seg_id) = segment_id {
1084 if let Some(growing_arc) = self.growing.read().as_ref() {
1085 let mut growing = growing_arc.write();
1086 if growing.id() == seg_id && growing.state().is_writable() {
1087 return growing.update_hot(
1088 entity.take().expect("entity already moved"),
1089 modified_columns,
1090 );
1091 }
1092 }
1093 }
1094
1095 if let Some(entity) = entity.take() {
1096 return self.rewrite_sealed_entity_into_growing(entity, None);
1097 }
1098
1099 Err(SegmentError::NotFound(entity_id))
1100 }
1101
1102 pub fn update_hot_with_metadata(
1105 &self,
1106 entity: UnifiedEntity,
1107 modified_columns: &[String],
1108 metadata: Option<&Metadata>,
1109 ) -> Result<(), SegmentError> {
1110 let entity_id = entity.id;
1111 let mut entity = Some(entity);
1112
1113 if let Some(growing_arc) = self.growing.read().as_ref() {
1114 let mut growing = growing_arc.write();
1115 if growing.contains(entity_id) && growing.state().is_writable() {
1116 growing.update_hot(
1117 entity.take().expect("entity already moved"),
1118 modified_columns,
1119 )?;
1120 if let Some(metadata) = metadata {
1121 growing.set_metadata(entity_id, metadata.clone())?;
1122 }
1123 return Ok(());
1124 }
1125 }
1126
1127 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_hot(
1133 entity.take().expect("entity already moved"),
1134 modified_columns,
1135 )?;
1136 if let Some(metadata) = metadata {
1137 growing.set_metadata(entity_id, metadata.clone())?;
1138 }
1139 return Ok(());
1140 }
1141 }
1142 }
1143
1144 if let Some(entity) = entity.take() {
1145 return self.rewrite_sealed_entity_into_growing(entity, metadata);
1146 }
1147
1148 Err(SegmentError::NotFound(entity_id))
1149 }
1150
1151 pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
1154 where
1155 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
1156 {
1157 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
1158 items.into_iter().collect();
1159 if items.is_empty() {
1160 return Ok(());
1161 }
1162
1163 if let Some(growing_arc) = self.growing.read().as_ref() {
1164 let mut growing = growing_arc.write();
1165 if growing.state().is_writable() {
1166 match growing.update_hot_batch_with_metadata(items.iter().copied()) {
1167 Ok(()) => return Ok(()),
1168 Err(SegmentError::NotFound(_)) => {}
1169 Err(other) => return Err(other),
1170 }
1171 }
1172 }
1173
1174 for (entity, modified_columns, metadata) in items {
1175 self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
1176 }
1177 Ok(())
1178 }
1179
1180 pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
1182 if let Some(growing_arc) = self.growing.read().as_ref() {
1185 let mut growing = growing_arc.write();
1186 if growing.contains(id) && growing.state().is_writable() {
1187 let seg_id = growing.id();
1188 let deleted = growing.delete(id)?;
1189 if deleted {
1190 self.entity_segment.write().remove(&id);
1191 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1192 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1193 }
1194 return Ok(deleted);
1195 }
1196 }
1197
1198 let segment_id = self.entity_segment.read().get(&id).copied();
1201
1202 if let Some(seg_id) = segment_id {
1203 if let Some(growing_arc) = self.growing.read().as_ref() {
1204 let mut growing = growing_arc.write();
1205 if growing.id() == seg_id && growing.state().is_writable() {
1206 let deleted = growing.delete(id)?;
1207 if deleted {
1208 self.entity_segment.write().remove(&id);
1209 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1210 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1211 }
1212 return Ok(deleted);
1213 }
1214 }
1215 }
1216
1217 {
1220 let sealed = self.sealed.read();
1221 for segment_arc in sealed.iter() {
1222 let mut seg = segment_arc.write();
1223 let seg_id = seg.id();
1224 if seg.contains(id) {
1225 let deleted = seg.force_delete(id);
1226 drop(seg);
1227 if deleted {
1228 self.entity_segment.write().remove(&id);
1229 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1230 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1231 }
1232 return Ok(deleted);
1233 }
1234 }
1235 }
1236
1237 Ok(false)
1238 }
1239
1240 pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1241 if ids.is_empty() {
1242 return Ok(Vec::new());
1243 }
1244
1245 let mut deleted_ids = Vec::with_capacity(ids.len());
1246
1247 if let Some(growing_arc) = self.growing.read().as_ref() {
1248 let mut growing = growing_arc.write();
1249 if growing.state().is_writable() {
1250 let seg_id = growing.id();
1251 let deleted = growing.delete_batch(ids)?;
1252 if !deleted.is_empty() {
1253 {
1254 let mut entity_segment = self.entity_segment.write();
1255 for id in &deleted {
1256 entity_segment.remove(id);
1257 }
1258 }
1259 self.total_entities_atomic
1260 .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
1261 for id in &deleted {
1262 self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1263 }
1264 deleted_ids.extend(deleted);
1265 }
1266 }
1267 }
1268
1269 if deleted_ids.len() == ids.len() {
1270 return Ok(deleted_ids);
1271 }
1272
1273 let deleted_set: std::collections::HashSet<EntityId> =
1274 deleted_ids.iter().copied().collect();
1275 for &id in ids {
1276 if deleted_set.contains(&id) {
1277 continue;
1278 }
1279 if self.delete(id)? {
1280 deleted_ids.push(id);
1281 }
1282 }
1283
1284 Ok(deleted_ids)
1285 }
1286
1287 pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1289 if let Some(growing_arc) = self.growing.read().as_ref() {
1291 let growing = growing_arc.read();
1292 if growing.contains(id) {
1293 return growing.get_metadata(id);
1294 }
1295 }
1296
1297 let segment_id = self.entity_segment.read().get(&id).copied();
1299
1300 if let Some(seg_id) = segment_id {
1301 if let Some(growing_arc) = self.growing.read().as_ref() {
1302 let growing = growing_arc.read();
1303 if growing.id() == seg_id {
1304 return growing.get_metadata(id);
1305 }
1306 }
1307
1308 let sealed = self.sealed.read();
1309 for segment in sealed.iter() {
1310 if segment.id() == seg_id {
1311 return segment.get_metadata(id);
1312 }
1313 }
1314 }
1315
1316 if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1317 return segment_arc.read().get_metadata(id);
1318 }
1319
1320 None
1321 }
1322
1323 pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1325 if let Some(growing_arc) = self.growing.read().as_ref() {
1328 let mut growing = growing_arc.write();
1329 if growing.contains(id) && growing.state().is_writable() {
1330 return growing.set_metadata(id, metadata);
1331 }
1332 }
1333
1334 let segment_id = self.entity_segment.read().get(&id).copied();
1336
1337 if let Some(seg_id) = segment_id {
1338 if let Some(growing_arc) = self.growing.read().as_ref() {
1339 let mut growing = growing_arc.write();
1340 if growing.id() == seg_id && growing.state().is_writable() {
1341 return growing.set_metadata(id, metadata);
1342 }
1343 }
1344 }
1345
1346 if let Some(entity) = self.get(id) {
1347 return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1348 }
1349
1350 Err(SegmentError::NotFound(id))
1351 }
1352
1353 fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1355 let should_seal = {
1356 let growing_opt = self.growing.read();
1357 if let Some(growing_arc) = growing_opt.as_ref() {
1358 let growing = growing_arc.read();
1359 growing.should_seal(&self.config.segment_config)
1360 || growing.idle_secs() >= self.config.idle_seal_secs
1361 } else {
1362 false
1363 }
1364 };
1365
1366 if should_seal {
1367 self.seal_current()?;
1368 }
1369
1370 Ok(())
1371 }
1372
1373 pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1375 let mut growing_slot = self.growing.write();
1378 if let Some(growing_arc) = growing_slot.take() {
1379 let mut growing = growing_arc.write();
1380 let seg_id = growing.id();
1381 let entity_count = growing.stats().entity_count as u64;
1382
1383 growing.seal()?;
1385
1386 drop(growing); self.sealed.write().push(growing_arc);
1392 drop(growing_slot);
1393
1394 self.mark_sealed_pages_visible(entity_count);
1396
1397 {
1399 let mut stats = self.stats.write();
1400 stats.growing_count = stats.growing_count.saturating_sub(1);
1401 stats.sealed_count += 1;
1402 stats.seal_ops += 1;
1403 }
1404
1405 self.emit(LifecycleEvent::SegmentSealed(seg_id));
1406
1407 return Ok(seg_id);
1408 }
1409
1410 Err(SegmentError::InvalidState(SegmentState::Sealed))
1411 }
1412
1413 pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1415 let has_growing = self.growing.read().is_some();
1416 if has_growing {
1417 Ok(Some(self.seal_current()?))
1418 } else {
1419 Ok(None)
1420 }
1421 }
1422
1423 pub fn all_visible_fraction(&self) -> f64 {
1433 const ROWS_PER_PAGE: u32 = 256;
1434 let sealed = self.sealed.read();
1435 if sealed.is_empty() {
1436 return 0.0;
1437 }
1438 let mut total_pages: u64 = 0;
1439 for seg_arc in sealed.iter() {
1440 let seg = seg_arc.read();
1441 let entity_count = seg.stats().entity_count as u64;
1442 let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1443 total_pages += pages;
1444 }
1445 if total_pages == 0 {
1446 return 0.0;
1447 }
1448 let visible = self.visibility_map.all_visible_count();
1449 (visible as f64 / total_pages as f64).min(1.0)
1450 }
1451
1452 fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1455 const ROWS_PER_PAGE: u32 = 256;
1456 let existing_visible = self.visibility_map.all_visible_count();
1457 let start_page = existing_visible as u32;
1459 let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1460 let end_page = start_page + new_pages as u32;
1461 self.visibility_map.mark_range_visible(start_page, end_page);
1462 }
1463
1464 #[doc(hidden)]
1468 pub fn find_unique_key_conflict(
1469 &self,
1470 columns: &[Vec<String>],
1471 key_index: usize,
1472 signatures: &[String],
1473 exclude_id: Option<EntityId>,
1474 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1475 ) -> Option<EntityId> {
1476 if let Some(growing) = self.growing.read().as_ref() {
1477 if let Some(id) = growing.read().find_unique_key_conflict(
1478 columns,
1479 key_index,
1480 signatures,
1481 exclude_id,
1482 reserves_key,
1483 ) {
1484 return Some(id);
1485 }
1486 }
1487 for segment in self.sealed.read().iter() {
1488 if let Some(id) = segment.read().find_unique_key_conflict(
1489 columns,
1490 key_index,
1491 signatures,
1492 exclude_id,
1493 reserves_key,
1494 ) {
1495 return Some(id);
1496 }
1497 }
1498 None
1499 }
1500
1501 pub fn for_each_entity<F>(&self, callback: F)
1507 where
1508 F: FnMut(&UnifiedEntity) -> bool,
1509 {
1510 self.for_each_entity_in_segments(|_| true, || true, callback);
1511 }
1512
1513 #[doc(hidden)]
1515 pub fn checkpoint_revisions(&self) -> Vec<(u64, u64)> {
1516 let growing = self.growing.read();
1517 let sealed = self.sealed.read();
1518 growing
1519 .iter()
1520 .chain(sealed.iter())
1521 .map(|segment| {
1522 let segment = segment.read();
1523 (segment.id(), segment.checkpoint_revision())
1524 })
1525 .collect()
1526 }
1527
1528 #[doc(hidden)]
1529 pub fn checkpoint_mutations_since(&self, before: &[(u64, u64)]) -> Option<Vec<EntityId>> {
1530 let growing = self.growing.read();
1531 let sealed = self.sealed.read();
1532 let previous: HashMap<_, _> = before.iter().copied().collect();
1537 if previous.len() != before.len() {
1538 return None;
1539 }
1540 let mut retained_sources = 0;
1541 let mut changed = std::collections::BTreeSet::new();
1542 for segment in growing.iter().chain(sealed.iter()) {
1543 let segment = segment.read();
1544 let revision = previous.get(&segment.id()).copied();
1545 retained_sources += usize::from(revision.is_some());
1546 let revision = revision.unwrap_or(0);
1549 changed.extend(segment.checkpoint_mutations_since(revision)?);
1550 if changed.len() > super::segment::CHECKPOINT_MUTATIONS_MAX {
1551 return None;
1552 }
1553 }
1554 (retained_sources == previous.len()).then(|| changed.into_iter().collect())
1557 }
1558
1559 #[doc(hidden)]
1560 pub fn checkpoint_tracking_bytes(&self) -> usize {
1561 let growing = self.growing.read();
1562 let sealed = self.sealed.read();
1563 growing
1564 .iter()
1565 .chain(sealed.iter())
1566 .map(|segment| segment.read().checkpoint_tracking_bytes())
1567 .sum()
1568 }
1569
1570 #[doc(hidden)]
1572 pub fn cached_physical_root(&self, seed: u64, compute: impl FnOnce() -> u64) -> u64 {
1573 let segments = self.checkpoint_revisions();
1574 let mut cache = self.physical_root_cache.lock();
1575 if let Some(previous) = cache.as_ref() {
1576 if previous.seed == seed && previous.segments == segments {
1577 return previous.root;
1578 }
1579 }
1580 let root = compute();
1581 *cache = Some(PhysicalRootCacheEntry {
1582 segments,
1583 seed,
1584 root,
1585 });
1586 root
1587 }
1588
1589 #[doc(hidden)]
1594 pub fn for_each_entity_in_id_order(&self, mut visit: impl FnMut(&UnifiedEntity) -> bool) {
1595 self.for_each_entity_source_in_id_order(|_, entity| visit(entity));
1596 }
1597
1598 #[doc(hidden)]
1600 pub fn page_in_id_order(&self, offset: usize, limit: usize) -> (usize, Vec<UnifiedEntity>) {
1601 let end = offset.saturating_add(limit);
1602 let mut total = 0;
1603 let mut items = Vec::with_capacity(limit.min(self.count()));
1604 self.for_each_entity_in_id_order(|entity| {
1605 if total >= offset && total < end {
1606 items.push(entity.clone());
1607 }
1608 total += 1;
1609 true
1610 });
1611 self.restore_owned_row_schemas(items.iter_mut());
1612 (total, items)
1613 }
1614
1615 #[doc(hidden)]
1616 pub fn for_each_entity_with_metadata_in_id_order(
1617 &self,
1618 mut visit: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1619 ) {
1620 self.for_each_entity_source_in_id_order(|segment, entity| {
1621 visit(entity, segment.get_metadata(entity.id))
1622 });
1623 }
1624
1625 fn for_each_entity_source_in_id_order(
1626 &self,
1627 mut visit: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
1628 ) {
1629 let segments: Vec<_> = {
1630 let growing = self.growing.read();
1631 let sealed = self.sealed.read();
1632 growing.iter().chain(sealed.iter()).cloned().collect()
1633 };
1634 let mut order = Vec::new();
1635 for (source, segment) in segments.iter().enumerate() {
1636 segment.read().for_each_fast(|entity| {
1637 order.push((entity.id, source));
1638 true
1639 });
1640 }
1641 order.sort_by_key(|(id, _)| id.raw());
1642 for (id, source) in order {
1643 let segment = segments[source].read();
1644 if let Some(entity) = segment.get(id) {
1645 if !visit(&segment, entity) {
1646 break;
1647 }
1648 }
1649 }
1650 }
1651
1652 #[doc(hidden)]
1655 pub fn for_each_entity_with_metadata(
1656 &self,
1657 mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1658 ) {
1659 if let Some(growing_arc) = self.growing.read().as_ref() {
1660 let growing = growing_arc.read();
1661 if !growing.for_each_fast(|entity| callback(entity, growing.get_metadata(entity.id))) {
1662 return;
1663 }
1664 }
1665 let sealed = self.sealed.read();
1666 for segment_arc in sealed.iter() {
1667 let segment = segment_arc.read();
1668 if !segment.for_each_fast(|entity| callback(entity, segment.get_metadata(entity.id))) {
1669 return;
1670 }
1671 }
1672 }
1673
1674 fn for_each_entity_in_segments<F>(
1675 &self,
1676 segment_filter: impl Fn(&GrowingSegment) -> bool,
1677 mut before_segment: impl FnMut() -> bool,
1678 mut callback: F,
1679 ) where
1680 F: FnMut(&UnifiedEntity) -> bool,
1681 {
1682 if let Some(growing_arc) = self.growing.read().as_ref() {
1686 if !before_segment() {
1687 return;
1688 }
1689 let growing = if let Some(g) = growing_arc.try_read() {
1690 g
1691 } else {
1692 growing_arc.read()
1693 };
1694 if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
1695 return;
1696 }
1697 }
1698
1699 let sealed = self.sealed.read();
1701 for segment_arc in sealed.iter() {
1702 if !before_segment() {
1703 return;
1704 }
1705 let segment = segment_arc.read();
1706 if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
1707 return;
1708 }
1709 }
1710 }
1711
1712 pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1728 &self,
1729 init: FInit,
1730 fold: FFold,
1731 reduce: FReduce,
1732 ) -> T
1733 where
1734 T: Send,
1735 FInit: Fn() -> T + Send + Sync,
1736 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1737 FReduce: Fn(T, T) -> T + Send + Sync,
1738 {
1739 #[cfg(not(target_arch = "wasm32"))]
1740 use rayon::prelude::*;
1741
1742 let mut acc = init();
1745 if let Some(growing_arc) = self.growing.read().as_ref() {
1746 let growing = if let Some(g) = growing_arc.try_read() {
1747 g
1748 } else {
1749 growing_arc.read()
1750 };
1751 growing.for_each_fast(|entity| {
1752 acc = fold(std::mem::replace(&mut acc, init()), entity);
1753 true
1754 });
1755 }
1756
1757 let segments: Vec<_> = {
1761 let sealed = self.sealed.read();
1762 sealed.iter().cloned().collect()
1763 };
1764
1765 if segments.len() <= 1 {
1766 for seg_arc in &segments {
1767 let seg = seg_arc.read();
1768 seg.for_each_fast(|entity| {
1769 acc = fold(std::mem::replace(&mut acc, init()), entity);
1770 true
1771 });
1772 }
1773 return acc;
1774 }
1775
1776 #[cfg(not(target_arch = "wasm32"))]
1777 let segments = segments.into_par_iter();
1778 #[cfg(target_arch = "wasm32")]
1779 let segments = segments.into_iter();
1780 let partials = segments.map(|seg_arc| {
1781 let mut local = init();
1782 let seg = seg_arc.read();
1783 seg.for_each_fast(|entity| {
1784 local = fold(std::mem::replace(&mut local, init()), entity);
1785 true
1786 });
1787 local
1788 });
1789 #[cfg(not(target_arch = "wasm32"))]
1790 let sealed_acc = partials.reduce(&init, &reduce);
1791 #[cfg(target_arch = "wasm32")]
1792 let sealed_acc = partials.fold(init(), &reduce);
1793
1794 reduce(acc, sealed_acc)
1795 }
1796
1797 pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
1806 where
1807 F: FnMut(&UnifiedEntity) -> bool,
1808 {
1809 let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
1810 }
1811
1812 pub fn for_each_entity_zoned_with_stats<F>(
1813 &self,
1814 zone_preds: &[(&str, ZoneColPred<'_>)],
1815 mut callback: F,
1816 ) -> SegmentScanStats
1817 where
1818 F: FnMut(&UnifiedEntity) -> bool,
1819 {
1820 let mut scan_stats = SegmentScanStats::default();
1821 if let Some(growing_arc) = self.growing.read().as_ref() {
1827 scan_stats.segments_total += 1;
1828 scan_stats.segments_scanned += 1;
1829 let growing = if let Some(g) = growing_arc.try_read() {
1830 g
1831 } else {
1832 growing_arc.read()
1833 };
1834 if !growing.for_each_fast(&mut callback) {
1835 return scan_stats;
1836 }
1837 }
1838
1839 let sealed = self.sealed.read();
1841 for segment_arc in sealed.iter() {
1842 scan_stats.segments_total += 1;
1843 let segment = segment_arc.read();
1844 if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
1845 scan_stats.segments_pruned += 1;
1846 continue; }
1848 scan_stats.segments_scanned += 1;
1849 if !segment.for_each_fast(&mut callback) {
1850 return scan_stats;
1851 }
1852 }
1853 scan_stats
1854 }
1855
1856 pub fn query_all_zoned<F>(
1866 &self,
1867 zone_preds: &[(&str, ZoneColPred<'_>)],
1868 filter: F,
1869 ) -> Vec<UnifiedEntity>
1870 where
1871 F: Fn(&UnifiedEntity) -> bool + Sync,
1872 {
1873 self.query_all_zoned_with_stats(zone_preds, filter).0
1874 }
1875
1876 pub fn query_all_zoned_with_stats<F>(
1877 &self,
1878 zone_preds: &[(&str, ZoneColPred<'_>)],
1879 filter: F,
1880 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
1881 where
1882 F: Fn(&UnifiedEntity) -> bool + Sync,
1883 {
1884 let mut results = Vec::new();
1885 let mut scan_stats = SegmentScanStats::default();
1886
1887 if let Some(growing_arc) = self.growing.read().as_ref() {
1890 scan_stats.segments_total += 1;
1891 scan_stats.segments_scanned += 1;
1892 let growing = if let Some(g) = growing_arc.try_read() {
1893 g
1894 } else {
1895 growing_arc.read()
1896 };
1897 results.extend(growing.iter().filter(|e| filter(e)).cloned());
1898 }
1899
1900 let sealed = self.sealed.read();
1902 let surviving: Vec<_> = sealed
1904 .iter()
1905 .filter(|seg_arc| {
1906 scan_stats.segments_total += 1;
1907 if zone_preds.is_empty() {
1908 scan_stats.segments_scanned += 1;
1909 return true;
1910 }
1911 let seg = seg_arc.read();
1912 if seg.can_skip_zone_preds(zone_preds) {
1913 scan_stats.segments_pruned += 1;
1914 false
1915 } else {
1916 scan_stats.segments_scanned += 1;
1917 true
1918 }
1919 })
1920 .collect();
1921
1922 let use_parallel = surviving.len() > 1 && crate::execution::should_parallelize();
1923
1924 if use_parallel {
1925 let filter_ref = &filter;
1926 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1927 surviving
1928 .iter()
1929 .map(|segment| {
1930 s.spawn(move || {
1931 segment
1932 .read()
1933 .iter()
1934 .filter(|e| filter_ref(e))
1935 .cloned()
1936 .collect::<Vec<_>>()
1937 })
1938 })
1939 .collect::<Vec<_>>()
1940 .into_iter()
1941 .map(|handle| handle.join().unwrap_or_default())
1942 .collect()
1943 });
1944 for batch in segment_results {
1945 results.extend(batch);
1946 }
1947 } else {
1948 for segment_arc in surviving {
1949 let seg = segment_arc.read();
1950 results.extend(seg.iter().filter(|e| filter(e)).cloned());
1951 }
1952 }
1953
1954 (results, scan_stats)
1955 }
1956
1957 fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
1963 match snapshot {
1964 Some(ctx) => entity_visible_with_context(Some(ctx), entity),
1965 None => {
1966 !crate::storage::visibility::entity_moderation_hidden(entity) && entity.xmax == 0
1967 }
1968 }
1969 }
1970
1971 pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
1978 where
1979 F: Fn(&UnifiedEntity) -> bool + Sync,
1980 {
1981 self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
1982 }
1983
1984 #[doc(hidden)]
1989 pub fn project_key_entity<T>(
1990 &self,
1991 key: &str,
1992 snapshot: Option<&SnapshotContext>,
1993 versioned: bool,
1994 mut project: impl FnMut(&UnifiedEntity) -> T,
1995 ) -> Option<T> {
1996 let mut probe = |segment: &GrowingSegment, best: &mut Option<(u64, T)>| {
1997 if let Some(entity) = segment.find_key_entity(key, versioned, &mut |entity| {
1998 !versioned || Self::scan_entity_visible(snapshot, entity)
1999 }) {
2000 if best.as_ref().is_none_or(|(xmin, _)| entity.xmin >= *xmin) {
2001 *best = Some((entity.xmin, project(entity)));
2002 }
2003 }
2004 };
2005 let mut best = None;
2006 if let Some(growing) = self.growing.read().as_ref() {
2007 probe(&growing.read(), &mut best);
2008 }
2009 for segment in self.sealed.read().iter() {
2010 if !versioned && best.is_some() {
2011 break;
2012 }
2013 probe(&segment.read(), &mut best);
2014 }
2015 best.map(|(_, projected)| projected)
2016 }
2017
2018 #[doc(hidden)]
2020 pub fn scan_graph_kind(
2021 &self,
2022 snapshot: Option<&SnapshotContext>,
2023 kind: GraphEntityKind,
2024 ) -> Vec<UnifiedEntity> {
2025 self.query_all_in_segments(
2026 |segment| segment.may_contain_graph_kind(kind),
2027 |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
2028 )
2029 }
2030
2031 #[doc(hidden)]
2036 pub fn scan_graph_candidates_for_each<F>(
2037 &self,
2038 snapshot: Option<&SnapshotContext>,
2039 kind: GraphEntityKind,
2040 before_segment: impl FnMut() -> bool,
2041 mut callback: F,
2042 ) where
2043 F: FnMut(&UnifiedEntity) -> bool,
2044 {
2045 self.for_each_entity_in_segments(
2046 |segment| segment.may_contain_graph_kind(kind),
2047 before_segment,
2048 |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
2049 );
2050 }
2051
2052 pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
2057 where
2058 F: FnMut(&UnifiedEntity) -> bool,
2059 {
2060 self.for_each_entity(|entity| {
2061 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2062 });
2063 }
2064
2065 #[doc(hidden)]
2066 pub fn physical_scan_cursor(&self) -> PhysicalScanCursor {
2067 let growing = self.growing.read();
2068 let sealed = self.sealed.read();
2069 PhysicalScanCursor {
2070 sources: growing
2071 .iter()
2072 .chain(sealed.iter())
2073 .map(|segment| (Arc::downgrade(segment), segment.read().scan_cursor()))
2074 .collect(),
2075 source: 0,
2076 }
2077 }
2078
2079 #[doc(hidden)]
2082 pub fn physical_scan_batch(
2083 &self,
2084 cursor: &mut PhysicalScanCursor,
2085 work_left: &mut usize,
2086 ) -> Vec<EntityId> {
2087 let mut ids = Vec::with_capacity(super::segment::SCAN_BATCH_SIZE.min(*work_left));
2088 while !cursor.finished() && *work_left > 0 {
2089 let (source, position) = &mut cursor.sources[cursor.source];
2090 if position.finished() {
2091 *work_left -= 1;
2092 cursor.source += 1;
2093 continue;
2094 }
2095 let Some(segment) = source.upgrade() else {
2096 *work_left -= 1;
2097 cursor.source += 1;
2098 continue;
2099 };
2100 segment.read().scan_batch(
2101 position,
2102 &mut || {
2103 if *work_left == 0 {
2104 return false;
2105 }
2106 *work_left -= 1;
2107 true
2108 },
2109 &mut |id| ids.push(id),
2110 );
2111 if position.finished() {
2112 cursor.source += 1;
2113 }
2114 if !ids.is_empty() {
2115 break;
2116 }
2117 }
2118 ids
2119 }
2120
2121 #[doc(hidden)]
2127 pub fn scan_batches(
2128 &self,
2129 snapshot: Option<&SnapshotContext>,
2130 mut before_work: impl FnMut() -> bool,
2131 mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
2132 ) -> bool {
2133 let segments = {
2134 let growing = self.growing.read();
2135 let sealed = self.sealed.read();
2136 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
2137 for segment in growing.iter().chain(sealed.iter()) {
2138 if !before_work() {
2139 return false;
2140 }
2141 let cursor = segment.read().scan_cursor();
2142 segments.push((Arc::clone(segment), cursor));
2143 }
2144 segments
2145 };
2146 let mut ids = Vec::new();
2147 let mut batch = Vec::new();
2148 for (segment, mut cursor) in segments {
2149 while !cursor.finished() {
2150 ids.clear();
2151 batch.clear();
2152 let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
2153 ids.reserve_exact(batch_size);
2154 batch.reserve_exact(batch_size);
2155 {
2156 let segment = segment.read();
2157 if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
2158 return false;
2159 }
2160 }
2161 if !before_work() {
2164 return false;
2165 }
2166 self.for_each_id(&ids, |_, entity| {
2167 if Self::scan_entity_visible(snapshot, entity) {
2168 batch.push(entity.clone());
2169 }
2170 });
2171 if !before_work() {
2172 return false;
2173 }
2174 if !consume(&batch) {
2175 return false;
2176 }
2177 }
2178 }
2179 true
2180 }
2181
2182 pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
2184 &self,
2185 snapshot: Option<&SnapshotContext>,
2186 init: FInit,
2187 fold: FFold,
2188 reduce: FReduce,
2189 ) -> T
2190 where
2191 T: Send,
2192 FInit: Fn() -> T + Send + Sync,
2193 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
2194 FReduce: Fn(T, T) -> T + Send + Sync,
2195 {
2196 self.fold_entities_parallel(
2197 init,
2198 |acc, entity| {
2199 if Self::scan_entity_visible(snapshot, entity) {
2200 fold(acc, entity)
2201 } else {
2202 acc
2203 }
2204 },
2205 reduce,
2206 )
2207 }
2208
2209 pub fn scan_for_each_zoned_with_stats<F>(
2211 &self,
2212 snapshot: Option<&SnapshotContext>,
2213 zone_preds: &[(&str, ZoneColPred<'_>)],
2214 mut callback: F,
2215 ) -> SegmentScanStats
2216 where
2217 F: FnMut(&UnifiedEntity) -> bool,
2218 {
2219 self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
2220 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2221 })
2222 }
2223
2224 pub fn scan_zoned_with_stats<F>(
2226 &self,
2227 snapshot: Option<&SnapshotContext>,
2228 zone_preds: &[(&str, ZoneColPred<'_>)],
2229 filter: F,
2230 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2231 where
2232 F: Fn(&UnifiedEntity) -> bool + Sync,
2233 {
2234 self.query_all_zoned_with_stats(zone_preds, |entity| {
2235 Self::scan_entity_visible(snapshot, entity) && filter(entity)
2236 })
2237 }
2238
2239 pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
2242 where
2243 F: Fn(&UnifiedEntity) -> bool + Sync,
2244 {
2245 let mut entities = self.query_all_in_segments(|_| true, filter);
2246 self.restore_owned_row_schemas(entities.iter_mut());
2247 entities
2248 }
2249
2250 fn query_all_in_segments<F>(
2251 &self,
2252 segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
2253 filter: F,
2254 ) -> Vec<UnifiedEntity>
2255 where
2256 F: Fn(&UnifiedEntity) -> bool + Sync,
2257 {
2258 let mut results = Vec::new();
2259
2260 if let Some(growing_arc) = self.growing.read().as_ref() {
2263 let growing = if let Some(g) = growing_arc.try_read() {
2264 g
2265 } else {
2266 growing_arc.read()
2267 };
2268 if segment_filter(&growing) {
2269 results.extend(growing.iter().filter(|e| filter(e)).cloned());
2270 }
2271 }
2272
2273 let sealed = self.sealed.read();
2275 let use_parallel = sealed.len() > 1 && crate::execution::should_parallelize();
2276 if use_parallel {
2277 let filter_ref = &filter;
2278 let segment_filter_ref = &segment_filter;
2279 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
2280 sealed
2281 .iter()
2282 .map(|segment| {
2283 s.spawn(move || {
2284 let segment = segment.read();
2285 if segment_filter_ref(&segment) {
2286 segment
2287 .iter()
2288 .filter(|e| filter_ref(e))
2289 .cloned()
2290 .collect::<Vec<_>>()
2291 } else {
2292 Vec::new()
2293 }
2294 })
2295 })
2296 .collect::<Vec<_>>()
2297 .into_iter()
2298 .map(|handle| handle.join().unwrap_or_default())
2299 .collect()
2300 });
2301 for batch in segment_results {
2302 results.extend(batch);
2303 }
2304 } else {
2305 for segment in sealed.iter() {
2306 let seg = segment.read();
2307 if segment_filter(&seg) {
2308 results.extend(seg.iter().filter(|e| filter(e)).cloned());
2309 }
2310 }
2311 }
2312
2313 results
2314 }
2315
2316 pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
2321 if let Some(growing_arc) = self.growing.read().as_ref() {
2322 let growing = growing_arc.read();
2323 if growing.may_contain_exact_key(key) {
2324 return true;
2325 }
2326 } else {
2327 return true;
2328 }
2329
2330 let sealed = self.sealed.read();
2331 for segment_arc in sealed.iter() {
2332 let segment = segment_arc.read();
2333 if segment.may_contain_exact_key(key) {
2334 return true;
2335 }
2336 }
2337
2338 false
2339 }
2340
2341 pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2343 let mut results = Vec::new();
2344
2345 if let Some(growing_arc) = self.growing.read().as_ref() {
2347 let growing = growing_arc.read();
2348 results.extend(growing.filter_metadata(filters));
2349 }
2350
2351 let sealed = self.sealed.read();
2353 for segment in sealed.iter() {
2354 results.extend(segment.filter_metadata(filters));
2355 }
2356
2357 results
2358 }
2359
2360 pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2362 let mut results = Vec::new();
2363
2364 if let Some(growing_arc) = self.growing.read().as_ref() {
2366 let growing = growing_arc.read();
2367 for entity in growing.iter_kind(kind) {
2368 results.push(entity.clone());
2369 }
2370 }
2371
2372 let sealed = self.sealed.read();
2377 for segment_arc in sealed.iter() {
2378 let segment = segment_arc.read();
2379 for entity in segment.iter_kind(kind) {
2380 results.push(entity.clone());
2381 }
2382 }
2383
2384 results
2385 }
2386
2387 pub fn count(&self) -> usize {
2389 self.total_entities_atomic.load(Ordering::Relaxed) as usize
2390 }
2391
2392 pub fn segment_ids(&self) -> Vec<SegmentId> {
2394 let mut ids = Vec::new();
2395
2396 if let Some(growing_arc) = self.growing.read().as_ref() {
2397 ids.push(growing_arc.read().id());
2398 }
2399
2400 let sealed = self.sealed.read();
2401 for segment in sealed.iter() {
2402 ids.push(segment.id());
2403 }
2404
2405 ids.extend(self.archived.read().iter().copied());
2406
2407 ids
2408 }
2409
2410 #[inline]
2423 #[allow(clippy::unused_self)]
2424 fn emit(&self, _event: LifecycleEvent) {}
2425
2426 pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2429 std::mem::take(&mut *self.events.write())
2430 }
2431
2432 pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2434 self.maybe_seal_growing()?;
2436
2437 if self.config.enable_consolidation {
2438 self.consolidation_tick();
2439 }
2440
2441 Ok(())
2442 }
2443
2444 fn consolidation_tick(&self) {
2455 let running = self.consolidation.read().is_some();
2458 if !running && !self.start_consolidation() {
2459 return;
2460 }
2461
2462 let mut guard = self.consolidation.write();
2465 let Some(run) = guard.as_mut() else {
2466 return;
2467 };
2468 self.copy_bounded(run);
2469 if run.cursor >= run.sources.len() {
2470 if let Some(run) = guard.take() {
2471 self.finish_consolidation(run);
2472 }
2473 }
2474 }
2475
2476 fn start_consolidation(&self) -> bool {
2480 let sources = self.select_candidates();
2481 self.start_consolidation_from_sources(sources)
2482 }
2483
2484 fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
2485 if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
2486 return false;
2487 }
2488
2489 let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
2490 let merged = GrowingSegment::new(merged_id, &self.collection);
2491
2492 *self.consolidation.write() = Some(Consolidation {
2493 sources: sources
2494 .into_iter()
2495 .map(|id| SourceProgress {
2496 id,
2497 copied: Vec::new(),
2498 tombstones_at_open: None,
2499 })
2500 .collect(),
2501 merged_id,
2502 merged,
2503 cursor: 0,
2504 pending_ids: Vec::new(),
2505 pending_cursor: 0,
2506 entities_copied: 0,
2507 });
2508
2509 self.stats.write().consolidation.runs_started += 1;
2510 true
2511 }
2512
2513 pub fn pressure_consolidation_tick(&self) -> bool {
2518 if !self.config.enable_consolidation {
2519 return false;
2520 }
2521
2522 let running = self.consolidation.read().is_some();
2523 if !running {
2524 let sources = self.dirty_sealed_segments();
2525 if !self.start_consolidation_from_sources(sources) {
2526 return false;
2527 }
2528 }
2529
2530 self.consolidation_tick();
2531 true
2532 }
2533
2534 fn select_candidates(&self) -> Vec<SegmentId> {
2549 let sealed = self.sealed.read();
2550 if sealed.is_empty() {
2551 return Vec::new();
2552 }
2553
2554 if sealed.len() > self.config.max_sealed_segments {
2555 return sealed.iter().map(|segment| segment.read().id()).collect();
2556 }
2557
2558 let mut live = 0usize;
2559 let mut tombstones = 0usize;
2560 let mut held = 0u64;
2561 let mut reclaimable = 0u64;
2562 let mut dirty = Vec::new();
2563
2564 for segment_arc in sealed.iter() {
2565 let segment = segment_arc.read();
2566 live += segment.live_entity_count();
2567 tombstones += segment.tombstone_count();
2568 held += segment.resident_bytes();
2569 reclaimable += segment.reclaimable_bytes();
2570 if segment.tombstone_count() > 0 {
2571 dirty.push(segment.id());
2572 }
2573 }
2574
2575 let total = live + tombstones;
2576 let tombstone_ratio = if total == 0 {
2577 0.0
2578 } else {
2579 tombstones as f64 / total as f64
2580 };
2581 let fragmentation = if held == 0 {
2582 0.0
2583 } else {
2584 reclaimable as f64 / held as f64
2585 };
2586
2587 if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
2588 || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
2589 {
2590 return dirty;
2591 }
2592
2593 Vec::new()
2594 }
2595
2596 fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
2597 self.sealed
2598 .read()
2599 .iter()
2600 .filter_map(|segment_arc| {
2601 let segment = segment_arc.read();
2602 (segment.reclaimable_bytes() > 0).then(|| segment.id())
2603 })
2604 .collect()
2605 }
2606
2607 fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
2610 if sources.is_empty() {
2611 return false;
2612 }
2613 let sealed = self.sealed.read();
2614 sealed.iter().any(|segment_arc| {
2615 let segment = segment_arc.read();
2616 sources.contains(&segment.id()) && segment.tombstone_count() > 0
2617 })
2618 }
2619
2620 fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
2621 let sealed = self.sealed.read();
2622 sealed
2623 .iter()
2624 .find(|segment| segment.read().id() == id)
2625 .map(Arc::clone)
2626 }
2627
2628 fn copy_bounded(&self, run: &mut Consolidation) {
2633 let mut budget = self.config.consolidation_entities_per_tick;
2634
2635 while budget > 0 && run.cursor < run.sources.len() {
2636 let source_id = run.sources[run.cursor].id;
2637 let Some(source_arc) = self.sealed_segment(source_id) else {
2638 run.cursor += 1;
2642 run.pending_ids.clear();
2643 run.pending_cursor = 0;
2644 continue;
2645 };
2646
2647 if run.sources[run.cursor].tombstones_at_open.is_none() {
2648 let source = source_arc.read();
2649 run.pending_ids = source.live_entity_ids();
2650 run.pending_cursor = 0;
2651 run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
2652 }
2653
2654 {
2655 let source = source_arc.read();
2656 while budget > 0 && run.pending_cursor < run.pending_ids.len() {
2657 let id = run.pending_ids[run.pending_cursor];
2658 run.pending_cursor += 1;
2659 budget -= 1;
2660
2661 let Some(entity) = source.get(id) else {
2664 continue;
2665 };
2666 let metadata = source.get_metadata(id);
2667 run.merged.adopt_entity(entity.clone(), metadata);
2668 run.sources[run.cursor].copied.push(id);
2669 run.entities_copied += 1;
2670 }
2671 }
2672
2673 if run.pending_cursor >= run.pending_ids.len() {
2674 run.cursor += 1;
2675 run.pending_ids.clear();
2676 run.pending_cursor = 0;
2677 }
2678 }
2679 }
2680
2681 fn finish_consolidation(&self, mut run: Consolidation) {
2687 let mut sealed = self.sealed.write();
2688
2689 let mut positions = Vec::with_capacity(run.sources.len());
2690 for source in &run.sources {
2691 if let Some(index) = sealed
2692 .iter()
2693 .position(|segment| segment.read().id() == source.id)
2694 {
2695 positions.push(index);
2696 }
2697 }
2698 if positions.is_empty() {
2699 return;
2700 }
2701 positions.sort_unstable();
2702
2703 let mut stale = Vec::new();
2707 for source in &run.sources {
2708 let Some(source_arc) = sealed
2709 .iter()
2710 .find(|segment| segment.read().id() == source.id)
2711 else {
2712 continue;
2713 };
2714 let segment = source_arc.read();
2715 if source.tombstones_at_open == Some(segment.tombstone_count()) {
2716 continue;
2717 }
2718 stale.extend(
2719 source
2720 .copied
2721 .iter()
2722 .copied()
2723 .filter(|id| !segment.contains(*id)),
2724 );
2725 }
2726 for id in stale {
2727 run.merged.evict_entity(id);
2728 }
2729
2730 let reclaimed_bytes: u64 = positions
2731 .iter()
2732 .map(|&index| sealed[index].read().resident_bytes())
2733 .sum();
2734 let reclaimed_tombstones: u64 = positions
2735 .iter()
2736 .map(|&index| sealed[index].read().tombstone_count() as u64)
2737 .sum();
2738
2739 if run.merged.seal().is_err() {
2740 return;
2741 }
2742 let merged_bytes = run.merged.resident_bytes();
2743 let merged_arc = Arc::new(RwLock::new(run.merged));
2744
2745 {
2746 let mut retired = self.retired.write();
2747 retired.retain(|source| source.strong_count() > 0);
2748 for &index in positions.iter().rev() {
2749 let source = sealed.remove(index);
2750 if Arc::strong_count(&source) > 1 {
2751 retired.push(Arc::downgrade(&source));
2752 }
2753 }
2754 }
2755 let insert_at = positions[0].min(sealed.len());
2756 sealed.insert(insert_at, Arc::clone(&merged_arc));
2757 drop(sealed);
2758
2759 let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
2761 {
2762 let mut entity_segment = self.entity_segment.write();
2763 for segment_id in entity_segment.values_mut() {
2764 if source_ids.contains(segment_id) {
2765 *segment_id = run.merged_id;
2766 }
2767 }
2768 }
2769
2770 {
2771 let mut stats = self.stats.write();
2772 stats.sealed_count = stats
2773 .sealed_count
2774 .saturating_sub(positions.len())
2775 .saturating_add(1);
2776 let consolidation = &mut stats.consolidation;
2777 consolidation.runs_completed += 1;
2778 consolidation.segments_merged += positions.len() as u64;
2779 consolidation.tombstones_reclaimed += reclaimed_tombstones;
2780 consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
2781 }
2782
2783 self.emit(LifecycleEvent::SegmentConsolidated {
2784 source: source_ids,
2785 target: run.merged_id,
2786 });
2787 }
2788
2789 #[cfg(test)]
2792 fn consolidation_progress(&self) -> Option<u64> {
2793 self.consolidation
2794 .read()
2795 .as_ref()
2796 .map(|run| run.entities_copied)
2797 }
2798}
2799
2800#[cfg(test)]
2801mod tests {
2802 use std::collections::HashSet;
2803
2804 use super::*;
2805 use crate::storage::entity::{EntityData, EntityKind, RowData};
2806 use crate::transaction::snapshot::{Snapshot, SnapshotManager};
2807 use reddb_types::Value;
2808
2809 #[test]
2810 fn physical_pages_keep_global_order_across_segments_and_deleted_slots() {
2811 let manager = SegmentManager::new("pages");
2812 for group in [[50, 10], [30, 20]] {
2813 for id in group {
2814 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "pages", id, vec![]);
2815 entity.data = EntityData::Row(RowData::with_names(
2816 vec![Value::Integer(id as i64)],
2817 vec!["id".into()],
2818 ));
2819 manager.insert(entity).unwrap();
2820 }
2821 manager.force_seal().unwrap();
2822 }
2823 manager
2824 .insert(UnifiedEntity::table_row(
2825 EntityId::new(40),
2826 "pages",
2827 40,
2828 vec![],
2829 ))
2830 .unwrap();
2831 let (total, page) = manager.page_in_id_order(1, 2);
2832 assert_eq!(total, 5);
2833 assert_eq!(
2834 page.iter()
2835 .map(|entity| entity.id.raw())
2836 .collect::<Vec<_>>(),
2837 [20, 30]
2838 );
2839 let EntityData::Row(row) = &page[0].data else {
2840 panic!("row")
2841 };
2842 assert_eq!(row.get_field("id"), Some(&Value::Integer(20)));
2843 manager.delete(EntityId::new(30)).unwrap();
2844 let (total, page) = manager.page_in_id_order(1, 2);
2845 assert_eq!(total, 4);
2846 assert_eq!(
2847 page.iter()
2848 .map(|entity| entity.id.raw())
2849 .collect::<Vec<_>>(),
2850 [20, 40]
2851 );
2852 for offset in [4, usize::MAX] {
2853 let (total, page) = manager.page_in_id_order(offset, 2);
2854 assert_eq!(total, 4);
2855 assert!(page.is_empty());
2856 }
2857 }
2858
2859 fn lookup_version(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
2860 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
2861 entity.data = EntityData::Row(RowData::with_names(
2862 vec![Value::text(key)],
2863 vec!["key".into()],
2864 ));
2865 entity.set_xmin(xmin);
2866 entity.set_xmax(xmax);
2867 entity
2868 }
2869
2870 #[test]
2871 fn key_lookup_matches_snapshot_scan_across_segments() {
2872 let manager = Arc::new(SegmentManager::new("keys"));
2873 for entity in [
2874 lookup_version(1, "changed", 3, 12),
2875 lookup_version(2, "deleted", 3, 7),
2876 ] {
2877 manager.insert(entity).expect("old version");
2878 }
2879 manager.force_seal().expect("seal old versions");
2880 for entity in [
2881 lookup_version(3, "changed", 12, 0),
2882 lookup_version(4, "pending", 4, 0),
2883 lookup_version(5, "own", 15, 0),
2884 lookup_version(6, "live", 3, 0),
2885 ] {
2886 manager.insert(entity).expect("new version");
2887 }
2888 std::thread::spawn(move || {
2889 let snapshot = SnapshotContext {
2890 snapshot: Snapshot {
2891 xid: 10,
2892 in_progress: HashSet::from([4]),
2893 },
2894 manager: Arc::new(SnapshotManager::new()),
2895 own_xids: HashSet::from([15]),
2896 requires_index_fallback: false,
2897 serializable_reader: None,
2898 };
2899 for snapshot in [Some(&snapshot), None] {
2900 for key in ["changed", "deleted", "pending", "own", "live", "absent"] {
2901 let expected = manager
2902 .scan(snapshot, |entity| {
2903 crate::storage::segment_key_lookup::SegmentKeyLookup::key(entity)
2904 == Some(key)
2905 })
2906 .into_iter()
2907 .max_by_key(|entity| entity.xmin)
2908 .map(|entity| entity.id);
2909 assert_eq!(
2910 manager.project_key_entity(key, snapshot, true, |entity| entity.id),
2911 expected,
2912 "key={key}"
2913 );
2914 }
2915 }
2916 assert_eq!(
2917 manager.project_key_entity("changed", Some(&snapshot), true, |entity| entity.id),
2918 Some(EntityId::new(1))
2919 );
2920 assert_eq!(
2921 manager.project_key_entity("changed", None, true, |entity| entity.id),
2922 Some(EntityId::new(3))
2923 );
2924 assert!(manager
2925 .project_key_entity("deleted", None, true, |entity| entity.id)
2926 .is_none());
2927 })
2928 .join()
2929 .expect("snapshot worker");
2930 }
2931
2932 #[test]
2933 fn key_lookup_survives_paced_consolidation() {
2934 let manager = consolidating_manager(1);
2935 manager
2936 .bulk_insert(
2937 (1..=10)
2938 .map(|id| lookup_version(id, &format!("key-{id}"), 1, 0))
2939 .collect(),
2940 )
2941 .expect("rows");
2942 manager.force_seal().expect("seal");
2943 for id in 1..=4 {
2944 manager.delete(EntityId::new(id)).expect("delete");
2945 }
2946 assert!(manager.pressure_consolidation_tick());
2947 drain_maintenance(&manager);
2948 for id in 1..=10 {
2949 assert_eq!(
2950 manager.project_key_entity(&format!("key-{id}"), None, false, |entity| entity.id),
2951 (id > 4).then_some(EntityId::new(id))
2952 );
2953 }
2954 }
2955
2956 #[test]
2957 #[ignore = "diagnostic of segment fragmentation and retained history; no timing assertions"]
2958 fn key_lookup_fragmentation_history_diagnostic() {
2959 use std::{hint::black_box, time::Instant};
2960 for segments in [1, 16, 64] {
2961 for versions in [1, 100, 1000] {
2962 let manager = SegmentManager::new("keys");
2963 let mut id = 1;
2964 for segment in 0..segments {
2965 let mut rows: Vec<_> = (0..128)
2966 .map(|_| {
2967 let entity = lookup_version(id, &format!("other-{id}"), 1, 0);
2968 id += 1;
2969 entity
2970 })
2971 .collect();
2972 for version in (segment..versions).step_by(segments) {
2973 rows.push(lookup_version(
2974 id,
2975 "history",
2976 version as u64 + 1,
2977 if version + 1 == versions {
2978 0
2979 } else {
2980 version as u64 + 2
2981 },
2982 ));
2983 id += 1;
2984 }
2985 manager.bulk_insert(rows).expect("rows");
2986 manager.force_seal().expect("seal");
2987 }
2988 let started = Instant::now();
2989 for _ in 0..2000 {
2990 let xmin =
2991 manager.project_key_entity(black_box("history"), None, true, |entity| {
2992 entity.xmin
2993 });
2994 assert_eq!(xmin, Some(versions as u64));
2995 }
2996 eprintln!("segments={segments} versions={versions} reads=2000 elapsed_us={} resident_bytes={}", started.elapsed().as_micros(), manager.resident_bytes());
2997 }
2998 }
2999 }
3000
3001 #[test]
3002 fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
3003 let manager = SegmentManager::new("fingerprint");
3004 let row = |id, label: &str| {
3005 UnifiedEntity::table_row(
3006 EntityId::new(id),
3007 "fingerprint",
3008 id,
3009 vec![Value::text(label)],
3010 )
3011 };
3012 manager
3013 .bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
3014 .expect("sealed rows");
3015 manager.force_seal().expect("seal");
3016 manager
3017 .bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
3018 .expect("growing rows");
3019
3020 let mut expected = manager.query_all(|_| true);
3021 expected.sort_by_key(|entity| entity.id.raw());
3022 let expected: Vec<_> = expected
3023 .iter()
3024 .map(|entity| format!("{entity:?}"))
3025 .collect();
3026 let mut stored_addresses = HashSet::new();
3027 manager.for_each_entity(|entity| {
3028 stored_addresses.insert(entity as *const UnifiedEntity as usize);
3029 true
3030 });
3031
3032 let mut actual = Vec::new();
3033 manager.for_each_entity_in_id_order(|entity| {
3034 assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
3035 assert!(
3036 manager.growing.try_write().is_some(),
3037 "topology guard retained"
3038 );
3039 assert!(
3040 manager.sealed.try_write().is_some(),
3041 "topology guard retained"
3042 );
3043 actual.push(format!("{entity:?}"));
3044 true
3045 });
3046 assert_eq!(
3047 actual, expected,
3048 "equal IDs retain their original segment order"
3049 );
3050 }
3051
3052 #[test]
3053 fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
3054 let manager = SegmentManager::new("fingerprint");
3055 for id in [3, 1, 2] {
3056 manager
3057 .insert(UnifiedEntity::table_row(
3058 EntityId::new(id),
3059 "fingerprint",
3060 id,
3061 vec![Value::Integer(id as i64)],
3062 ))
3063 .expect("row");
3064 }
3065 let source = manager.growing.read().as_ref().expect("source").clone();
3066 let handles = Arc::strong_count(&source);
3067 let mut seen = Vec::new();
3068 manager.for_each_entity_in_id_order(|entity| {
3069 seen.push(entity.id.raw());
3070 false
3071 });
3072 assert_eq!(seen, vec![1]);
3073 assert_eq!(Arc::strong_count(&source), handles);
3074 assert!(
3075 source.try_write().is_some(),
3076 "segment guard retained after early stop"
3077 );
3078 }
3079
3080 #[test]
3081 fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
3082 let manager = SegmentManager::with_config(
3083 "vectors",
3084 ManagerConfig {
3085 max_sealed_segments: 1,
3086 consolidation_entities_per_tick: 10_000,
3087 ..Default::default()
3088 },
3089 );
3090 for part in 0..3 {
3091 manager
3092 .bulk_insert(
3093 (1..=300)
3094 .map(|offset| {
3095 UnifiedEntity::vector(
3096 EntityId::new(part * 300 + offset),
3097 "vectors",
3098 vec![1.0, 0.0],
3099 )
3100 })
3101 .collect(),
3102 )
3103 .expect("batch");
3104 if part < 2 {
3105 manager.force_seal().expect("seal");
3106 }
3107 }
3108 let original_growing = manager.growing.read().as_ref().expect("growing").clone();
3109 let mut changed = false;
3110 let mut seen = Vec::new();
3111 assert!(manager.scan_batches(
3112 None,
3113 || true,
3114 |batch| {
3115 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3116 assert!(
3117 manager.growing.try_write().is_some(),
3118 "consumer holds no topology guard"
3119 );
3120 assert!(
3121 original_growing.try_write().is_some(),
3122 "consumer holds no segment guard"
3123 );
3124 for entity in batch {
3125 assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
3126 seen.push(entity.id.raw());
3127 }
3128 if !changed {
3129 changed = true;
3130 manager.force_seal().expect("seal during scan");
3131 manager
3132 .insert(UnifiedEntity::vector(
3133 EntityId::new(5000),
3134 "vectors",
3135 vec![1.0, 0.0],
3136 ))
3137 .expect("new write");
3138 drain_maintenance(&manager);
3139 assert!(manager.stats().consolidation.runs_completed > 0);
3140 assert!(manager
3141 .delete(EntityId::new(1))
3142 .expect("delete after source retirement"));
3143 assert!(manager.get(EntityId::new(1)).is_none());
3144 }
3145 true
3146 }
3147 ));
3148 seen.sort_unstable();
3149 assert_eq!(
3150 seen,
3151 (2..=900).collect::<Vec<_>>(),
3152 "no duplicate, missing or late candidates"
3153 );
3154 let mut after = Vec::new();
3155 assert!(manager.scan_batches(
3156 None,
3157 || true,
3158 |batch| {
3159 after.extend(batch.iter().map(|entity| entity.id.raw()));
3160 true
3161 }
3162 ));
3163 after.sort_unstable();
3164 let mut expected: Vec<_> = (2..=900).collect();
3165 expected.push(5000);
3166 assert_eq!(after, expected, "adopted entities have scan positions");
3167 }
3168
3169 #[test]
3170 fn bounded_scan_applies_snapshot_to_relocated_versions() {
3171 for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
3172 let manager = SegmentManager::with_config(
3173 "vectors",
3174 ManagerConfig {
3175 max_sealed_segments: 1,
3176 ..Default::default()
3177 },
3178 );
3179 for ids in [[1, 2], [3, 4]] {
3180 let entities = ids.map(|id| {
3181 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3182 entity.set_xmin(3);
3183 entity
3184 });
3185 if bulk {
3186 manager
3187 .bulk_insert(entities.into())
3188 .expect("initial versions");
3189 } else {
3190 for entity in entities {
3191 manager.insert(entity).expect("initial version");
3192 }
3193 }
3194 manager.force_seal().expect("seal source");
3195 }
3196 let snapshot = SnapshotContext {
3197 snapshot: Snapshot {
3198 xid: 10,
3199 in_progress: HashSet::new(),
3200 },
3201 manager: Arc::new(SnapshotManager::new()),
3202 own_xids: HashSet::new(),
3203 requires_index_fallback: false,
3204 serializable_reader: None,
3205 };
3206 let mut changed = false;
3207 let mut seen = Vec::new();
3208 assert!(manager.scan_batches(
3209 Some(&snapshot),
3210 || true,
3211 |batch| {
3212 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3213 if !changed {
3214 changed = true;
3215 if consolidate {
3216 drain_maintenance(&manager);
3217 assert!(manager.stats().consolidation.runs_completed > 0);
3218 }
3219 for (id, xmax) in [(3, 8), (4, 12)] {
3220 let mut entity =
3221 manager.get(EntityId::new(id)).expect("retained version");
3222 entity.set_xmax(xmax);
3223 manager
3224 .update(entity)
3225 .expect("update after source retirement");
3226 }
3227 }
3228 true
3229 }
3230 ));
3231 seen.sort_unstable();
3232 assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
3233 let mut again = Vec::new();
3234 assert!(manager.scan_batches(
3235 Some(&snapshot),
3236 || true,
3237 |batch| {
3238 again.extend(batch.iter().map(|entity| entity.id.raw()));
3239 true
3240 }
3241 ));
3242 again.sort_unstable();
3243 assert_eq!(
3244 again,
3245 vec![1, 2, 4],
3246 "old slots do not duplicate relocated versions"
3247 );
3248 let mut current = Vec::new();
3249 assert!(manager.scan_batches(
3250 None,
3251 || true,
3252 |batch| {
3253 current.extend(batch.iter().map(|entity| entity.id.raw()));
3254 true
3255 }
3256 ));
3257 current.sort_unstable();
3258 assert_eq!(
3259 current,
3260 vec![1, 2],
3261 "current view excludes both deleted versions"
3262 );
3263 }
3264 }
3265
3266 #[test]
3267 fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
3268 let manager = SegmentManager::new("mixed");
3269 let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
3270 manager
3271 .bulk_insert((1..=300).map(vector).collect())
3272 .expect("flat");
3273 manager.insert(vector(500)).expect("hash insert");
3274 manager
3275 .bulk_insert(vec![vector(502), vector(505)])
3276 .expect("gapped bulk");
3277 manager.delete(EntityId::new(2)).expect("flat delete");
3278 manager.delete(EntityId::new(505)).expect("hash delete");
3279 let mut changed = false;
3280 let mut seen = Vec::new();
3281 assert!(manager.scan_batches(
3282 None,
3283 || true,
3284 |batch| {
3285 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3286 if !changed {
3287 changed = true;
3288 manager
3289 .insert(vector(700))
3290 .expect("append after cursor boundary");
3291 manager
3292 .update(UnifiedEntity::graph_node(
3293 EntityId::new(500),
3294 "mixed",
3295 "Node",
3296 HashMap::new(),
3297 ))
3298 .expect("structural update");
3299 }
3300 true
3301 }
3302 ));
3303 seen.sort_unstable();
3304 let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
3305 expected.extend([500, 502]);
3306 assert_eq!(seen, expected);
3307 }
3308
3309 #[test]
3310 fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
3311 let manager = SegmentManager::new("vectors");
3312 let mut entities = Vec::new();
3313 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
3314 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3315 entity.set_xmin(xmin);
3316 entity.set_xmax(xmax);
3317 entities.push(entity);
3318 }
3319 manager.bulk_insert(entities).expect("versions");
3320 manager.force_seal().expect("seal");
3321 let snapshot = SnapshotContext {
3322 snapshot: Snapshot {
3323 xid: 10,
3324 in_progress: HashSet::from([4]),
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 work = 0;
3332 let mut callbacks = 0;
3333 assert!(!manager.scan_batches(
3334 Some(&snapshot),
3335 || {
3336 work += 1;
3337 work <= 2
3338 },
3339 |_| {
3340 callbacks += 1;
3341 true
3342 }
3343 ));
3344 assert_eq!(callbacks, 0, "a stopped batch is not consumed");
3345 for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
3346 let mut seen = Vec::new();
3347 assert!(manager.scan_batches(
3348 view,
3349 || true,
3350 |batch| {
3351 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3352 true
3353 }
3354 ));
3355 seen.sort_unstable();
3356 assert_eq!(seen, expected);
3357 }
3358 }
3359
3360 #[test]
3361 fn test_manager_basic() {
3362 let manager = SegmentManager::new("test_collection");
3363
3364 let entity = UnifiedEntity::table_row(
3365 manager.next_entity_id(),
3366 "users",
3367 1,
3368 vec![Value::text("Alice".to_string())],
3369 );
3370
3371 let id = manager.insert(entity).unwrap();
3372 assert!(manager.get(id).is_some());
3373 assert_eq!(manager.count(), 1);
3374 }
3375
3376 #[test]
3377 fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
3378 let manager = Arc::new(SegmentManager::new("mixed"));
3379 manager
3380 .bulk_insert(
3381 (1..=128)
3382 .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
3383 .collect(),
3384 )
3385 .expect("vector-only segment");
3386 manager.force_seal().expect("seal vectors");
3387 let mut old =
3388 UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
3389 old.set_xmin(3);
3390 old.set_xmax(12);
3391 let mut new =
3392 UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
3393 new.set_xmin(12);
3394 let edge = UnifiedEntity::graph_edge(
3395 EntityId::new(202),
3396 "mixed",
3397 "a",
3398 "b",
3399 1.0,
3400 Default::default(),
3401 );
3402 manager
3403 .bulk_insert(vec![old, new, edge])
3404 .expect("physical graph versions");
3405 manager.force_seal().expect("seal graph");
3406 manager
3407 .insert(UnifiedEntity::graph_node(
3408 EntityId::new(203),
3409 "mixed",
3410 "person",
3411 Default::default(),
3412 ))
3413 .expect("growing node");
3414 let snapshot = SnapshotContext {
3415 snapshot: Snapshot {
3416 xid: 10,
3417 in_progress: HashSet::new(),
3418 },
3419 manager: Arc::new(SnapshotManager::new()),
3420 own_xids: HashSet::new(),
3421 requires_index_fallback: false,
3422 serializable_reader: None,
3423 };
3424 std::thread::spawn(move || {
3425 for view in [None, Some(&snapshot)] {
3426 for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
3427 let expected: HashSet<_> = manager
3428 .scan(view, |item| kind.matches(&item.kind))
3429 .into_iter()
3430 .map(|item| item.id)
3431 .collect();
3432 let actual: HashSet<_> = manager
3433 .scan_graph_kind(view, kind)
3434 .into_iter()
3435 .map(|item| item.id)
3436 .collect();
3437 assert_eq!(actual, expected);
3438 let mut visited = 0;
3439 let mut candidates = HashSet::new();
3440 manager.scan_graph_candidates_for_each(
3441 view,
3442 kind,
3443 || true,
3444 |item| {
3445 visited += 1;
3446 if kind.matches(&item.kind) {
3447 candidates.insert(item.id);
3448 }
3449 true
3450 },
3451 );
3452 assert_eq!(candidates, expected);
3453 assert_eq!(
3454 visited,
3455 if matches!(kind, GraphEntityKind::Node) {
3456 3
3457 } else {
3458 2
3459 }
3460 );
3461 }
3462 }
3463 let ids: HashSet<_> = manager
3464 .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
3465 .into_iter()
3466 .map(|item| item.id.raw())
3467 .collect();
3468 assert_eq!(ids, HashSet::from([200, 203]));
3469 let ids: HashSet<_> = manager
3470 .scan_graph_kind(None, GraphEntityKind::Node)
3471 .into_iter()
3472 .map(|item| item.id.raw())
3473 .collect();
3474 assert_eq!(ids, HashSet::from([201, 203]));
3475 let mut visits = 0;
3476 manager.scan_graph_candidates_for_each(
3477 None,
3478 GraphEntityKind::Node,
3479 || true,
3480 |_| {
3481 visits += 1;
3482 false
3483 },
3484 );
3485 assert_eq!(
3486 visits, 1,
3487 "early termination must also stop subsequent segments"
3488 );
3489 })
3490 .join()
3491 .expect("snapshot worker");
3492 }
3493
3494 #[test]
3495 fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
3496 let manager = Arc::new(SegmentManager::new("accounts"));
3497 let cases = [
3498 ("legacy row", 0, 0, true),
3499 ("committed creator", 3, 0, true),
3500 ("future creator", 11, 0, false),
3501 ("in-progress creator", 4, 0, false),
3502 ("committed deleter", 3, 7, false),
3503 ("future deleter", 3, 11, true),
3504 ("in-progress deleter", 3, 4, true),
3505 ];
3506
3507 for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
3508 let row_id = index as u64 + 1;
3509 let mut entity = UnifiedEntity::table_row(
3510 EntityId::new(row_id),
3511 "accounts",
3512 row_id,
3513 vec![Value::Integer(row_id as i64)],
3514 );
3515 entity.set_xmin(*xmin);
3516 entity.set_xmax(*xmax);
3517 manager.insert(entity).unwrap();
3518 }
3519
3520 let snapshot = SnapshotContext {
3521 snapshot: Snapshot {
3522 xid: 10,
3523 in_progress: HashSet::from([4]),
3524 },
3525 manager: Arc::new(SnapshotManager::new()),
3526 own_xids: HashSet::new(),
3527 requires_index_fallback: false,
3528 serializable_reader: None,
3529 };
3530
3531 let raw_visible_row_ids: HashSet<u64> = manager
3532 .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
3533 .into_iter()
3534 .filter_map(|entity| match entity.kind {
3535 EntityKind::TableRow { row_id, .. } => Some(row_id),
3536 _ => None,
3537 })
3538 .collect();
3539
3540 let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
3541 .join()
3542 .unwrap();
3543 let visible_row_ids: HashSet<u64> = scanned
3544 .into_iter()
3545 .filter_map(|entity| match entity.kind {
3546 EntityKind::TableRow { row_id, .. } => Some(row_id),
3547 _ => None,
3548 })
3549 .collect();
3550 assert_eq!(visible_row_ids, raw_visible_row_ids);
3551
3552 for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
3553 let row_id = index as u64 + 1;
3554 assert_eq!(
3555 visible_row_ids.contains(&row_id),
3556 *expected_visible,
3557 "{name} visibility mismatch"
3558 );
3559 }
3560 }
3561
3562 #[test]
3563 fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
3564 let manager = SegmentManager::new("accounts");
3568 let mut live =
3569 UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
3570 live.set_xmin(2);
3571 live.set_xmax(0);
3572 manager.insert(live).unwrap();
3573
3574 let mut superseded =
3575 UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
3576 superseded.set_xmin(2);
3577 superseded.set_xmax(9);
3578 manager.insert(superseded).unwrap();
3579
3580 let rows: Vec<u64> = manager
3581 .scan(None, |_| true)
3582 .into_iter()
3583 .filter_map(|entity| match entity.kind {
3584 EntityKind::TableRow { row_id, .. } => Some(row_id),
3585 _ => None,
3586 })
3587 .collect();
3588 assert_eq!(
3589 rows,
3590 vec![1],
3591 "frameless scan must hide superseded versions"
3592 );
3593 }
3594
3595 #[test]
3596 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
3597 let manager = SegmentManager::new("test_collection");
3598
3599 let entity = UnifiedEntity::table_row(
3600 manager.next_entity_id(),
3601 "users",
3602 1,
3603 vec![Value::text("Alice".to_string())],
3604 );
3605
3606 let id = manager.insert(entity).unwrap();
3607
3608 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3609 assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
3610 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
3611
3612 assert!(manager.delete(id).unwrap());
3613 assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3614 }
3615
3616 #[test]
3617 fn zoned_scan_reports_pruned_segments_and_preserves_results() {
3618 let config = ManagerConfig {
3619 segment_config: SegmentConfig {
3620 max_entities: 2,
3621 ..Default::default()
3622 },
3623 ..Default::default()
3624 };
3625 let manager = SegmentManager::with_config("events", config);
3626
3627 for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
3628 let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
3629 let entity = UnifiedEntity::new(
3630 manager.next_entity_id(),
3631 EntityKind::TableRow {
3632 table: "events".into(),
3633 row_id: (idx + 1) as u64,
3634 },
3635 EntityData::Row(row),
3636 );
3637 manager.insert(entity).unwrap();
3638 }
3639
3640 let full = manager.query_all(|entity| match &entity.data {
3641 EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
3642 _ => false,
3643 });
3644 let probe = Value::Integer(100);
3645 let (pruned, stats) =
3646 manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
3647 match &entity.data {
3648 EntityData::Row(row) => {
3649 matches!(row.get_field("ts"), Some(Value::Integer(100)))
3650 }
3651 _ => false,
3652 }
3653 });
3654
3655 assert_eq!(pruned.len(), full.len());
3656 assert_eq!(pruned[0].id, full[0].id);
3657 assert_eq!(stats.segments_total, 3);
3658 assert_eq!(stats.segments_scanned, 2);
3659 assert_eq!(stats.segments_pruned, 1);
3660 }
3661
3662 #[test]
3663 fn test_manager_auto_seal() {
3664 let config = ManagerConfig {
3665 segment_config: SegmentConfig {
3666 max_entities: 2,
3667 ..Default::default()
3668 },
3669 ..Default::default()
3670 };
3671
3672 let manager = SegmentManager::with_config("test", config);
3673
3674 manager
3676 .insert(UnifiedEntity::vector(
3677 manager.next_entity_id(),
3678 "v",
3679 vec![0.1],
3680 ))
3681 .unwrap();
3682
3683 manager
3685 .insert(UnifiedEntity::vector(
3686 manager.next_entity_id(),
3687 "v",
3688 vec![0.2],
3689 ))
3690 .unwrap();
3691
3692 manager
3694 .insert(UnifiedEntity::vector(
3695 manager.next_entity_id(),
3696 "v",
3697 vec![0.3],
3698 ))
3699 .unwrap();
3700
3701 let stats = manager.stats();
3702 assert_eq!(stats.total_entities, 3);
3703 }
3704
3705 #[test]
3706 fn test_manager_delete() {
3707 let manager = SegmentManager::new("test");
3708
3709 let id = manager
3710 .insert(UnifiedEntity::vector(
3711 manager.next_entity_id(),
3712 "v",
3713 vec![0.1],
3714 ))
3715 .unwrap();
3716
3717 assert!(manager.get(id).is_some());
3718 assert!(manager.delete(id).unwrap());
3719 assert!(manager.get(id).is_none());
3720 }
3721
3722 #[test]
3723 fn test_manager_metadata() {
3724 let manager = SegmentManager::new("test");
3725
3726 let id = manager
3727 .insert(UnifiedEntity::table_row(
3728 manager.next_entity_id(),
3729 "hosts",
3730 1,
3731 vec![Value::text("192.168.1.1".to_string())],
3732 ))
3733 .unwrap();
3734
3735 let mut meta = Metadata::new();
3736 meta.set(
3737 "os",
3738 super::super::metadata::MetadataValue::String("linux".to_string()),
3739 );
3740
3741 manager.set_metadata(id, meta).unwrap();
3742
3743 let retrieved = manager.get_metadata(id).unwrap();
3744 assert!(retrieved.has("os"));
3745 }
3746
3747 #[test]
3748 fn test_manager_query_by_kind() {
3749 let manager = SegmentManager::new("test");
3750
3751 manager
3752 .insert(UnifiedEntity::table_row(
3753 manager.next_entity_id(),
3754 "hosts",
3755 1,
3756 vec![],
3757 ))
3758 .unwrap();
3759
3760 manager
3761 .insert(UnifiedEntity::vector(
3762 manager.next_entity_id(),
3763 "embeddings",
3764 vec![0.1],
3765 ))
3766 .unwrap();
3767
3768 manager
3769 .insert(UnifiedEntity::table_row(
3770 manager.next_entity_id(),
3771 "hosts",
3772 2,
3773 vec![],
3774 ))
3775 .unwrap();
3776
3777 let rows = manager.get_by_kind("table");
3778 assert_eq!(rows.len(), 2);
3779
3780 let vectors = manager.get_by_kind("vector");
3781 assert_eq!(vectors.len(), 1);
3782 }
3783
3784 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
3791 SegmentManager::with_config(
3792 "test",
3793 ManagerConfig {
3794 max_sealed_segments: 1_000,
3795 consolidation_entities_per_tick: entities_per_tick,
3796 ..Default::default()
3797 },
3798 )
3799 }
3800
3801 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
3805 let entities: Vec<UnifiedEntity> = (0..count)
3806 .map(|i| {
3807 UnifiedEntity::table_row(
3808 manager.next_entity_id(),
3809 "users",
3810 i as u64,
3811 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
3812 )
3813 })
3814 .collect();
3815 let ids = manager.bulk_insert(entities).expect("bulk insert");
3816 manager.force_seal().expect("seal");
3817 ids
3818 }
3819
3820 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
3821 let mut ids: Vec<u64> = manager
3822 .query_all(|_| true)
3823 .into_iter()
3824 .map(|entity| entity.id.raw())
3825 .collect();
3826 ids.sort_unstable();
3827 ids
3828 }
3829
3830 fn drain_maintenance(manager: &SegmentManager) -> usize {
3831 let mut ticks = 0;
3832 while manager.consolidation.read().is_some() || ticks == 0 {
3833 manager.run_maintenance().expect("maintenance");
3834 ticks += 1;
3835 assert!(ticks < 1_000, "consolidation failed to converge");
3836 }
3837 ticks
3838 }
3839
3840 #[test]
3841 fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
3842 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3843 let ids = manager
3844 .bulk_insert(
3845 (1..=1024)
3846 .map(|id| {
3847 UnifiedEntity::graph_edge(
3848 EntityId::new(id),
3849 "step",
3850 "2000",
3851 "3000",
3852 1.0,
3853 HashMap::new(),
3854 )
3855 })
3856 .collect(),
3857 )
3858 .expect("edges");
3859 manager.force_seal().expect("seal");
3860 for id in ids.iter().take(300) {
3861 manager.delete(*id).expect("fragment source");
3862 }
3863 let source = Arc::downgrade(&manager.sealed.read()[0]);
3864 let mut seen = Vec::new();
3865 let mut changed = false;
3866 let mut metadata_bytes = 0;
3867 assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
3868 || true,
3869 |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
3870 |batch| {
3871 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3872 assert!(manager.growing.try_write().is_some(), "no growing topology guard");
3873 assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
3874 assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
3875 if !changed {
3876 changed = true;
3877 drain_maintenance(&manager);
3878 assert!(manager.stats().consolidation.runs_completed > 0);
3879 assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
3880 let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
3881 if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
3882 manager.update(updated).expect("update authoritative merged copy");
3883 }
3884 for id in batch {
3885 let entity = manager.get(*id).expect("hydrate current manager");
3886 if id.raw() == 1024 {
3887 assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
3888 }
3889 seen.push(id.raw());
3890 }
3891 Ok(true)
3892 }).expect("cursor"));
3893 assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
3894 assert!(
3895 source.upgrade().is_none(),
3896 "finished cursor releases retired source"
3897 );
3898 assert!(
3899 metadata_bytes < 4096,
3900 "fixed ID buffer plus source handles: {metadata_bytes}"
3901 );
3902 }
3903
3904 #[test]
3905 fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
3906 let manager = SegmentManager::new("graph");
3907 manager
3908 .bulk_insert(
3909 (1..=1024)
3910 .map(|id| {
3911 UnifiedEntity::graph_edge(
3912 EntityId::new(id),
3913 "step",
3914 "2000",
3915 "3000",
3916 1.0,
3917 HashMap::new(),
3918 )
3919 })
3920 .collect(),
3921 )
3922 .expect("edges");
3923 let mut delivered = 0;
3924 let denied = manager.visit_graph_index_batches(
3925 GraphEntityKind::Edge,
3926 &[EntityId::new(2000)],
3927 || true,
3928 |_| Err("out of memory"),
3929 |ids| {
3930 delivered += ids.len();
3931 Ok(true)
3932 },
3933 );
3934 assert_eq!(denied, Err("out of memory"));
3935 assert_eq!(delivered, 0);
3936 let mut work = 0;
3937 assert!(!manager
3938 .visit_graph_index_batches(
3939 GraphEntityKind::Edge,
3940 &[EntityId::new(2000)],
3941 || {
3942 work += 1;
3943 work < 100
3944 },
3945 |_| Ok::<_, ()>(()),
3946 |ids| {
3947 delivered += ids.len();
3948 Ok(true)
3949 }
3950 )
3951 .expect("cancelled cursor"));
3952 assert_eq!(
3953 delivered, 0,
3954 "cancelled partial batch cannot reach consumer"
3955 );
3956 assert!(!manager
3957 .visit_graph_index_batches(
3958 GraphEntityKind::Edge,
3959 &[EntityId::new(2000)],
3960 || true,
3961 |_| Ok::<_, ()>(()),
3962 |ids| {
3963 delivered += ids.len();
3964 Ok(false)
3965 }
3966 )
3967 .expect("consumer stop"));
3968 assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
3969 }
3970
3971 #[test]
3972 fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
3973 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3974 let ids = seal_bulk_segment(&manager, 100);
3975 for id in ids.iter().take(40) {
3976 manager.delete(*id).expect("delete before scan");
3977 }
3978 let source_bytes = manager.resident_bytes();
3979 let source = Arc::downgrade(&manager.sealed.read()[0]);
3980 std::thread::scope(|scope| {
3981 let (ready, waiting) = std::sync::mpsc::channel();
3982 let (release_first, first_wait) = std::sync::mpsc::channel();
3983 let (release_last, last_wait) = std::sync::mpsc::channel();
3984 let reader = |resume: std::sync::mpsc::Receiver<()>,
3985 ready: std::sync::mpsc::Sender<()>| {
3986 assert!(manager.scan_batches(
3987 None,
3988 || true,
3989 |_| {
3990 ready.send(()).expect("signal captured cursor");
3991 let _ = resume.recv();
3992 true
3993 }
3994 ));
3995 };
3996 let ready_first = ready.clone();
3997 let first = scope.spawn(move || reader(first_wait, ready_first));
3998 let last = scope.spawn(move || reader(last_wait, ready));
3999 for _ in 0..2 {
4000 waiting
4001 .recv_timeout(std::time::Duration::from_secs(10))
4002 .expect("reader ready");
4003 }
4004 drain_maintenance(&manager);
4005 assert!(
4006 source.upgrade().is_some(),
4007 "both cursors pin the retired source"
4008 );
4009 let active_bytes = manager.sealed.read()[0].read().resident_bytes();
4010 let held = manager.resident_bytes();
4011 assert!(
4012 held >= active_bytes + source_bytes,
4013 "retired source remains resident"
4014 );
4015 release_first.send(()).expect("release first reader");
4016 first.join().expect("first reader");
4017 assert_eq!(
4018 manager.resident_bytes(),
4019 held,
4020 "one reader still pins the source"
4021 );
4022 release_last.send(()).expect("release last reader");
4023 last.join().expect("last reader");
4024 assert!(
4025 source.upgrade().is_none(),
4026 "accounting must not retain payloads"
4027 );
4028 assert_eq!(held - manager.resident_bytes(), source_bytes);
4029 assert_eq!(
4030 manager.stats().total_memory_bytes as u64,
4031 manager.resident_bytes()
4032 );
4033 });
4034 }
4035
4036 #[test]
4037 fn retained_memory_includes_an_inflight_consolidation_copy() {
4038 let manager = consolidating_manager(1);
4039 let ids = seal_bulk_segment(&manager, 100);
4040 for id in ids.iter().take(40) {
4041 manager.delete(*id).expect("delete");
4042 }
4043 let before = manager.resident_bytes();
4044 assert!(manager.pressure_consolidation_tick());
4045 let copied = manager
4046 .consolidation
4047 .read()
4048 .as_ref()
4049 .expect("paced merge")
4050 .merged
4051 .resident_bytes();
4052 assert!(copied > 0);
4053 assert!(
4054 manager.resident_bytes() >= before + copied,
4055 "unpublished copy consumes memory"
4056 );
4057 drain_maintenance(&manager);
4058 assert!(
4059 manager.resident_bytes() < before,
4060 "completed merge releases the old storage"
4061 );
4062 }
4063
4064 #[test]
4065 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
4066 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4067 let ids = seal_bulk_segment(&manager, 100);
4068
4069 let bytes_before_delete = manager.resident_bytes();
4070
4071 for id in ids.iter().take(40) {
4073 assert!(manager.delete(*id).expect("delete"));
4074 }
4075
4076 let bytes_with_tombstones = manager.resident_bytes();
4079 assert!(bytes_with_tombstones > bytes_before_delete);
4080
4081 let survivors = live_ids(&manager);
4082 assert_eq!(survivors.len(), 60);
4083
4084 drain_maintenance(&manager);
4085
4086 let stats = manager.stats();
4087 assert_eq!(stats.consolidation.runs_started, 1);
4088 assert_eq!(stats.consolidation.runs_completed, 1);
4089 assert_eq!(stats.consolidation.segments_merged, 1);
4090 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
4091 assert!(stats.consolidation.bytes_reclaimed > 0);
4092
4093 let bytes_after = manager.resident_bytes();
4095 assert!(
4096 bytes_after <= bytes_before_delete,
4097 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
4098 );
4099 assert_eq!(
4100 bytes_with_tombstones - bytes_after,
4101 stats.consolidation.bytes_reclaimed
4102 );
4103 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
4104
4105 assert_eq!(live_ids(&manager), survivors);
4106 assert_eq!(manager.sealed.read().len(), 1);
4107 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
4108 }
4109
4110 #[test]
4111 fn query_results_are_identical_before_during_and_after_consolidation() {
4112 let manager = consolidating_manager(3);
4113 let ids = seal_bulk_segment(&manager, 20);
4114 seal_bulk_segment(&manager, 20);
4115 for id in ids.iter().take(10) {
4116 assert!(manager.delete(*id).expect("delete"));
4117 }
4118
4119 let before = live_ids(&manager);
4120 let survivor = *ids.last().expect("at least one row");
4121 let tombstoned = ids[0];
4122
4123 let mut mid_tick_observations = 0;
4126 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
4127 manager.run_maintenance().expect("maintenance");
4128 assert_eq!(live_ids(&manager), before);
4129 assert!(manager.get(survivor).is_some());
4130 assert!(manager.get(tombstoned).is_none());
4131 assert_eq!(manager.get_by_kind("table").len(), before.len());
4132 mid_tick_observations += 1;
4133 assert!(mid_tick_observations < 1_000);
4134 }
4135
4136 assert!(
4137 mid_tick_observations > 2,
4138 "expected the run to span several ticks, saw {mid_tick_observations}"
4139 );
4140 assert_eq!(live_ids(&manager), before);
4141 assert!(manager.get(survivor).is_some());
4142 assert!(manager.get(tombstoned).is_none());
4143 assert_eq!(manager.get_by_kind("table").len(), before.len());
4144 }
4145
4146 #[test]
4147 fn get_by_kind_reaches_sealed_segments() {
4148 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4149 let ids = seal_bulk_segment(&manager, 5);
4150 assert!(manager.delete(ids[0]).expect("delete"));
4151
4152 assert_eq!(manager.get_by_kind("table").len(), 4);
4156 assert_eq!(manager.get_by_kind("vector").len(), 0);
4157 }
4158
4159 #[test]
4160 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
4161 const PER_TICK: usize = 4;
4162 let manager = consolidating_manager(PER_TICK);
4163 let ids = seal_bulk_segment(&manager, 60);
4164 for id in ids.iter().take(30) {
4165 assert!(manager.delete(*id).expect("delete"));
4166 }
4167
4168 manager.run_maintenance().expect("maintenance");
4170 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
4171
4172 let mut previous = PER_TICK as u64;
4175 let mut ticks = 1;
4176 while let Some(copied) = manager.consolidation_progress() {
4177 assert!(
4178 copied - previous <= PER_TICK as u64,
4179 "tick copied {} entities, bound is {PER_TICK}",
4180 copied - previous
4181 );
4182 previous = copied;
4183 manager.run_maintenance().expect("maintenance");
4184 ticks += 1;
4185 assert!(ticks < 100);
4186 }
4187
4188 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
4191 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4192 assert_eq!(live_ids(&manager).len(), 30);
4193 }
4194
4195 #[test]
4196 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
4197 let manager = SegmentManager::with_config(
4198 "test",
4199 ManagerConfig {
4200 max_sealed_segments: 3,
4201 ..Default::default()
4202 },
4203 );
4204 for _ in 0..4 {
4205 seal_bulk_segment(&manager, 5);
4206 }
4207 assert_eq!(manager.sealed.read().len(), 4);
4208
4209 let before = live_ids(&manager);
4210 drain_maintenance(&manager);
4211
4212 assert_eq!(manager.sealed.read().len(), 1);
4215 assert_eq!(manager.stats().consolidation.segments_merged, 4);
4216 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
4217 assert_eq!(live_ids(&manager), before);
4218 }
4219
4220 #[test]
4221 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
4222 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4223 seal_bulk_segment(&manager, 10);
4224 seal_bulk_segment(&manager, 10);
4225
4226 for _ in 0..5 {
4227 manager.run_maintenance().expect("maintenance");
4228 }
4229
4230 assert_eq!(manager.stats().consolidation.runs_started, 0);
4231 assert_eq!(manager.sealed.read().len(), 2);
4232 }
4233
4234 #[test]
4235 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
4236 use std::sync::atomic::AtomicBool;
4237
4238 let manager = Arc::new(consolidating_manager(2));
4239 let ids = seal_bulk_segment(&manager, 40);
4240 seal_bulk_segment(&manager, 40);
4241 for id in ids.iter().take(20) {
4242 assert!(manager.delete(*id).expect("delete"));
4243 }
4244
4245 let expected = live_ids(&manager);
4246 let stop = Arc::new(AtomicBool::new(false));
4247
4248 let readers: Vec<_> = (0..4)
4249 .map(|_| {
4250 let manager = Arc::clone(&manager);
4251 let stop = Arc::clone(&stop);
4252 let expected = expected.clone();
4253 std::thread::spawn(move || {
4254 let mut scans = 0u64;
4255 loop {
4258 let mut seen: Vec<u64> = manager
4259 .query_all(|_| true)
4260 .into_iter()
4261 .map(|entity| entity.id.raw())
4262 .collect();
4263 seen.sort_unstable();
4264 assert_eq!(seen, expected, "reader saw a torn sealed set");
4267 scans += 1;
4268 if stop.load(Ordering::Relaxed) {
4269 break;
4270 }
4271 }
4272 scans
4273 })
4274 })
4275 .collect();
4276
4277 drain_maintenance(&manager);
4278 stop.store(true, Ordering::Relaxed);
4279
4280 for reader in readers {
4281 let scans = reader.join().expect("reader thread");
4282 assert!(scans > 0, "reader never got to scan");
4283 }
4284
4285 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4286 assert_eq!(live_ids(&manager), expected);
4287 }
4288
4289 #[test]
4290 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
4291 let manager = consolidating_manager(4);
4292 let ids = seal_bulk_segment(&manager, 20);
4293 for id in ids.iter().take(6) {
4294 assert!(manager.delete(*id).expect("delete"));
4295 }
4296
4297 manager.run_maintenance().expect("maintenance");
4299 assert_eq!(manager.consolidation_progress(), Some(4));
4300 let doomed = ids[6];
4301 assert!(manager.delete(doomed).expect("delete"));
4302
4303 let expected = live_ids(&manager);
4304 drain_maintenance(&manager);
4305
4306 assert!(manager.get(doomed).is_none(), "deleted row came back");
4307 assert_eq!(live_ids(&manager), expected);
4308 }
4309
4310 #[test]
4311 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
4312 let manager = consolidating_manager(3);
4313 let ids = seal_bulk_segment(&manager, 30);
4314 for id in ids.iter().take(12) {
4315 assert!(manager.delete(*id).expect("delete"));
4316 }
4317
4318 let before = live_ids(&manager);
4319 let sealed_before: Vec<SegmentId> = manager
4320 .sealed
4321 .read()
4322 .iter()
4323 .map(|segment| segment.read().id())
4324 .collect();
4325
4326 manager.run_maintenance().expect("maintenance");
4327 manager.run_maintenance().expect("maintenance");
4328 assert!(manager.consolidation.read().is_some());
4329
4330 *manager.consolidation.write() = None;
4334
4335 assert_eq!(live_ids(&manager), before);
4336 let sealed_after: Vec<SegmentId> = manager
4337 .sealed
4338 .read()
4339 .iter()
4340 .map(|segment| segment.read().id())
4341 .collect();
4342 assert_eq!(sealed_after, sealed_before);
4343 assert_eq!(manager.stats().consolidation.runs_completed, 0);
4344
4345 drain_maintenance(&manager);
4346 assert_eq!(live_ids(&manager), before);
4347 assert_eq!(manager.stats().consolidation.runs_completed, 1);
4348 }
4349
4350 #[test]
4351 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
4352 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4353
4354 let entities: Vec<UnifiedEntity> = (0..10)
4358 .map(|i| {
4359 UnifiedEntity::table_row(
4360 manager.next_entity_id(),
4361 "users",
4362 i as u64,
4363 vec![Value::Integer(i as i64)],
4364 )
4365 })
4366 .collect();
4367 let ids = manager.bulk_insert(entities).expect("bulk insert");
4368
4369 let mut meta = Metadata::new();
4370 meta.set(
4371 "os",
4372 super::super::metadata::MetadataValue::String("linux".to_string()),
4373 );
4374 manager.set_metadata(ids[9], meta).expect("set metadata");
4375 manager.force_seal().expect("seal");
4376
4377 assert!(manager.delete(ids[0]).expect("delete"));
4378 assert!(manager.delete(ids[1]).expect("delete"));
4379 assert!(manager.delete(ids[2]).expect("delete"));
4380
4381 let before: Vec<UnifiedEntity> = {
4382 let mut entities = manager.query_all(|_| true);
4383 entities.sort_by_key(|entity| entity.id.raw());
4384 entities
4385 };
4386
4387 drain_maintenance(&manager);
4388
4389 let after: Vec<UnifiedEntity> = {
4390 let mut entities = manager.query_all(|_| true);
4391 entities.sort_by_key(|entity| entity.id.raw());
4392 entities
4393 };
4394
4395 let columns = |entity: &UnifiedEntity| match &entity.data {
4396 super::super::entity::EntityData::Row(row) => row.columns.clone(),
4397 _ => Vec::new(),
4398 };
4399
4400 assert_eq!(before.len(), after.len());
4401 for (old, new) in before.iter().zip(after.iter()) {
4402 assert_eq!(old.id, new.id);
4403 assert_eq!(columns(old), columns(new));
4404 assert_eq!(
4405 old.sequence_id, new.sequence_id,
4406 "merge must not re-issue sequence ids"
4407 );
4408 }
4409 assert!(manager
4410 .get_metadata(ids[9])
4411 .expect("metadata survives the merge")
4412 .has("os"));
4413 }
4414
4415 #[test]
4416 #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
4417 fn test_lifecycle_events() {
4418 let manager = SegmentManager::new("test");
4419
4420 manager
4421 .insert(UnifiedEntity::vector(
4422 manager.next_entity_id(),
4423 "v",
4424 vec![0.1],
4425 ))
4426 .unwrap();
4427
4428 let events = manager.drain_events();
4429
4430 assert!(events
4432 .iter()
4433 .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
4434 assert!(events
4435 .iter()
4436 .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
4437 }
4438}