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 fn ordered_fingerprint_scan_matches_legacy_sort_without_cloning_rows() {
3539 let manager = SegmentManager::new("fingerprint");
3540 let row = |id, label: &str| {
3541 UnifiedEntity::table_row(
3542 EntityId::new(id),
3543 "fingerprint",
3544 id,
3545 vec![Value::text(label)],
3546 )
3547 };
3548 manager
3549 .bulk_insert(vec![row(8, "eight"), row(2, "sealed-two"), row(5, "five")])
3550 .expect("sealed rows");
3551 manager.force_seal().expect("seal");
3552 manager
3553 .bulk_insert(vec![row(7, "seven"), row(2, "growing-two"), row(1, "one")])
3554 .expect("growing rows");
3555
3556 let mut expected = manager.query_all(|_| true);
3557 expected.sort_by_key(|entity| entity.id.raw());
3558 let expected: Vec<_> = expected
3559 .iter()
3560 .map(|entity| format!("{entity:?}"))
3561 .collect();
3562 let mut stored_addresses = HashSet::new();
3563 manager.for_each_entity(|entity| {
3564 stored_addresses.insert(entity as *const UnifiedEntity as usize);
3565 true
3566 });
3567
3568 let mut actual = Vec::new();
3569 manager.for_each_entity_in_id_order(|entity| {
3570 assert!(stored_addresses.contains(&(entity as *const UnifiedEntity as usize)));
3571 assert!(
3572 manager.growing.try_write().is_some(),
3573 "topology guard retained"
3574 );
3575 assert!(
3576 manager.sealed.try_write().is_some(),
3577 "topology guard retained"
3578 );
3579 actual.push(format!("{entity:?}"));
3580 true
3581 });
3582 assert_eq!(
3583 actual, expected,
3584 "equal IDs retain their original segment order"
3585 );
3586 }
3587
3588 #[test]
3589 fn ordered_fingerprint_scan_stops_early_and_releases_sources() {
3590 let manager = SegmentManager::new("fingerprint");
3591 for id in [3, 1, 2] {
3592 manager
3593 .insert(UnifiedEntity::table_row(
3594 EntityId::new(id),
3595 "fingerprint",
3596 id,
3597 vec![Value::Integer(id as i64)],
3598 ))
3599 .expect("row");
3600 }
3601 let source = manager.growing.read().as_ref().expect("source").clone();
3602 let handles = Arc::strong_count(&source);
3603 let mut seen = Vec::new();
3604 manager.for_each_entity_in_id_order(|entity| {
3605 seen.push(entity.id.raw());
3606 false
3607 });
3608 assert_eq!(seen, vec![1]);
3609 assert_eq!(Arc::strong_count(&source), handles);
3610 assert!(
3611 source.try_write().is_some(),
3612 "segment guard retained after early stop"
3613 );
3614 }
3615
3616 #[test]
3617 fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
3618 let manager = SegmentManager::with_config(
3619 "vectors",
3620 ManagerConfig {
3621 max_sealed_segments: 1,
3622 consolidation_entities_per_tick: 10_000,
3623 ..Default::default()
3624 },
3625 );
3626 for part in 0..3 {
3627 manager
3628 .bulk_insert(
3629 (1..=300)
3630 .map(|offset| {
3631 UnifiedEntity::vector(
3632 EntityId::new(part * 300 + offset),
3633 "vectors",
3634 vec![1.0, 0.0],
3635 )
3636 })
3637 .collect(),
3638 )
3639 .expect("batch");
3640 if part < 2 {
3641 manager.force_seal().expect("seal");
3642 }
3643 }
3644 let original_growing = manager.growing.read().as_ref().expect("growing").clone();
3645 let mut changed = false;
3646 let mut seen = Vec::new();
3647 assert!(manager.scan_batches(
3648 None,
3649 || true,
3650 |batch| {
3651 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3652 assert!(
3653 manager.growing.try_write().is_some(),
3654 "consumer holds no topology guard"
3655 );
3656 assert!(
3657 original_growing.try_write().is_some(),
3658 "consumer holds no segment guard"
3659 );
3660 for entity in batch {
3661 assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
3662 seen.push(entity.id.raw());
3663 }
3664 if !changed {
3665 changed = true;
3666 manager.force_seal().expect("seal during scan");
3667 manager
3668 .insert(UnifiedEntity::vector(
3669 EntityId::new(5000),
3670 "vectors",
3671 vec![1.0, 0.0],
3672 ))
3673 .expect("new write");
3674 drain_maintenance(&manager);
3675 assert!(manager.stats().consolidation.runs_completed > 0);
3676 assert!(manager
3677 .delete(EntityId::new(1))
3678 .expect("delete after source retirement"));
3679 assert!(manager.get(EntityId::new(1)).is_none());
3680 }
3681 true
3682 }
3683 ));
3684 seen.sort_unstable();
3685 assert_eq!(
3686 seen,
3687 (2..=900).collect::<Vec<_>>(),
3688 "no duplicate, missing or late candidates"
3689 );
3690 let mut after = Vec::new();
3691 assert!(manager.scan_batches(
3692 None,
3693 || true,
3694 |batch| {
3695 after.extend(batch.iter().map(|entity| entity.id.raw()));
3696 true
3697 }
3698 ));
3699 after.sort_unstable();
3700 let mut expected: Vec<_> = (2..=900).collect();
3701 expected.push(5000);
3702 assert_eq!(after, expected, "adopted entities have scan positions");
3703 }
3704
3705 #[test]
3706 fn bounded_scan_applies_snapshot_to_relocated_versions() {
3707 for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
3708 let manager = SegmentManager::with_config(
3709 "vectors",
3710 ManagerConfig {
3711 max_sealed_segments: 1,
3712 ..Default::default()
3713 },
3714 );
3715 for ids in [[1, 2], [3, 4]] {
3716 let entities = ids.map(|id| {
3717 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3718 entity.set_xmin(3);
3719 entity
3720 });
3721 if bulk {
3722 manager
3723 .bulk_insert(entities.into())
3724 .expect("initial versions");
3725 } else {
3726 for entity in entities {
3727 manager.insert(entity).expect("initial version");
3728 }
3729 }
3730 manager.force_seal().expect("seal source");
3731 }
3732 let snapshot = SnapshotContext {
3733 snapshot: Snapshot {
3734 xid: 10,
3735 in_progress: HashSet::new(),
3736 },
3737 manager: Arc::new(SnapshotManager::new()),
3738 own_xids: HashSet::new(),
3739 requires_index_fallback: false,
3740 serializable_reader: None,
3741 };
3742 let mut changed = false;
3743 let mut seen = Vec::new();
3744 assert!(manager.scan_batches(
3745 Some(&snapshot),
3746 || true,
3747 |batch| {
3748 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3749 if !changed {
3750 changed = true;
3751 if consolidate {
3752 drain_maintenance(&manager);
3753 assert!(manager.stats().consolidation.runs_completed > 0);
3754 }
3755 for (id, xmax) in [(3, 8), (4, 12)] {
3756 let mut entity =
3757 manager.get(EntityId::new(id)).expect("retained version");
3758 entity.set_xmax(xmax);
3759 manager
3760 .update(entity)
3761 .expect("update after source retirement");
3762 }
3763 }
3764 true
3765 }
3766 ));
3767 seen.sort_unstable();
3768 assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
3769 let mut again = Vec::new();
3770 assert!(manager.scan_batches(
3771 Some(&snapshot),
3772 || true,
3773 |batch| {
3774 again.extend(batch.iter().map(|entity| entity.id.raw()));
3775 true
3776 }
3777 ));
3778 again.sort_unstable();
3779 assert_eq!(
3780 again,
3781 vec![1, 2, 4],
3782 "old slots do not duplicate relocated versions"
3783 );
3784 let mut current = Vec::new();
3785 assert!(manager.scan_batches(
3786 None,
3787 || true,
3788 |batch| {
3789 current.extend(batch.iter().map(|entity| entity.id.raw()));
3790 true
3791 }
3792 ));
3793 current.sort_unstable();
3794 assert_eq!(
3795 current,
3796 vec![1, 2],
3797 "current view excludes both deleted versions"
3798 );
3799 }
3800 }
3801
3802 #[test]
3803 fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
3804 let manager = SegmentManager::new("mixed");
3805 let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
3806 manager
3807 .bulk_insert((1..=300).map(vector).collect())
3808 .expect("flat");
3809 manager.insert(vector(500)).expect("hash insert");
3810 manager
3811 .bulk_insert(vec![vector(502), vector(505)])
3812 .expect("gapped bulk");
3813 manager.delete(EntityId::new(2)).expect("flat delete");
3814 manager.delete(EntityId::new(505)).expect("hash delete");
3815 let mut changed = false;
3816 let mut seen = Vec::new();
3817 assert!(manager.scan_batches(
3818 None,
3819 || true,
3820 |batch| {
3821 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3822 if !changed {
3823 changed = true;
3824 manager
3825 .insert(vector(700))
3826 .expect("append after cursor boundary");
3827 manager
3828 .update(UnifiedEntity::graph_node(
3829 EntityId::new(500),
3830 "mixed",
3831 "Node",
3832 HashMap::new(),
3833 ))
3834 .expect("structural update");
3835 }
3836 true
3837 }
3838 ));
3839 seen.sort_unstable();
3840 let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
3841 expected.extend([500, 502]);
3842 assert_eq!(seen, expected);
3843 }
3844
3845 #[test]
3846 fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
3847 let manager = SegmentManager::new("vectors");
3848 let mut entities = Vec::new();
3849 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
3850 let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
3851 entity.set_xmin(xmin);
3852 entity.set_xmax(xmax);
3853 entities.push(entity);
3854 }
3855 manager.bulk_insert(entities).expect("versions");
3856 manager.force_seal().expect("seal");
3857 let snapshot = SnapshotContext {
3858 snapshot: Snapshot {
3859 xid: 10,
3860 in_progress: HashSet::from([4]),
3861 },
3862 manager: Arc::new(SnapshotManager::new()),
3863 own_xids: HashSet::new(),
3864 requires_index_fallback: false,
3865 serializable_reader: None,
3866 };
3867 let mut work = 0;
3868 let mut callbacks = 0;
3869 assert!(!manager.scan_batches(
3870 Some(&snapshot),
3871 || {
3872 work += 1;
3873 work <= 2
3874 },
3875 |_| {
3876 callbacks += 1;
3877 true
3878 }
3879 ));
3880 assert_eq!(callbacks, 0, "a stopped batch is not consumed");
3881 for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
3882 let mut seen = Vec::new();
3883 assert!(manager.scan_batches(
3884 view,
3885 || true,
3886 |batch| {
3887 seen.extend(batch.iter().map(|entity| entity.id.raw()));
3888 true
3889 }
3890 ));
3891 seen.sort_unstable();
3892 assert_eq!(seen, expected);
3893 }
3894 }
3895
3896 #[test]
3897 fn test_manager_basic() {
3898 let manager = SegmentManager::new("test_collection");
3899
3900 let entity = UnifiedEntity::table_row(
3901 manager.next_entity_id(),
3902 "users",
3903 1,
3904 vec![Value::text("Alice".to_string())],
3905 );
3906
3907 let id = manager.insert(entity).unwrap();
3908 assert!(manager.get(id).is_some());
3909 assert_eq!(manager.count(), 1);
3910 }
3911
3912 #[test]
3913 fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
3914 let manager = Arc::new(SegmentManager::new("mixed"));
3915 manager
3916 .bulk_insert(
3917 (1..=128)
3918 .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
3919 .collect(),
3920 )
3921 .expect("vector-only segment");
3922 manager.force_seal().expect("seal vectors");
3923 let mut old =
3924 UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
3925 old.set_xmin(3);
3926 old.set_xmax(12);
3927 let mut new =
3928 UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
3929 new.set_xmin(12);
3930 let edge = UnifiedEntity::graph_edge(
3931 EntityId::new(202),
3932 "mixed",
3933 "a",
3934 "b",
3935 1.0,
3936 Default::default(),
3937 );
3938 manager
3939 .bulk_insert(vec![old, new, edge])
3940 .expect("physical graph versions");
3941 manager.force_seal().expect("seal graph");
3942 manager
3943 .insert(UnifiedEntity::graph_node(
3944 EntityId::new(203),
3945 "mixed",
3946 "person",
3947 Default::default(),
3948 ))
3949 .expect("growing node");
3950 let snapshot = SnapshotContext {
3951 snapshot: Snapshot {
3952 xid: 10,
3953 in_progress: HashSet::new(),
3954 },
3955 manager: Arc::new(SnapshotManager::new()),
3956 own_xids: HashSet::new(),
3957 requires_index_fallback: false,
3958 serializable_reader: None,
3959 };
3960 std::thread::spawn(move || {
3961 for view in [None, Some(&snapshot)] {
3962 for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
3963 let expected: HashSet<_> = manager
3964 .scan(view, |item| kind.matches(&item.kind))
3965 .into_iter()
3966 .map(|item| item.id)
3967 .collect();
3968 let actual: HashSet<_> = manager
3969 .scan_graph_kind(view, kind)
3970 .into_iter()
3971 .map(|item| item.id)
3972 .collect();
3973 assert_eq!(actual, expected);
3974 let mut batched = HashSet::new();
3975 assert!(manager.scan_graph_batches(
3976 view,
3977 kind,
3978 || true,
3979 |batch| {
3980 batched.extend(batch.into_iter().map(|item| item.id));
3981 true
3982 }
3983 ));
3984 assert_eq!(batched, expected);
3985 let mut visited = 0;
3986 let mut candidates = HashSet::new();
3987 manager.scan_graph_candidates_for_each(
3988 view,
3989 kind,
3990 || true,
3991 |item| {
3992 visited += 1;
3993 if kind.matches(&item.kind) {
3994 candidates.insert(item.id);
3995 }
3996 true
3997 },
3998 );
3999 assert_eq!(candidates, expected);
4000 assert_eq!(
4001 visited,
4002 if matches!(kind, GraphEntityKind::Node) {
4003 3
4004 } else {
4005 2
4006 }
4007 );
4008 }
4009 }
4010 let ids: HashSet<_> = manager
4011 .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
4012 .into_iter()
4013 .map(|item| item.id.raw())
4014 .collect();
4015 assert_eq!(ids, HashSet::from([200, 203]));
4016 let ids: HashSet<_> = manager
4017 .scan_graph_kind(None, GraphEntityKind::Node)
4018 .into_iter()
4019 .map(|item| item.id.raw())
4020 .collect();
4021 assert_eq!(ids, HashSet::from([201, 203]));
4022 let mut visits = 0;
4023 manager.scan_graph_candidates_for_each(
4024 None,
4025 GraphEntityKind::Node,
4026 || true,
4027 |_| {
4028 visits += 1;
4029 false
4030 },
4031 );
4032 assert_eq!(
4033 visits, 1,
4034 "early termination must also stop subsequent segments"
4035 );
4036 })
4037 .join()
4038 .expect("snapshot worker");
4039 }
4040
4041 #[test]
4042 fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
4043 let manager = Arc::new(SegmentManager::new("accounts"));
4044 let cases = [
4045 ("legacy row", 0, 0, true),
4046 ("committed creator", 3, 0, true),
4047 ("future creator", 11, 0, false),
4048 ("in-progress creator", 4, 0, false),
4049 ("committed deleter", 3, 7, false),
4050 ("future deleter", 3, 11, true),
4051 ("in-progress deleter", 3, 4, true),
4052 ];
4053
4054 for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
4055 let row_id = index as u64 + 1;
4056 let mut entity = UnifiedEntity::table_row(
4057 EntityId::new(row_id),
4058 "accounts",
4059 row_id,
4060 vec![Value::Integer(row_id as i64)],
4061 );
4062 entity.set_xmin(*xmin);
4063 entity.set_xmax(*xmax);
4064 manager.insert(entity).unwrap();
4065 }
4066
4067 let snapshot = SnapshotContext {
4068 snapshot: Snapshot {
4069 xid: 10,
4070 in_progress: HashSet::from([4]),
4071 },
4072 manager: Arc::new(SnapshotManager::new()),
4073 own_xids: HashSet::new(),
4074 requires_index_fallback: false,
4075 serializable_reader: None,
4076 };
4077
4078 let raw_visible_row_ids: HashSet<u64> = manager
4079 .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
4080 .into_iter()
4081 .filter_map(|entity| match entity.kind {
4082 EntityKind::TableRow { row_id, .. } => Some(row_id),
4083 _ => None,
4084 })
4085 .collect();
4086
4087 let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
4088 .join()
4089 .unwrap();
4090 let visible_row_ids: HashSet<u64> = scanned
4091 .into_iter()
4092 .filter_map(|entity| match entity.kind {
4093 EntityKind::TableRow { row_id, .. } => Some(row_id),
4094 _ => None,
4095 })
4096 .collect();
4097 assert_eq!(visible_row_ids, raw_visible_row_ids);
4098
4099 for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
4100 let row_id = index as u64 + 1;
4101 assert_eq!(
4102 visible_row_ids.contains(&row_id),
4103 *expected_visible,
4104 "{name} visibility mismatch"
4105 );
4106 }
4107 }
4108
4109 #[test]
4110 fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
4111 let manager = SegmentManager::new("accounts");
4115 let mut live =
4116 UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
4117 live.set_xmin(2);
4118 live.set_xmax(0);
4119 manager.insert(live).unwrap();
4120
4121 let mut superseded =
4122 UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
4123 superseded.set_xmin(2);
4124 superseded.set_xmax(9);
4125 manager.insert(superseded).unwrap();
4126
4127 let rows: Vec<u64> = manager
4128 .scan(None, |_| true)
4129 .into_iter()
4130 .filter_map(|entity| match entity.kind {
4131 EntityKind::TableRow { row_id, .. } => Some(row_id),
4132 _ => None,
4133 })
4134 .collect();
4135 assert_eq!(
4136 rows,
4137 vec![1],
4138 "frameless scan must hide superseded versions"
4139 );
4140 }
4141
4142 #[test]
4143 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
4144 let manager = SegmentManager::new("test_collection");
4145
4146 let entity = UnifiedEntity::table_row(
4147 manager.next_entity_id(),
4148 "users",
4149 1,
4150 vec![Value::text("Alice".to_string())],
4151 );
4152
4153 let id = manager.insert(entity).unwrap();
4154
4155 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
4156 assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
4157 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
4158
4159 assert!(manager.delete(id).unwrap());
4160 assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
4161 }
4162
4163 #[test]
4164 fn zoned_scan_reports_pruned_segments_and_preserves_results() {
4165 let config = ManagerConfig {
4166 segment_config: SegmentConfig {
4167 max_entities: 2,
4168 ..Default::default()
4169 },
4170 ..Default::default()
4171 };
4172 let manager = SegmentManager::with_config("events", config);
4173
4174 for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
4175 let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
4176 let entity = UnifiedEntity::new(
4177 manager.next_entity_id(),
4178 EntityKind::TableRow {
4179 table: "events".into(),
4180 row_id: (idx + 1) as u64,
4181 },
4182 EntityData::Row(row),
4183 );
4184 manager.insert(entity).unwrap();
4185 }
4186
4187 let full = manager.query_all(|entity| match &entity.data {
4188 EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
4189 _ => false,
4190 });
4191 let probe = Value::Integer(100);
4192 let (pruned, stats) =
4193 manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
4194 match &entity.data {
4195 EntityData::Row(row) => {
4196 matches!(row.get_field("ts"), Some(Value::Integer(100)))
4197 }
4198 _ => false,
4199 }
4200 });
4201
4202 assert_eq!(pruned.len(), full.len());
4203 assert_eq!(pruned[0].id, full[0].id);
4204 assert_eq!(stats.segments_total, 3);
4205 assert_eq!(stats.segments_scanned, 2);
4206 assert_eq!(stats.segments_pruned, 1);
4207 }
4208
4209 #[test]
4210 fn projected_zoned_scan_preserves_segment_order_and_pruning() {
4211 let manager = SegmentManager::with_config(
4212 "events",
4213 ManagerConfig {
4214 segment_config: SegmentConfig {
4215 max_entities: 2,
4216 ..Default::default()
4217 },
4218 ..Default::default()
4219 },
4220 );
4221 for (index, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
4222 let mut entity = UnifiedEntity::table_row(
4223 EntityId::new(index as u64 + 1),
4224 "events",
4225 index as u64 + 1,
4226 vec![],
4227 );
4228 entity.data = EntityData::Row(RowData::with_names(
4229 vec![Value::Integer(value)],
4230 vec!["ts".into()],
4231 ));
4232 manager.insert(entity).expect("test row");
4233 }
4234 assert_eq!(manager.stats().sealed_count, 2);
4235 let project = |entity: &UnifiedEntity| match entity.data.as_row()?.get_field("ts")? {
4236 Value::Integer(value) => Some(*value),
4237 _ => None,
4238 };
4239 let (all, stats) = manager.scan_project_zoned_with_stats(None, &[], project);
4240 let (legacy, _) = manager.scan_zoned_with_stats(None, &[], |_| true);
4241 assert_eq!(all, legacy.iter().filter_map(project).collect::<Vec<_>>());
4242 assert_eq!(all[0], 200, "growing segment first");
4243 assert!(all[1..3].iter().all(|value| [10, 20].contains(value)));
4244 assert!(all[3..5].iter().all(|value| [100, 110].contains(value)));
4245 assert_eq!(stats.segments_total, 3);
4246 assert_eq!(stats.segments_scanned, 3);
4247 assert_eq!(stats.segments_pruned, 0);
4248 let probe = Value::Integer(100);
4249 let (pruned, stats) = manager.scan_project_zoned_with_stats(
4250 None,
4251 &[("ts", ZoneColPred::Gte(&probe))],
4252 |entity| project(entity).filter(|value| *value >= 100),
4253 );
4254 assert_eq!(
4255 pruned,
4256 all.into_iter()
4257 .filter(|value| *value >= 100)
4258 .collect::<Vec<_>>()
4259 );
4260 assert_eq!(stats.segments_total, 3);
4261 assert_eq!(stats.segments_scanned, 2);
4262 assert_eq!(stats.segments_pruned, 1);
4263 let (empty, _) = manager.scan_project_zoned_with_stats(None, &[], |_| None::<u64>);
4264 assert!(empty.is_empty(), "projection may reject every candidate");
4265
4266 let failure = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4267 manager.scan_project_zoned_with_stats(None, &[], |entity| {
4268 let value = project(entity)?;
4269 assert_ne!(value, 10, "injected projection worker failure");
4270 Some(value)
4271 })
4272 }));
4273 assert!(
4274 failure.is_err(),
4275 "a failed worker must not return a successful partial scan"
4276 );
4277 }
4278
4279 #[test]
4280 fn projected_zoned_scan_checks_visibility_before_projection_on_workers() {
4281 let manager = SegmentManager::with_config(
4282 "accounts",
4283 ManagerConfig {
4284 segment_config: SegmentConfig {
4285 max_entities: 2,
4286 ..Default::default()
4287 },
4288 ..Default::default()
4289 },
4290 );
4291 for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4), (5, 3, 0)] {
4292 let mut entity = UnifiedEntity::table_row(EntityId::new(id), "accounts", id, vec![]);
4293 entity.set_xmin(xmin);
4294 entity.set_xmax(xmax);
4295 manager.insert(entity).expect("versioned test row");
4296 }
4297 let mut hidden = UnifiedEntity::table_row(EntityId::new(6), "accounts", 6, vec![]);
4298 hidden.data = EntityData::Row(RowData::with_names(
4299 vec![Value::text("pending")],
4300 vec!["__moderation_status".into()],
4301 ));
4302 manager.insert(hidden).expect("moderated row");
4303 assert_eq!(manager.stats().sealed_count, 2);
4304 let snapshot = SnapshotContext {
4305 snapshot: Snapshot {
4306 xid: 10,
4307 in_progress: HashSet::from([4]),
4308 },
4309 manager: Arc::new(SnapshotManager::new()),
4310 own_xids: HashSet::new(),
4311 requires_index_fallback: false,
4312 serializable_reader: None,
4313 };
4314 std::thread::spawn(move || {
4315 for (view, expected) in [(Some(&snapshot), vec![5, 1, 3, 4]), (None, vec![5, 2, 3])] {
4316 let called = std::sync::Mutex::new(Vec::new());
4317 let (rows, stats) = manager.scan_project_zoned_with_stats(view, &[], |entity| {
4318 let id = entity.id.raw();
4319 assert!(
4320 expected.contains(&id),
4321 "invisible entity reached projection: {id}"
4322 );
4323 called.lock().expect("callback log").push(id);
4324 Some(id)
4325 });
4326 let (legacy, _) = manager.scan_zoned_with_stats(view, &[], |_| true);
4327 assert_eq!(
4328 rows,
4329 legacy
4330 .iter()
4331 .map(|entity| entity.id.raw())
4332 .collect::<Vec<_>>()
4333 );
4334 assert_eq!(rows[0], 5, "visible growing row first");
4335 let mut called = called.into_inner().expect("callback log");
4336 called.sort_unstable();
4337 let mut expected = expected;
4338 expected.sort_unstable();
4339 assert_eq!(called, expected, "each visible row projected exactly once");
4340 assert_eq!(stats.segments_total, 3);
4341 assert_eq!(stats.segments_scanned, 3);
4342 }
4343 })
4344 .join()
4345 .expect("projected scan worker");
4346 }
4347
4348 #[test]
4349 fn test_manager_auto_seal() {
4350 let config = ManagerConfig {
4351 segment_config: SegmentConfig {
4352 max_entities: 2,
4353 ..Default::default()
4354 },
4355 ..Default::default()
4356 };
4357
4358 let manager = SegmentManager::with_config("test", config);
4359
4360 manager
4362 .insert(UnifiedEntity::vector(
4363 manager.next_entity_id(),
4364 "v",
4365 vec![0.1],
4366 ))
4367 .unwrap();
4368
4369 manager
4371 .insert(UnifiedEntity::vector(
4372 manager.next_entity_id(),
4373 "v",
4374 vec![0.2],
4375 ))
4376 .unwrap();
4377
4378 manager
4380 .insert(UnifiedEntity::vector(
4381 manager.next_entity_id(),
4382 "v",
4383 vec![0.3],
4384 ))
4385 .unwrap();
4386
4387 let stats = manager.stats();
4388 assert_eq!(stats.total_entities, 3);
4389 }
4390
4391 #[test]
4392 fn test_manager_delete() {
4393 let manager = SegmentManager::new("test");
4394
4395 let id = manager
4396 .insert(UnifiedEntity::vector(
4397 manager.next_entity_id(),
4398 "v",
4399 vec![0.1],
4400 ))
4401 .unwrap();
4402
4403 assert!(manager.get(id).is_some());
4404 assert!(manager.delete(id).unwrap());
4405 assert!(manager.get(id).is_none());
4406 }
4407
4408 #[test]
4409 fn test_manager_metadata() {
4410 let manager = SegmentManager::new("test");
4411
4412 let id = manager
4413 .insert(UnifiedEntity::table_row(
4414 manager.next_entity_id(),
4415 "hosts",
4416 1,
4417 vec![Value::text("192.168.1.1".to_string())],
4418 ))
4419 .unwrap();
4420
4421 let mut meta = Metadata::new();
4422 meta.set(
4423 "os",
4424 super::super::metadata::MetadataValue::String("linux".to_string()),
4425 );
4426
4427 manager.set_metadata(id, meta).unwrap();
4428
4429 let retrieved = manager.get_metadata(id).unwrap();
4430 assert!(retrieved.has("os"));
4431 }
4432
4433 #[test]
4434 fn metadata_expiry_index_is_bounded_rekeyed_and_removed_with_entity() {
4435 let manager = SegmentManager::new("sessions");
4436 let mut first = UnifiedEntity::table_row(EntityId::new(1), "sessions", 1, vec![]);
4437 first.created_at = 10;
4438 let mut second = UnifiedEntity::table_row(EntityId::new(2), "sessions", 2, vec![]);
4439 second.created_at = 20;
4440 manager.insert(first).expect("first session");
4441 manager.insert(second).expect("second session");
4442
4443 let mut relative = Metadata::new();
4444 relative.set("_ttl_ms", MetadataValue::Int(5_000));
4445 manager
4446 .set_metadata(EntityId::new(1), relative)
4447 .expect("relative TTL");
4448 let mut absolute = Metadata::new();
4449 absolute.set("_expires_at", MetadataValue::Timestamp(30_000));
4450 manager
4451 .set_metadata(EntityId::new(2), absolute)
4452 .expect("absolute expiry");
4453
4454 assert_eq!(manager.scheduled_metadata_expiry_count(), 2);
4455 assert_eq!(
4456 manager.expired_metadata_ids(30_000, 1),
4457 vec![EntityId::new(1)],
4458 "the deadline index must stop at the requested batch"
4459 );
4460 assert_eq!(
4461 manager.expired_metadata_ids(30_000, 2),
4462 vec![EntityId::new(1), EntityId::new(2)]
4463 );
4464
4465 let mut postponed = Metadata::new();
4466 postponed.set("_expires_at", MetadataValue::Timestamp(60_000));
4467 manager
4468 .set_metadata(EntityId::new(1), postponed)
4469 .expect("postpone expiry");
4470 assert_eq!(
4471 manager.expired_metadata_ids(30_000, 2),
4472 vec![EntityId::new(2)],
4473 "metadata replacement must remove the stale deadline"
4474 );
4475
4476 assert!(manager.delete(EntityId::new(2)).expect("delete session"));
4477 assert!(manager
4478 .expired_metadata_ids(u64::MAX, 2)
4479 .contains(&EntityId::new(1)));
4480 assert!(!manager
4481 .expired_metadata_ids(u64::MAX, 2)
4482 .contains(&EntityId::new(2)));
4483 assert_eq!(manager.scheduled_metadata_expiry_count(), 1);
4484 }
4485
4486 #[test]
4487 fn test_manager_query_by_kind() {
4488 let manager = SegmentManager::new("test");
4489
4490 manager
4491 .insert(UnifiedEntity::table_row(
4492 manager.next_entity_id(),
4493 "hosts",
4494 1,
4495 vec![],
4496 ))
4497 .unwrap();
4498
4499 manager
4500 .insert(UnifiedEntity::vector(
4501 manager.next_entity_id(),
4502 "embeddings",
4503 vec![0.1],
4504 ))
4505 .unwrap();
4506
4507 manager
4508 .insert(UnifiedEntity::table_row(
4509 manager.next_entity_id(),
4510 "hosts",
4511 2,
4512 vec![],
4513 ))
4514 .unwrap();
4515
4516 let rows = manager.get_by_kind("table");
4517 assert_eq!(rows.len(), 2);
4518
4519 let vectors = manager.get_by_kind("vector");
4520 assert_eq!(vectors.len(), 1);
4521 }
4522
4523 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
4530 SegmentManager::with_config(
4531 "test",
4532 ManagerConfig {
4533 max_sealed_segments: 1_000,
4534 consolidation_entities_per_tick: entities_per_tick,
4535 ..Default::default()
4536 },
4537 )
4538 }
4539
4540 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
4544 let entities: Vec<UnifiedEntity> = (0..count)
4545 .map(|i| {
4546 UnifiedEntity::table_row(
4547 manager.next_entity_id(),
4548 "users",
4549 i as u64,
4550 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
4551 )
4552 })
4553 .collect();
4554 let ids = manager.bulk_insert(entities).expect("bulk insert");
4555 manager.force_seal().expect("seal");
4556 ids
4557 }
4558
4559 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
4560 let mut ids: Vec<u64> = manager
4561 .query_all(|_| true)
4562 .into_iter()
4563 .map(|entity| entity.id.raw())
4564 .collect();
4565 ids.sort_unstable();
4566 ids
4567 }
4568
4569 fn drain_maintenance(manager: &SegmentManager) -> usize {
4570 let mut ticks = 0;
4571 while manager.consolidation.read().is_some() || ticks == 0 {
4572 manager.run_maintenance().expect("maintenance");
4573 ticks += 1;
4574 assert!(ticks < 1_000, "consolidation failed to converge");
4575 }
4576 ticks
4577 }
4578
4579 #[test]
4580 fn graph_batches_bound_payloads_and_hydrate_current_entities_after_sealing() {
4581 let manager = SegmentManager::new("graph");
4582 let node =
4583 |id| UnifiedEntity::graph_node(EntityId::new(id), "graph", "node", HashMap::new());
4584 manager
4585 .bulk_insert((1..=1024).map(node).collect())
4586 .expect("nodes");
4587 let mut hidden = node(1025);
4588 hidden.set_xmax(5);
4589 manager.insert(hidden).expect("hidden version");
4590 manager
4591 .insert(UnifiedEntity::vector(EntityId::new(1026), "graph", vec![]))
4592 .expect("mixed kind");
4593 let mut work = 0;
4594 let mut delivered = 0;
4595 assert!(!manager.scan_graph_batches(
4596 None,
4597 GraphEntityKind::Node,
4598 || {
4599 work += 1;
4600 work < 100
4601 },
4602 |batch| {
4603 delivered += batch.len();
4604 true
4605 }
4606 ));
4607 assert_eq!(delivered, 0, "cancelled partial batch must not escape");
4608 let mut seen = Vec::new();
4609 assert!(manager.scan_graph_batches(None, GraphEntityKind::Node, || true, |batch| {
4610 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4611 assert!(manager.growing.try_write().is_some(), "no topology guard");
4612 assert!(manager.sealed.try_write().is_some(), "no sealed guard");
4613 if seen.is_empty() {
4614 manager.force_seal().expect("seal during delivery");
4615 manager.insert(node(2000)).expect("append beyond captured boundary");
4616 let mut updated = node(900);
4617 if let EntityKind::GraphNode(node) = &mut updated.kind {
4618 node.label = "updated".into();
4619 }
4620 manager.update(updated).expect("update unread candidate");
4621 }
4622 for entity in batch {
4623 if entity.id.raw() == 900 {
4624 assert!(matches!(&entity.kind, EntityKind::GraphNode(node) if node.label == "updated"));
4625 }
4626 seen.push(entity.id.raw());
4627 }
4628 true
4629 }));
4630 assert_eq!(seen, (1..=1024).collect::<Vec<_>>());
4631 }
4632
4633 #[test]
4634 fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
4635 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4636 let ids = manager
4637 .bulk_insert(
4638 (1..=1024)
4639 .map(|id| {
4640 UnifiedEntity::graph_edge(
4641 EntityId::new(id),
4642 "step",
4643 "2000",
4644 "3000",
4645 1.0,
4646 HashMap::new(),
4647 )
4648 })
4649 .collect(),
4650 )
4651 .expect("edges");
4652 manager.force_seal().expect("seal");
4653 for id in ids.iter().take(300) {
4654 manager.delete(*id).expect("fragment source");
4655 }
4656 let source = Arc::downgrade(&manager.sealed.read()[0]);
4657 let mut seen = Vec::new();
4658 let mut changed = false;
4659 let mut metadata_bytes = 0;
4660 assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
4661 || true,
4662 |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
4663 |batch| {
4664 assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
4665 assert!(manager.growing.try_write().is_some(), "no growing topology guard");
4666 assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
4667 assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
4668 if !changed {
4669 changed = true;
4670 drain_maintenance(&manager);
4671 assert!(manager.stats().consolidation.runs_completed > 0);
4672 assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
4673 let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
4674 if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
4675 manager.update(updated).expect("update authoritative merged copy");
4676 }
4677 for id in batch {
4678 let entity = manager.get(*id).expect("hydrate current manager");
4679 if id.raw() == 1024 {
4680 assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
4681 }
4682 seen.push(id.raw());
4683 }
4684 Ok(true)
4685 }).expect("cursor"));
4686 assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
4687 assert!(
4688 source.upgrade().is_none(),
4689 "finished cursor releases retired source"
4690 );
4691 assert!(
4692 metadata_bytes < 4096,
4693 "fixed ID buffer plus source handles: {metadata_bytes}"
4694 );
4695 }
4696
4697 #[test]
4698 fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
4699 let manager = SegmentManager::new("graph");
4700 manager
4701 .bulk_insert(
4702 (1..=1024)
4703 .map(|id| {
4704 UnifiedEntity::graph_edge(
4705 EntityId::new(id),
4706 "step",
4707 "2000",
4708 "3000",
4709 1.0,
4710 HashMap::new(),
4711 )
4712 })
4713 .collect(),
4714 )
4715 .expect("edges");
4716 let mut delivered = 0;
4717 let denied = manager.visit_graph_index_batches(
4718 GraphEntityKind::Edge,
4719 &[EntityId::new(2000)],
4720 || true,
4721 |_| Err("out of memory"),
4722 |ids| {
4723 delivered += ids.len();
4724 Ok(true)
4725 },
4726 );
4727 assert_eq!(denied, Err("out of memory"));
4728 assert_eq!(delivered, 0);
4729 let mut work = 0;
4730 assert!(!manager
4731 .visit_graph_index_batches(
4732 GraphEntityKind::Edge,
4733 &[EntityId::new(2000)],
4734 || {
4735 work += 1;
4736 work < 100
4737 },
4738 |_| Ok::<_, ()>(()),
4739 |ids| {
4740 delivered += ids.len();
4741 Ok(true)
4742 }
4743 )
4744 .expect("cancelled cursor"));
4745 assert_eq!(
4746 delivered, 0,
4747 "cancelled partial batch cannot reach consumer"
4748 );
4749 assert!(!manager
4750 .visit_graph_index_batches(
4751 GraphEntityKind::Edge,
4752 &[EntityId::new(2000)],
4753 || true,
4754 |_| Ok::<_, ()>(()),
4755 |ids| {
4756 delivered += ids.len();
4757 Ok(false)
4758 }
4759 )
4760 .expect("consumer stop"));
4761 assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
4762 }
4763
4764 #[test]
4765 fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
4766 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4767 let ids = seal_bulk_segment(&manager, 100);
4768 for id in ids.iter().take(40) {
4769 manager.delete(*id).expect("delete before scan");
4770 }
4771 let source_bytes = manager.resident_bytes();
4772 let source = Arc::downgrade(&manager.sealed.read()[0]);
4773 std::thread::scope(|scope| {
4774 let (ready, waiting) = std::sync::mpsc::channel();
4775 let (release_first, first_wait) = std::sync::mpsc::channel();
4776 let (release_last, last_wait) = std::sync::mpsc::channel();
4777 let reader = |resume: std::sync::mpsc::Receiver<()>,
4778 ready: std::sync::mpsc::Sender<()>| {
4779 assert!(manager.scan_batches(
4780 None,
4781 || true,
4782 |_| {
4783 ready.send(()).expect("signal captured cursor");
4784 let _ = resume.recv();
4785 true
4786 }
4787 ));
4788 };
4789 let ready_first = ready.clone();
4790 let first = scope.spawn(move || reader(first_wait, ready_first));
4791 let last = scope.spawn(move || reader(last_wait, ready));
4792 for _ in 0..2 {
4793 waiting
4794 .recv_timeout(std::time::Duration::from_secs(10))
4795 .expect("reader ready");
4796 }
4797 drain_maintenance(&manager);
4798 assert!(
4799 source.upgrade().is_some(),
4800 "both cursors pin the retired source"
4801 );
4802 let active_bytes = manager.sealed.read()[0].read().resident_bytes();
4803 let held = manager.resident_bytes();
4804 assert!(
4805 held >= active_bytes + source_bytes,
4806 "retired source remains resident"
4807 );
4808 release_first.send(()).expect("release first reader");
4809 first.join().expect("first reader");
4810 assert_eq!(
4811 manager.resident_bytes(),
4812 held,
4813 "one reader still pins the source"
4814 );
4815 release_last.send(()).expect("release last reader");
4816 last.join().expect("last reader");
4817 assert!(
4818 source.upgrade().is_none(),
4819 "accounting must not retain payloads"
4820 );
4821 assert_eq!(held - manager.resident_bytes(), source_bytes);
4822 assert_eq!(
4823 manager.stats().total_memory_bytes as u64,
4824 manager.resident_bytes()
4825 );
4826 });
4827 }
4828
4829 #[test]
4830 fn physical_cursor_tracks_individual_and_mixed_insertions() {
4831 for mode in 0..3 {
4832 let manager = SegmentManager::new("physical_rows");
4833 let row = |id| {
4834 UnifiedEntity::table_row(
4835 EntityId::new(id),
4836 "physical_rows",
4837 id,
4838 vec![Value::Integer(id as i64)],
4839 )
4840 };
4841 match mode {
4842 0 => {
4843 for id in [10, 30, 20] {
4844 manager.insert(row(id)).expect("individual row");
4845 }
4846 }
4847 1 => {
4848 manager.insert(row(10)).expect("individual prefix");
4849 manager
4850 .bulk_insert(vec![row(30), row(31)])
4851 .expect("batch suffix");
4852 }
4853 _ => {
4854 manager
4855 .bulk_insert(vec![row(30), row(31)])
4856 .expect("batch prefix");
4857 manager.insert(row(10)).expect("individual suffix");
4858 }
4859 }
4860 let mut cursor = manager.physical_scan_cursor();
4861 let mut ids = Vec::new();
4862 for _ in 0..32 {
4863 if cursor.finished() {
4864 break;
4865 }
4866 let mut work = 1;
4867 let batch = manager.physical_scan_batch(&mut cursor, &mut work);
4868 assert!(batch.len() <= 1);
4869 ids.extend(batch.into_iter().map(|id| id.raw()));
4870 }
4871 assert!(cursor.finished());
4872 ids.sort_unstable();
4873 assert_eq!(
4874 ids,
4875 if mode == 0 {
4876 vec![10, 20, 30]
4877 } else {
4878 vec![10, 30, 31]
4879 },
4880 "mode={mode}"
4881 );
4882 }
4883 }
4884
4885 fn growing_tombstones_fixture() -> SegmentManager {
4886 let manager = consolidating_manager(1);
4887 let ids = manager
4888 .bulk_insert(
4889 (0..64)
4890 .map(|id| {
4891 UnifiedEntity::table_row(
4892 manager.next_entity_id(),
4893 "users",
4894 id,
4895 vec![Value::Blob(vec![7; 2048])],
4896 )
4897 })
4898 .collect(),
4899 )
4900 .expect("growing rows");
4901 for id in ids.iter().take(48) {
4902 manager.delete(*id).expect("physical tombstone");
4903 }
4904 assert_eq!(manager.stats().sealed_count, 0);
4905 manager
4906 }
4907
4908 #[test]
4909 fn pressure_consolidation_reclaims_unsealed_tombstones() {
4910 let manager = growing_tombstones_fixture();
4911 let before = manager.resident_bytes();
4912 assert!(manager.reclaimable_bytes() > 0);
4913 assert!(manager.pressure_consolidation_tick());
4914 drain_maintenance(&manager);
4915 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
4916 assert!(manager.resident_bytes() < before);
4917 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 48);
4918 }
4919
4920 #[test]
4921 fn pressure_consolidation_preserves_pinned_growing_source_accounting() {
4922 let manager = growing_tombstones_fixture();
4923 let source = Arc::clone(manager.growing.read().as_ref().expect("growing source"));
4924 let weak = Arc::downgrade(&source);
4925 assert!(manager.pressure_consolidation_tick());
4926 drain_maintenance(&manager);
4927 assert_eq!(live_ids(&manager), (49..=64).collect::<Vec<_>>());
4928 let source_bytes = source.read().resident_bytes();
4929 let held = manager.resident_bytes();
4930 assert!(held >= source_bytes);
4931 drop(source);
4932 assert!(weak.upgrade().is_none());
4933 assert_eq!(held - manager.resident_bytes(), source_bytes);
4934 }
4935
4936 #[test]
4937 fn pressure_consolidation_keeps_clean_and_disabled_growing_sources_open() {
4938 let clean = consolidating_manager(1);
4939 clean
4940 .insert(UnifiedEntity::table_row(
4941 clean.next_entity_id(),
4942 "users",
4943 1,
4944 vec![Value::Integer(1)],
4945 ))
4946 .expect("clean row");
4947 assert!(!clean.pressure_consolidation_tick());
4948 assert_eq!(clean.stats().sealed_count, 0);
4949 let mut disabled = growing_tombstones_fixture();
4950 disabled.config.enable_consolidation = false;
4951 assert!(!disabled.pressure_consolidation_tick());
4952 assert_eq!(disabled.stats().sealed_count, 0);
4953 assert_eq!(live_ids(&disabled), (49..=64).collect::<Vec<_>>());
4954 }
4955
4956 #[test]
4957 fn retained_memory_includes_an_inflight_consolidation_copy() {
4958 let manager = consolidating_manager(1);
4959 let ids = seal_bulk_segment(&manager, 100);
4960 for id in ids.iter().take(40) {
4961 manager.delete(*id).expect("delete");
4962 }
4963 let before = manager.resident_bytes();
4964 assert!(manager.pressure_consolidation_tick());
4965 let copied = manager
4966 .consolidation
4967 .read()
4968 .as_ref()
4969 .expect("paced merge")
4970 .merged
4971 .resident_bytes();
4972 assert!(copied > 0);
4973 assert!(
4974 manager.resident_bytes() >= before + copied,
4975 "unpublished copy consumes memory"
4976 );
4977 drain_maintenance(&manager);
4978 assert!(
4979 manager.resident_bytes() < before,
4980 "completed merge releases the old storage"
4981 );
4982 }
4983
4984 #[test]
4985 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
4986 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4987 let ids = seal_bulk_segment(&manager, 100);
4988
4989 let bytes_before_delete = manager.resident_bytes();
4990
4991 for id in ids.iter().take(40) {
4993 assert!(manager.delete(*id).expect("delete"));
4994 }
4995
4996 let bytes_with_tombstones = manager.resident_bytes();
4999 assert!(bytes_with_tombstones > bytes_before_delete);
5000
5001 let survivors = live_ids(&manager);
5002 assert_eq!(survivors.len(), 60);
5003
5004 drain_maintenance(&manager);
5005
5006 let stats = manager.stats();
5007 assert_eq!(stats.consolidation.runs_started, 1);
5008 assert_eq!(stats.consolidation.runs_completed, 1);
5009 assert_eq!(stats.consolidation.segments_merged, 1);
5010 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
5011 assert!(stats.consolidation.bytes_reclaimed > 0);
5012
5013 let bytes_after = manager.resident_bytes();
5015 assert!(
5016 bytes_after <= bytes_before_delete,
5017 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
5018 );
5019 assert_eq!(
5020 bytes_with_tombstones - bytes_after,
5021 stats.consolidation.bytes_reclaimed
5022 );
5023 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
5024
5025 assert_eq!(live_ids(&manager), survivors);
5026 assert_eq!(manager.sealed.read().len(), 1);
5027 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
5028 }
5029
5030 #[test]
5031 fn query_results_are_identical_before_during_and_after_consolidation() {
5032 let manager = consolidating_manager(3);
5033 let ids = seal_bulk_segment(&manager, 20);
5034 seal_bulk_segment(&manager, 20);
5035 for id in ids.iter().take(10) {
5036 assert!(manager.delete(*id).expect("delete"));
5037 }
5038
5039 let before = live_ids(&manager);
5040 let survivor = *ids.last().expect("at least one row");
5041 let tombstoned = ids[0];
5042
5043 let mut mid_tick_observations = 0;
5046 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
5047 manager.run_maintenance().expect("maintenance");
5048 assert_eq!(live_ids(&manager), before);
5049 assert!(manager.get(survivor).is_some());
5050 assert!(manager.get(tombstoned).is_none());
5051 assert_eq!(manager.get_by_kind("table").len(), before.len());
5052 mid_tick_observations += 1;
5053 assert!(mid_tick_observations < 1_000);
5054 }
5055
5056 assert!(
5057 mid_tick_observations > 2,
5058 "expected the run to span several ticks, saw {mid_tick_observations}"
5059 );
5060 assert_eq!(live_ids(&manager), before);
5061 assert!(manager.get(survivor).is_some());
5062 assert!(manager.get(tombstoned).is_none());
5063 assert_eq!(manager.get_by_kind("table").len(), before.len());
5064 }
5065
5066 #[test]
5067 fn get_by_kind_reaches_sealed_segments() {
5068 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5069 let ids = seal_bulk_segment(&manager, 5);
5070 assert!(manager.delete(ids[0]).expect("delete"));
5071
5072 assert_eq!(manager.get_by_kind("table").len(), 4);
5076 assert_eq!(manager.get_by_kind("vector").len(), 0);
5077 }
5078
5079 #[test]
5080 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
5081 const PER_TICK: usize = 4;
5082 let manager = consolidating_manager(PER_TICK);
5083 let ids = seal_bulk_segment(&manager, 60);
5084 for id in ids.iter().take(30) {
5085 assert!(manager.delete(*id).expect("delete"));
5086 }
5087
5088 manager.run_maintenance().expect("maintenance");
5090 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
5091
5092 let mut previous = PER_TICK as u64;
5095 let mut ticks = 1;
5096 while let Some(copied) = manager.consolidation_progress() {
5097 assert!(
5098 copied - previous <= PER_TICK as u64,
5099 "tick copied {} entities, bound is {PER_TICK}",
5100 copied - previous
5101 );
5102 previous = copied;
5103 manager.run_maintenance().expect("maintenance");
5104 ticks += 1;
5105 assert!(ticks < 100);
5106 }
5107
5108 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
5111 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5112 assert_eq!(live_ids(&manager).len(), 30);
5113 }
5114
5115 #[test]
5116 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
5117 let manager = SegmentManager::with_config(
5118 "test",
5119 ManagerConfig {
5120 max_sealed_segments: 3,
5121 ..Default::default()
5122 },
5123 );
5124 for _ in 0..4 {
5125 seal_bulk_segment(&manager, 5);
5126 }
5127 assert_eq!(manager.sealed.read().len(), 4);
5128
5129 let before = live_ids(&manager);
5130 drain_maintenance(&manager);
5131
5132 assert_eq!(manager.sealed.read().len(), 1);
5135 assert_eq!(manager.stats().consolidation.segments_merged, 4);
5136 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
5137 assert_eq!(live_ids(&manager), before);
5138 }
5139
5140 #[test]
5141 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
5142 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5143 seal_bulk_segment(&manager, 10);
5144 seal_bulk_segment(&manager, 10);
5145
5146 for _ in 0..5 {
5147 manager.run_maintenance().expect("maintenance");
5148 }
5149
5150 assert_eq!(manager.stats().consolidation.runs_started, 0);
5151 assert_eq!(manager.sealed.read().len(), 2);
5152 }
5153
5154 #[test]
5155 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
5156 use std::sync::atomic::AtomicBool;
5157
5158 let manager = Arc::new(consolidating_manager(2));
5159 let ids = seal_bulk_segment(&manager, 40);
5160 seal_bulk_segment(&manager, 40);
5161 for id in ids.iter().take(20) {
5162 assert!(manager.delete(*id).expect("delete"));
5163 }
5164
5165 let expected = live_ids(&manager);
5166 let stop = Arc::new(AtomicBool::new(false));
5167
5168 let readers: Vec<_> = (0..4)
5169 .map(|_| {
5170 let manager = Arc::clone(&manager);
5171 let stop = Arc::clone(&stop);
5172 let expected = expected.clone();
5173 std::thread::spawn(move || {
5174 let mut scans = 0u64;
5175 loop {
5178 let mut seen: Vec<u64> = manager
5179 .query_all(|_| true)
5180 .into_iter()
5181 .map(|entity| entity.id.raw())
5182 .collect();
5183 seen.sort_unstable();
5184 assert_eq!(seen, expected, "reader saw a torn sealed set");
5187 scans += 1;
5188 if stop.load(Ordering::Relaxed) {
5189 break;
5190 }
5191 }
5192 scans
5193 })
5194 })
5195 .collect();
5196
5197 drain_maintenance(&manager);
5198 stop.store(true, Ordering::Relaxed);
5199
5200 for reader in readers {
5201 let scans = reader.join().expect("reader thread");
5202 assert!(scans > 0, "reader never got to scan");
5203 }
5204
5205 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5206 assert_eq!(live_ids(&manager), expected);
5207 }
5208
5209 #[test]
5210 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
5211 let manager = consolidating_manager(4);
5212 let ids = seal_bulk_segment(&manager, 20);
5213 for id in ids.iter().take(6) {
5214 assert!(manager.delete(*id).expect("delete"));
5215 }
5216
5217 manager.run_maintenance().expect("maintenance");
5219 assert_eq!(manager.consolidation_progress(), Some(4));
5220 let doomed = ids[6];
5221 assert!(manager.delete(doomed).expect("delete"));
5222
5223 let expected = live_ids(&manager);
5224 drain_maintenance(&manager);
5225
5226 assert!(manager.get(doomed).is_none(), "deleted row came back");
5227 assert_eq!(live_ids(&manager), expected);
5228 }
5229
5230 #[test]
5231 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
5232 let manager = consolidating_manager(3);
5233 let ids = seal_bulk_segment(&manager, 30);
5234 for id in ids.iter().take(12) {
5235 assert!(manager.delete(*id).expect("delete"));
5236 }
5237
5238 let before = live_ids(&manager);
5239 let sealed_before: Vec<SegmentId> = manager
5240 .sealed
5241 .read()
5242 .iter()
5243 .map(|segment| segment.read().id())
5244 .collect();
5245
5246 manager.run_maintenance().expect("maintenance");
5247 manager.run_maintenance().expect("maintenance");
5248 assert!(manager.consolidation.read().is_some());
5249
5250 *manager.consolidation.write() = None;
5254
5255 assert_eq!(live_ids(&manager), before);
5256 let sealed_after: Vec<SegmentId> = manager
5257 .sealed
5258 .read()
5259 .iter()
5260 .map(|segment| segment.read().id())
5261 .collect();
5262 assert_eq!(sealed_after, sealed_before);
5263 assert_eq!(manager.stats().consolidation.runs_completed, 0);
5264
5265 drain_maintenance(&manager);
5266 assert_eq!(live_ids(&manager), before);
5267 assert_eq!(manager.stats().consolidation.runs_completed, 1);
5268 }
5269
5270 #[test]
5271 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
5272 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
5273
5274 let entities: Vec<UnifiedEntity> = (0..10)
5278 .map(|i| {
5279 UnifiedEntity::table_row(
5280 manager.next_entity_id(),
5281 "users",
5282 i as u64,
5283 vec![Value::Integer(i as i64)],
5284 )
5285 })
5286 .collect();
5287 let ids = manager.bulk_insert(entities).expect("bulk insert");
5288
5289 let mut meta = Metadata::new();
5290 meta.set(
5291 "os",
5292 super::super::metadata::MetadataValue::String("linux".to_string()),
5293 );
5294 manager.set_metadata(ids[9], meta).expect("set metadata");
5295 manager.force_seal().expect("seal");
5296
5297 assert!(manager.delete(ids[0]).expect("delete"));
5298 assert!(manager.delete(ids[1]).expect("delete"));
5299 assert!(manager.delete(ids[2]).expect("delete"));
5300
5301 let before: Vec<UnifiedEntity> = {
5302 let mut entities = manager.query_all(|_| true);
5303 entities.sort_by_key(|entity| entity.id.raw());
5304 entities
5305 };
5306
5307 drain_maintenance(&manager);
5308
5309 let after: Vec<UnifiedEntity> = {
5310 let mut entities = manager.query_all(|_| true);
5311 entities.sort_by_key(|entity| entity.id.raw());
5312 entities
5313 };
5314
5315 let columns = |entity: &UnifiedEntity| match &entity.data {
5316 super::super::entity::EntityData::Row(row) => row.columns.clone(),
5317 _ => Vec::new(),
5318 };
5319
5320 assert_eq!(before.len(), after.len());
5321 for (old, new) in before.iter().zip(after.iter()) {
5322 assert_eq!(old.id, new.id);
5323 assert_eq!(columns(old), columns(new));
5324 assert_eq!(
5325 old.sequence_id, new.sequence_id,
5326 "merge must not re-issue sequence ids"
5327 );
5328 }
5329 assert!(manager
5330 .get_metadata(ids[9])
5331 .expect("metadata survives the merge")
5332 .has("os"));
5333 }
5334}