1#![allow(clippy::too_many_arguments)] mod backing;
18pub use backing::PagedEntityCursor;
19mod graph_index_merge;
20
21use parking_lot::RwLock;
22use std::collections::{BTreeSet, HashMap};
23use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
24use std::sync::{Arc, Weak};
25
26use super::entity::{EntityId, EntityKind, UnifiedEntity};
27use super::metadata::{Metadata, MetadataFilter, MetadataValue};
28use super::segment::{
29 GraphEntityKind, GrowingSegment, SegmentConfig, SegmentError, SegmentId, SegmentScanCursor,
30 SegmentState, UnifiedSegment, ZoneColPred,
31};
32use super::visibility_map::VisibilityMap;
33use crate::storage::visibility::entity_visible_with_context;
34use crate::transaction::read::SnapshotContext;
35
36pub const CONSOLIDATION_TOMBSTONE_RATIO: f64 = 0.20;
43
44pub const CONSOLIDATION_FRAGMENTATION_RATIO: f64 = 0.30;
49
50pub const CONSOLIDATION_ENTITIES_PER_TICK: usize = 4_096;
56
57#[derive(Debug, Clone)]
59pub struct ManagerConfig {
60 pub segment_config: SegmentConfig,
62 pub max_sealed_segments: usize,
64 pub idle_seal_secs: u64,
66 pub enable_consolidation: bool,
68 pub consolidation_entities_per_tick: usize,
70 pub enable_archival: bool,
72 pub archive_age_secs: u64,
74}
75
76impl Default for ManagerConfig {
77 fn default() -> Self {
78 Self {
79 segment_config: SegmentConfig::default(),
80 max_sealed_segments: 10,
81 idle_seal_secs: 300, enable_consolidation: true,
83 consolidation_entities_per_tick: CONSOLIDATION_ENTITIES_PER_TICK,
84 enable_archival: true,
85 archive_age_secs: 86400 * 7, }
87 }
88}
89
90#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
95pub struct ConsolidationStats {
96 pub runs_started: u64,
98 pub runs_completed: u64,
100 pub segments_merged: u64,
102 pub tombstones_reclaimed: u64,
104 pub bytes_reclaimed: u64,
106}
107
108#[derive(Debug, Clone, Default)]
110pub struct ManagerStats {
111 pub total_entities: usize,
113 pub growing_count: usize,
115 pub sealed_count: usize,
117 pub archived_count: usize,
119 pub total_memory_bytes: usize,
121 pub seal_ops: u64,
123 pub consolidation: ConsolidationStats,
125}
126
127#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
128pub struct SegmentScanStats {
129 pub segments_total: u64,
130 pub segments_scanned: u64,
131 pub segments_pruned: u64,
132}
133
134#[derive(Debug, Clone)]
136pub enum LifecycleEvent {
137 SegmentCreated(SegmentId),
138 SegmentSealed(SegmentId),
139 SegmentConsolidated {
140 source: Vec<SegmentId>,
141 target: SegmentId,
142 },
143 SegmentArchived(SegmentId),
144 EntityInserted(EntityId, SegmentId),
145 EntityDeleted(EntityId, SegmentId),
146}
147
148struct Consolidation {
155 sources: Vec<SourceProgress>,
157 merged_id: SegmentId,
159 merged: GrowingSegment,
161 cursor: usize,
163 pending_ids: Vec<EntityId>,
165 pending_cursor: usize,
167 entities_copied: u64,
169}
170
171struct SourceProgress {
173 id: SegmentId,
174 copied: Vec<EntityId>,
176 tombstones_at_open: Option<usize>,
185}
186
187struct PhysicalRootCacheEntry {
188 segments: Vec<(u64, u64)>,
189 seed: u64,
190 root: u64,
191}
192
193#[doc(hidden)]
196pub struct PhysicalScanCursor {
197 sources: Vec<(Weak<RwLock<GrowingSegment>>, SegmentScanCursor)>,
198 source: usize,
199}
200
201#[derive(Default)]
202struct MetadataExpiryIndex {
203 by_deadline: BTreeSet<(u64, EntityId)>,
204 by_entity: HashMap<EntityId, u64>,
205}
206
207impl MetadataExpiryIndex {
208 fn replace(&mut self, entity_id: EntityId, deadline_ms: Option<u64>) {
209 if let Some(previous_ms) = self.by_entity.remove(&entity_id) {
210 self.by_deadline.remove(&(previous_ms, entity_id));
211 }
212 if let Some(deadline_ms) = deadline_ms {
213 self.by_deadline.insert((deadline_ms, entity_id));
214 self.by_entity.insert(entity_id, deadline_ms);
215 }
216 }
217
218 fn expired(&self, now_ms: u64, limit: usize) -> Vec<EntityId> {
219 if limit == 0 {
220 return Vec::new();
221 }
222 self.by_deadline
223 .range(..=(now_ms, EntityId::new(u64::MAX)))
224 .map(|(_, entity_id)| *entity_id)
225 .take(limit)
226 .collect()
227 }
228
229 fn scheduled_count(&self) -> usize {
230 self.by_entity.len()
231 }
232
233 fn entries(&self) -> Vec<(u64, EntityId)> {
234 self.by_deadline.iter().copied().collect()
235 }
236
237 fn replace_all(&mut self, entries: &[(u64, EntityId)]) {
238 self.by_deadline.clear();
239 self.by_entity.clear();
240 for &(deadline_ms, entity_id) in entries {
241 self.by_deadline.insert((deadline_ms, entity_id));
242 self.by_entity.insert(entity_id, deadline_ms);
243 }
244 }
245}
246
247fn metadata_u64(value: &MetadataValue) -> Option<u64> {
248 match value {
249 MetadataValue::Int(value) if *value >= 0 => Some(*value as u64),
250 MetadataValue::Timestamp(value) => Some(*value),
251 MetadataValue::Float(value)
252 if value.is_finite()
253 && value.is_sign_positive()
254 && value.fract().abs() < f64::EPSILON
255 && *value <= u64::MAX as f64 =>
256 {
257 Some(value.trunc() as u64)
258 }
259 MetadataValue::String(value) => value.parse().ok(),
260 _ => None,
261 }
262}
263
264fn metadata_expiry_deadline_ms(entity: &UnifiedEntity, metadata: &Metadata) -> Option<u64> {
265 let absolute_ms = metadata.get("_expires_at").and_then(metadata_u64);
266 let ttl_ms = metadata.get("_ttl_ms").and_then(metadata_u64).or_else(|| {
267 metadata
268 .get("_ttl")
269 .and_then(metadata_u64)
270 .and_then(|seconds| seconds.checked_mul(1_000))
271 });
272 let relative_ms = ttl_ms.map(|ttl_ms| {
273 entity
274 .created_at
275 .saturating_mul(1_000)
276 .saturating_add(ttl_ms)
277 });
278 match (absolute_ms, relative_ms) {
279 (Some(absolute_ms), Some(relative_ms)) => Some(absolute_ms.min(relative_ms)),
280 (Some(deadline_ms), None) | (None, Some(deadline_ms)) => Some(deadline_ms),
281 (None, None) => None,
282 }
283}
284
285impl PhysicalScanCursor {
286 #[doc(hidden)]
287 pub fn finished(&self) -> bool {
288 self.source == self.sources.len()
289 }
290}
291
292pub struct SegmentManager {
294 row_constraint_lock: Arc<parking_lot::ReentrantMutex<()>>,
297 collection: String,
299 config: ManagerConfig,
301 next_segment_id: AtomicU64,
303 next_entity_id: AtomicU64,
305 next_sequence_id: AtomicU64,
306 next_row_id: AtomicU64,
308 total_entities_atomic: AtomicU64,
312 growing: RwLock<Option<Arc<RwLock<GrowingSegment>>>>,
314 sealed: RwLock<Vec<Arc<RwLock<GrowingSegment>>>>,
316 retired: RwLock<Vec<Weak<RwLock<GrowingSegment>>>>,
318 archived: RwLock<Vec<SegmentId>>,
320 entity_segment: RwLock<HashMap<EntityId, SegmentId>>,
324 metadata_expiry: RwLock<MetadataExpiryIndex>,
328 metadata_expiry_projection_ready: AtomicBool,
332 column_schema: RwLock<Option<Arc<Vec<String>>>>,
335 recovery_row_schemas: RwLock<Vec<Arc<Vec<String>>>>,
338 stats: RwLock<ManagerStats>,
340 consolidation: RwLock<Option<Consolidation>>,
343 events: RwLock<Vec<LifecycleEvent>>,
345 visibility_map: VisibilityMap,
349 physical_root_cache: parking_lot::Mutex<Option<PhysicalRootCacheEntry>>,
350 backing: RwLock<Option<backing::BackingRead>>,
351}
352
353impl SegmentManager {
354 pub fn new(collection: impl Into<String>) -> Self {
356 Self::with_config(collection, ManagerConfig::default())
357 }
358
359 pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
361 Self {
362 row_constraint_lock: Arc::new(parking_lot::ReentrantMutex::new(())),
363 physical_root_cache: parking_lot::Mutex::new(None),
364 backing: RwLock::new(None),
365 collection: collection.into(),
366 config,
367 next_segment_id: AtomicU64::new(1),
368 next_entity_id: AtomicU64::new(1),
369 next_sequence_id: AtomicU64::new(0),
370 next_row_id: AtomicU64::new(1),
371 total_entities_atomic: AtomicU64::new(0),
372 growing: RwLock::new(None),
373 sealed: RwLock::new(Vec::new()),
374 retired: RwLock::new(Vec::new()),
375 archived: RwLock::new(Vec::new()),
376 entity_segment: RwLock::new(HashMap::new()),
377 metadata_expiry: RwLock::new(MetadataExpiryIndex::default()),
378 metadata_expiry_projection_ready: AtomicBool::new(true),
379 column_schema: RwLock::new(None),
380 recovery_row_schemas: RwLock::new(Vec::new()),
381 stats: RwLock::new(ManagerStats::default()),
382 consolidation: RwLock::new(None),
383 events: RwLock::new(Vec::new()),
384 visibility_map: VisibilityMap::new(),
385 }
386 }
387
388 #[doc(hidden)]
389 pub fn row_constraint_lock(&self) -> Arc<parking_lot::ReentrantMutex<()>> {
390 Arc::clone(&self.row_constraint_lock)
391 }
392
393 pub fn get_or_init_schema(
395 &self,
396 named: &HashMap<String, reddb_types::Value>,
397 ) -> Arc<Vec<String>> {
398 {
399 let schema = self.column_schema.read();
400 if let Some(ref s) = *schema {
401 return Arc::clone(s);
402 }
403 }
404 let cols: Vec<String> = named.keys().cloned().collect();
405 let arc = Arc::new(cols);
406 *self.column_schema.write() = Some(Arc::clone(&arc));
407 arc
408 }
409
410 pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
412 self.column_schema.read().clone()
413 }
414
415 #[doc(hidden)]
419 pub fn restore_owned_row_schemas<'a>(
420 &self,
421 entities: impl Iterator<Item = &'a mut UnifiedEntity>,
422 ) {
423 let schema = self.column_schema();
424 for entity in entities {
425 if let super::entity::EntityData::Row(row) = &mut entity.data {
426 if row.named.is_none() && row.schema.is_none() {
427 row.schema = schema.clone();
428 }
429 }
430 }
431 }
432
433 #[doc(hidden)]
434 pub fn set_column_schema_if_empty(&self, columns: Vec<String>) {
435 if columns.is_empty()
436 || columns
437 .iter()
438 .collect::<std::collections::HashSet<_>>()
439 .len()
440 != columns.len()
441 {
442 return;
443 }
444 let mut schema = self.column_schema.write();
445 if schema.is_none() {
446 *schema = Some(Arc::new(columns));
447 }
448 }
449
450 #[doc(hidden)]
454 pub fn compact_recovered_rows(&self) -> usize {
455 if self.recovery_preserves_named_rows() {
456 return 0;
457 }
458 let Some(schema) = self.column_schema() else {
459 return 0;
460 };
461 if schema
462 .iter()
463 .collect::<std::collections::HashSet<_>>()
464 .len()
465 != schema.len()
466 {
467 return 0;
468 }
469 let mut growing = self.growing.write();
470 let mut sealed = self.sealed.write();
471 let mut compacted = 0;
472 for segment in growing.iter_mut().chain(sealed.iter_mut()) {
473 if let Some(segment) = Arc::get_mut(segment) {
474 compacted += segment.get_mut().compact_recovered_rows(&schema);
475 }
476 }
477 compacted
478 }
479
480 pub fn collection(&self) -> &str {
482 &self.collection
483 }
484
485 pub fn config(&self) -> &ManagerConfig {
487 &self.config
488 }
489
490 pub fn stats(&self) -> ManagerStats {
493 let mut s = self.stats.read().clone();
494 s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
495 s.total_memory_bytes = self.resident_bytes() as usize;
496 s
497 }
498
499 pub fn resident_bytes(&self) -> u64 {
502 self.memory_usage().1
503 }
504
505 pub fn memory_bytes(&self) -> u64 {
508 self.memory_usage().0
509 }
510
511 fn memory_usage(&self) -> (u64, u64) {
515 let consolidation = self.consolidation.read();
516 let growing = self.growing.read();
517 let sealed = self.sealed.read();
518 let mut retired = self.retired.write();
519 let mut payload = 0u64;
520 let mut resident = 0u64;
521 {
522 let mut add = |segment: &GrowingSegment| {
523 payload = payload.saturating_add(segment.memory_bytes());
524 resident = resident.saturating_add(segment.resident_bytes());
525 };
526 for segment in growing.iter().chain(sealed.iter()) {
527 add(&segment.read());
528 }
529 retired.retain(|source| {
530 if let Some(source) = source.upgrade() {
531 add(&source.read());
532 true
533 } else {
534 false
535 }
536 });
537 if let Some(run) = consolidation.as_ref() {
538 add(&run.merged);
539 }
540 }
541 let mut bookkeeping =
542 (retired.capacity() * std::mem::size_of::<Weak<RwLock<GrowingSegment>>>()) as u64;
543 if let Some(run) = consolidation.as_ref() {
544 bookkeeping = bookkeeping
545 .saturating_add(
546 (run.sources.capacity() * std::mem::size_of::<SourceProgress>()) as u64,
547 )
548 .saturating_add(
549 (run.pending_ids.capacity() * std::mem::size_of::<EntityId>()) as u64,
550 );
551 for source in &run.sources {
552 bookkeeping = bookkeeping.saturating_add(
553 (source.copied.capacity() * std::mem::size_of::<EntityId>()) as u64,
554 );
555 }
556 }
557 (
558 payload.saturating_add(bookkeeping),
559 resident.saturating_add(bookkeeping),
560 )
561 }
562
563 pub fn reclaimable_bytes(&self) -> u64 {
565 let growing = self.growing.read();
566 let sealed = self.sealed.read();
567 growing
568 .iter()
569 .chain(sealed.iter())
570 .map(|segment| segment.read().reclaimable_bytes())
571 .fold(0, u64::saturating_add)
572 }
573
574 pub fn next_entity_id(&self) -> EntityId {
576 EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
577 }
578
579 pub fn identity_counters(&self) -> (u64, u64) {
581 (
582 self.next_entity_id.load(Ordering::SeqCst),
583 self.next_row_id.load(Ordering::SeqCst),
584 )
585 }
586
587 pub fn restore_identity_counters(&self, entity_id_next: u64, row_id_next: u64) {
588 assert!(
589 entity_id_next > 0 && row_id_next > 0,
590 "invariant: next identities are positive"
591 );
592 self.next_entity_id
593 .fetch_max(entity_id_next, Ordering::SeqCst);
594 self.next_row_id.fetch_max(row_id_next, Ordering::SeqCst);
595 }
596
597 pub fn sequence_id_next(&self) -> u64 {
598 let mut next = self.next_sequence_id.load(Ordering::SeqCst);
599 let growing = self.growing.read();
600 let sealed = self.sealed.read();
601 for segment in growing.iter().chain(sealed.iter()) {
602 next = next.max(segment.read().sequence_id_next());
603 }
604 self.next_sequence_id.fetch_max(next, Ordering::SeqCst);
605 next
606 }
607
608 pub fn restore_sequence_id_next(&self, next: u64) {
609 self.next_sequence_id.fetch_max(next, Ordering::SeqCst);
610 }
611
612 pub fn next_row_id(&self) -> u64 {
614 self.next_row_id.fetch_add(1, Ordering::SeqCst)
615 }
616
617 pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
621 let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
622 start..start + n
623 }
624
625 pub fn register_row_id(&self, id: u64) {
628 let candidate = id.saturating_add(1);
629 let mut current = self.next_row_id.load(Ordering::SeqCst);
630 while candidate > current {
631 match self.next_row_id.compare_exchange(
632 current,
633 candidate,
634 Ordering::SeqCst,
635 Ordering::SeqCst,
636 ) {
637 Ok(_) => break,
638 Err(updated) => current = updated,
639 }
640 }
641 }
642
643 fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
650 {
652 let growing = self.growing.read();
653 if let Some(segment) = growing.as_ref() {
654 return Arc::clone(segment);
655 }
656 }
657
658 let mut growing = self.growing.write();
660 if let Some(segment) = growing.as_ref() {
662 return Arc::clone(segment);
663 }
664
665 let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
666 let segment = GrowingSegment::new(id, &self.collection);
667 segment.restore_sequence_id_next(self.next_sequence_id.load(Ordering::SeqCst));
668 let segment_arc = Arc::new(RwLock::new(segment));
669 *growing = Some(Arc::clone(&segment_arc));
670
671 self.emit(LifecycleEvent::SegmentCreated(id));
672
673 self.stats.write().growing_count += 1;
676
677 segment_arc
678 }
679
680 pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
682 if entity.id.raw() == 0 {
683 entity.id = self.next_entity_id();
684 }
685 self.reject_backing_duplicate(entity.id)?;
686 if self.has_backing() && self.get_with(entity.id, |_| ()).is_some() {
687 return Err(SegmentError::AlreadyExists(entity.id));
688 }
689 self.next_entity_id
690 .fetch_max(entity.id.raw().saturating_add(1), Ordering::SeqCst);
691 self.maybe_seal_growing()?;
693
694 let segment_arc = self.get_or_create_growing();
695 let mut segment = segment_arc.write();
696
697 let entity_id = entity.id;
698 let segment_id = segment.id();
699
700 segment.insert(entity)?;
701
702 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
704
705 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
714
715 Ok(entity_id)
716 }
717
718 fn with_restoration_segment<T>(
722 &self,
723 id: EntityId,
724 restore: impl FnOnce(&mut GrowingSegment) -> Result<T, SegmentError>,
725 ) -> Result<T, SegmentError> {
726 loop {
727 let segment = self.get_or_create_growing();
728 let mut growing = segment.write();
729 if !growing.state().is_writable() {
730 continue;
731 }
732 if !growing.is_tombstoned(id) {
733 return restore(&mut growing);
734 }
735 drop(growing);
736 self.force_seal()?;
737 }
738 }
739
740 fn recovery_preserves_named_rows(&self) -> bool {
741 self.collection.starts_with("red_") || self.collection.starts_with("red.")
745 }
746
747 fn recovery_schema(
748 &self,
749 row: &super::entity::RowData,
750 column_schema: Option<Arc<Vec<String>>>,
751 ) -> Option<Arc<Vec<String>>> {
752 if self.recovery_preserves_named_rows() {
753 return None;
754 }
755 let named = row.named.as_ref()?;
756 let matches = |schema: &Vec<String>| {
757 schema.len() == named.len() && schema.iter().all(|name| named.contains_key(name))
758 };
759 if let Some(schema) = column_schema.filter(|schema| matches(schema)) {
760 return Some(schema);
761 }
762 {
763 let schemas = self.recovery_row_schemas.read();
764 if let Some(schema) = schemas.iter().find(|schema| matches(schema)) {
765 return Some(Arc::clone(schema));
766 }
767 }
768 let mut schemas = self.recovery_row_schemas.write();
769 if let Some(schema) = schemas.iter().find(|schema| matches(schema)) {
770 return Some(Arc::clone(schema));
771 }
772 const MAX_RECOVERY_ROW_SCHEMAS: usize = 16;
773 if schemas.len() >= MAX_RECOVERY_ROW_SCHEMAS {
774 return None;
775 }
776 let mut names: Vec<_> = named.keys().cloned().collect();
777 names.sort_unstable();
778 let schema = Arc::new(names);
779 schemas.push(Arc::clone(&schema));
780 Some(schema)
781 }
782
783 #[doc(hidden)]
787 pub fn restore_snapshot_row_layout(&self, entity: &mut UnifiedEntity) {
788 self.restore_recovery_row_layout(entity, self.column_schema());
789 }
790
791 fn restore_recovery_row_layout(
792 &self,
793 entity: &mut UnifiedEntity,
794 column_schema: Option<Arc<Vec<String>>>,
795 ) {
796 if let super::entity::EntityData::Row(row) = &mut entity.data {
797 if let Some(schema) = self.recovery_schema(row, column_schema) {
798 row.compact_with_schema(&schema);
799 }
800 }
801 }
802
803 #[doc(hidden)]
805 pub fn restore_snapshot_entity(
806 &self,
807 mut entity: UnifiedEntity,
808 ) -> Result<EntityId, SegmentError> {
809 self.restore_snapshot_row_layout(&mut entity);
810 self.maybe_seal_growing()?;
811 if self.get(entity.id).is_some() {
812 return Err(SegmentError::AlreadyExists(entity.id));
813 }
814 let entity_id = entity.id;
815 self.next_entity_id
816 .fetch_max(entity_id.raw().saturating_add(1), Ordering::SeqCst);
817 self.with_restoration_segment(entity_id, |segment| {
818 let segment_id = segment.id();
819 segment.adopt_entity(entity, None);
821 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
822 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
823 Ok(entity_id)
824 })
825 }
826
827 pub fn insert_batch(
829 &self,
830 entities: Vec<UnifiedEntity>,
831 ) -> Result<Vec<EntityId>, SegmentError> {
832 let mut ids = Vec::with_capacity(entities.len());
833 for entity in entities {
834 ids.push(self.insert(entity)?);
835 }
836 Ok(ids)
837 }
838
839 pub fn bulk_insert(
842 &self,
843 mut entities: Vec<UnifiedEntity>,
844 ) -> Result<Vec<EntityId>, SegmentError> {
845 if self.has_backing() {
847 let mut identities = std::collections::HashSet::with_capacity(entities.len());
848 for entity in &entities {
849 if entity.id.raw() != 0 {
850 self.reject_backing_duplicate(entity.id)?;
851 if self.get_with(entity.id, |_| ()).is_some() || !identities.insert(entity.id) {
852 return Err(SegmentError::AlreadyExists(entity.id));
853 }
854 }
855 }
856 }
857 for entity in &mut entities {
859 if entity.id.raw() == 0 {
860 entity.id = self.next_entity_id();
861 }
862 self.next_entity_id
863 .fetch_max(entity.id.raw().saturating_add(1), Ordering::SeqCst);
864 if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
865 if *row_id == 0 {
866 *row_id = self.next_row_id();
867 } else {
868 self.register_row_id(*row_id);
869 }
870 }
871 }
872
873 if let Some(first_row) = entities.first() {
876 if let super::entity::EntityData::Row(ref row) = first_row.data {
877 if let Some(ref named) = row.named {
878 let schema = self.get_or_init_schema(named);
879 for entity in &mut entities {
880 if let super::entity::EntityData::Row(ref mut row) = entity.data {
881 if let Some(named) = row.named.take() {
882 let mut cols = Vec::with_capacity(schema.len());
883 for col_name in schema.iter() {
884 cols.push(
885 named
886 .get(col_name)
887 .cloned()
888 .unwrap_or(reddb_types::Value::Null),
889 );
890 }
891 row.columns = cols;
892 row.schema = Some(Arc::clone(&schema));
893 }
894 }
895 }
896 }
897 }
898 }
899
900 let segment_arc = self.get_or_create_growing();
901 let mut segment = segment_arc.write();
902 let ids = segment.bulk_insert(entities)?;
904
905 self.total_entities_atomic
910 .fetch_add(ids.len() as u64, Ordering::Relaxed);
911
912 Ok(ids)
913 }
914
915 pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
917 let mut entity = self.get_with(id, Clone::clone)?;
918 self.restore_owned_row_schemas(std::iter::once(&mut entity));
919 Some(entity)
920 }
921
922 #[doc(hidden)]
925 pub fn get_with<T>(&self, id: EntityId, visit: impl FnOnce(&UnifiedEntity) -> T) -> Option<T> {
926 if let Some(growing_arc) = self.growing.read().as_ref() {
928 let growing = growing_arc.read();
929 if let Some(entity) = growing.get(id) {
930 return Some(visit(entity));
931 }
932 }
933
934 let sealed = self.sealed.read();
936 for segment in sealed.iter() {
937 let seg = segment.read();
938 if let Some(entity) = seg.get(id) {
939 return Some(visit(entity));
940 }
941 }
942
943 None
944 }
945
946 #[doc(hidden)]
950 pub fn graph_node_logical_id(
951 &self,
952 id: EntityId,
953 mut before_segment: impl FnMut() -> bool,
954 ) -> Option<EntityId> {
955 if let Some(growing_arc) = self.growing.read().as_ref() {
956 if !before_segment() {
957 return None;
958 }
959 let growing = growing_arc.read();
960 if let Some(entity) = growing.get(id) {
961 return matches!(entity.kind, EntityKind::GraphNode(_))
962 .then(|| entity.logical_id());
963 }
964 }
965 let sealed = self.sealed.read();
966 for segment in sealed.iter() {
967 if !before_segment() {
968 return None;
969 }
970 let segment = segment.read();
971 if let Some(entity) = segment.get(id) {
972 return matches!(entity.kind, EntityKind::GraphNode(_))
973 .then(|| entity.logical_id());
974 }
975 }
976 None
977 }
978
979 #[doc(hidden)]
984 pub fn visit_graph_index_batches<E>(
985 &self,
986 kind: GraphEntityKind,
987 keys: &[EntityId],
988 mut before_work: impl FnMut() -> bool,
989 mut admit: impl FnMut(usize) -> Result<(), E>,
990 mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
991 ) -> Result<bool, E> {
992 use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
993 admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
994 let mut admitted = 0;
995 let sources = loop {
996 let count = {
997 let growing = self.growing.read();
998 let sealed = self.sealed.read();
999 sealed.len() + usize::from(growing.is_some())
1000 };
1001 let bytes = count.saturating_mul(std::mem::size_of::<(
1002 Arc<RwLock<GrowingSegment>>,
1003 GraphIndexCursor,
1004 )>());
1005 admit(bytes.saturating_sub(admitted))?;
1006 admitted = admitted.max(bytes);
1007 let mut sources = Vec::with_capacity(count);
1008 let growing = self.growing.read();
1009 let sealed = self.sealed.read();
1010 if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
1011 continue;
1012 }
1013 for segment in growing.iter().chain(sealed.iter()) {
1014 if !before_work() {
1015 return Ok(false);
1016 }
1017 let view = segment.read();
1018 if !view.may_contain_graph_kind(kind) {
1019 continue;
1020 }
1021 let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
1022 return Ok(false);
1023 };
1024 sources.push((Arc::clone(segment), cursor));
1025 }
1026 break sources;
1027 };
1028 let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
1029 for (source, mut cursor) in sources {
1030 while !cursor.finished(keys) {
1031 let mut count = 0;
1032 {
1033 let segment = source.read();
1034 if !segment.may_contain_graph_kind(kind) {
1035 break;
1036 }
1037 if !segment.graph_index_batch(
1038 kind,
1039 keys,
1040 &mut cursor,
1041 &mut before_work,
1042 &mut |id| {
1043 ids[count] = id;
1044 count += 1;
1045 },
1046 ) {
1047 return Ok(false);
1048 }
1049 }
1050 if !before_work() || !consume(&ids[..count])? {
1051 return Ok(false);
1052 }
1053 }
1054 }
1055 Ok(true)
1056 }
1057
1058 pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
1063 let schema = self.column_schema();
1064 let clone_row = |entity: &UnifiedEntity| {
1065 let mut owned = entity.clone();
1066 if let super::entity::EntityData::Row(row) = &mut owned.data {
1067 if row.schema.is_none() {
1068 row.schema = schema.clone();
1069 }
1070 }
1071 owned
1072 };
1073 let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
1074 let mut remaining: Vec<usize> = (0..ids.len()).collect(); if let Some(growing_arc) = self.growing.read().as_ref() {
1079 let growing = if let Some(g) = growing_arc.try_read() {
1080 g
1081 } else {
1082 growing_arc.read()
1083 };
1084 remaining.retain(|&i| {
1085 if let Some(entity) = growing.get(ids[i]) {
1086 out[i] = Some(clone_row(entity));
1087 false } else {
1089 true }
1091 });
1092 }
1093
1094 if remaining.is_empty() {
1095 self.restore_owned_row_schemas(out.iter_mut().flatten());
1096 return out;
1097 }
1098
1099 let sealed = self.sealed.read();
1101 for segment in sealed.iter() {
1102 if remaining.is_empty() {
1103 break;
1104 }
1105 let seg = segment.read();
1106 remaining.retain(|&i| {
1107 if let Some(entity) = seg.get(ids[i]) {
1108 out[i] = Some(clone_row(entity));
1109 false
1110 } else {
1111 true
1112 }
1113 });
1114 }
1115
1116 drop(sealed);
1117 self.restore_owned_row_schemas(out.iter_mut().flatten());
1118 out
1119 }
1120
1121 pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
1134 where
1135 F: FnMut(usize, &UnifiedEntity),
1136 {
1137 thread_local! {
1146 static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
1147 const { std::cell::Cell::new(Vec::new()) };
1148 }
1149
1150 let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
1151 remaining.clear();
1152
1153 if let Some(growing_arc) = self.growing.read().as_ref() {
1154 let growing = if let Some(g) = growing_arc.try_read() {
1155 g
1156 } else {
1157 growing_arc.read()
1158 };
1159 for (i, id) in ids.iter().enumerate() {
1160 if let Some(entity) = growing.get(*id) {
1161 f(i, entity);
1162 } else {
1163 remaining.push(i);
1164 }
1165 }
1166 } else {
1167 remaining.reserve(ids.len());
1168 remaining.extend(0..ids.len());
1169 }
1170
1171 if !remaining.is_empty() {
1172 let sealed = self.sealed.read();
1173 for segment in sealed.iter() {
1174 if remaining.is_empty() {
1175 break;
1176 }
1177 let seg = segment.read();
1178 remaining.retain(|&i| {
1179 if let Some(entity) = seg.get(ids[i]) {
1180 f(i, entity);
1181 false
1182 } else {
1183 true
1184 }
1185 });
1186 }
1187 }
1188
1189 REMAINING_SCRATCH.with(|cell| cell.set(remaining));
1190 }
1191
1192 #[allow(dead_code)] fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
1195 if let Some(growing_arc) = self.growing.read().as_ref() {
1197 let growing = growing_arc.read();
1198 if let Some(entity) = growing.get(id) {
1199 return Some(entity.clone());
1200 }
1201 }
1202
1203 let sealed = self.sealed.read();
1205 for segment in sealed.iter() {
1206 if let Some(entity) = segment.get(id) {
1207 return Some(entity.clone());
1208 }
1209 }
1210
1211 None
1212 }
1213
1214 fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
1215 let sealed = self.sealed.read();
1216 sealed
1217 .iter()
1218 .find(|segment_arc| segment_arc.read().contains(id))
1219 .map(Arc::clone)
1220 }
1221
1222 fn rewrite_sealed_entity_into_growing(
1223 &self,
1224 entity: UnifiedEntity,
1225 metadata: Option<&Metadata>,
1226 ) -> Result<(), SegmentError> {
1227 let entity_id = entity.id;
1228 let sealed_arc = self
1229 .find_sealed_segment_arc(entity_id)
1230 .ok_or(SegmentError::NotFound(entity_id))?;
1231
1232 let metadata_to_apply = {
1233 let mut sealed = sealed_arc.write();
1234 let existing_metadata = sealed.get_metadata(entity_id);
1235 if !sealed.force_delete(entity_id) {
1236 return Err(SegmentError::NotFound(entity_id));
1237 }
1238 metadata.cloned().or(existing_metadata)
1239 };
1240
1241 let growing_id = self.with_restoration_segment(entity_id, |growing| {
1242 growing.insert_preserving_sequence(entity)?;
1243 if let Some(metadata) = metadata_to_apply {
1244 growing.set_metadata(entity_id, metadata)?;
1245 }
1246 Ok(growing.id())
1247 })?;
1248
1249 self.entity_segment.write().insert(entity_id, growing_id);
1250 Ok(())
1251 }
1252
1253 pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
1255 self.materialize_backing_entity(entity.id)?;
1256 let entity_id = entity.id;
1257 let mut entity = Some(entity);
1258
1259 if let Some(growing_arc) = self.growing.read().as_ref() {
1261 let mut growing = growing_arc.write();
1262 if growing.contains(entity_id) && growing.state().is_writable() {
1263 return growing.update(entity.take().expect("entity already moved"));
1264 }
1265 }
1266
1267 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1269 if let Some(seg_id) = segment_id {
1270 if let Some(growing_arc) = self.growing.read().as_ref() {
1271 let mut growing = growing_arc.write();
1272 if growing.id() == seg_id && growing.state().is_writable() {
1273 return growing.update(entity.take().expect("entity already moved"));
1274 }
1275 }
1276 }
1277
1278 if let Some(entity) = entity.take() {
1279 return self.rewrite_sealed_entity_into_growing(entity, None);
1280 }
1281
1282 Err(SegmentError::NotFound(entity_id))
1283 }
1284
1285 pub fn update_with_metadata(
1288 &self,
1289 entity: UnifiedEntity,
1290 metadata: Option<&Metadata>,
1291 ) -> Result<(), SegmentError> {
1292 self.materialize_backing_entity(entity.id)?;
1293 let entity_id = entity.id;
1294 let deadline_ms =
1295 metadata.and_then(|metadata| metadata_expiry_deadline_ms(&entity, metadata));
1296 let mut entity = Some(entity);
1297
1298 if let Some(growing_arc) = self.growing.read().as_ref() {
1300 let mut growing = growing_arc.write();
1301 if growing.contains(entity_id) && growing.state().is_writable() {
1302 growing.update(entity.take().expect("entity already moved"))?;
1303 if let Some(metadata) = metadata {
1304 growing.set_metadata(entity_id, metadata.clone())?;
1305 self.metadata_expiry.write().replace(entity_id, deadline_ms);
1306 }
1307 return Ok(());
1308 }
1309 }
1310
1311 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1313 if let Some(seg_id) = segment_id {
1314 if let Some(growing_arc) = self.growing.read().as_ref() {
1315 let mut growing = growing_arc.write();
1316 if growing.id() == seg_id && growing.state().is_writable() {
1317 growing.update(entity.take().expect("entity already moved"))?;
1318 if let Some(metadata) = metadata {
1319 growing.set_metadata(entity_id, metadata.clone())?;
1320 self.metadata_expiry.write().replace(entity_id, deadline_ms);
1321 }
1322 return Ok(());
1323 }
1324 }
1325 }
1326
1327 if let Some(entity) = entity.take() {
1328 self.rewrite_sealed_entity_into_growing(entity, metadata)?;
1329 if metadata.is_some() {
1330 self.metadata_expiry.write().replace(entity_id, deadline_ms);
1331 }
1332 return Ok(());
1333 }
1334
1335 Err(SegmentError::NotFound(entity_id))
1336 }
1337
1338 pub fn update_hot(
1342 &self,
1343 entity: UnifiedEntity,
1344 modified_columns: &[String],
1345 ) -> Result<(), SegmentError> {
1346 self.materialize_backing_entity(entity.id)?;
1347 let entity_id = entity.id;
1348 let mut entity = Some(entity);
1349
1350 if let Some(growing_arc) = self.growing.read().as_ref() {
1351 let mut growing = growing_arc.write();
1352 if growing.contains(entity_id) && growing.state().is_writable() {
1353 return growing.update_hot(
1354 entity.take().expect("entity already moved"),
1355 modified_columns,
1356 );
1357 }
1358 }
1359
1360 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1361 if let Some(seg_id) = segment_id {
1362 if let Some(growing_arc) = self.growing.read().as_ref() {
1363 let mut growing = growing_arc.write();
1364 if growing.id() == seg_id && growing.state().is_writable() {
1365 return growing.update_hot(
1366 entity.take().expect("entity already moved"),
1367 modified_columns,
1368 );
1369 }
1370 }
1371 }
1372
1373 if let Some(entity) = entity.take() {
1374 return self.rewrite_sealed_entity_into_growing(entity, None);
1375 }
1376
1377 Err(SegmentError::NotFound(entity_id))
1378 }
1379
1380 pub fn update_hot_with_metadata(
1383 &self,
1384 entity: UnifiedEntity,
1385 modified_columns: &[String],
1386 metadata: Option<&Metadata>,
1387 ) -> Result<(), SegmentError> {
1388 self.materialize_backing_entity(entity.id)?;
1389 let entity_id = entity.id;
1390 let deadline_ms =
1391 metadata.and_then(|metadata| metadata_expiry_deadline_ms(&entity, metadata));
1392 let mut entity = Some(entity);
1393
1394 if let Some(growing_arc) = self.growing.read().as_ref() {
1395 let mut growing = growing_arc.write();
1396 if growing.contains(entity_id) && growing.state().is_writable() {
1397 growing.update_hot(
1398 entity.take().expect("entity already moved"),
1399 modified_columns,
1400 )?;
1401 if let Some(metadata) = metadata {
1402 growing.set_metadata(entity_id, metadata.clone())?;
1403 self.metadata_expiry.write().replace(entity_id, deadline_ms);
1404 }
1405 return Ok(());
1406 }
1407 }
1408
1409 let segment_id = self.entity_segment.read().get(&entity_id).copied();
1410 if let Some(seg_id) = segment_id {
1411 if let Some(growing_arc) = self.growing.read().as_ref() {
1412 let mut growing = growing_arc.write();
1413 if growing.id() == seg_id && growing.state().is_writable() {
1414 growing.update_hot(
1415 entity.take().expect("entity already moved"),
1416 modified_columns,
1417 )?;
1418 if let Some(metadata) = metadata {
1419 growing.set_metadata(entity_id, metadata.clone())?;
1420 self.metadata_expiry.write().replace(entity_id, deadline_ms);
1421 }
1422 return Ok(());
1423 }
1424 }
1425 }
1426
1427 if let Some(entity) = entity.take() {
1428 self.rewrite_sealed_entity_into_growing(entity, metadata)?;
1429 if metadata.is_some() {
1430 self.metadata_expiry.write().replace(entity_id, deadline_ms);
1431 }
1432 return Ok(());
1433 }
1434
1435 Err(SegmentError::NotFound(entity_id))
1436 }
1437
1438 pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
1441 where
1442 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
1443 {
1444 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
1445 items.into_iter().collect();
1446 if items.is_empty() {
1447 return Ok(());
1448 }
1449
1450 if let Some(growing_arc) = self.growing.read().as_ref() {
1451 let mut growing = growing_arc.write();
1452 if growing.state().is_writable() {
1453 match growing.update_hot_batch_with_metadata(items.iter().copied()) {
1454 Ok(()) => {
1455 let mut expiry = self.metadata_expiry.write();
1456 for (entity, _, metadata) in &items {
1457 if let Some(metadata) = metadata {
1458 expiry.replace(
1459 entity.id,
1460 metadata_expiry_deadline_ms(entity, metadata),
1461 );
1462 }
1463 }
1464 return Ok(());
1465 }
1466 Err(SegmentError::NotFound(_)) => {}
1467 Err(other) => return Err(other),
1468 }
1469 }
1470 }
1471
1472 for (entity, modified_columns, metadata) in items {
1473 self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
1474 }
1475 Ok(())
1476 }
1477
1478 pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
1480 self.materialize_backing_entity(id)?;
1481 let deleted = self.delete_resident(id)?;
1482 if deleted {
1483 self.remember_backing_delete(id);
1484 self.metadata_expiry.write().replace(id, None);
1485 }
1486 Ok(deleted)
1487 }
1488
1489 fn delete_resident(&self, id: EntityId) -> Result<bool, SegmentError> {
1490 if let Some(growing_arc) = self.growing.read().as_ref() {
1493 let mut growing = growing_arc.write();
1494 if growing.contains(id) && growing.state().is_writable() {
1495 let seg_id = growing.id();
1496 let deleted = growing.delete(id)?;
1497 if deleted {
1498 self.entity_segment.write().remove(&id);
1499 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1500 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1501 }
1502 return Ok(deleted);
1503 }
1504 }
1505
1506 let segment_id = self.entity_segment.read().get(&id).copied();
1509
1510 if let Some(seg_id) = segment_id {
1511 if let Some(growing_arc) = self.growing.read().as_ref() {
1512 let mut growing = growing_arc.write();
1513 if growing.id() == seg_id && growing.state().is_writable() {
1514 let deleted = growing.delete(id)?;
1515 if deleted {
1516 self.entity_segment.write().remove(&id);
1517 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1518 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1519 }
1520 return Ok(deleted);
1521 }
1522 }
1523 }
1524
1525 {
1528 let sealed = self.sealed.read();
1529 for segment_arc in sealed.iter() {
1530 let mut seg = segment_arc.write();
1531 let seg_id = seg.id();
1532 if seg.contains(id) {
1533 let deleted = seg.force_delete(id);
1534 drop(seg);
1535 if deleted {
1536 self.entity_segment.write().remove(&id);
1537 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1538 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1539 }
1540 return Ok(deleted);
1541 }
1542 }
1543 }
1544
1545 Ok(false)
1546 }
1547
1548 pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1549 if self.has_backing() {
1550 let mut deleted = Vec::new();
1551 for id in ids {
1552 if self.delete(*id)? {
1553 deleted.push(*id);
1554 }
1555 }
1556 return Ok(deleted);
1557 }
1558 if ids.is_empty() {
1559 return Ok(Vec::new());
1560 }
1561
1562 let mut deleted_ids = Vec::with_capacity(ids.len());
1563
1564 if let Some(growing_arc) = self.growing.read().as_ref() {
1565 let mut growing = growing_arc.write();
1566 if growing.state().is_writable() {
1567 let seg_id = growing.id();
1568 let deleted = growing.delete_batch(ids)?;
1569 if !deleted.is_empty() {
1570 {
1571 let mut entity_segment = self.entity_segment.write();
1572 for id in &deleted {
1573 entity_segment.remove(id);
1574 }
1575 }
1576 self.total_entities_atomic
1577 .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
1578 for id in &deleted {
1579 self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1580 }
1581 deleted_ids.extend(deleted);
1582 }
1583 }
1584 }
1585
1586 if deleted_ids.len() == ids.len() {
1587 let mut expiry = self.metadata_expiry.write();
1588 for id in &deleted_ids {
1589 expiry.replace(*id, None);
1590 }
1591 return Ok(deleted_ids);
1592 }
1593
1594 let deleted_set: std::collections::HashSet<EntityId> =
1595 deleted_ids.iter().copied().collect();
1596 for &id in ids {
1597 if deleted_set.contains(&id) {
1598 continue;
1599 }
1600 if self.delete(id)? {
1601 deleted_ids.push(id);
1602 }
1603 }
1604
1605 if !deleted_ids.is_empty() {
1606 let mut expiry = self.metadata_expiry.write();
1607 for id in &deleted_ids {
1608 expiry.replace(*id, None);
1609 }
1610 }
1611
1612 Ok(deleted_ids)
1613 }
1614
1615 pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1617 if let Some(growing_arc) = self.growing.read().as_ref() {
1619 let growing = growing_arc.read();
1620 if growing.contains(id) {
1621 return growing.get_metadata(id);
1622 }
1623 }
1624
1625 let segment_id = self.entity_segment.read().get(&id).copied();
1627
1628 if let Some(seg_id) = segment_id {
1629 if let Some(growing_arc) = self.growing.read().as_ref() {
1630 let growing = growing_arc.read();
1631 if growing.id() == seg_id {
1632 return growing.get_metadata(id);
1633 }
1634 }
1635
1636 let sealed = self.sealed.read();
1637 for segment in sealed.iter() {
1638 if segment.id() == seg_id {
1639 return segment.get_metadata(id);
1640 }
1641 }
1642 }
1643
1644 if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1645 return segment_arc.read().get_metadata(id);
1646 }
1647
1648 None
1649 }
1650
1651 pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1653 self.materialize_backing_entity(id)?;
1654 let entity = self.get(id).ok_or(SegmentError::NotFound(id))?;
1655 let deadline_ms = metadata_expiry_deadline_ms(&entity, &metadata);
1656 self.set_metadata_resident(id, metadata)?;
1657 self.metadata_expiry.write().replace(id, deadline_ms);
1658 Ok(())
1659 }
1660
1661 fn set_metadata_resident(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1662 if let Some(growing_arc) = self.growing.read().as_ref() {
1665 let mut growing = growing_arc.write();
1666 if growing.contains(id) && growing.state().is_writable() {
1667 return growing.set_metadata(id, metadata);
1668 }
1669 }
1670
1671 let segment_id = self.entity_segment.read().get(&id).copied();
1673
1674 if let Some(seg_id) = segment_id {
1675 if let Some(growing_arc) = self.growing.read().as_ref() {
1676 let mut growing = growing_arc.write();
1677 if growing.id() == seg_id && growing.state().is_writable() {
1678 return growing.set_metadata(id, metadata);
1679 }
1680 }
1681 }
1682
1683 if let Some(entity) = self.get(id) {
1684 return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1685 }
1686
1687 Err(SegmentError::NotFound(id))
1688 }
1689
1690 pub fn expired_metadata_ids(&self, now_ms: u64, limit: usize) -> Vec<EntityId> {
1693 self.metadata_expiry.read().expired(now_ms, limit)
1694 }
1695
1696 pub fn scheduled_metadata_expiry_count(&self) -> usize {
1698 self.metadata_expiry.read().scheduled_count()
1699 }
1700
1701 pub fn metadata_expiry_entries(&self) -> Vec<(u64, EntityId)> {
1702 self.metadata_expiry.read().entries()
1703 }
1704
1705 pub fn restore_metadata_expiry_projection(&self, entries: &[(u64, EntityId)], ready: bool) {
1706 self.metadata_expiry.write().replace_all(entries);
1707 self.metadata_expiry_projection_ready
1708 .store(ready, Ordering::Release);
1709 }
1710
1711 pub fn metadata_expiry_projection_ready(&self) -> bool {
1712 self.metadata_expiry_projection_ready
1713 .load(Ordering::Acquire)
1714 }
1715
1716 pub fn mark_metadata_expiry_projection_ready(&self) {
1717 self.metadata_expiry_projection_ready
1718 .store(true, Ordering::Release);
1719 }
1720
1721 pub fn remove_metadata_expiry(&self, entity_id: EntityId) {
1722 self.metadata_expiry.write().replace(entity_id, None);
1723 }
1724
1725 pub fn restore_metadata_expiry(&self, entity: &UnifiedEntity, metadata: Option<&Metadata>) {
1728 let deadline_ms = (entity.xmax == 0)
1729 .then(|| metadata.and_then(|value| metadata_expiry_deadline_ms(entity, value)))
1730 .flatten();
1731 self.metadata_expiry.write().replace(entity.id, deadline_ms);
1732 }
1733
1734 fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1736 let should_seal = {
1737 let growing_opt = self.growing.read();
1738 if let Some(growing_arc) = growing_opt.as_ref() {
1739 let growing = growing_arc.read();
1740 growing.should_seal(&self.config.segment_config)
1741 || growing.idle_secs() >= self.config.idle_seal_secs
1742 } else {
1743 false
1744 }
1745 };
1746
1747 if should_seal {
1748 self.seal_current()?;
1749 }
1750
1751 Ok(())
1752 }
1753
1754 pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1756 self.seal_current_if(|_| true)?
1757 .ok_or(SegmentError::InvalidState(SegmentState::Sealed))
1758 }
1759
1760 fn seal_current_if(
1761 &self,
1762 should_seal: impl FnOnce(&GrowingSegment) -> bool,
1763 ) -> Result<Option<SegmentId>, SegmentError> {
1764 let disk_residency = self.has_backing();
1765 let mut growing_slot = self.growing.write();
1768 if let Some(growing_arc) = growing_slot.as_ref() {
1769 let mut growing = growing_arc.write();
1770 if !should_seal(&growing) {
1771 return Ok(None);
1772 }
1773 if disk_residency {
1774 self.next_sequence_id
1775 .fetch_max(growing.sequence_id_next(), Ordering::SeqCst);
1776 }
1777 let seg_id = growing.id();
1778 let entity_count = growing.stats().entity_count as u64;
1779
1780 growing.seal()?;
1782
1783 drop(growing); let growing_arc = growing_slot
1789 .take()
1790 .expect("sealed source remains published");
1791 self.sealed.write().push(growing_arc);
1792 drop(growing_slot);
1793
1794 self.mark_sealed_pages_visible(entity_count);
1796
1797 {
1799 let mut stats = self.stats.write();
1800 stats.growing_count = stats.growing_count.saturating_sub(1);
1801 stats.sealed_count += 1;
1802 stats.seal_ops += 1;
1803 }
1804
1805 self.emit(LifecycleEvent::SegmentSealed(seg_id));
1806
1807 return Ok(Some(seg_id));
1808 }
1809
1810 Ok(None)
1811 }
1812
1813 pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1815 let has_growing = self.growing.read().is_some();
1816 if has_growing {
1817 Ok(Some(self.seal_current()?))
1818 } else {
1819 Ok(None)
1820 }
1821 }
1822
1823 pub fn all_visible_fraction(&self) -> f64 {
1833 const ROWS_PER_PAGE: u32 = 256;
1834 let sealed = self.sealed.read();
1835 if sealed.is_empty() {
1836 return 0.0;
1837 }
1838 let mut total_pages: u64 = 0;
1839 for seg_arc in sealed.iter() {
1840 let seg = seg_arc.read();
1841 let entity_count = seg.stats().entity_count as u64;
1842 let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1843 total_pages += pages;
1844 }
1845 if total_pages == 0 {
1846 return 0.0;
1847 }
1848 let visible = self.visibility_map.all_visible_count();
1849 (visible as f64 / total_pages as f64).min(1.0)
1850 }
1851
1852 fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1855 const ROWS_PER_PAGE: u32 = 256;
1856 let existing_visible = self.visibility_map.all_visible_count();
1857 let start_page = existing_visible as u32;
1859 let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1860 let end_page = start_page + new_pages as u32;
1861 self.visibility_map.mark_range_visible(start_page, end_page);
1862 }
1863
1864 #[doc(hidden)]
1868 pub fn find_unique_key_conflict(
1869 &self,
1870 columns: &[Vec<String>],
1871 key_index: usize,
1872 signatures: &[String],
1873 exclude_id: Option<EntityId>,
1874 reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1875 ) -> Option<EntityId> {
1876 if let Some(growing) = self.growing.read().as_ref() {
1877 if let Some(id) = growing.read().find_unique_key_conflict(
1878 columns,
1879 key_index,
1880 signatures,
1881 exclude_id,
1882 reserves_key,
1883 ) {
1884 return Some(id);
1885 }
1886 }
1887 for segment in self.sealed.read().iter() {
1888 if let Some(id) = segment.read().find_unique_key_conflict(
1889 columns,
1890 key_index,
1891 signatures,
1892 exclude_id,
1893 reserves_key,
1894 ) {
1895 return Some(id);
1896 }
1897 }
1898 None
1899 }
1900
1901 pub fn for_each_entity<F>(&self, callback: F)
1907 where
1908 F: FnMut(&UnifiedEntity) -> bool,
1909 {
1910 self.for_each_entity_in_segments(|_| true, || true, callback);
1911 }
1912
1913 #[doc(hidden)]
1915 pub fn checkpoint_revisions(&self) -> Vec<(u64, u64)> {
1916 let growing = self.growing.read();
1917 let sealed = self.sealed.read();
1918 growing
1919 .iter()
1920 .chain(sealed.iter())
1921 .map(|segment| {
1922 let segment = segment.read();
1923 (segment.id(), segment.checkpoint_revision())
1924 })
1925 .collect()
1926 }
1927
1928 #[doc(hidden)]
1929 pub fn checkpoint_mutations_since(&self, before: &[(u64, u64)]) -> Option<Vec<EntityId>> {
1930 let growing = self.growing.read();
1931 let sealed = self.sealed.read();
1932 let previous: HashMap<_, _> = before.iter().copied().collect();
1937 if previous.len() != before.len() {
1938 return None;
1939 }
1940 let mut retained_sources = 0;
1941 let mut changed = std::collections::BTreeSet::new();
1942 for segment in growing.iter().chain(sealed.iter()) {
1943 let segment = segment.read();
1944 let revision = previous.get(&segment.id()).copied();
1945 retained_sources += usize::from(revision.is_some());
1946 let revision = revision.unwrap_or(0);
1949 changed.extend(segment.checkpoint_mutations_since(revision)?);
1950 if changed.len() > super::segment::CHECKPOINT_MUTATIONS_MAX {
1951 return None;
1952 }
1953 }
1954 (retained_sources == previous.len()).then(|| changed.into_iter().collect())
1957 }
1958
1959 #[doc(hidden)]
1960 pub fn checkpoint_tracking_bytes(&self) -> usize {
1961 let growing = self.growing.read();
1962 let sealed = self.sealed.read();
1963 growing
1964 .iter()
1965 .chain(sealed.iter())
1966 .map(|segment| segment.read().checkpoint_tracking_bytes())
1967 .sum()
1968 }
1969
1970 #[doc(hidden)]
1972 pub fn cached_physical_root(&self, seed: u64, compute: impl FnOnce() -> u64) -> u64 {
1973 let segments = self.checkpoint_revisions();
1974 let mut cache = self.physical_root_cache.lock();
1975 if let Some(previous) = cache.as_ref() {
1976 if previous.seed == seed && previous.segments == segments {
1977 return previous.root;
1978 }
1979 }
1980 let root = compute();
1981 *cache = Some(PhysicalRootCacheEntry {
1982 segments,
1983 seed,
1984 root,
1985 });
1986 root
1987 }
1988
1989 #[doc(hidden)]
1994 pub fn for_each_entity_in_id_order(&self, mut visit: impl FnMut(&UnifiedEntity) -> bool) {
1995 self.for_each_entity_source_in_id_order(|_, entity| visit(entity));
1996 }
1997
1998 #[doc(hidden)]
2000 pub fn page_in_id_order(&self, offset: usize, limit: usize) -> (usize, Vec<UnifiedEntity>) {
2001 let end = offset.saturating_add(limit);
2002 let mut total = 0;
2003 let mut items = Vec::with_capacity(limit.min(self.count()));
2004 self.for_each_entity_in_id_order(|entity| {
2005 if total >= offset && total < end {
2006 items.push(entity.clone());
2007 }
2008 total += 1;
2009 true
2010 });
2011 self.restore_owned_row_schemas(items.iter_mut());
2012 (total, items)
2013 }
2014
2015 #[doc(hidden)]
2016 pub fn for_each_entity_with_metadata_in_id_order(
2017 &self,
2018 mut visit: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
2019 ) {
2020 self.for_each_entity_source_in_id_order(|segment, entity| {
2021 visit(entity, segment.get_metadata(entity.id))
2022 });
2023 }
2024
2025 fn for_each_entity_source_in_id_order(
2026 &self,
2027 mut visit: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
2028 ) {
2029 let segments: Vec<_> = {
2030 let growing = self.growing.read();
2031 let sealed = self.sealed.read();
2032 growing.iter().chain(sealed.iter()).cloned().collect()
2033 };
2034 let mut order = Vec::new();
2035 for (source, segment) in segments.iter().enumerate() {
2036 segment.read().for_each_fast(|entity| {
2037 order.push((entity.id, source));
2038 true
2039 });
2040 }
2041 order.sort_by_key(|(id, _)| id.raw());
2042 for (id, source) in order {
2043 let segment = segments[source].read();
2044 if let Some(entity) = segment.get(id) {
2045 if !visit(&segment, entity) {
2046 break;
2047 }
2048 }
2049 }
2050 }
2051
2052 #[doc(hidden)]
2055 pub fn for_each_entity_with_metadata(
2056 &self,
2057 mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
2058 ) {
2059 self.for_each_entity_source(|segment, entity| {
2060 callback(entity, segment.get_metadata(entity.id))
2061 });
2062 }
2063
2064 fn for_each_entity_source(
2065 &self,
2066 mut callback: impl FnMut(&GrowingSegment, &UnifiedEntity) -> bool,
2067 ) {
2068 if let Some(growing_arc) = self.growing.read().as_ref() {
2069 let growing = growing_arc.read();
2070 if !growing.for_each_fast(|entity| callback(&growing, entity)) {
2071 return;
2072 }
2073 }
2074 let sealed = self.sealed.read();
2075 for segment_arc in sealed.iter() {
2076 let segment = segment_arc.read();
2077 if !segment.for_each_fast(|entity| callback(&segment, entity)) {
2078 return;
2079 }
2080 }
2081 }
2082
2083 fn for_each_entity_in_segments<F>(
2084 &self,
2085 segment_filter: impl Fn(&GrowingSegment) -> bool,
2086 mut before_segment: impl FnMut() -> bool,
2087 mut callback: F,
2088 ) where
2089 F: FnMut(&UnifiedEntity) -> bool,
2090 {
2091 if let Some(growing_arc) = self.growing.read().as_ref() {
2095 if !before_segment() {
2096 return;
2097 }
2098 let growing = if let Some(g) = growing_arc.try_read() {
2099 g
2100 } else {
2101 growing_arc.read()
2102 };
2103 if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
2104 return;
2105 }
2106 }
2107
2108 let sealed = self.sealed.read();
2110 for segment_arc in sealed.iter() {
2111 if !before_segment() {
2112 return;
2113 }
2114 let segment = segment_arc.read();
2115 if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
2116 return;
2117 }
2118 }
2119 }
2120
2121 pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
2137 &self,
2138 init: FInit,
2139 fold: FFold,
2140 reduce: FReduce,
2141 ) -> T
2142 where
2143 T: Send,
2144 FInit: Fn() -> T + Send + Sync,
2145 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
2146 FReduce: Fn(T, T) -> T + Send + Sync,
2147 {
2148 #[cfg(not(target_arch = "wasm32"))]
2149 use rayon::prelude::*;
2150
2151 let mut acc = init();
2154 if let Some(growing_arc) = self.growing.read().as_ref() {
2155 let growing = if let Some(g) = growing_arc.try_read() {
2156 g
2157 } else {
2158 growing_arc.read()
2159 };
2160 growing.for_each_fast(|entity| {
2161 acc = fold(std::mem::replace(&mut acc, init()), entity);
2162 true
2163 });
2164 }
2165
2166 let segments: Vec<_> = {
2170 let sealed = self.sealed.read();
2171 sealed.iter().cloned().collect()
2172 };
2173
2174 if segments.len() <= 1 {
2175 for seg_arc in &segments {
2176 let seg = seg_arc.read();
2177 seg.for_each_fast(|entity| {
2178 acc = fold(std::mem::replace(&mut acc, init()), entity);
2179 true
2180 });
2181 }
2182 return acc;
2183 }
2184
2185 #[cfg(not(target_arch = "wasm32"))]
2186 let segments = segments.into_par_iter();
2187 #[cfg(target_arch = "wasm32")]
2188 let segments = segments.into_iter();
2189 let partials = segments.map(|seg_arc| {
2190 let mut local = init();
2191 let seg = seg_arc.read();
2192 seg.for_each_fast(|entity| {
2193 local = fold(std::mem::replace(&mut local, init()), entity);
2194 true
2195 });
2196 local
2197 });
2198 #[cfg(not(target_arch = "wasm32"))]
2199 let sealed_acc = partials.reduce(&init, &reduce);
2200 #[cfg(target_arch = "wasm32")]
2201 let sealed_acc = partials.fold(init(), &reduce);
2202
2203 reduce(acc, sealed_acc)
2204 }
2205
2206 pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
2215 where
2216 F: FnMut(&UnifiedEntity) -> bool,
2217 {
2218 let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
2219 }
2220
2221 pub fn for_each_entity_zoned_with_stats<F>(
2222 &self,
2223 zone_preds: &[(&str, ZoneColPred<'_>)],
2224 mut callback: F,
2225 ) -> SegmentScanStats
2226 where
2227 F: FnMut(&UnifiedEntity) -> bool,
2228 {
2229 let mut scan_stats = SegmentScanStats::default();
2230 if let Some(growing_arc) = self.growing.read().as_ref() {
2236 scan_stats.segments_total += 1;
2237 scan_stats.segments_scanned += 1;
2238 let growing = if let Some(g) = growing_arc.try_read() {
2239 g
2240 } else {
2241 growing_arc.read()
2242 };
2243 if !growing.for_each_fast(&mut callback) {
2244 return scan_stats;
2245 }
2246 }
2247
2248 let sealed = self.sealed.read();
2250 for segment_arc in sealed.iter() {
2251 scan_stats.segments_total += 1;
2252 let segment = segment_arc.read();
2253 if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
2254 scan_stats.segments_pruned += 1;
2255 continue; }
2257 scan_stats.segments_scanned += 1;
2258 if !segment.for_each_fast(&mut callback) {
2259 return scan_stats;
2260 }
2261 }
2262 scan_stats
2263 }
2264
2265 pub fn query_all_zoned<F>(
2275 &self,
2276 zone_preds: &[(&str, ZoneColPred<'_>)],
2277 filter: F,
2278 ) -> Vec<UnifiedEntity>
2279 where
2280 F: Fn(&UnifiedEntity) -> bool + Sync,
2281 {
2282 self.query_all_zoned_with_stats(zone_preds, filter).0
2283 }
2284
2285 pub fn query_all_zoned_with_stats<F>(
2286 &self,
2287 zone_preds: &[(&str, ZoneColPred<'_>)],
2288 filter: F,
2289 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2290 where
2291 F: Fn(&UnifiedEntity) -> bool + Sync,
2292 {
2293 self.project_all_zoned_with_stats(zone_preds, |entity| {
2294 filter(entity).then(|| entity.clone())
2295 })
2296 }
2297
2298 fn project_all_zoned_with_stats<T, F>(
2299 &self,
2300 zone_preds: &[(&str, ZoneColPred<'_>)],
2301 project: F,
2302 ) -> (Vec<T>, SegmentScanStats)
2303 where
2304 T: Send,
2305 F: Fn(&UnifiedEntity) -> Option<T> + Sync,
2306 {
2307 let mut results = Vec::new();
2308 let mut scan_stats = SegmentScanStats::default();
2309
2310 if let Some(growing_arc) = self.growing.read().as_ref() {
2313 scan_stats.segments_total += 1;
2314 scan_stats.segments_scanned += 1;
2315 let growing = if let Some(g) = growing_arc.try_read() {
2316 g
2317 } else {
2318 growing_arc.read()
2319 };
2320 results.extend(growing.iter().filter_map(&project));
2321 }
2322
2323 let sealed = self.sealed.read();
2325 let surviving: Vec<_> = sealed
2327 .iter()
2328 .filter(|seg_arc| {
2329 scan_stats.segments_total += 1;
2330 if zone_preds.is_empty() {
2331 scan_stats.segments_scanned += 1;
2332 return true;
2333 }
2334 let seg = seg_arc.read();
2335 if seg.can_skip_zone_preds(zone_preds) {
2336 scan_stats.segments_pruned += 1;
2337 false
2338 } else {
2339 scan_stats.segments_scanned += 1;
2340 true
2341 }
2342 })
2343 .collect();
2344
2345 let use_parallel = surviving.len() > 1 && crate::execution::should_parallelize();
2346
2347 if use_parallel {
2348 let project_ref = &project;
2349 let segment_results: Vec<Vec<T>> = std::thread::scope(|s| {
2350 surviving
2351 .iter()
2352 .map(|segment| {
2353 s.spawn(move || {
2354 segment
2355 .read()
2356 .iter()
2357 .filter_map(project_ref)
2358 .collect::<Vec<_>>()
2359 })
2360 })
2361 .collect::<Vec<_>>()
2362 .into_iter()
2363 .map(|handle| match handle.join() {
2364 Ok(batch) => batch,
2365 Err(panic) => std::panic::resume_unwind(panic),
2366 })
2367 .collect()
2368 });
2369 for batch in segment_results {
2370 results.extend(batch);
2371 }
2372 } else {
2373 for segment_arc in surviving {
2374 let seg = segment_arc.read();
2375 results.extend(seg.iter().filter_map(&project));
2376 }
2377 }
2378
2379 (results, scan_stats)
2380 }
2381
2382 pub fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
2388 match snapshot {
2389 Some(ctx) => entity_visible_with_context(Some(ctx), entity),
2390 None => {
2391 !crate::storage::visibility::entity_moderation_hidden(entity) && entity.xmax == 0
2392 }
2393 }
2394 }
2395
2396 pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
2403 where
2404 F: Fn(&UnifiedEntity) -> bool + Sync,
2405 {
2406 self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
2407 }
2408
2409 #[doc(hidden)]
2414 pub fn project_key_entity<T>(
2415 &self,
2416 key: &str,
2417 snapshot: Option<&SnapshotContext>,
2418 versioned: bool,
2419 mut project: impl FnMut(&UnifiedEntity) -> T,
2420 ) -> Option<T> {
2421 let mut probe = |segment: &GrowingSegment, best: &mut Option<(u64, T)>| {
2422 if let Some(entity) = segment.find_key_entity(key, versioned, &mut |entity| {
2423 !versioned || Self::scan_entity_visible(snapshot, entity)
2424 }) {
2425 if best.as_ref().is_none_or(|(xmin, _)| entity.xmin >= *xmin) {
2426 *best = Some((entity.xmin, project(entity)));
2427 }
2428 }
2429 };
2430 let mut best = None;
2431 if let Some(growing) = self.growing.read().as_ref() {
2432 probe(&growing.read(), &mut best);
2433 }
2434 for segment in self.sealed.read().iter() {
2435 if !versioned && best.is_some() {
2436 break;
2437 }
2438 probe(&segment.read(), &mut best);
2439 }
2440 best.map(|(_, projected)| projected)
2441 }
2442
2443 #[doc(hidden)]
2446 pub fn visit_queue_dedup(
2447 &self,
2448 queue: &str,
2449 key: Option<&str>,
2450 now: u64,
2451 snapshot: Option<&SnapshotContext>,
2452 mut visit: impl FnMut(&UnifiedEntity) -> bool,
2453 ) {
2454 let mut visible = |entity: &UnifiedEntity| {
2455 (key.is_some() && !Self::scan_entity_visible(snapshot, entity)) || visit(entity)
2456 };
2457 if let Some(growing) = self.growing.read().as_ref() {
2458 if !growing
2459 .read()
2460 .visit_queue_dedup(queue, key, now, &mut visible)
2461 {
2462 return;
2463 }
2464 }
2465 for segment in self.sealed.read().iter() {
2466 if !segment
2467 .read()
2468 .visit_queue_dedup(queue, key, now, &mut visible)
2469 {
2470 return;
2471 }
2472 }
2473 }
2474
2475 #[doc(hidden)]
2477 pub fn scan_graph_kind(
2478 &self,
2479 snapshot: Option<&SnapshotContext>,
2480 kind: GraphEntityKind,
2481 ) -> Vec<UnifiedEntity> {
2482 self.query_all_in_segments(
2483 |segment| segment.may_contain_graph_kind(kind),
2484 |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
2485 )
2486 }
2487
2488 #[doc(hidden)]
2492 pub fn scan_graph_batches(
2493 &self,
2494 snapshot: Option<&SnapshotContext>,
2495 kind: GraphEntityKind,
2496 mut before_work: impl FnMut() -> bool,
2497 mut consume: impl FnMut(Vec<UnifiedEntity>) -> bool,
2498 ) -> bool {
2499 let segments = {
2500 let growing = self.growing.read();
2501 let sealed = self.sealed.read();
2502 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
2503 for segment in growing.iter().chain(sealed.iter()) {
2504 if !before_work() {
2505 return false;
2506 }
2507 let view = segment.read();
2508 if view.may_contain_graph_kind(kind) {
2509 segments.push((Arc::clone(segment), view.scan_cursor()));
2510 }
2511 }
2512 segments
2513 };
2514 let mut ids = [EntityId::new(0); super::segment::SCAN_BATCH_SIZE];
2515 for (segment, mut cursor) in segments {
2516 while !cursor.finished() {
2517 let mut count = 0;
2518 {
2519 let view = segment.read();
2520 if !view.scan_batch(&mut cursor, &mut before_work, &mut |id| {
2521 ids[count] = id;
2522 count += 1;
2523 }) {
2524 return false;
2525 }
2526 }
2527 if !before_work() {
2528 return false;
2529 }
2530 let mut batch = vec![None; count];
2533 self.for_each_id(&ids[..count], |position, entity| {
2534 if kind.matches(&entity.kind) && Self::scan_entity_visible(snapshot, entity) {
2535 batch[position] = Some(entity.clone());
2536 }
2537 });
2538 if !before_work() || !consume(batch.into_iter().flatten().collect()) {
2539 return false;
2540 }
2541 }
2542 }
2543 true
2544 }
2545
2546 #[doc(hidden)]
2551 pub fn scan_graph_candidates_for_each<F>(
2552 &self,
2553 snapshot: Option<&SnapshotContext>,
2554 kind: GraphEntityKind,
2555 before_segment: impl FnMut() -> bool,
2556 mut callback: F,
2557 ) where
2558 F: FnMut(&UnifiedEntity) -> bool,
2559 {
2560 self.for_each_entity_in_segments(
2561 |segment| segment.may_contain_graph_kind(kind),
2562 before_segment,
2563 |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
2564 );
2565 }
2566
2567 pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
2572 where
2573 F: FnMut(&UnifiedEntity) -> bool,
2574 {
2575 self.for_each_entity(|entity| {
2576 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2577 });
2578 }
2579
2580 #[doc(hidden)]
2581 pub fn physical_scan_cursor(&self) -> PhysicalScanCursor {
2582 let growing = self.growing.read();
2583 let sealed = self.sealed.read();
2584 PhysicalScanCursor {
2585 sources: growing
2586 .iter()
2587 .chain(sealed.iter())
2588 .map(|segment| (Arc::downgrade(segment), segment.read().scan_cursor()))
2589 .collect(),
2590 source: 0,
2591 }
2592 }
2593
2594 #[doc(hidden)]
2597 pub fn physical_scan_batch(
2598 &self,
2599 cursor: &mut PhysicalScanCursor,
2600 work_left: &mut usize,
2601 ) -> Vec<EntityId> {
2602 let mut ids = Vec::with_capacity(super::segment::SCAN_BATCH_SIZE.min(*work_left));
2603 while !cursor.finished() && *work_left > 0 {
2604 let (source, position) = &mut cursor.sources[cursor.source];
2605 if position.finished() {
2606 *work_left -= 1;
2607 cursor.source += 1;
2608 continue;
2609 }
2610 let Some(segment) = source.upgrade() else {
2611 *work_left -= 1;
2612 cursor.source += 1;
2613 continue;
2614 };
2615 segment.read().scan_batch(
2616 position,
2617 &mut || {
2618 if *work_left == 0 {
2619 return false;
2620 }
2621 *work_left -= 1;
2622 true
2623 },
2624 &mut |id| ids.push(id),
2625 );
2626 if position.finished() {
2627 cursor.source += 1;
2628 }
2629 if !ids.is_empty() {
2630 break;
2631 }
2632 }
2633 ids
2634 }
2635
2636 #[doc(hidden)]
2642 pub fn scan_batches(
2643 &self,
2644 snapshot: Option<&SnapshotContext>,
2645 mut before_work: impl FnMut() -> bool,
2646 mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
2647 ) -> bool {
2648 let segments = {
2649 let growing = self.growing.read();
2650 let sealed = self.sealed.read();
2651 let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
2652 for segment in growing.iter().chain(sealed.iter()) {
2653 if !before_work() {
2654 return false;
2655 }
2656 let cursor = segment.read().scan_cursor();
2657 segments.push((Arc::clone(segment), cursor));
2658 }
2659 segments
2660 };
2661 let mut ids = Vec::new();
2662 let mut batch = Vec::new();
2663 for (segment, mut cursor) in segments {
2664 while !cursor.finished() {
2665 ids.clear();
2666 batch.clear();
2667 let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
2668 ids.reserve_exact(batch_size);
2669 batch.reserve_exact(batch_size);
2670 {
2671 let segment = segment.read();
2672 if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
2673 return false;
2674 }
2675 }
2676 if !before_work() {
2679 return false;
2680 }
2681 self.for_each_id(&ids, |_, entity| {
2682 if Self::scan_entity_visible(snapshot, entity) {
2683 batch.push(entity.clone());
2684 }
2685 });
2686 if !before_work() {
2687 return false;
2688 }
2689 if !consume(&batch) {
2690 return false;
2691 }
2692 }
2693 }
2694 true
2695 }
2696
2697 pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
2699 &self,
2700 snapshot: Option<&SnapshotContext>,
2701 init: FInit,
2702 fold: FFold,
2703 reduce: FReduce,
2704 ) -> T
2705 where
2706 T: Send,
2707 FInit: Fn() -> T + Send + Sync,
2708 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
2709 FReduce: Fn(T, T) -> T + Send + Sync,
2710 {
2711 self.fold_entities_parallel(
2712 init,
2713 |acc, entity| {
2714 if Self::scan_entity_visible(snapshot, entity) {
2715 fold(acc, entity)
2716 } else {
2717 acc
2718 }
2719 },
2720 reduce,
2721 )
2722 }
2723
2724 pub fn scan_for_each_zoned_with_stats<F>(
2726 &self,
2727 snapshot: Option<&SnapshotContext>,
2728 zone_preds: &[(&str, ZoneColPred<'_>)],
2729 mut callback: F,
2730 ) -> SegmentScanStats
2731 where
2732 F: FnMut(&UnifiedEntity) -> bool,
2733 {
2734 self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
2735 !Self::scan_entity_visible(snapshot, entity) || callback(entity)
2736 })
2737 }
2738
2739 pub fn scan_zoned_with_stats<F>(
2741 &self,
2742 snapshot: Option<&SnapshotContext>,
2743 zone_preds: &[(&str, ZoneColPred<'_>)],
2744 filter: F,
2745 ) -> (Vec<UnifiedEntity>, SegmentScanStats)
2746 where
2747 F: Fn(&UnifiedEntity) -> bool + Sync,
2748 {
2749 self.query_all_zoned_with_stats(zone_preds, |entity| {
2750 Self::scan_entity_visible(snapshot, entity) && filter(entity)
2751 })
2752 }
2753
2754 pub fn scan_project_zoned_with_stats<T, F>(
2758 &self,
2759 snapshot: Option<&SnapshotContext>,
2760 zone_preds: &[(&str, ZoneColPred<'_>)],
2761 project: F,
2762 ) -> (Vec<T>, SegmentScanStats)
2763 where
2764 T: Send,
2765 F: Fn(&UnifiedEntity) -> Option<T> + Sync,
2766 {
2767 self.project_all_zoned_with_stats(zone_preds, |entity| {
2768 if Self::scan_entity_visible(snapshot, entity) {
2769 project(entity)
2770 } else {
2771 None
2772 }
2773 })
2774 }
2775
2776 pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
2779 where
2780 F: Fn(&UnifiedEntity) -> bool + Sync,
2781 {
2782 let mut entities = self.query_all_in_segments(|_| true, filter);
2783 self.restore_owned_row_schemas(entities.iter_mut());
2784 entities
2785 }
2786
2787 fn query_all_in_segments<F>(
2788 &self,
2789 segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
2790 filter: F,
2791 ) -> Vec<UnifiedEntity>
2792 where
2793 F: Fn(&UnifiedEntity) -> bool + Sync,
2794 {
2795 let mut results = Vec::new();
2796
2797 if let Some(growing_arc) = self.growing.read().as_ref() {
2800 let growing = if let Some(g) = growing_arc.try_read() {
2801 g
2802 } else {
2803 growing_arc.read()
2804 };
2805 if segment_filter(&growing) {
2806 results.extend(growing.iter().filter(|e| filter(e)).cloned());
2807 }
2808 }
2809
2810 let sealed = self.sealed.read();
2812 let use_parallel = sealed.len() > 1 && crate::execution::should_parallelize();
2813 if use_parallel {
2814 let filter_ref = &filter;
2815 let segment_filter_ref = &segment_filter;
2816 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
2817 sealed
2818 .iter()
2819 .map(|segment| {
2820 s.spawn(move || {
2821 let segment = segment.read();
2822 if segment_filter_ref(&segment) {
2823 segment
2824 .iter()
2825 .filter(|e| filter_ref(e))
2826 .cloned()
2827 .collect::<Vec<_>>()
2828 } else {
2829 Vec::new()
2830 }
2831 })
2832 })
2833 .collect::<Vec<_>>()
2834 .into_iter()
2835 .map(|handle| handle.join().unwrap_or_default())
2836 .collect()
2837 });
2838 for batch in segment_results {
2839 results.extend(batch);
2840 }
2841 } else {
2842 for segment in sealed.iter() {
2843 let seg = segment.read();
2844 if segment_filter(&seg) {
2845 results.extend(seg.iter().filter(|e| filter(e)).cloned());
2846 }
2847 }
2848 }
2849
2850 results
2851 }
2852
2853 pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
2858 if let Some(growing_arc) = self.growing.read().as_ref() {
2859 let growing = growing_arc.read();
2860 if growing.may_contain_exact_key(key) {
2861 return true;
2862 }
2863 } else {
2864 return true;
2865 }
2866
2867 let sealed = self.sealed.read();
2868 for segment_arc in sealed.iter() {
2869 let segment = segment_arc.read();
2870 if segment.may_contain_exact_key(key) {
2871 return true;
2872 }
2873 }
2874
2875 false
2876 }
2877
2878 pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2880 let mut results = Vec::new();
2881
2882 if let Some(growing_arc) = self.growing.read().as_ref() {
2884 let growing = growing_arc.read();
2885 results.extend(growing.filter_metadata(filters));
2886 }
2887
2888 let sealed = self.sealed.read();
2890 for segment in sealed.iter() {
2891 results.extend(segment.filter_metadata(filters));
2892 }
2893
2894 results
2895 }
2896
2897 pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2899 let mut results = Vec::new();
2900
2901 if let Some(growing_arc) = self.growing.read().as_ref() {
2903 let growing = growing_arc.read();
2904 for entity in growing.iter_kind(kind) {
2905 results.push(entity.clone());
2906 }
2907 }
2908
2909 let sealed = self.sealed.read();
2914 for segment_arc in sealed.iter() {
2915 let segment = segment_arc.read();
2916 for entity in segment.iter_kind(kind) {
2917 results.push(entity.clone());
2918 }
2919 }
2920
2921 results
2922 }
2923
2924 pub fn count(&self) -> usize {
2926 self.total_entities_atomic.load(Ordering::Relaxed) as usize
2927 }
2928
2929 pub fn segment_ids(&self) -> Vec<SegmentId> {
2931 let mut ids = Vec::new();
2932
2933 if let Some(growing_arc) = self.growing.read().as_ref() {
2934 ids.push(growing_arc.read().id());
2935 }
2936
2937 let sealed = self.sealed.read();
2938 for segment in sealed.iter() {
2939 ids.push(segment.id());
2940 }
2941
2942 ids.extend(self.archived.read().iter().copied());
2943
2944 ids
2945 }
2946
2947 #[inline]
2960 #[allow(clippy::unused_self)]
2961 fn emit(&self, _event: LifecycleEvent) {}
2962
2963 pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2966 std::mem::take(&mut *self.events.write())
2967 }
2968
2969 pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2971 self.maybe_seal_growing()?;
2973
2974 if self.config.enable_consolidation {
2975 self.consolidation_tick();
2976 }
2977
2978 Ok(())
2979 }
2980
2981 fn consolidation_tick(&self) {
2992 let running = self.consolidation.read().is_some();
2995 if !running && !self.start_consolidation() {
2996 return;
2997 }
2998
2999 let mut guard = self.consolidation.write();
3002 let Some(run) = guard.as_mut() else {
3003 return;
3004 };
3005 self.copy_bounded(run);
3006 if run.cursor >= run.sources.len() {
3007 if let Some(run) = guard.take() {
3008 self.finish_consolidation(run);
3009 }
3010 }
3011 }
3012
3013 fn start_consolidation(&self) -> bool {
3017 let sources = self.select_candidates();
3018 self.start_consolidation_from_sources(sources)
3019 }
3020
3021 fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
3022 if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
3023 return false;
3024 }
3025
3026 let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
3027 let merged = GrowingSegment::new(merged_id, &self.collection);
3028
3029 *self.consolidation.write() = Some(Consolidation {
3030 sources: sources
3031 .into_iter()
3032 .map(|id| SourceProgress {
3033 id,
3034 copied: Vec::new(),
3035 tombstones_at_open: None,
3036 })
3037 .collect(),
3038 merged_id,
3039 merged,
3040 cursor: 0,
3041 pending_ids: Vec::new(),
3042 pending_cursor: 0,
3043 entities_copied: 0,
3044 });
3045
3046 self.stats.write().consolidation.runs_started += 1;
3047 true
3048 }
3049
3050 pub fn pressure_consolidation_tick(&self) -> bool {
3055 if !self.config.enable_consolidation {
3056 return false;
3057 }
3058
3059 let running = self.consolidation.read().is_some();
3060 if !running {
3061 if self
3065 .seal_current_if(|segment| segment.reclaimable_bytes() > 0)
3066 .is_err()
3067 {
3068 return false;
3069 }
3070 let sources = self.dirty_sealed_segments();
3071 if !self.start_consolidation_from_sources(sources) {
3072 return false;
3073 }
3074 }
3075
3076 self.consolidation_tick();
3077 true
3078 }
3079
3080 fn select_candidates(&self) -> Vec<SegmentId> {
3095 let sealed = self.sealed.read();
3096 if sealed.is_empty() {
3097 return Vec::new();
3098 }
3099
3100 if sealed.len() > self.config.max_sealed_segments {
3101 return sealed.iter().map(|segment| segment.read().id()).collect();
3102 }
3103
3104 let mut live = 0usize;
3105 let mut tombstones = 0usize;
3106 let mut held = 0u64;
3107 let mut reclaimable = 0u64;
3108 let mut dirty = Vec::new();
3109
3110 for segment_arc in sealed.iter() {
3111 let segment = segment_arc.read();
3112 live += segment.live_entity_count();
3113 tombstones += segment.tombstone_count();
3114 held += segment.resident_bytes();
3115 reclaimable += segment.reclaimable_bytes();
3116 if segment.tombstone_count() > 0 {
3117 dirty.push(segment.id());
3118 }
3119 }
3120
3121 let total = live + tombstones;
3122 let tombstone_ratio = if total == 0 {
3123 0.0
3124 } else {
3125 tombstones as f64 / total as f64
3126 };
3127 let fragmentation = if held == 0 {
3128 0.0
3129 } else {
3130 reclaimable as f64 / held as f64
3131 };
3132
3133 if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
3134 || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
3135 {
3136 return dirty;
3137 }
3138
3139 Vec::new()
3140 }
3141
3142 fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
3143 self.sealed
3144 .read()
3145 .iter()
3146 .filter_map(|segment_arc| {
3147 let segment = segment_arc.read();
3148 (segment.reclaimable_bytes() > 0).then(|| segment.id())
3149 })
3150 .collect()
3151 }
3152
3153 fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
3156 if sources.is_empty() {
3157 return false;
3158 }
3159 let sealed = self.sealed.read();
3160 sealed.iter().any(|segment_arc| {
3161 let segment = segment_arc.read();
3162 sources.contains(&segment.id()) && segment.tombstone_count() > 0
3163 })
3164 }
3165
3166 fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
3167 let sealed = self.sealed.read();
3168 sealed
3169 .iter()
3170 .find(|segment| segment.read().id() == id)
3171 .map(Arc::clone)
3172 }
3173
3174 fn copy_bounded(&self, run: &mut Consolidation) {
3179 let mut budget = self.config.consolidation_entities_per_tick;
3180
3181 while budget > 0 && run.cursor < run.sources.len() {
3182 let source_id = run.sources[run.cursor].id;
3183 let Some(source_arc) = self.sealed_segment(source_id) else {
3184 run.cursor += 1;
3188 run.pending_ids.clear();
3189 run.pending_cursor = 0;
3190 continue;
3191 };
3192
3193 if run.sources[run.cursor].tombstones_at_open.is_none() {
3194 let source = source_arc.read();
3195 run.pending_ids = source.live_entity_ids();
3196 run.pending_cursor = 0;
3197 run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
3198 }
3199
3200 {
3201 let source = source_arc.read();
3202 while budget > 0 && run.pending_cursor < run.pending_ids.len() {
3203 let id = run.pending_ids[run.pending_cursor];
3204 run.pending_cursor += 1;
3205 budget -= 1;
3206
3207 let Some(entity) = source.get(id) else {
3210 continue;
3211 };
3212 let metadata = source.get_metadata(id);
3213 run.merged.adopt_entity(entity.clone(), metadata);
3214 run.sources[run.cursor].copied.push(id);
3215 run.entities_copied += 1;
3216 }
3217 }
3218
3219 if run.pending_cursor >= run.pending_ids.len() {
3220 run.cursor += 1;
3221 run.pending_ids.clear();
3222 run.pending_cursor = 0;
3223 }
3224 }
3225 }
3226
3227 fn finish_consolidation(&self, mut run: Consolidation) {
3233 let mut sealed = self.sealed.write();
3234
3235 let mut positions = Vec::with_capacity(run.sources.len());
3236 for source in &run.sources {
3237 if let Some(index) = sealed
3238 .iter()
3239 .position(|segment| segment.read().id() == source.id)
3240 {
3241 positions.push(index);
3242 }
3243 }
3244 if positions.is_empty() {
3245 return;
3246 }
3247 positions.sort_unstable();
3248
3249 let mut stale = Vec::new();
3253 for source in &run.sources {
3254 let Some(source_arc) = sealed
3255 .iter()
3256 .find(|segment| segment.read().id() == source.id)
3257 else {
3258 continue;
3259 };
3260 let segment = source_arc.read();
3261 if source.tombstones_at_open == Some(segment.tombstone_count()) {
3262 continue;
3263 }
3264 stale.extend(
3265 source
3266 .copied
3267 .iter()
3268 .copied()
3269 .filter(|id| !segment.contains(*id)),
3270 );
3271 }
3272 for id in stale {
3273 run.merged.evict_entity(id);
3274 }
3275
3276 let reclaimed_bytes: u64 = positions
3277 .iter()
3278 .map(|&index| sealed[index].read().resident_bytes())
3279 .sum();
3280 let reclaimed_tombstones: u64 = positions
3281 .iter()
3282 .map(|&index| sealed[index].read().tombstone_count() as u64)
3283 .sum();
3284
3285 if run.merged.seal().is_err() {
3286 return;
3287 }
3288 let merged_bytes = run.merged.resident_bytes();
3289 let merged_arc = Arc::new(RwLock::new(run.merged));
3290
3291 {
3292 let mut retired = self.retired.write();
3293 retired.retain(|source| source.strong_count() > 0);
3294 for &index in positions.iter().rev() {
3295 let source = sealed.remove(index);
3296 if Arc::strong_count(&source) > 1 {
3297 retired.push(Arc::downgrade(&source));
3298 }
3299 }
3300 }
3301 let insert_at = positions[0].min(sealed.len());
3302 sealed.insert(insert_at, Arc::clone(&merged_arc));
3303 drop(sealed);
3304
3305 let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
3307 {
3308 let mut entity_segment = self.entity_segment.write();
3309 for segment_id in entity_segment.values_mut() {
3310 if source_ids.contains(segment_id) {
3311 *segment_id = run.merged_id;
3312 }
3313 }
3314 }
3315
3316 {
3317 let mut stats = self.stats.write();
3318 stats.sealed_count = stats
3319 .sealed_count
3320 .saturating_sub(positions.len())
3321 .saturating_add(1);
3322 let consolidation = &mut stats.consolidation;
3323 consolidation.runs_completed += 1;
3324 consolidation.segments_merged += positions.len() as u64;
3325 consolidation.tombstones_reclaimed += reclaimed_tombstones;
3326 consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
3327 }
3328
3329 self.emit(LifecycleEvent::SegmentConsolidated {
3330 source: source_ids,
3331 target: run.merged_id,
3332 });
3333 }
3334
3335 #[cfg(test)]
3338 fn consolidation_progress(&self) -> Option<u64> {
3339 self.consolidation
3340 .read()
3341 .as_ref()
3342 .map(|run| run.entities_copied)
3343 }
3344}
3345
3346#[cfg(test)]
3347mod tests {
3348 use std::collections::HashSet;
3349
3350 use super::*;
3351 use crate::storage::entity::{EntityData, EntityKind, RowData};
3352 use crate::transaction::snapshot::{Snapshot, SnapshotManager};
3353 use reddb_types::Value;
3354
3355 #[test]
3356 fn recovery_view_projection_keeps_schema_bounds_without_retaining_rows_or_identities() {
3357 for collection in ["rows", "red_internal"] {
3358 let manager = SegmentManager::new(collection);
3359 let identities = manager.identity_counters();
3360 for id in 1..=18 {
3361 let field = format!("field_{id}");
3362 let mut entity =
3363 UnifiedEntity::table_row(EntityId::new(id), collection, id, vec![]);
3364 entity.data = EntityData::Row(RowData {
3365 columns: Vec::new(),
3366 named: Some(HashMap::from([
3367 (field.clone(), Value::Integer(id as i64)),
3368 ("id".into(), Value::Integer(id as i64)),
3369 ])),
3370 schema: None,
3371 });
3372 manager.restore_snapshot_row_layout(&mut entity);
3373 let row = entity.data.as_row().expect("row");
3374 assert_eq!(row.get_field(&field), Some(&Value::Integer(id as i64)));
3375 assert_eq!(row.get_field("id"), Some(&Value::Integer(id as i64)));
3376 assert_eq!(
3377 row.named.is_some(),
3378 collection.starts_with("red_") || id > 16
3379 );
3380 }
3381 assert_eq!(manager.count(), 0, "a projection never adopts payloads");
3382 assert_eq!(manager.identity_counters(), identities);
3383 assert_eq!(
3384 manager.recovery_row_schemas.read().len(),
3385 if collection == "rows" { 16 } else { 0 }
3386 );
3387 }
3388 }
3389
3390 #[test]
3391 fn physical_pages_keep_global_order_across_segments_and_deleted_slots() {
3392 let manager = SegmentManager::new("pages");
3393 for group in [[50, 10], [30, 20]] {
3394 for id in group {
3395 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "pages", id, vec![]);
3396 entity.data = EntityData::Row(RowData::with_names(
3397 vec![Value::Integer(id as i64)],
3398 vec!["id".into()],
3399 ));
3400 manager.insert(entity).unwrap();
3401 }
3402 manager.force_seal().unwrap();
3403 }
3404 manager
3405 .insert(UnifiedEntity::table_row(
3406 EntityId::new(40),
3407 "pages",
3408 40,
3409 vec![],
3410 ))
3411 .unwrap();
3412 let (total, page) = manager.page_in_id_order(1, 2);
3413 assert_eq!(total, 5);
3414 assert_eq!(
3415 page.iter()
3416 .map(|entity| entity.id.raw())
3417 .collect::<Vec<_>>(),
3418 [20, 30]
3419 );
3420 let EntityData::Row(row) = &page[0].data else {
3421 panic!("row")
3422 };
3423 assert_eq!(row.get_field("id"), Some(&Value::Integer(20)));
3424 manager.delete(EntityId::new(30)).unwrap();
3425 let (total, page) = manager.page_in_id_order(1, 2);
3426 assert_eq!(total, 4);
3427 assert_eq!(
3428 page.iter()
3429 .map(|entity| entity.id.raw())
3430 .collect::<Vec<_>>(),
3431 [20, 40]
3432 );
3433 for offset in [4, usize::MAX] {
3434 let (total, page) = manager.page_in_id_order(offset, 2);
3435 assert_eq!(total, 4);
3436 assert!(page.is_empty());
3437 }
3438 }
3439
3440 fn lookup_version(id: u64, key: &str, xmin: u64, xmax: u64) -> UnifiedEntity {
3441 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "keys", id, vec![]);
3442 entity.data = EntityData::Row(RowData::with_names(
3443 vec![Value::text(key)],
3444 vec!["key".into()],
3445 ));
3446 entity.set_xmin(xmin);
3447 entity.set_xmax(xmax);
3448 entity
3449 }
3450
3451 #[test]
3452 fn key_lookup_matches_snapshot_scan_across_segments() {
3453 let manager = Arc::new(SegmentManager::new("keys"));
3454 for entity in [
3455 lookup_version(1, "changed", 3, 12),
3456 lookup_version(2, "deleted", 3, 7),
3457 ] {
3458 manager.insert(entity).expect("old version");
3459 }
3460 manager.force_seal().expect("seal old versions");
3461 for entity in [
3462 lookup_version(3, "changed", 12, 0),
3463 lookup_version(4, "pending", 4, 0),
3464 lookup_version(5, "own", 15, 0),
3465 lookup_version(6, "live", 3, 0),
3466 ] {
3467 manager.insert(entity).expect("new version");
3468 }
3469 std::thread::spawn(move || {
3470 let snapshot = SnapshotContext {
3471 snapshot: Snapshot {
3472 xid: 10,
3473 in_progress: HashSet::from([4]),
3474 },
3475 manager: Arc::new(SnapshotManager::new()),
3476 own_xids: HashSet::from([15]),
3477 requires_index_fallback: false,
3478 serializable_reader: None,
3479 };
3480 for snapshot in [Some(&snapshot), None] {
3481 for key in ["changed", "deleted", "pending", "own", "live", "absent"] {
3482 let expected = manager
3483 .scan(snapshot, |entity| {
3484 crate::storage::segment_key_lookup::SegmentKeyLookup::key(entity)
3485 == Some(key)
3486 })
3487 .into_iter()
3488 .max_by_key(|entity| entity.xmin)
3489 .map(|entity| entity.id);
3490 assert_eq!(
3491 manager.project_key_entity(key, snapshot, true, |entity| entity.id),
3492 expected,
3493 "key={key}"
3494 );
3495 }
3496 }
3497 assert_eq!(
3498 manager.project_key_entity("changed", Some(&snapshot), true, |entity| entity.id),
3499 Some(EntityId::new(1))
3500 );
3501 assert_eq!(
3502 manager.project_key_entity("changed", None, true, |entity| entity.id),
3503 Some(EntityId::new(3))
3504 );
3505 assert!(manager
3506 .project_key_entity("deleted", None, true, |entity| entity.id)
3507 .is_none());
3508 })
3509 .join()
3510 .expect("snapshot worker");
3511 }
3512
3513 #[test]
3514 fn key_lookup_survives_paced_consolidation() {
3515 let manager = consolidating_manager(1);
3516 manager
3517 .bulk_insert(
3518 (1..=10)
3519 .map(|id| lookup_version(id, &format!("key-{id}"), 1, 0))
3520 .collect(),
3521 )
3522 .expect("rows");
3523 manager.force_seal().expect("seal");
3524 for id in 1..=4 {
3525 manager.delete(EntityId::new(id)).expect("delete");
3526 }
3527 assert!(manager.pressure_consolidation_tick());
3528 drain_maintenance(&manager);
3529 for id in 1..=10 {
3530 assert_eq!(
3531 manager.project_key_entity(&format!("key-{id}"), None, false, |entity| entity.id),
3532 (id > 4).then_some(EntityId::new(id))
3533 );
3534 }
3535 }
3536
3537 #[test]
3538 #[ignore = "diagnostic of segment fragmentation and retained history; no timing assertions"]
3539 fn key_lookup_fragmentation_history_diagnostic() {
3540 use std::{hint::black_box, time::Instant};
3541 for segments in [1, 16, 64] {
3542 for versions in [1, 100, 1000] {
3543 let manager = SegmentManager::new("keys");
3544 let mut id = 1;
3545 for segment in 0..segments {
3546 let mut rows: Vec<_> = (0..128)
3547 .map(|_| {
3548 let entity = lookup_version(id, &format!("other-{id}"), 1, 0);
3549 id += 1;
3550 entity
3551 })
3552 .collect();
3553 for version in (segment..versions).step_by(segments) {
3554 rows.push(lookup_version(
3555 id,
3556 "history",
3557 version as u64 + 1,
3558 if version + 1 == versions {
3559 0
3560 } else {
3561 version as u64 + 2
3562 },
3563 ));
3564 id += 1;
3565 }
3566 manager.bulk_insert(rows).expect("rows");
3567 manager.force_seal().expect("seal");
3568 }
3569 let started = Instant::now();
3570 for _ in 0..2000 {
3571 let xmin =
3572 manager.project_key_entity(black_box("history"), None, true, |entity| {
3573 entity.xmin
3574 });
3575 assert_eq!(xmin, Some(versions as u64));
3576 }
3577 eprintln!("segments={segments} versions={versions} reads=2000 elapsed_us={} resident_bytes={}", started.elapsed().as_micros(), manager.resident_bytes());
3578 }
3579 }
3580 }
3581
3582 #[test]
3583 fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
3584 let manager = SegmentManager::new("fingerprint");
3585 let row = |id, label: &str| {
3586 UnifiedEntity::table_row(
3587 EntityId::new(id),
3588 "fingerprint",
3589 id,
3590 vec![Value::text(label)],
3591 )
3592 };
3593 manager
3594 .bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
3595 .expect("sealed rows");
3596 manager.force_seal().expect("seal");
3597 manager
3598 .bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
3599 .expect("growing rows");
3600
3601 let mut expected = manager.query_all(|_| true);
3602 expected.sort_by_key(|entity| entity.id.raw());
3603 let expected: Vec<_> = expected
3604 .iter()
3605 .map(|entity| format!("{entity:?}"))
3606 .collect();
3607 let mut stored_addresses = HashSet::new();
3608 manager.for_each_entity(|entity| {
3609 stored_addresses.insert(entity as *const UnifiedEntity as usize);
3610 true
3611 });
3612
3613 let mut actual = Vec::new();
3614 manager.for_each_entity_in_id_order(|entity| {
3615 assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
3616 assert!(
3617 manager.growing.try_write().is_some(),
3618 "topology guard retained"
3619 );
3620 assert!(
3621 manager.sealed.try_write().is_some(),
3622 "topology guard retained"
3623 );
3624 actual.push(format!("{entity:?}"));
3625 true
3626 });
3627 assert_eq!(
3628 actual, expected,
3629 "equal IDs retain their original segment order"
3630 );
3631 }
3632
3633 #[test]
3634 fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
3635 let manager = SegmentManager::new("fingerprint");
3636 for id in [3, 1, 2] {
3637 manager
3638 .insert(UnifiedEntity::table_row(
3639 EntityId::new(id),
3640 "fingerprint",
3641 id,
3642 vec![Value::Integer(id as i64)],
3643 ))
3644 .expect("row");
3645 }
3646 let source = manager.growing.read().as_ref().expect("source").clone();
3647 let handles = Arc::strong_count(&source);
3648 let mut seen = Vec::new();
3649 manager.for_each_entity_in_id_order(|entity| {
3650 seen.push(entity.id.raw());
3651 false
3652 });
3653 assert_eq!(seen, vec![1]);
3654 assert_eq!(Arc::strong_count(&source), handles);
3655 assert!(
3656 source.try_write().is_some(),
3657 "segment guard retained after early stop"
3658 );
3659 }
3660
3661 #[test]
3662 fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
3663 let manager = SegmentManager::with_config(
3664 "vectors",
3665 ManagerConfig {
3666 max_sealed_segments: 1,
3667 consolidation_entities_per_tick: 10_000,
3668 ..Default::default()
3669 },
3670 );
3671 for part in 0..3 {
3672 manager
3673 .bulk_insert(
3674 (1..=300)
3675 .map(|offset| {
3676 UnifiedEntity::vector(
3677 EntityId::new(part * 300 + offset),
3678 "vectors",
3679 vec![1.0, 0.0],
3680 )
3681 })
3682 .collect(),
3683 )
3684 .expect("batch");
3685 if part < 2 {
3686 manager.force_seal().expect("seal");
3687 }
3688 }
3689 let original_growing = manager.growing.read().as_ref().expect("growing").clone();
3690 let mut changed = false;
3691 let mut seen = Vec::new();
3692 assert!(manager.scan_batches(
3693 None,
3694 || true,
3695 |batch| {
3696 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3697 assert!(
3698 manager.growing.try_write().is_some(),
3699 "consumer holds no topology guard"
3700 );
3701 assert!(
3702 original_growing.try_write().is_some(),
3703 "consumer holds no segment guard"
3704 );
3705 for entity in batch {
3706 assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
3707 seen.push(entity.id.raw());
3708 }
3709 if !changed {
3710 changed = true;
3711 manager.force_seal().expect("seal during scan");
3712 manager
3713 .insert(UnifiedEntity::vector(
3714 EntityId::new(5000),
3715 "vectors",
3716 vec![1.0, 0.0],
3717 ))
3718 .expect("new write");
3719 drain_maintenance(&manager);
3720 assert!(manager.stats().consolidation.runs_completed > 0);
3721 assert!(manager
3722 .delete(EntityId::new(1))
3723 .expect("delete after source retirement"));
3724 assert!(manager.get(EntityId::new(1)).is_none());
3725 }
3726 true
3727 }
3728 ));
3729 seen.sort_unstable();
3730 assert_eq!(
3731 seen,
3732 (2..=900).collect::<Vec<_>>(),
3733 "no duplicate, missing or late candidates"
3734 );
3735 let mut after = Vec::new();
3736 assert!(manager.scan_batches(
3737 None,
3738 || true,
3739 |batch| {
3740 after.extend(batch.iter().map(|entity| entity.id.raw()));
3741 true
3742 }
3743 ));
3744 after.sort_unstable();
3745 let mut expected: Vec<_> = (2..=900).collect();
3746 expected.push(5000);
3747 assert_eq!(after, expected, "adopted entities have scan positions");
3748 }
3749
3750 #[test]
3751 fn bounded_scan_applies_snapshot_to_relocated_versions() {
3752 for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
3753 let manager = SegmentManager::with_config(
3754 "vectors",
3755 ManagerConfig {
3756 max_sealed_segments: 1,
3757 ..Default::default()
3758 },
3759 );
3760 for ids in [[1, 2], [3, 4]] {
3761 let entities = ids.map(|id| {
3762 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3763 entity.set_xmin(3);
3764 entity
3765 });
3766 if bulk {
3767 manager
3768 .bulk_insert(entities.into())
3769 .expect("initial versions");
3770 } else {
3771 for entity in entities {
3772 manager.insert(entity).expect("initial version");
3773 }
3774 }
3775 manager.force_seal().expect("seal source");
3776 }
3777 let snapshot = SnapshotContext {
3778 snapshot: Snapshot {
3779 xid: 10,
3780 in_progress: HashSet::new(),
3781 },
3782 manager: Arc::new(SnapshotManager::new()),
3783 own_xids: HashSet::new(),
3784 requires_index_fallback: false,
3785 serializable_reader: None,
3786 };
3787 let mut changed = false;
3788 let mut seen = Vec::new();
3789 assert!(manager.scan_batches(
3790 Some(&snapshot),
3791 || true,
3792 |batch| {
3793 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3794 if !changed {
3795 changed = true;
3796 if consolidate {
3797 drain_maintenance(&manager);
3798 assert!(manager.stats().consolidation.runs_completed > 0);
3799 }
3800 for (id, xmax) in [(3, 8), (4, 12)] {
3801 let mut entity =
3802 manager.get(EntityId::new(id)).expect("retained version");
3803 entity.set_xmax(xmax);
3804 manager
3805 .update(entity)
3806 .expect("update after source retirement");
3807 }
3808 }
3809 true
3810 }
3811 ));
3812 seen.sort_unstable();
3813 assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
3814 let mut again = Vec::new();
3815 assert!(manager.scan_batches(
3816 Some(&snapshot),
3817 || true,
3818 |batch| {
3819 again.extend(batch.iter().map(|entity| entity.id.raw()));
3820 true
3821 }
3822 ));
3823 again.sort_unstable();
3824 assert_eq!(
3825 again,
3826 vec![1, 2, 4],
3827 "old slots do not duplicate relocated versions"
3828 );
3829 let mut current = Vec::new();
3830 assert!(manager.scan_batches(
3831 None,
3832 || true,
3833 |batch| {
3834 current.extend(batch.iter().map(|entity| entity.id.raw()));
3835 true
3836 }
3837 ));
3838 current.sort_unstable();
3839 assert_eq!(
3840 current,
3841 vec![1, 2],
3842 "current view excludes both deleted versions"
3843 );
3844 }
3845 }
3846
3847 #[test]
3848 fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
3849 let manager = SegmentManager::new("mixed");
3850 let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
3851 manager
3852 .bulk_insert((1..=300).map(vector).collect())
3853 .expect("flat");
3854 manager.insert(vector(500)).expect("hash insert");
3855 manager
3856 .bulk_insert(vec![vector(502), vector(505)])
3857 .expect("gapped bulk");
3858 manager.delete(EntityId::new(2)).expect("flat delete");
3859 manager.delete(EntityId::new(505)).expect("hash delete");
3860 let mut changed = false;
3861 let mut seen = Vec::new();
3862 assert!(manager.scan_batches(
3863 None,
3864 || true,
3865 |batch| {
3866 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3867 if !changed {
3868 changed = true;
3869 manager
3870 .insert(vector(700))
3871 .expect("append after cursor boundary");
3872 manager
3873 .update(UnifiedEntity::graph_node(
3874 EntityId::new(500),
3875 "mixed",
3876 "Node",
3877 HashMap::new(),
3878 ))
3879 .expect("structural update");
3880 }
3881 true
3882 }
3883 ));
3884 seen.sort_unstable();
3885 let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
3886 expected.extend([500, 502]);
3887 assert_eq!(seen, expected);
3888 }
3889
3890 #[test]
3891 fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
3892 let manager = SegmentManager::new("vectors");
3893 let mut entities = Vec::new();
3894 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
3895 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3896 entity.set_xmin(xmin);
3897 entity.set_xmax(xmax);
3898 entities.push(entity);
3899 }
3900 manager.bulk_insert(entities).expect("versions");
3901 manager.force_seal().expect("seal");
3902 let snapshot = SnapshotContext {
3903 snapshot: Snapshot {
3904 xid: 10,
3905 in_progress: HashSet::from([4]),
3906 },
3907 manager: Arc::new(SnapshotManager::new()),
3908 own_xids: HashSet::new(),
3909 requires_index_fallback: false,
3910 serializable_reader: None,
3911 };
3912 let mut work = 0;
3913 let mut callbacks = 0;
3914 assert!(!manager.scan_batches(
3915 Some(&snapshot),
3916 || {
3917 work += 1;
3918 work <= 2
3919 },
3920 |_| {
3921 callbacks += 1;
3922 true
3923 }
3924 ));
3925 assert_eq!(callbacks, 0, "a stopped batch is not consumed");
3926 for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
3927 let mut seen = Vec::new();
3928 assert!(manager.scan_batches(
3929 view,
3930 || true,
3931 |batch| {
3932 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3933 true
3934 }
3935 ));
3936 seen.sort_unstable();
3937 assert_eq!(seen, expected);
3938 }
3939 }
3940
3941 #[test]
3942 fn test_manager_basic() {
3943 let manager = SegmentManager::new("test_collection");
3944
3945 let entity = UnifiedEntity::table_row(
3946 manager.next_entity_id(),
3947 "users",
3948 1,
3949 vec![Value::text("Alice".to_string())],
3950 );
3951
3952 let id = manager.insert(entity).unwrap();
3953 assert!(manager.get(id).is_some());
3954 assert_eq!(manager.count(), 1);
3955 }
3956
3957 #[test]
3958 fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
3959 let manager = Arc::new(SegmentManager::new("mixed"));
3960 manager
3961 .bulk_insert(
3962 (1..=128)
3963 .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
3964 .collect(),
3965 )
3966 .expect("vector-only segment");
3967 manager.force_seal().expect("seal vectors");
3968 let mut old =
3969 UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
3970 old.set_xmin(3);
3971 old.set_xmax(12);
3972 let mut new =
3973 UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
3974 new.set_xmin(12);
3975 let edge = UnifiedEntity::graph_edge(
3976 EntityId::new(202),
3977 "mixed",
3978 "a",
3979 "b",
3980 1.0,
3981 Default::default(),
3982 );
3983 manager
3984 .bulk_insert(vec![old, new, edge])
3985 .expect("physical graph versions");
3986 manager.force_seal().expect("seal graph");
3987 manager
3988 .insert(UnifiedEntity::graph_node(
3989 EntityId::new(203),
3990 "mixed",
3991 "person",
3992 Default::default(),
3993 ))
3994 .expect("growing node");
3995 let snapshot = SnapshotContext {
3996 snapshot: Snapshot {
3997 xid: 10,
3998 in_progress: HashSet::new(),
3999 },
4000 manager: Arc::new(SnapshotManager::new()),
4001 own_xids: HashSet::new(),
4002 requires_index_fallback: false,
4003 serializable_reader: None,
4004 };
4005 std::thread::spawn(move || {
4006 for view in [None, Some(&snapshot)] {
4007 for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
4008 let expected: HashSet<_> = manager
4009 .scan(view, |item| kind.matches(&item.kind))
4010 .into_iter()
4011 .map(|item| item.id)
4012 .collect();
4013 let actual: HashSet<_> = manager
4014 .scan_graph_kind(view, kind)
4015 .into_iter()
4016 .map(|item| item.id)
4017 .collect();
4018 assert_eq!(actual, expected);
4019 let mut batched = HashSet::new();
4020 assert!(manager.scan_graph_batches(
4021 view,
4022 kind,
4023 || true,
4024 |batch| {
4025 batched.extend(batch.into_iter().map(|item| item.id));
4026 true
4027 }
4028 ));
4029 assert_eq!(batched, expected);
4030 let mut visited = 0;
4031 let mut candidates = HashSet::new();
4032 manager.scan_graph_candidates_for_each(
4033 view,
4034 kind,
4035 || true,
4036 |item| {
4037 visited += 1;
4038 if kind.matches(&item.kind) {
4039 candidates.insert(item.id);
4040 }
4041 true
4042 },
4043 );
4044 assert_eq!(candidates, expected);
4045 assert_eq!(
4046 visited,
4047 if matches!(kind, GraphEntityKind::Node) {
4048 3
4049 } else {
4050 2
4051 }
4052 );
4053 }
4054 }
4055 let ids: HashSet<_> = manager
4056 .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
4057 .into_iter()
4058 .map(|item| item.id.raw())
4059 .collect();
4060 assert_eq!(ids, HashSet::from([200, 203]));
4061 let ids: HashSet<_> = manager
4062 .scan_graph_kind(None, GraphEntityKind::Node)
4063 .into_iter()
4064 .map(|item| item.id.raw())
4065 .collect();
4066 assert_eq!(ids, HashSet::from([201, 203]));
4067 let mut visits = 0;
4068 manager.scan_graph_candidates_for_each(
4069 None,
4070 GraphEntityKind::Node,
4071 || true,
4072 |_| {
4073 visits += 1;
4074 false
4075 },
4076 );
4077 assert_eq!(
4078 visits, 1,
4079 "early termination must also stop subsequent segments"
4080 );
4081 })
4082 .join()
4083 .expect("snapshot worker");
4084 }
4085
4086 #[test]
4087 fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
4088 let manager = Arc::new(SegmentManager::new("accounts"));
4089 let cases = [
4090 ("legacy row", 0, 0, true),
4091 ("committed creator", 3, 0, true),
4092 ("future creator", 11, 0, false),
4093 ("in-progress creator", 4, 0, false),
4094 ("committed deleter", 3, 7, false),
4095 ("future deleter", 3, 11, true),
4096 ("in-progress deleter", 3, 4, true),
4097 ];
4098
4099 for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
4100 let row_id = index as u64 + 1;
4101 let mut entity = UnifiedEntity::table_row(
4102 EntityId::new(row_id),
4103 "accounts",
4104 row_id,
4105 vec![Value::Integer(row_id as i64)],
4106 );
4107 entity.set_xmin(*xmin);
4108 entity.set_xmax(*xmax);
4109 manager.insert(entity).unwrap();
4110 }
4111
4112 let snapshot = SnapshotContext {
4113 snapshot: Snapshot {
4114 xid: 10,
4115 in_progress: HashSet::from([4]),
4116 },
4117 manager: Arc::new(SnapshotManager::new()),
4118 own_xids: HashSet::new(),
4119 requires_index_fallback: false,
4120 serializable_reader: None,
4121 };
4122
4123 let raw_visible_row_ids: HashSet<u64> = manager
4124 .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
4125 .into_iter()
4126 .filter_map(|entity| match entity.kind {
4127 EntityKind::TableRow { row_id, .. } => Some(row_id),
4128 _ => None,
4129 })
4130 .collect();
4131
4132 let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
4133 .join()
4134 .unwrap();
4135 let visible_row_ids: HashSet<u64> = scanned
4136 .into_iter()
4137 .filter_map(|entity| match entity.kind {
4138 EntityKind::TableRow { row_id, .. } => Some(row_id),
4139 _ => None,
4140 })
4141 .collect();
4142 assert_eq!(visible_row_ids, raw_visible_row_ids);
4143
4144 for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
4145 let row_id = index as u64 + 1;
4146 assert_eq!(
4147 visible_row_ids.contains(&row_id),
4148 *expected_visible,
4149 "{name} visibility mismatch"
4150 );
4151 }
4152 }
4153
4154 #[test]
4155 fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
4156 let manager = SegmentManager::new("accounts");
4160 let mut live =
4161 UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
4162 live.set_xmin(2);
4163 live.set_xmax(0);
4164 manager.insert(live).unwrap();
4165
4166 let mut superseded =
4167 UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
4168 superseded.set_xmin(2);
4169 superseded.set_xmax(9);
4170 manager.insert(superseded).unwrap();
4171
4172 let rows: Vec<u64> = manager
4173 .scan(None, |_| true)
4174 .into_iter()
4175 .filter_map(|entity| match entity.kind {
4176 EntityKind::TableRow { row_id, .. } => Some(row_id),
4177 _ => None,
4178 })
4179 .collect();
4180 assert_eq!(
4181 rows,
4182 vec![1],
4183 "frameless scan must hide superseded versions"
4184 );
4185 }
4186
4187 #[test]
4188 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
4189 let manager = SegmentManager::new("test_collection");
4190
4191 let entity = UnifiedEntity::table_row(
4192 manager.next_entity_id(),
4193 "users",
4194 1,
4195 vec![Value::text("Alice".to_string())],
4196 );
4197
4198 let id = manager.insert(entity).unwrap();
4199
4200 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
4201 assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
4202 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
4203
4204 assert!(manager.delete(id).unwrap());
4205 assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
4206 }
4207
4208 #[test]
4209 fn zoned_scan_reports_pruned_segments_and_preserves_results() {
4210 let config = ManagerConfig {
4211 segment_config: SegmentConfig {
4212 max_entities: 2,
4213 ..Default::default()
4214 },
4215 ..Default::default()
4216 };
4217 let manager = SegmentManager::with_config("events", config);
4218
4219 for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
4220 let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
4221 let entity = UnifiedEntity::new(
4222 manager.next_entity_id(),
4223 EntityKind::TableRow {
4224 table: "events".into(),
4225 row_id: (idx + 1) as u64,
4226 },
4227 EntityData::Row(row),
4228 );
4229 manager.insert(entity).unwrap();
4230 }
4231
4232 let full = manager.query_all(|entity| match &entity.data {
4233 EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
4234 _ => false,
4235 });
4236 let probe = Value::Integer(100);
4237 let (pruned, stats) =
4238 manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
4239 match &entity.data {
4240 EntityData::Row(row) => {
4241 matches!(row.get_field("ts"), Some(Value::Integer(100)))
4242 }
4243 _ => false,
4244 }
4245 });
4246
4247 assert_eq!(pruned.len(), full.len());
4248 assert_eq!(pruned[0].id, full[0].id);
4249 assert_eq!(stats.segments_total, 3);
4250 assert_eq!(stats.segments_scanned, 2);
4251 assert_eq!(stats.segments_pruned, 1);
4252 }
4253
4254 #[test]
4255 fn projected_zoned_scan_preserves_segment_order_and_pruning() {
4256 let manager = SegmentManager::with_config(
4257 "events",
4258 ManagerConfig {
4259 segment_config: SegmentConfig {
4260 max_entities: 2,
4261 ..Default::default()
4262 },
4263 ..Default::default()
4264 },
4265 );
4266 for (index, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
4267 let mut entity = UnifiedEntity::table_row(
4268 EntityId::new(index as u64 + 1),
4269 "events",
4270 index as u64 + 1,
4271 vec![],
4272 );
4273 entity.data = EntityData::Row(RowData::with_names(
4274 vec![Value::Integer(value)],
4275 vec!["ts".into()],
4276 ));
4277 manager.insert(entity).expect("test row");
4278 }
4279 assert_eq!(manager.stats().sealed_count, 2);
4280 let project = |entity: &UnifiedEntity| match entity.data.as_row()?.get_field("ts")? {
4281 Value::Integer(value) => Some(*value),
4282 _ => None,
4283 };
4284 let (all, stats) = manager.scan_project_zoned_with_stats(None, &[], project);
4285 let (legacy, _) = manager.scan_zoned_with_stats(None, &[], |_| true);
4286 assert_eq!(all, legacy.iter().filter_map(project).collect::<Vec<_>>());
4287 assert_eq!(all[0], 200, "growing segment first");
4288 assert!(all[1..3].iter().all(|value| [10, 20].contains(value)));
4289 assert!(all[3..5].iter().all(|value| [100, 110].contains(value)));
4290 assert_eq!(stats.segments_total, 3);
4291 assert_eq!(stats.segments_scanned, 3);
4292 assert_eq!(stats.segments_pruned, 0);
4293 let probe = Value::Integer(100);
4294 let (pruned, stats) = manager.scan_project_zoned_with_stats(
4295 None,
4296 &[("ts", ZoneColPred::Gte(&probe))],
4297 |entity| project(entity).filter(|value| *value >= 100),
4298 );
4299 assert_eq!(
4300 pruned,
4301 all.into_iter()
4302 .filter(|value| *value >= 100)
4303 .collect::<Vec<_>>()
4304 );
4305 assert_eq!(stats.segments_total, 3);
4306 assert_eq!(stats.segments_scanned, 2);
4307 assert_eq!(stats.segments_pruned, 1);
4308 let (empty, _) = manager.scan_project_zoned_with_stats(None, &[], |_| None::<u64>);
4309 assert!(empty.is_empty(), "projection may reject every candidate");
4310
4311 let failure = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4312 manager.scan_project_zoned_with_stats(None, &[], |entity| {
4313 let value = project(entity)?;
4314 assert_ne!(value, 10, "injected projection worker failure");
4315 Some(value)
4316 })
4317 }));
4318 assert!(
4319 failure.is_err(),
4320 "a failed worker must not return a successful partial scan"
4321 );
4322 }
4323
4324 #[test]
4325 fn projected_zoned_scan_checks_visibility_before_projection_on_workers() {
4326 let manager = SegmentManager::with_config(
4327 "accounts",
4328 ManagerConfig {
4329 segment_config: SegmentConfig {
4330 max_entities: 2,
4331 ..Default::default()
4332 },
4333 ..Default::default()
4334 },
4335 );
4336 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4), (5, 3, 0)] {
4337 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "accounts", id, vec![]);
4338 entity.set_xmin(xmin);
4339 entity.set_xmax(xmax);
4340 manager.insert(entity).expect("versioned test row");
4341 }
4342 let mut hidden = UnifiedEntity::table_row(EntityId::new(6), "accounts", 6, vec![]);
4343 hidden.data = EntityData::Row(RowData::with_names(
4344 vec![Value::text("pending")],
4345 vec!["__moderation_status".into()],
4346 ));
4347 manager.insert(hidden).expect("moderated row");
4348 assert_eq!(manager.stats().sealed_count, 2);
4349 let snapshot = SnapshotContext {
4350 snapshot: Snapshot {
4351 xid: 10,
4352 in_progress: HashSet::from([4]),
4353 },
4354 manager: Arc::new(SnapshotManager::new()),
4355 own_xids: HashSet::new(),
4356 requires_index_fallback: false,
4357 serializable_reader: None,
4358 };
4359 std::thread::spawn(move || {
4360 for (view, expected) in [(Some(&snapshot), vec![5, 1, 3, 4]), (None, vec![5, 2, 3])] {
4361 let called = std::sync::Mutex::new(Vec::new());
4362 let (rows, stats) = manager.scan_project_zoned_with_stats(view, &[], |entity| {
4363 let id = entity.id.raw();
4364 assert!(
4365 expected.contains(&id),
4366 "invisible entity reached projection: {id}"
4367 );
4368 called.lock().expect("callback log").push(id);
4369 Some(id)
4370 });
4371 let (legacy, _) = manager.scan_zoned_with_stats(view, &[], |_| true);
4372 assert_eq!(
4373 rows,
4374 legacy
4375 .iter()
4376 .map(|entity| entity.id.raw())
4377 .collect::<Vec<_>>()
4378 );
4379 assert_eq!(rows[0], 5, "visible growing row first");
4380 let mut called = called.into_inner().expect("callback log");
4381 called.sort_unstable();
4382 let mut expected = expected;
4383 expected.sort_unstable();
4384 assert_eq!(called, expected, "each visible row projected exactly once");
4385 assert_eq!(stats.segments_total, 3);
4386 assert_eq!(stats.segments_scanned, 3);
4387 }
4388 })
4389 .join()
4390 .expect("projected scan worker");
4391 }
4392
4393 #[test]
4394 fn test_manager_auto_seal() {
4395 let config = ManagerConfig {
4396 segment_config: SegmentConfig {
4397 max_entities: 2,
4398 ..Default::default()
4399 },
4400 ..Default::default()
4401 };
4402
4403 let manager = SegmentManager::with_config("test", config);
4404
4405 manager
4407 .insert(UnifiedEntity::vector(
4408 manager.next_entity_id(),
4409 "v",
4410 vec![0.1],
4411 ))
4412 .unwrap();
4413
4414 manager
4416 .insert(UnifiedEntity::vector(
4417 manager.next_entity_id(),
4418 "v",
4419 vec![0.2],
4420 ))
4421 .unwrap();
4422
4423 manager
4425 .insert(UnifiedEntity::vector(
4426 manager.next_entity_id(),
4427 "v",
4428 vec![0.3],
4429 ))
4430 .unwrap();
4431
4432 let stats = manager.stats();
4433 assert_eq!(stats.total_entities, 3);
4434 }
4435
4436 #[test]
4437 fn test_manager_delete() {
4438 let manager = SegmentManager::new("test");
4439
4440 let id = manager
4441 .insert(UnifiedEntity::vector(
4442 manager.next_entity_id(),
4443 "v",
4444 vec![0.1],
4445 ))
4446 .unwrap();
4447
4448 assert!(manager.get(id).is_some());
4449 assert!(manager.delete(id).unwrap());
4450 assert!(manager.get(id).is_none());
4451 }
4452
4453 #[test]
4454 fn test_manager_metadata() {
4455 let manager = SegmentManager::new("test");
4456
4457 let id = manager
4458 .insert(UnifiedEntity::table_row(
4459 manager.next_entity_id(),
4460 "hosts",
4461 1,
4462 vec![Value::text("192.168.1.1".to_string())],
4463 ))
4464 .unwrap();
4465
4466 let mut meta = Metadata::new();
4467 meta.set(
4468 "os",
4469 super::super::metadata::MetadataValue::String("linux".to_string()),
4470 );
4471
4472 manager.set_metadata(id, meta).unwrap();
4473
4474 let retrieved = manager.get_metadata(id).unwrap();
4475 assert!(retrieved.has("os"));
4476 }
4477
4478 #[test]
4479 fn metadata_expiry_index_is_bounded_rekeyed_and_removed_with_entity() {
4480 let manager = SegmentManager::new("sessions");
4481 let mut first = UnifiedEntity::table_row(EntityId::new(1), "sessions", 1, vec![]);
4482 first.created_at = 10;
4483 let mut second = UnifiedEntity::table_row(EntityId::new(2), "sessions", 2, vec![]);
4484 second.created_at = 20;
4485 manager.insert(first).expect("first session");
4486 manager.insert(second).expect("second session");
4487
4488 let mut relative = Metadata::new();
4489 relative.set("_ttl_ms", MetadataValue::Int(5_000));
4490 manager
4491 .set_metadata(EntityId::new(1), relative)
4492 .expect("relative TTL");
4493 let mut absolute = Metadata::new();
4494 absolute.set("_expires_at", MetadataValue::Timestamp(30_000));
4495 manager
4496 .set_metadata(EntityId::new(2), absolute)
4497 .expect("absolute expiry");
4498
4499 assert_eq!(manager.scheduled_metadata_expiry_count(), 2);
4500 assert_eq!(
4501 manager.expired_metadata_ids(30_000, 1),
4502 vec![EntityId::new(1)],
4503 "the deadline index must stop at the requested batch"
4504 );
4505 assert_eq!(
4506 manager.expired_metadata_ids(30_000, 2),
4507 vec![EntityId::new(1), EntityId::new(2)]
4508 );
4509
4510 let mut postponed = Metadata::new();
4511 postponed.set("_expires_at", MetadataValue::Timestamp(60_000));
4512 manager
4513 .set_metadata(EntityId::new(1), postponed)
4514 .expect("postpone expiry");
4515 assert_eq!(
4516 manager.expired_metadata_ids(30_000, 2),
4517 vec![EntityId::new(2)],
4518 "metadata replacement must remove the stale deadline"
4519 );
4520
4521 assert!(manager.delete(EntityId::new(2)).expect("delete session"));
4522 assert!(manager
4523 .expired_metadata_ids(u64::MAX, 2)
4524 .contains(&EntityId::new(1)));
4525 assert!(!manager
4526 .expired_metadata_ids(u64::MAX, 2)
4527 .contains(&EntityId::new(2)));
4528 assert_eq!(manager.scheduled_metadata_expiry_count(), 1);
4529 }
4530
4531 #[test]
4532 #[ignore = "qualification: one million TTL deadlines"]
4533 fn qualification_million_ttl_backlog_drains_in_bounded_batches() {
4534 const ROWS: u64 = 1_000_000;
4535 const BATCH: usize = 1_000;
4536 const LOAD_BATCH: u64 = 10_000;
4537 let manager = SegmentManager::new("ttl_million");
4538 for first in (1..=ROWS).step_by(LOAD_BATCH as usize) {
4539 let last = first.saturating_add(LOAD_BATCH - 1).min(ROWS);
4540 manager
4541 .bulk_insert(
4542 (first..=last)
4543 .map(|id| {
4544 UnifiedEntity::table_row(EntityId::new(id), "ttl_million", id, vec![])
4545 })
4546 .collect(),
4547 )
4548 .expect("million-row TTL fixture");
4549 }
4550 assert_eq!(manager.count(), ROWS as usize);
4551 let entries = (1..=ROWS)
4552 .map(|id| (1, EntityId::new(id)))
4553 .collect::<Vec<_>>();
4554 manager.restore_metadata_expiry_projection(&entries, true);
4555 assert_eq!(manager.scheduled_metadata_expiry_count(), ROWS as usize);
4556
4557 let started = std::time::Instant::now();
4558 let mut drained = 0u64;
4559 loop {
4560 let batch = manager.expired_metadata_ids(1, BATCH);
4561 assert!(batch.len() <= BATCH);
4562 if batch.is_empty() {
4563 break;
4564 }
4565 drained += batch.len() as u64;
4566 for entity_id in batch {
4567 manager.remove_metadata_expiry(entity_id);
4568 }
4569 }
4570
4571 assert_eq!(drained, ROWS);
4572 assert_eq!(manager.scheduled_metadata_expiry_count(), 0);
4573 eprintln!(
4574 "ttl_deadlines={ROWS} batch={BATCH} drain_ms={}",
4575 started.elapsed().as_millis()
4576 );
4577 }
4578
4579 #[test]
4580 fn test_manager_query_by_kind() {
4581 let manager = SegmentManager::new("test");
4582
4583 manager
4584 .insert(UnifiedEntity::table_row(
4585 manager.next_entity_id(),
4586 "hosts",
4587 1,
4588 vec![],
4589 ))
4590 .unwrap();
4591
4592 manager
4593 .insert(UnifiedEntity::vector(
4594 manager.next_entity_id(),
4595 "embeddings",
4596 vec![0.1],
4597 ))
4598 .unwrap();
4599
4600 manager
4601 .insert(UnifiedEntity::table_row(
4602 manager.next_entity_id(),
4603 "hosts",
4604 2,
4605 vec![],
4606 ))
4607 .unwrap();
4608
4609 let rows = manager.get_by_kind("table");
4610 assert_eq!(rows.len(), 2);
4611
4612 let vectors = manager.get_by_kind("vector");
4613 assert_eq!(vectors.len(), 1);
4614 }
4615
4616 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
4623 SegmentManager::with_config(
4624 "test",
4625 ManagerConfig {
4626 max_sealed_segments: 1_000,
4627 consolidation_entities_per_tick: entities_per_tick,
4628 ..Default::default()
4629 },
4630 )
4631 }
4632
4633 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
4637 let entities: Vec<UnifiedEntity> = (0..count)
4638 .map(|i| {
4639 UnifiedEntity::table_row(
4640 manager.next_entity_id(),
4641 "users",
4642 i as u64,
4643 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
4644 )
4645 })
4646 .collect();
4647 let ids = manager.bulk_insert(entities).expect("bulk insert");
4648 manager.force_seal().expect("seal");
4649 ids
4650 }
4651
4652 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
4653 let mut ids: Vec<u64> = manager
4654 .query_all(|_| true)
4655 .into_iter()
4656 .map(|entity| entity.id.raw())
4657 .collect();
4658 ids.sort_unstable();
4659 ids
4660 }
4661
4662 fn drain_maintenance(manager: &SegmentManager) -> usize {
4663 let mut ticks = 0;
4664 while manager.consolidation.read().is_some() || ticks == 0 {
4665 manager.run_maintenance().expect("maintenance");
4666 ticks += 1;
4667 assert!(ticks < 1_000, "consolidation failed to converge");
4668 }
4669 ticks
4670 }
4671
4672 #[test]
4673 fn graph_batches_bound_payloads_and_hydrate_current_entities_after_sealing() {
4674 let manager = SegmentManager::new("graph");
4675 let node =
4676 |id| UnifiedEntity::graph_node(EntityId::new(id), "graph", "node", HashMap::new());
4677 manager
4678 .bulk_insert((1..=1024).map(node).collect())
4679 .expect("nodes");
4680 let mut hidden = node(1025);
4681 hidden.set_xmax(5);
4682 manager.insert(hidden).expect("hidden version");
4683 manager
4684 .insert(UnifiedEntity::vector(EntityId::new(1026), "graph", vec![]))
4685 .expect("mixed kind");
4686 let mut work = 0;
4687 let mut delivered = 0;
4688 assert!(!manager.scan_graph_batches(
4689 None,
4690 GraphEntityKind::Node,
4691 || {
4692 work += 1;
4693 work < 100
4694 },
4695 |batch| {
4696 delivered += batch.len();
4697 true
4698 }
4699 ));
4700 assert_eq!(delivered, 0, "cancelled partial batch must not escape");
4701 let mut seen = Vec::new();
4702 assert!(manager.scan_graph_batches(None, GraphEntityKind::Node, || true, |batch| {
4703 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4704 assert!(manager.growing.try_write().is_some(), "no topology guard");
4705 assert!(manager.sealed.try_write().is_some(), "no sealed guard");
4706 if seen.is_empty() {
4707 manager.force_seal().expect("seal during delivery");
4708 manager.insert(node(2000)).expect("append beyond captured boundary");
4709 let mut updated = node(900);
4710 if let EntityKind::GraphNode(node) = &mut updated.kind {
4711 node.label = "updated".into();
4712 }
4713 manager.update(updated).expect("update unread candidate");
4714 }
4715 for entity in batch {
4716 if entity.id.raw() == 900 {
4717 assert!(matches!(&entity.kind, EntityKind::GraphNode(node) if node.label == "updated"));
4718 }
4719 seen.push(entity.id.raw());
4720 }
4721 true
4722 }));
4723 assert_eq!(seen, (1..=1024).collect::<Vec<_>>());
4724 }
4725
4726 #[test]
4727 fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
4728 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4729 let ids = manager
4730 .bulk_insert(
4731 (1..=1024)
4732 .map(|id| {
4733 UnifiedEntity::graph_edge(
4734 EntityId::new(id),
4735 "step",
4736 "2000",
4737 "3000",
4738 1.0,
4739 HashMap::new(),
4740 )
4741 })
4742 .collect(),
4743 )
4744 .expect("edges");
4745 manager.force_seal().expect("seal");
4746 for id in ids.iter().take(300) {
4747 manager.delete(*id).expect("fragment source");
4748 }
4749 let source = Arc::downgrade(&manager.sealed.read()[0]);
4750 let mut seen = Vec::new();
4751 let mut changed = false;
4752 let mut metadata_bytes = 0;
4753 assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
4754 || true,
4755 |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
4756 |batch| {
4757 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4758 assert!(manager.growing.try_write().is_some(), "no growing topology guard");
4759 assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
4760 assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
4761 if !changed {
4762 changed = true;
4763 drain_maintenance(&manager);
4764 assert!(manager.stats().consolidation.runs_completed > 0);
4765 assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
4766 let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
4767 if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
4768 manager.update(updated).expect("update authoritative merged copy");
4769 }
4770 for id in batch {
4771 let entity = manager.get(*id).expect("hydrate current manager");
4772 if id.raw() == 1024 {
4773 assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
4774 }
4775 seen.push(id.raw());
4776 }
4777 Ok(true)
4778 }).expect("cursor"));
4779 assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
4780 assert!(
4781 source.upgrade().is_none(),
4782 "finished cursor releases retired source"
4783 );
4784 assert!(
4785 metadata_bytes < 4096,
4786 "fixed ID buffer plus source handles: {metadata_bytes}"
4787 );
4788 }
4789
4790 #[test]
4791 fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
4792 let manager = SegmentManager::new("graph");
4793 manager
4794 .bulk_insert(
4795 (1..=1024)
4796 .map(|id| {
4797 UnifiedEntity::graph_edge(
4798 EntityId::new(id),
4799 "step",
4800 "2000",
4801 "3000",
4802 1.0,
4803 HashMap::new(),
4804 )
4805 })
4806 .collect(),
4807 )
4808 .expect("edges");
4809 let mut delivered = 0;
4810 let denied = manager.visit_graph_index_batches(
4811 GraphEntityKind::Edge,
4812 &[EntityId::new(2000)],
4813 || true,
4814 |_| Err("out of memory"),
4815 |ids| {
4816 delivered += ids.len();
4817 Ok(true)
4818 },
4819 );
4820 assert_eq!(denied, Err("out of memory"));
4821 assert_eq!(delivered, 0);
4822 let mut work = 0;
4823 assert!(!manager
4824 .visit_graph_index_batches(
4825 GraphEntityKind::Edge,
4826 &[EntityId::new(2000)],
4827 || {
4828 work += 1;
4829 work < 100
4830 },
4831 |_| Ok::<_, ()>(()),
4832 |ids| {
4833 delivered += ids.len();
4834 Ok(true)
4835 }
4836 )
4837 .expect("cancelled cursor"));
4838 assert_eq!(
4839 delivered, 0,
4840 "cancelled partial batch cannot reach consumer"
4841 );
4842 assert!(!manager
4843 .visit_graph_index_batches(
4844 GraphEntityKind::Edge,
4845 &[EntityId::new(2000)],
4846 || true,
4847 |_| Ok::<_, ()>(()),
4848 |ids| {
4849 delivered += ids.len();
4850 Ok(false)
4851 }
4852 )
4853 .expect("consumer stop"));
4854 assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
4855 }
4856
4857 #[test]
4858 fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
4859 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4860 let ids = seal_bulk_segment(&manager, 100);
4861 for id in ids.iter().take(40) {
4862 manager.delete(*id).expect("delete before scan");
4863 }
4864 let source_bytes = manager.resident_bytes();
4865 let source = Arc::downgrade(&manager.sealed.read()[0]);
4866 std::thread::scope(|scope| {
4867 let (ready, waiting) = std::sync::mpsc::channel();
4868 let (release_first, first_wait) = std::sync::mpsc::channel();
4869 let (release_last, last_wait) = std::sync::mpsc::channel();
4870 let reader = |resume: std::sync::mpsc::Receiver<()>,
4871 ready: std::sync::mpsc::Sender<()>| {
4872 assert!(manager.scan_batches(
4873 None,
4874 || true,
4875 |_| {
4876 ready.send(()).expect("signal captured cursor");
4877 let _ = resume.recv();
4878 true
4879 }
4880 ));
4881 };
4882 let ready_first = ready.clone();
4883 let first = scope.spawn(move || reader(first_wait, ready_first));
4884 let last = scope.spawn(move || reader(last_wait, ready));
4885 for _ in 0..2 {
4886 waiting
4887 .recv_timeout(std::time::Duration::from_secs(10))
4888 .expect("reader ready");
4889 }
4890 drain_maintenance(&manager);
4891 assert!(
4892 source.upgrade().is_some(),
4893 "both cursors pin the retired source"
4894 );
4895 let active_bytes = manager.sealed.read()[0].read().resident_bytes();
4896 let held = manager.resident_bytes();
4897 assert!(
4898 held >= active_bytes + source_bytes,
4899 "retired source remains resident"
4900 );
4901 release_first.send(()).expect("release first reader");
4902 first.join().expect("first reader");
4903 assert_eq!(
4904 manager.resident_bytes(),
4905 held,
4906 "one reader still pins the source"
4907 );
4908 release_last.send(()).expect("release last reader");
4909 last.join().expect("last reader");
4910 assert!(
4911 source.upgrade().is_none(),
4912 "accounting must not retain payloads"
4913 );
4914 assert_eq!(held - manager.resident_bytes(), source_bytes);
4915 assert_eq!(
4916 manager.stats().total_memory_bytes as u64,
4917 manager.resident_bytes()
4918 );
4919 });
4920 }
4921
4922 #[test]
4923 fn physical_cursor_tracks_individual_and_mixed_insertions() {
4924 for mode in 0..3 {
4925 let manager = SegmentManager::new("physical_rows");
4926 let row = |id| {
4927 UnifiedEntity::table_row(
4928 EntityId::new(id),
4929 "physical_rows",
4930 id,
4931 vec![Value::Integer(id as i64)],
4932 )
4933 };
4934 match mode {
4935 0 => {
4936 for id in [10, 30, 20] {
4937 manager.insert(row(id)).expect("individual row");
4938 }
4939 }
4940 1 => {
4941 manager.insert(row(10)).expect("individual prefix");
4942 manager
4943 .bulk_insert(vec![row(30), row(31)])
4944 .expect("batch suffix");
4945 }
4946 _ => {
4947 manager
4948 .bulk_insert(vec![row(30), row(31)])
4949 .expect("batch prefix");
4950 manager.insert(row(10)).expect("individual suffix");
4951 }
4952 }
4953 let mut cursor = manager.physical_scan_cursor();
4954 let mut ids = Vec::new();
4955 for _ in 0..32 {
4956 if cursor.finished() {
4957 break;
4958 }
4959 let mut work = 1;
4960 let batch = manager.physical_scan_batch(&mut cursor, &mut work);
4961 assert!(batch.len() <= 1);
4962 ids.extend(batch.into_iter().map(|id| id.raw()));
4963 }
4964 assert!(cursor.finished());
4965 ids.sort_unstable();
4966 assert_eq!(
4967 ids,
4968 if mode == 0 {
4969 vec![10, 20, 30]
4970 } else {
4971 vec![10, 30, 31]
4972 },
4973 "mode={mode}"
4974 );
4975 }
4976 }
4977
4978 fn growing_tombstones_fixture() -> SegmentManager {
4979 let manager = consolidating_manager(1);
4980 let ids = manager
4981 .bulk_insert(
4982 (0..64)
4983 .map(|id| {
4984 UnifiedEntity::table_row(
4985 manager.next_entity_id(),
4986 "users",
4987 id,
4988 vec![Value::Blob(vec![7; 2048])],
4989 )
4990 })
4991 .collect(),
4992 )
4993 .expect("growing rows");
4994 for id in ids.iter().take(48) {
4995 manager.delete(*id).expect("physical tombstone");
4996 }
4997 assert_eq!(manager.stats().sealed_count, 0);
4998 manager
4999 }
5000
5001 #[test]
5002 fn pressure_consolidation_reclaims_unsealed_tombstones() {
5003 let manager = growing_tombstones_fixture();
5004 let before = manager.resident_bytes();
5005 assert!(manager.reclaimable_bytes() > 0);
5006 assert!(manager.pressure_consolidation_tick());
5007 drain_maintenance(&manager);
5008 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
5009 assert!(manager.resident_bytes() < before);
5010 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 48);
5011 }
5012
5013 #[test]
5014 fn pressure_consolidation_preserves_pinned_growing_source_accounting() {
5015 let manager = growing_tombstones_fixture();
5016 let source = Arc::clone(manager.growing.read().as_ref().expect("growing source"));
5017 let weak = Arc::downgrade(&source);
5018 assert!(manager.pressure_consolidation_tick());
5019 drain_maintenance(&manager);
5020 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
5021 let source_bytes = source.read().resident_bytes();
5022 let held = manager.resident_bytes();
5023 assert!(held >= source_bytes);
5024 drop(source);
5025 assert!(weak.upgrade().is_none());
5026 assert_eq!(held - manager.resident_bytes(), source_bytes);
5027 }
5028
5029 #[test]
5030 fn pressure_consolidation_keeps_clean_and_disabled_growing_sources_open() {
5031 let clean = consolidating_manager(1);
5032 clean
5033 .insert(UnifiedEntity::table_row(
5034 clean.next_entity_id(),
5035 "users",
5036 1,
5037 vec![Value::Integer(1)],
5038 ))
5039 .expect("clean row");
5040 assert!(!clean.pressure_consolidation_tick());
5041 assert_eq!(clean.stats().sealed_count, 0);
5042 let mut disabled = growing_tombstones_fixture();
5043 disabled.config.enable_consolidation = false;
5044 assert!(!disabled.pressure_consolidation_tick());
5045 assert_eq!(disabled.stats().sealed_count, 0);
5046 assert_eq!(live_ids(&disabled), (49..=64).collect::<Vec<_>>());
5047 }
5048
5049 #[test]
5050 fn retained_memory_includes_an_inflight_consolidation_copy() {
5051 let manager = consolidating_manager(1);
5052 let ids = seal_bulk_segment(&manager, 100);
5053 for id in ids.iter().take(40) {
5054 manager.delete(*id).expect("delete");
5055 }
5056 let before = manager.resident_bytes();
5057 assert!(manager.pressure_consolidation_tick());
5058 let copied = manager
5059 .consolidation
5060 .read()
5061 .as_ref()
5062 .expect("paced merge")
5063 .merged
5064 .resident_bytes();
5065 assert!(copied > 0);
5066 assert!(
5067 manager.resident_bytes() >= before + copied,
5068 "unpublished copy consumes memory"
5069 );
5070 drain_maintenance(&manager);
5071 assert!(
5072 manager.resident_bytes() < before,
5073 "completed merge releases the old storage"
5074 );
5075 }
5076
5077 #[test]
5078 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
5079 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5080 let ids = seal_bulk_segment(&manager, 100);
5081
5082 let bytes_before_delete = manager.resident_bytes();
5083
5084 for id in ids.iter().take(40) {
5086 assert!(manager.delete(*id).expect("delete"));
5087 }
5088
5089 let bytes_with_tombstones = manager.resident_bytes();
5092 assert!(bytes_with_tombstones > bytes_before_delete);
5093
5094 let survivors = live_ids(&manager);
5095 assert_eq!(survivors.len(), 60);
5096
5097 drain_maintenance(&manager);
5098
5099 let stats = manager.stats();
5100 assert_eq!(stats.consolidation.runs_started, 1);
5101 assert_eq!(stats.consolidation.runs_completed, 1);
5102 assert_eq!(stats.consolidation.segments_merged, 1);
5103 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
5104 assert!(stats.consolidation.bytes_reclaimed > 0);
5105
5106 let bytes_after = manager.resident_bytes();
5108 assert!(
5109 bytes_after <= bytes_before_delete,
5110 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
5111 );
5112 assert_eq!(
5113 bytes_with_tombstones - bytes_after,
5114 stats.consolidation.bytes_reclaimed
5115 );
5116 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
5117
5118 assert_eq!(live_ids(&manager), survivors);
5119 assert_eq!(manager.sealed.read().len(), 1);
5120 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
5121 }
5122
5123 #[test]
5124 fn query_results_are_identical_before_during_and_after_consolidation() {
5125 let manager = consolidating_manager(3);
5126 let ids = seal_bulk_segment(&manager, 20);
5127 seal_bulk_segment(&manager, 20);
5128 for id in ids.iter().take(10) {
5129 assert!(manager.delete(*id).expect("delete"));
5130 }
5131
5132 let before = live_ids(&manager);
5133 let survivor = *ids.last().expect("at least one row");
5134 let tombstoned = ids[0];
5135
5136 let mut mid_tick_observations = 0;
5139 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
5140 manager.run_maintenance().expect("maintenance");
5141 assert_eq!(live_ids(&manager), before);
5142 assert!(manager.get(survivor).is_some());
5143 assert!(manager.get(tombstoned).is_none());
5144 assert_eq!(manager.get_by_kind("table").len(), before.len());
5145 mid_tick_observations += 1;
5146 assert!(mid_tick_observations < 1_000);
5147 }
5148
5149 assert!(
5150 mid_tick_observations > 2,
5151 "expected the run to span several ticks, saw {mid_tick_observations}"
5152 );
5153 assert_eq!(live_ids(&manager), before);
5154 assert!(manager.get(survivor).is_some());
5155 assert!(manager.get(tombstoned).is_none());
5156 assert_eq!(manager.get_by_kind("table").len(), before.len());
5157 }
5158
5159 #[test]
5160 fn get_by_kind_reaches_sealed_segments() {
5161 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5162 let ids = seal_bulk_segment(&manager, 5);
5163 assert!(manager.delete(ids[0]).expect("delete"));
5164
5165 assert_eq!(manager.get_by_kind("table").len(), 4);
5169 assert_eq!(manager.get_by_kind("vector").len(), 0);
5170 }
5171
5172 #[test]
5173 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
5174 const PER_TICK: usize = 4;
5175 let manager = consolidating_manager(PER_TICK);
5176 let ids = seal_bulk_segment(&manager, 60);
5177 for id in ids.iter().take(30) {
5178 assert!(manager.delete(*id).expect("delete"));
5179 }
5180
5181 manager.run_maintenance().expect("maintenance");
5183 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
5184
5185 let mut previous = PER_TICK as u64;
5188 let mut ticks = 1;
5189 while let Some(copied) = manager.consolidation_progress() {
5190 assert!(
5191 copied - previous <= PER_TICK as u64,
5192 "tick copied {} entities, bound is {PER_TICK}",
5193 copied - previous
5194 );
5195 previous = copied;
5196 manager.run_maintenance().expect("maintenance");
5197 ticks += 1;
5198 assert!(ticks < 100);
5199 }
5200
5201 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
5204 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5205 assert_eq!(live_ids(&manager).len(), 30);
5206 }
5207
5208 #[test]
5209 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
5210 let manager = SegmentManager::with_config(
5211 "test",
5212 ManagerConfig {
5213 max_sealed_segments: 3,
5214 ..Default::default()
5215 },
5216 );
5217 for _ in 0..4 {
5218 seal_bulk_segment(&manager, 5);
5219 }
5220 assert_eq!(manager.sealed.read().len(), 4);
5221
5222 let before = live_ids(&manager);
5223 drain_maintenance(&manager);
5224
5225 assert_eq!(manager.sealed.read().len(), 1);
5228 assert_eq!(manager.stats().consolidation.segments_merged, 4);
5229 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
5230 assert_eq!(live_ids(&manager), before);
5231 }
5232
5233 #[test]
5234 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
5235 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5236 seal_bulk_segment(&manager, 10);
5237 seal_bulk_segment(&manager, 10);
5238
5239 for _ in 0..5 {
5240 manager.run_maintenance().expect("maintenance");
5241 }
5242
5243 assert_eq!(manager.stats().consolidation.runs_started, 0);
5244 assert_eq!(manager.sealed.read().len(), 2);
5245 }
5246
5247 #[test]
5248 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
5249 use std::sync::atomic::AtomicBool;
5250
5251 let manager = Arc::new(consolidating_manager(2));
5252 let ids = seal_bulk_segment(&manager, 40);
5253 seal_bulk_segment(&manager, 40);
5254 for id in ids.iter().take(20) {
5255 assert!(manager.delete(*id).expect("delete"));
5256 }
5257
5258 let expected = live_ids(&manager);
5259 let stop = Arc::new(AtomicBool::new(false));
5260
5261 let readers: Vec<_> = (0..4)
5262 .map(|_| {
5263 let manager = Arc::clone(&manager);
5264 let stop = Arc::clone(&stop);
5265 let expected = expected.clone();
5266 std::thread::spawn(move || {
5267 let mut scans = 0u64;
5268 loop {
5271 let mut seen: Vec<u64> = manager
5272 .query_all(|_| true)
5273 .into_iter()
5274 .map(|entity| entity.id.raw())
5275 .collect();
5276 seen.sort_unstable();
5277 assert_eq!(seen, expected, "reader saw a torn sealed set");
5280 scans += 1;
5281 if stop.load(Ordering::Relaxed) {
5282 break;
5283 }
5284 }
5285 scans
5286 })
5287 })
5288 .collect();
5289
5290 drain_maintenance(&manager);
5291 stop.store(true, Ordering::Relaxed);
5292
5293 for reader in readers {
5294 let scans = reader.join().expect("reader thread");
5295 assert!(scans > 0, "reader never got to scan");
5296 }
5297
5298 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5299 assert_eq!(live_ids(&manager), expected);
5300 }
5301
5302 #[test]
5303 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
5304 let manager = consolidating_manager(4);
5305 let ids = seal_bulk_segment(&manager, 20);
5306 for id in ids.iter().take(6) {
5307 assert!(manager.delete(*id).expect("delete"));
5308 }
5309
5310 manager.run_maintenance().expect("maintenance");
5312 assert_eq!(manager.consolidation_progress(), Some(4));
5313 let doomed = ids[6];
5314 assert!(manager.delete(doomed).expect("delete"));
5315
5316 let expected = live_ids(&manager);
5317 drain_maintenance(&manager);
5318
5319 assert!(manager.get(doomed).is_none(), "deleted row came back");
5320 assert_eq!(live_ids(&manager), expected);
5321 }
5322
5323 #[test]
5324 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
5325 let manager = consolidating_manager(3);
5326 let ids = seal_bulk_segment(&manager, 30);
5327 for id in ids.iter().take(12) {
5328 assert!(manager.delete(*id).expect("delete"));
5329 }
5330
5331 let before = live_ids(&manager);
5332 let sealed_before: Vec<SegmentId> = manager
5333 .sealed
5334 .read()
5335 .iter()
5336 .map(|segment| segment.read().id())
5337 .collect();
5338
5339 manager.run_maintenance().expect("maintenance");
5340 manager.run_maintenance().expect("maintenance");
5341 assert!(manager.consolidation.read().is_some());
5342
5343 *manager.consolidation.write() = None;
5347
5348 assert_eq!(live_ids(&manager), before);
5349 let sealed_after: Vec<SegmentId> = manager
5350 .sealed
5351 .read()
5352 .iter()
5353 .map(|segment| segment.read().id())
5354 .collect();
5355 assert_eq!(sealed_after, sealed_before);
5356 assert_eq!(manager.stats().consolidation.runs_completed, 0);
5357
5358 drain_maintenance(&manager);
5359 assert_eq!(live_ids(&manager), before);
5360 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5361 }
5362
5363 #[test]
5364 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
5365 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5366
5367 let entities: Vec<UnifiedEntity> = (0..10)
5371 .map(|i| {
5372 UnifiedEntity::table_row(
5373 manager.next_entity_id(),
5374 "users",
5375 i as u64,
5376 vec![Value::Integer(i as i64)],
5377 )
5378 })
5379 .collect();
5380 let ids = manager.bulk_insert(entities).expect("bulk insert");
5381
5382 let mut meta = Metadata::new();
5383 meta.set(
5384 "os",
5385 super::super::metadata::MetadataValue::String("linux".to_string()),
5386 );
5387 manager.set_metadata(ids[9], meta).expect("set metadata");
5388 manager.force_seal().expect("seal");
5389
5390 assert!(manager.delete(ids[0]).expect("delete"));
5391 assert!(manager.delete(ids[1]).expect("delete"));
5392 assert!(manager.delete(ids[2]).expect("delete"));
5393
5394 let before: Vec<UnifiedEntity> = {
5395 let mut entities = manager.query_all(|_| true);
5396 entities.sort_by_key(|entity| entity.id.raw());
5397 entities
5398 };
5399
5400 drain_maintenance(&manager);
5401
5402 let after: Vec<UnifiedEntity> = {
5403 let mut entities = manager.query_all(|_| true);
5404 entities.sort_by_key(|entity| entity.id.raw());
5405 entities
5406 };
5407
5408 let columns = |entity: &UnifiedEntity| match &entity.data {
5409 super::super::entity::EntityData::Row(row) => row.columns.clone(),
5410 _ => Vec::new(),
5411 };
5412
5413 assert_eq!(before.len(), after.len());
5414 for (old, new) in before.iter().zip(after.iter()) {
5415 assert_eq!(old.id, new.id);
5416 assert_eq!(columns(old), columns(new));
5417 assert_eq!(
5418 old.sequence_id, new.sequence_id,
5419 "merge must not re-issue sequence ids"
5420 );
5421 }
5422 assert!(manager
5423 .get_metadata(ids[9])
5424 .expect("metadata survives the merge")
5425 .has("os"));
5426 }
5427
5428 #[test]
5429 #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
5430 fn test_lifecycle_events() {
5431 let manager = SegmentManager::new("test");
5432
5433 manager
5434 .insert(UnifiedEntity::vector(
5435 manager.next_entity_id(),
5436 "v",
5437 vec![0.1],
5438 ))
5439 .unwrap();
5440
5441 let events = manager.drain_events();
5442
5443 assert!(events
5445 .iter()
5446 .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
5447 assert!(events
5448 .iter()
5449 .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
5450 }
5451}