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