1use parking_lot::RwLock;
16use std::collections::HashMap;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::Arc;
19
20use super::entity::{EntityId, UnifiedEntity};
21use super::metadata::{Metadata, MetadataFilter};
22use super::segment::{
23 GrowingSegment, SegmentConfig, SegmentError, SegmentId, SegmentState, SegmentStats,
24 UnifiedSegment, ZoneColPred, ZoneColPredKind,
25};
26use crate::storage::btree::visibility_map::VisibilityMap;
27
28pub const CONSOLIDATION_TOMBSTONE_RATIO: f64 = 0.20;
35
36pub const CONSOLIDATION_FRAGMENTATION_RATIO: f64 = 0.30;
41
42pub const CONSOLIDATION_ENTITIES_PER_TICK: usize = 4_096;
48
49#[derive(Debug, Clone)]
51pub struct ManagerConfig {
52 pub segment_config: SegmentConfig,
54 pub max_sealed_segments: usize,
56 pub idle_seal_secs: u64,
58 pub enable_consolidation: bool,
60 pub consolidation_entities_per_tick: usize,
62 pub enable_archival: bool,
64 pub archive_age_secs: u64,
66}
67
68impl Default for ManagerConfig {
69 fn default() -> Self {
70 Self {
71 segment_config: SegmentConfig::default(),
72 max_sealed_segments: 10,
73 idle_seal_secs: 300, enable_consolidation: true,
75 consolidation_entities_per_tick: CONSOLIDATION_ENTITIES_PER_TICK,
76 enable_archival: true,
77 archive_age_secs: 86400 * 7, }
79 }
80}
81
82#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
87pub struct ConsolidationStats {
88 pub runs_started: u64,
90 pub runs_completed: u64,
92 pub segments_merged: u64,
94 pub tombstones_reclaimed: u64,
96 pub bytes_reclaimed: u64,
98}
99
100#[derive(Debug, Clone, Default)]
102pub struct ManagerStats {
103 pub total_entities: usize,
105 pub growing_count: usize,
107 pub sealed_count: usize,
109 pub archived_count: usize,
111 pub total_memory_bytes: usize,
113 pub seal_ops: u64,
115 pub consolidation: ConsolidationStats,
117}
118
119#[derive(Debug, Clone)]
121pub enum LifecycleEvent {
122 SegmentCreated(SegmentId),
123 SegmentSealed(SegmentId),
124 SegmentConsolidated {
125 source: Vec<SegmentId>,
126 target: SegmentId,
127 },
128 SegmentArchived(SegmentId),
129 EntityInserted(EntityId, SegmentId),
130 EntityDeleted(EntityId, SegmentId),
131}
132
133struct Consolidation {
140 sources: Vec<SourceProgress>,
142 merged_id: SegmentId,
144 merged: GrowingSegment,
146 cursor: usize,
148 pending_ids: Vec<EntityId>,
150 pending_cursor: usize,
152 entities_copied: u64,
154}
155
156struct SourceProgress {
158 id: SegmentId,
159 copied: Vec<EntityId>,
161 tombstones_at_open: Option<usize>,
170}
171
172pub struct SegmentManager {
174 collection: String,
176 config: ManagerConfig,
178 next_segment_id: AtomicU64,
180 next_entity_id: AtomicU64,
182 next_row_id: AtomicU64,
184 total_entities_atomic: AtomicU64,
188 growing: RwLock<Option<Arc<RwLock<GrowingSegment>>>>,
190 sealed: RwLock<Vec<Arc<RwLock<GrowingSegment>>>>,
192 archived: RwLock<Vec<SegmentId>>,
194 entity_segment: RwLock<HashMap<EntityId, SegmentId>>,
198 column_schema: RwLock<Option<Arc<Vec<String>>>>,
201 stats: RwLock<ManagerStats>,
203 consolidation: RwLock<Option<Consolidation>>,
206 events: RwLock<Vec<LifecycleEvent>>,
208 visibility_map: VisibilityMap,
212}
213
214impl SegmentManager {
215 pub fn new(collection: impl Into<String>) -> Self {
217 Self::with_config(collection, ManagerConfig::default())
218 }
219
220 pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
222 Self {
223 collection: collection.into(),
224 config,
225 next_segment_id: AtomicU64::new(1),
226 next_entity_id: AtomicU64::new(1),
227 next_row_id: AtomicU64::new(1),
228 total_entities_atomic: AtomicU64::new(0),
229 growing: RwLock::new(None),
230 sealed: RwLock::new(Vec::new()),
231 archived: RwLock::new(Vec::new()),
232 entity_segment: RwLock::new(HashMap::new()),
233 column_schema: RwLock::new(None),
234 stats: RwLock::new(ManagerStats::default()),
235 consolidation: RwLock::new(None),
236 events: RwLock::new(Vec::new()),
237 visibility_map: VisibilityMap::new(),
238 }
239 }
240
241 pub fn get_or_init_schema(
243 &self,
244 named: &HashMap<String, crate::storage::schema::Value>,
245 ) -> Arc<Vec<String>> {
246 {
247 let schema = self.column_schema.read();
248 if let Some(ref s) = *schema {
249 return Arc::clone(s);
250 }
251 }
252 let cols: Vec<String> = named.keys().cloned().collect();
253 let arc = Arc::new(cols);
254 *self.column_schema.write() = Some(Arc::clone(&arc));
255 arc
256 }
257
258 pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
260 self.column_schema.read().clone()
261 }
262
263 pub(crate) fn set_column_schema_if_empty(&self, columns: Vec<String>) {
264 if columns.is_empty() {
265 return;
266 }
267 let mut schema = self.column_schema.write();
268 if schema.is_none() {
269 *schema = Some(Arc::new(columns));
270 }
271 }
272
273 pub fn collection(&self) -> &str {
275 &self.collection
276 }
277
278 pub fn config(&self) -> &ManagerConfig {
280 &self.config
281 }
282
283 pub fn stats(&self) -> ManagerStats {
288 let mut s = self.stats.read().clone();
289 s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
290 s.total_memory_bytes = self.resident_bytes() as usize;
291 s
292 }
293
294 pub fn resident_bytes(&self) -> u64 {
298 let mut bytes = 0;
299 if let Some(growing_arc) = self.growing.read().as_ref() {
300 bytes += growing_arc.read().resident_bytes();
301 }
302 for segment in self.sealed.read().iter() {
303 bytes += segment.read().resident_bytes();
304 }
305 bytes
306 }
307
308 pub fn memory_bytes(&self) -> u64 {
315 let growing = self
316 .growing
317 .read()
318 .as_ref()
319 .map_or(0, |segment| segment.read().memory_bytes());
320
321 self.sealed
322 .read()
323 .iter()
324 .map(|segment| segment.read().memory_bytes())
325 .fold(growing, u64::saturating_add)
326 }
327
328 pub fn reclaimable_bytes(&self) -> u64 {
330 self.sealed
331 .read()
332 .iter()
333 .map(|segment| segment.read().reclaimable_bytes())
334 .fold(0, u64::saturating_add)
335 }
336
337 pub fn next_entity_id(&self) -> EntityId {
339 EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
340 }
341
342 pub fn next_row_id(&self) -> u64 {
344 self.next_row_id.fetch_add(1, Ordering::SeqCst)
345 }
346
347 pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
351 let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
352 start..start + n
353 }
354
355 pub fn register_row_id(&self, id: u64) {
358 let candidate = id.saturating_add(1);
359 let mut current = self.next_row_id.load(Ordering::SeqCst);
360 while candidate > current {
361 match self.next_row_id.compare_exchange(
362 current,
363 candidate,
364 Ordering::SeqCst,
365 Ordering::SeqCst,
366 ) {
367 Ok(_) => break,
368 Err(updated) => current = updated,
369 }
370 }
371 }
372
373 fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
380 {
382 let growing = self.growing.read();
383 if let Some(segment) = growing.as_ref() {
384 return Arc::clone(segment);
385 }
386 }
387
388 let mut growing = self.growing.write();
390 if let Some(segment) = growing.as_ref() {
392 return Arc::clone(segment);
393 }
394
395 let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
396 let segment = GrowingSegment::new(id, &self.collection);
397 let segment_arc = Arc::new(RwLock::new(segment));
398 *growing = Some(Arc::clone(&segment_arc));
399
400 self.emit(LifecycleEvent::SegmentCreated(id));
401
402 self.stats.write().growing_count += 1;
405
406 segment_arc
407 }
408
409 pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
411 self.maybe_seal_growing()?;
413
414 let segment_arc = self.get_or_create_growing();
415 let mut segment = segment_arc.write();
416
417 if entity.id.raw() == 0 {
419 entity.id = self.next_entity_id();
420 }
421
422 let entity_id = entity.id;
423 let segment_id = segment.id();
424
425 segment.insert(entity)?;
426
427 self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
429
430 self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
439
440 Ok(entity_id)
441 }
442
443 pub fn insert_batch(
445 &self,
446 entities: Vec<UnifiedEntity>,
447 ) -> Result<Vec<EntityId>, SegmentError> {
448 let mut ids = Vec::with_capacity(entities.len());
449 for entity in entities {
450 ids.push(self.insert(entity)?);
451 }
452 Ok(ids)
453 }
454
455 pub fn bulk_insert(
458 &self,
459 mut entities: Vec<UnifiedEntity>,
460 ) -> Result<Vec<EntityId>, SegmentError> {
461 for entity in &mut entities {
463 if entity.id.raw() == 0 {
464 entity.id = self.next_entity_id();
465 }
466 if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
467 if *row_id == 0 {
468 *row_id = self.next_row_id();
469 } else {
470 self.register_row_id(*row_id);
471 }
472 }
473 }
474
475 if let Some(first_row) = entities.first() {
478 if let super::entity::EntityData::Row(ref row) = first_row.data {
479 if let Some(ref named) = row.named {
480 let schema = self.get_or_init_schema(named);
481 for entity in &mut entities {
482 if let super::entity::EntityData::Row(ref mut row) = entity.data {
483 if let Some(named) = row.named.take() {
484 let mut cols = Vec::with_capacity(schema.len());
485 for col_name in schema.iter() {
486 cols.push(
487 named
488 .get(col_name)
489 .cloned()
490 .unwrap_or(crate::storage::schema::Value::Null),
491 );
492 }
493 row.columns = cols;
494 row.schema = Some(Arc::clone(&schema));
495 }
496 }
497 }
498 }
499 }
500 }
501
502 let segment_arc = self.get_or_create_growing();
503 let mut segment = segment_arc.write();
504 let segment_id = segment.id();
505
506 let ids = segment.bulk_insert(entities)?;
508
509 self.total_entities_atomic
514 .fetch_add(ids.len() as u64, Ordering::Relaxed);
515
516 Ok(ids)
517 }
518
519 pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
521 if let Some(growing_arc) = self.growing.read().as_ref() {
523 let growing = growing_arc.read();
524 if let Some(entity) = growing.get(id) {
525 return Some(entity.clone());
526 }
527 }
528
529 let sealed = self.sealed.read();
531 for segment in sealed.iter() {
532 let seg = segment.read();
533 if let Some(entity) = seg.get(id) {
534 return Some(entity.clone());
535 }
536 }
537
538 None
539 }
540
541 pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
546 let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
547 let mut remaining: Vec<usize> = (0..ids.len()).collect(); if let Some(growing_arc) = self.growing.read().as_ref() {
552 let growing = if let Some(g) = growing_arc.try_read() {
553 g
554 } else {
555 growing_arc.read()
556 };
557 remaining.retain(|&i| {
558 if let Some(entity) = growing.get(ids[i]) {
559 out[i] = Some(entity.clone());
560 false } else {
562 true }
564 });
565 }
566
567 if remaining.is_empty() {
568 return out;
569 }
570
571 let sealed = self.sealed.read();
573 for segment in sealed.iter() {
574 if remaining.is_empty() {
575 break;
576 }
577 let seg = segment.read();
578 remaining.retain(|&i| {
579 if let Some(entity) = seg.get(ids[i]) {
580 out[i] = Some(entity.clone());
581 false
582 } else {
583 true
584 }
585 });
586 }
587
588 out
589 }
590
591 pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
604 where
605 F: FnMut(usize, &UnifiedEntity),
606 {
607 thread_local! {
616 static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
617 const { std::cell::Cell::new(Vec::new()) };
618 }
619
620 let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
621 remaining.clear();
622
623 if let Some(growing_arc) = self.growing.read().as_ref() {
624 let growing = if let Some(g) = growing_arc.try_read() {
625 g
626 } else {
627 growing_arc.read()
628 };
629 for (i, id) in ids.iter().enumerate() {
630 if let Some(entity) = growing.get(*id) {
631 f(i, entity);
632 } else {
633 remaining.push(i);
634 }
635 }
636 } else {
637 remaining.reserve(ids.len());
638 remaining.extend(0..ids.len());
639 }
640
641 if !remaining.is_empty() {
642 let sealed = self.sealed.read();
643 for segment in sealed.iter() {
644 if remaining.is_empty() {
645 break;
646 }
647 let seg = segment.read();
648 remaining.retain(|&i| {
649 if let Some(entity) = seg.get(ids[i]) {
650 f(i, entity);
651 false
652 } else {
653 true
654 }
655 });
656 }
657 }
658
659 REMAINING_SCRATCH.with(|cell| cell.set(remaining));
660 }
661
662 fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
664 if let Some(growing_arc) = self.growing.read().as_ref() {
666 let growing = growing_arc.read();
667 if let Some(entity) = growing.get(id) {
668 return Some(entity.clone());
669 }
670 }
671
672 let sealed = self.sealed.read();
674 for segment in sealed.iter() {
675 if let Some(entity) = segment.get(id) {
676 return Some(entity.clone());
677 }
678 }
679
680 None
681 }
682
683 fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
684 let sealed = self.sealed.read();
685 sealed
686 .iter()
687 .find(|segment_arc| segment_arc.read().contains(id))
688 .map(Arc::clone)
689 }
690
691 fn rewrite_sealed_entity_into_growing(
692 &self,
693 entity: UnifiedEntity,
694 metadata: Option<&Metadata>,
695 ) -> Result<(), SegmentError> {
696 let entity_id = entity.id;
697 let sealed_arc = self
698 .find_sealed_segment_arc(entity_id)
699 .ok_or(SegmentError::NotFound(entity_id))?;
700
701 let metadata_to_apply = {
702 let mut sealed = sealed_arc.write();
703 let existing_metadata = sealed.get_metadata(entity_id);
704 if !sealed.force_delete(entity_id) {
705 return Err(SegmentError::NotFound(entity_id));
706 }
707 metadata.cloned().or(existing_metadata)
708 };
709
710 let growing_arc = self.get_or_create_growing();
711 let growing_id = {
712 let mut growing = growing_arc.write();
713 growing.insert(entity)?;
714 if let Some(metadata) = metadata_to_apply {
715 growing.set_metadata(entity_id, metadata)?;
716 }
717 growing.id()
718 };
719
720 self.entity_segment.write().insert(entity_id, growing_id);
721 Ok(())
722 }
723
724 pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
726 let entity_id = entity.id;
727 let mut entity = Some(entity);
728
729 if let Some(growing_arc) = self.growing.read().as_ref() {
731 let mut growing = growing_arc.write();
732 if growing.contains(entity_id) && growing.state().is_writable() {
733 return growing.update(entity.take().expect("entity already moved"));
734 }
735 }
736
737 let segment_id = self.entity_segment.read().get(&entity_id).copied();
739 if let Some(seg_id) = segment_id {
740 if let Some(growing_arc) = self.growing.read().as_ref() {
741 let mut growing = growing_arc.write();
742 if growing.id() == seg_id && growing.state().is_writable() {
743 return growing.update(entity.take().expect("entity already moved"));
744 }
745 }
746 }
747
748 if let Some(entity) = entity.take() {
749 return self.rewrite_sealed_entity_into_growing(entity, None);
750 }
751
752 Err(SegmentError::NotFound(entity_id))
753 }
754
755 pub fn update_with_metadata(
758 &self,
759 entity: UnifiedEntity,
760 metadata: Option<&Metadata>,
761 ) -> Result<(), SegmentError> {
762 let entity_id = entity.id;
763 let mut entity = Some(entity);
764
765 if let Some(growing_arc) = self.growing.read().as_ref() {
767 let mut growing = growing_arc.write();
768 if growing.contains(entity_id) && growing.state().is_writable() {
769 growing.update(entity.take().expect("entity already moved"))?;
770 if let Some(metadata) = metadata {
771 growing.set_metadata(entity_id, metadata.clone())?;
772 }
773 return Ok(());
774 }
775 }
776
777 let segment_id = self.entity_segment.read().get(&entity_id).copied();
779 if let Some(seg_id) = segment_id {
780 if let Some(growing_arc) = self.growing.read().as_ref() {
781 let mut growing = growing_arc.write();
782 if growing.id() == seg_id && growing.state().is_writable() {
783 growing.update(entity.take().expect("entity already moved"))?;
784 if let Some(metadata) = metadata {
785 growing.set_metadata(entity_id, metadata.clone())?;
786 }
787 return Ok(());
788 }
789 }
790 }
791
792 if let Some(entity) = entity.take() {
793 return self.rewrite_sealed_entity_into_growing(entity, metadata);
794 }
795
796 Err(SegmentError::NotFound(entity_id))
797 }
798
799 pub fn update_hot(
803 &self,
804 entity: UnifiedEntity,
805 modified_columns: &[String],
806 ) -> Result<(), SegmentError> {
807 let entity_id = entity.id;
808 let mut entity = Some(entity);
809
810 if let Some(growing_arc) = self.growing.read().as_ref() {
811 let mut growing = growing_arc.write();
812 if growing.contains(entity_id) && growing.state().is_writable() {
813 return growing.update_hot(
814 entity.take().expect("entity already moved"),
815 modified_columns,
816 );
817 }
818 }
819
820 let segment_id = self.entity_segment.read().get(&entity_id).copied();
821 if let Some(seg_id) = segment_id {
822 if let Some(growing_arc) = self.growing.read().as_ref() {
823 let mut growing = growing_arc.write();
824 if growing.id() == seg_id && growing.state().is_writable() {
825 return growing.update_hot(
826 entity.take().expect("entity already moved"),
827 modified_columns,
828 );
829 }
830 }
831 }
832
833 if let Some(entity) = entity.take() {
834 return self.rewrite_sealed_entity_into_growing(entity, None);
835 }
836
837 Err(SegmentError::NotFound(entity_id))
838 }
839
840 pub fn update_hot_with_metadata(
843 &self,
844 entity: UnifiedEntity,
845 modified_columns: &[String],
846 metadata: Option<&Metadata>,
847 ) -> Result<(), SegmentError> {
848 let entity_id = entity.id;
849 let mut entity = Some(entity);
850
851 if let Some(growing_arc) = self.growing.read().as_ref() {
852 let mut growing = growing_arc.write();
853 if growing.contains(entity_id) && growing.state().is_writable() {
854 growing.update_hot(
855 entity.take().expect("entity already moved"),
856 modified_columns,
857 )?;
858 if let Some(metadata) = metadata {
859 growing.set_metadata(entity_id, metadata.clone())?;
860 }
861 return Ok(());
862 }
863 }
864
865 let segment_id = self.entity_segment.read().get(&entity_id).copied();
866 if let Some(seg_id) = segment_id {
867 if let Some(growing_arc) = self.growing.read().as_ref() {
868 let mut growing = growing_arc.write();
869 if growing.id() == seg_id && growing.state().is_writable() {
870 growing.update_hot(
871 entity.take().expect("entity already moved"),
872 modified_columns,
873 )?;
874 if let Some(metadata) = metadata {
875 growing.set_metadata(entity_id, metadata.clone())?;
876 }
877 return Ok(());
878 }
879 }
880 }
881
882 if let Some(entity) = entity.take() {
883 return self.rewrite_sealed_entity_into_growing(entity, metadata);
884 }
885
886 Err(SegmentError::NotFound(entity_id))
887 }
888
889 pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
892 where
893 I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
894 {
895 let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
896 items.into_iter().collect();
897 if items.is_empty() {
898 return Ok(());
899 }
900
901 if let Some(growing_arc) = self.growing.read().as_ref() {
902 let mut growing = growing_arc.write();
903 if growing.state().is_writable() {
904 match growing.update_hot_batch_with_metadata(items.iter().copied()) {
905 Ok(()) => return Ok(()),
906 Err(SegmentError::NotFound(_)) => {}
907 Err(other) => return Err(other),
908 }
909 }
910 }
911
912 for (entity, modified_columns, metadata) in items {
913 self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
914 }
915 Ok(())
916 }
917
918 pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
920 if let Some(growing_arc) = self.growing.read().as_ref() {
923 let mut growing = growing_arc.write();
924 if growing.contains(id) && growing.state().is_writable() {
925 let seg_id = growing.id();
926 let deleted = growing.delete(id)?;
927 if deleted {
928 self.entity_segment.write().remove(&id);
929 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
930 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
931 }
932 return Ok(deleted);
933 }
934 }
935
936 let segment_id = self.entity_segment.read().get(&id).copied();
939
940 if let Some(seg_id) = segment_id {
941 if let Some(growing_arc) = self.growing.read().as_ref() {
942 let mut growing = growing_arc.write();
943 if growing.id() == seg_id && growing.state().is_writable() {
944 let deleted = growing.delete(id)?;
945 if deleted {
946 self.entity_segment.write().remove(&id);
947 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
948 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
949 }
950 return Ok(deleted);
951 }
952 }
953 }
954
955 {
958 let sealed = self.sealed.read();
959 for segment_arc in sealed.iter() {
960 let mut seg = segment_arc.write();
961 let seg_id = seg.id();
962 if seg.contains(id) {
963 let deleted = seg.force_delete(id);
964 drop(seg);
965 if deleted {
966 self.entity_segment.write().remove(&id);
967 self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
968 self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
969 }
970 return Ok(deleted);
971 }
972 }
973 }
974
975 Ok(false)
976 }
977
978 pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
979 if ids.is_empty() {
980 return Ok(Vec::new());
981 }
982
983 let mut deleted_ids = Vec::with_capacity(ids.len());
984
985 if let Some(growing_arc) = self.growing.read().as_ref() {
986 let mut growing = growing_arc.write();
987 if growing.state().is_writable() {
988 let seg_id = growing.id();
989 let deleted = growing.delete_batch(ids)?;
990 if !deleted.is_empty() {
991 {
992 let mut entity_segment = self.entity_segment.write();
993 for id in &deleted {
994 entity_segment.remove(id);
995 }
996 }
997 self.total_entities_atomic
998 .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
999 for id in &deleted {
1000 self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1001 }
1002 deleted_ids.extend(deleted);
1003 }
1004 }
1005 }
1006
1007 if deleted_ids.len() == ids.len() {
1008 return Ok(deleted_ids);
1009 }
1010
1011 let deleted_set: std::collections::HashSet<EntityId> =
1012 deleted_ids.iter().copied().collect();
1013 for &id in ids {
1014 if deleted_set.contains(&id) {
1015 continue;
1016 }
1017 if self.delete(id)? {
1018 deleted_ids.push(id);
1019 }
1020 }
1021
1022 Ok(deleted_ids)
1023 }
1024
1025 pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1027 if let Some(growing_arc) = self.growing.read().as_ref() {
1029 let growing = growing_arc.read();
1030 if growing.contains(id) {
1031 return growing.get_metadata(id);
1032 }
1033 }
1034
1035 let segment_id = self.entity_segment.read().get(&id).copied();
1037
1038 if let Some(seg_id) = segment_id {
1039 if let Some(growing_arc) = self.growing.read().as_ref() {
1040 let growing = growing_arc.read();
1041 if growing.id() == seg_id {
1042 return growing.get_metadata(id);
1043 }
1044 }
1045
1046 let sealed = self.sealed.read();
1047 for segment in sealed.iter() {
1048 if segment.id() == seg_id {
1049 return segment.get_metadata(id);
1050 }
1051 }
1052 }
1053
1054 if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1055 return segment_arc.read().get_metadata(id);
1056 }
1057
1058 None
1059 }
1060
1061 pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1063 if let Some(growing_arc) = self.growing.read().as_ref() {
1066 let mut growing = growing_arc.write();
1067 if growing.contains(id) && growing.state().is_writable() {
1068 return growing.set_metadata(id, metadata);
1069 }
1070 }
1071
1072 let segment_id = self.entity_segment.read().get(&id).copied();
1074
1075 if let Some(seg_id) = segment_id {
1076 if let Some(growing_arc) = self.growing.read().as_ref() {
1077 let mut growing = growing_arc.write();
1078 if growing.id() == seg_id && growing.state().is_writable() {
1079 return growing.set_metadata(id, metadata);
1080 }
1081 }
1082 }
1083
1084 if let Some(entity) = self.get(id) {
1085 return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1086 }
1087
1088 Err(SegmentError::NotFound(id))
1089 }
1090
1091 fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1093 let should_seal = {
1094 let growing_opt = self.growing.read();
1095 if let Some(growing_arc) = growing_opt.as_ref() {
1096 let growing = growing_arc.read();
1097 growing.should_seal(&self.config.segment_config)
1098 || growing.idle_secs() >= self.config.idle_seal_secs
1099 } else {
1100 false
1101 }
1102 };
1103
1104 if should_seal {
1105 self.seal_current()?;
1106 }
1107
1108 Ok(())
1109 }
1110
1111 pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1113 let growing_opt = self.growing.write().take();
1114
1115 if let Some(growing_arc) = growing_opt {
1116 let mut growing = growing_arc.write();
1117 let seg_id = growing.id();
1118 let entity_count = growing.stats().entity_count as u64;
1119
1120 growing.seal()?;
1122
1123 drop(growing); self.sealed.write().push(growing_arc);
1129
1130 self.mark_sealed_pages_visible(entity_count);
1132
1133 {
1135 let mut stats = self.stats.write();
1136 stats.growing_count = stats.growing_count.saturating_sub(1);
1137 stats.sealed_count += 1;
1138 stats.seal_ops += 1;
1139 }
1140
1141 self.emit(LifecycleEvent::SegmentSealed(seg_id));
1142
1143 return Ok(seg_id);
1144 }
1145
1146 Err(SegmentError::InvalidState(SegmentState::Sealed))
1147 }
1148
1149 pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1151 let has_growing = self.growing.read().is_some();
1152 if has_growing {
1153 Ok(Some(self.seal_current()?))
1154 } else {
1155 Ok(None)
1156 }
1157 }
1158
1159 pub fn all_visible_fraction(&self) -> f64 {
1169 const ROWS_PER_PAGE: u32 = 256;
1170 let sealed = self.sealed.read();
1171 if sealed.is_empty() {
1172 return 0.0;
1173 }
1174 let mut total_pages: u64 = 0;
1175 for seg_arc in sealed.iter() {
1176 let seg = seg_arc.read();
1177 let entity_count = seg.stats().entity_count as u64;
1178 let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1179 total_pages += pages;
1180 }
1181 if total_pages == 0 {
1182 return 0.0;
1183 }
1184 let visible = self.visibility_map.all_visible_count();
1185 (visible as f64 / total_pages as f64).min(1.0)
1186 }
1187
1188 fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1191 const ROWS_PER_PAGE: u32 = 256;
1192 let existing_visible = self.visibility_map.all_visible_count();
1193 let start_page = existing_visible as u32;
1195 let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1196 let end_page = start_page + new_pages as u32;
1197 self.visibility_map.mark_range_visible(start_page, end_page);
1198 }
1199
1200 pub fn for_each_entity<F>(&self, mut callback: F)
1206 where
1207 F: FnMut(&UnifiedEntity) -> bool,
1208 {
1209 if let Some(growing_arc) = self.growing.read().as_ref() {
1213 let growing = if let Some(g) = growing_arc.try_read() {
1214 g
1215 } else {
1216 growing_arc.read()
1217 };
1218 if !growing.for_each_fast(&mut callback) {
1219 return;
1220 }
1221 }
1222
1223 let sealed = self.sealed.read();
1225 for segment_arc in sealed.iter() {
1226 let segment = segment_arc.read();
1227 if !segment.for_each_fast(&mut callback) {
1228 return;
1229 }
1230 }
1231 }
1232
1233 pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1249 &self,
1250 init: FInit,
1251 fold: FFold,
1252 reduce: FReduce,
1253 ) -> T
1254 where
1255 T: Send,
1256 FInit: Fn() -> T + Send + Sync,
1257 FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1258 FReduce: Fn(T, T) -> T + Send + Sync,
1259 {
1260 use rayon::prelude::*;
1261
1262 let mut acc = init();
1265 if let Some(growing_arc) = self.growing.read().as_ref() {
1266 let growing = if let Some(g) = growing_arc.try_read() {
1267 g
1268 } else {
1269 growing_arc.read()
1270 };
1271 growing.for_each_fast(|entity| {
1272 acc = fold(std::mem::replace(&mut acc, init()), entity);
1273 true
1274 });
1275 }
1276
1277 let segments: Vec<_> = {
1281 let sealed = self.sealed.read();
1282 sealed.iter().cloned().collect()
1283 };
1284
1285 if segments.len() <= 1 {
1286 for seg_arc in &segments {
1287 let seg = seg_arc.read();
1288 seg.for_each_fast(|entity| {
1289 acc = fold(std::mem::replace(&mut acc, init()), entity);
1290 true
1291 });
1292 }
1293 return acc;
1294 }
1295
1296 let sealed_acc = segments
1297 .into_par_iter()
1298 .map(|seg_arc| {
1299 let mut local = init();
1300 let seg = seg_arc.read();
1301 seg.for_each_fast(|entity| {
1302 local = fold(std::mem::replace(&mut local, init()), entity);
1303 true
1304 });
1305 local
1306 })
1307 .reduce(&init, &reduce);
1308
1309 reduce(acc, sealed_acc)
1310 }
1311
1312 pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
1321 where
1322 F: FnMut(&UnifiedEntity) -> bool,
1323 {
1324 if let Some(growing_arc) = self.growing.read().as_ref() {
1330 let growing = if let Some(g) = growing_arc.try_read() {
1331 g
1332 } else {
1333 growing_arc.read()
1334 };
1335 if !growing.for_each_fast(&mut callback) {
1336 return;
1337 }
1338 }
1339
1340 let sealed = self.sealed.read();
1342 for segment_arc in sealed.iter() {
1343 let segment = segment_arc.read();
1344 if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
1345 continue; }
1347 if !segment.for_each_fast(&mut callback) {
1348 return;
1349 }
1350 }
1351 }
1352
1353 pub fn query_all_zoned<F>(
1363 &self,
1364 zone_preds: &[(&str, ZoneColPred<'_>)],
1365 filter: F,
1366 ) -> Vec<UnifiedEntity>
1367 where
1368 F: Fn(&UnifiedEntity) -> bool + Sync,
1369 {
1370 let mut results = Vec::new();
1371
1372 if let Some(growing_arc) = self.growing.read().as_ref() {
1375 let growing = if let Some(g) = growing_arc.try_read() {
1376 g
1377 } else {
1378 growing_arc.read()
1379 };
1380 results.extend(growing.iter().filter(|e| filter(e)).cloned());
1381 }
1382
1383 let sealed = self.sealed.read();
1385 let surviving: Vec<_> = sealed
1387 .iter()
1388 .filter(|seg_arc| {
1389 if zone_preds.is_empty() {
1390 return true;
1391 }
1392 let seg = seg_arc.read();
1393 !seg.can_skip_zone_preds(zone_preds)
1394 })
1395 .collect();
1396
1397 let use_parallel = surviving.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1398
1399 if use_parallel {
1400 let filter_ref = &filter;
1401 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1402 surviving
1403 .iter()
1404 .map(|segment| {
1405 s.spawn(move || {
1406 segment
1407 .read()
1408 .iter()
1409 .filter(|e| filter_ref(e))
1410 .cloned()
1411 .collect::<Vec<_>>()
1412 })
1413 })
1414 .collect::<Vec<_>>()
1415 .into_iter()
1416 .map(|handle| handle.join().unwrap_or_default())
1417 .collect()
1418 });
1419 for batch in segment_results {
1420 results.extend(batch);
1421 }
1422 } else {
1423 for segment_arc in surviving {
1424 let seg = segment_arc.read();
1425 results.extend(seg.iter().filter(|e| filter(e)).cloned());
1426 }
1427 }
1428
1429 results
1430 }
1431
1432 pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
1435 where
1436 F: Fn(&UnifiedEntity) -> bool + Sync,
1437 {
1438 let mut results = Vec::new();
1439
1440 if let Some(growing_arc) = self.growing.read().as_ref() {
1443 let growing = if let Some(g) = growing_arc.try_read() {
1444 g
1445 } else {
1446 growing_arc.read()
1447 };
1448 results.extend(growing.iter().filter(|e| filter(e)).cloned());
1449 }
1450
1451 let sealed = self.sealed.read();
1453 let use_parallel = sealed.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1454 if use_parallel {
1455 let filter_ref = &filter;
1456 let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1457 sealed
1458 .iter()
1459 .map(|segment| {
1460 s.spawn(move || {
1461 segment
1462 .read()
1463 .iter()
1464 .filter(|e| filter_ref(e))
1465 .cloned()
1466 .collect::<Vec<_>>()
1467 })
1468 })
1469 .collect::<Vec<_>>()
1470 .into_iter()
1471 .map(|handle| handle.join().unwrap_or_default())
1472 .collect()
1473 });
1474 for batch in segment_results {
1475 results.extend(batch);
1476 }
1477 } else {
1478 for segment in sealed.iter() {
1479 let seg = segment.read();
1480 results.extend(seg.iter().filter(|e| filter(e)).cloned());
1481 }
1482 }
1483
1484 results
1485 }
1486
1487 pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
1492 if let Some(growing_arc) = self.growing.read().as_ref() {
1493 let growing = growing_arc.read();
1494 if growing.bloom_might_contain_key(key) {
1495 return true;
1496 }
1497 } else {
1498 return true;
1499 }
1500
1501 let sealed = self.sealed.read();
1502 for segment_arc in sealed.iter() {
1503 let segment = segment_arc.read();
1504 if segment.bloom_might_contain_key(key) {
1505 return true;
1506 }
1507 }
1508
1509 false
1510 }
1511
1512 pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
1514 let mut results = Vec::new();
1515
1516 if let Some(growing_arc) = self.growing.read().as_ref() {
1518 let growing = growing_arc.read();
1519 results.extend(growing.filter_metadata(filters));
1520 }
1521
1522 let sealed = self.sealed.read();
1524 for segment in sealed.iter() {
1525 results.extend(segment.filter_metadata(filters));
1526 }
1527
1528 results
1529 }
1530
1531 pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
1533 let mut results = Vec::new();
1534
1535 if let Some(growing_arc) = self.growing.read().as_ref() {
1537 let growing = growing_arc.read();
1538 for entity in growing.iter_kind(kind) {
1539 results.push(entity.clone());
1540 }
1541 }
1542
1543 let sealed = self.sealed.read();
1548 for segment_arc in sealed.iter() {
1549 let segment = segment_arc.read();
1550 for entity in segment.iter_kind(kind) {
1551 results.push(entity.clone());
1552 }
1553 }
1554
1555 results
1556 }
1557
1558 pub fn count(&self) -> usize {
1560 self.total_entities_atomic.load(Ordering::Relaxed) as usize
1561 }
1562
1563 pub fn segment_ids(&self) -> Vec<SegmentId> {
1565 let mut ids = Vec::new();
1566
1567 if let Some(growing_arc) = self.growing.read().as_ref() {
1568 ids.push(growing_arc.read().id());
1569 }
1570
1571 let sealed = self.sealed.read();
1572 for segment in sealed.iter() {
1573 ids.push(segment.id());
1574 }
1575
1576 ids.extend(self.archived.read().iter().copied());
1577
1578 ids
1579 }
1580
1581 #[inline]
1594 #[allow(clippy::unused_self)]
1595 fn emit(&self, _event: LifecycleEvent) {}
1596
1597 pub fn drain_events(&self) -> Vec<LifecycleEvent> {
1600 std::mem::take(&mut *self.events.write())
1601 }
1602
1603 pub fn run_maintenance(&self) -> Result<(), SegmentError> {
1605 self.maybe_seal_growing()?;
1607
1608 if self.config.enable_consolidation {
1609 self.consolidation_tick();
1610 }
1611
1612 Ok(())
1613 }
1614
1615 fn consolidation_tick(&self) {
1626 let running = self.consolidation.read().is_some();
1629 if !running && !self.start_consolidation() {
1630 return;
1631 }
1632
1633 let finished = {
1634 let mut guard = self.consolidation.write();
1635 let Some(run) = guard.as_mut() else {
1636 return;
1637 };
1638 self.copy_bounded(run);
1639 run.cursor >= run.sources.len()
1640 };
1641
1642 if finished {
1643 let run = self.consolidation.write().take();
1644 if let Some(run) = run {
1645 self.finish_consolidation(run);
1646 }
1647 }
1648 }
1649
1650 fn start_consolidation(&self) -> bool {
1654 let sources = self.select_candidates();
1655 self.start_consolidation_from_sources(sources)
1656 }
1657
1658 fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
1659 if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
1660 return false;
1661 }
1662
1663 let live_estimate: usize = {
1664 let sealed = self.sealed.read();
1665 sealed
1666 .iter()
1667 .filter_map(|segment_arc| {
1668 let segment = segment_arc.read();
1669 if sources.contains(&segment.id()) {
1670 Some(segment.live_entity_count())
1671 } else {
1672 None
1673 }
1674 })
1675 .sum()
1676 };
1677
1678 let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
1679 let merged =
1680 GrowingSegment::with_bloom_capacity(merged_id, &self.collection, live_estimate);
1681
1682 *self.consolidation.write() = Some(Consolidation {
1683 sources: sources
1684 .into_iter()
1685 .map(|id| SourceProgress {
1686 id,
1687 copied: Vec::new(),
1688 tombstones_at_open: None,
1689 })
1690 .collect(),
1691 merged_id,
1692 merged,
1693 cursor: 0,
1694 pending_ids: Vec::new(),
1695 pending_cursor: 0,
1696 entities_copied: 0,
1697 });
1698
1699 self.stats.write().consolidation.runs_started += 1;
1700 true
1701 }
1702
1703 pub fn pressure_consolidation_tick(&self) -> bool {
1708 if !self.config.enable_consolidation {
1709 return false;
1710 }
1711
1712 let running = self.consolidation.read().is_some();
1713 if !running {
1714 let sources = self.dirty_sealed_segments();
1715 if !self.start_consolidation_from_sources(sources) {
1716 return false;
1717 }
1718 }
1719
1720 self.consolidation_tick();
1721 true
1722 }
1723
1724 fn select_candidates(&self) -> Vec<SegmentId> {
1739 let sealed = self.sealed.read();
1740 if sealed.is_empty() {
1741 return Vec::new();
1742 }
1743
1744 if sealed.len() > self.config.max_sealed_segments {
1745 return sealed.iter().map(|segment| segment.read().id()).collect();
1746 }
1747
1748 let mut live = 0usize;
1749 let mut tombstones = 0usize;
1750 let mut held = 0u64;
1751 let mut reclaimable = 0u64;
1752 let mut dirty = Vec::new();
1753
1754 for segment_arc in sealed.iter() {
1755 let segment = segment_arc.read();
1756 live += segment.live_entity_count();
1757 tombstones += segment.tombstone_count();
1758 held += segment.resident_bytes();
1759 reclaimable += segment.reclaimable_bytes();
1760 if segment.tombstone_count() > 0 {
1761 dirty.push(segment.id());
1762 }
1763 }
1764
1765 let total = live + tombstones;
1766 let tombstone_ratio = if total == 0 {
1767 0.0
1768 } else {
1769 tombstones as f64 / total as f64
1770 };
1771 let fragmentation = if held == 0 {
1772 0.0
1773 } else {
1774 reclaimable as f64 / held as f64
1775 };
1776
1777 if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
1778 || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
1779 {
1780 return dirty;
1781 }
1782
1783 Vec::new()
1784 }
1785
1786 fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
1787 self.sealed
1788 .read()
1789 .iter()
1790 .filter_map(|segment_arc| {
1791 let segment = segment_arc.read();
1792 (segment.reclaimable_bytes() > 0).then(|| segment.id())
1793 })
1794 .collect()
1795 }
1796
1797 fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
1800 if sources.is_empty() {
1801 return false;
1802 }
1803 let sealed = self.sealed.read();
1804 sealed.iter().any(|segment_arc| {
1805 let segment = segment_arc.read();
1806 sources.contains(&segment.id()) && segment.tombstone_count() > 0
1807 })
1808 }
1809
1810 fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
1811 let sealed = self.sealed.read();
1812 sealed
1813 .iter()
1814 .find(|segment| segment.read().id() == id)
1815 .map(Arc::clone)
1816 }
1817
1818 fn copy_bounded(&self, run: &mut Consolidation) {
1823 let mut budget = self.config.consolidation_entities_per_tick;
1824
1825 while budget > 0 && run.cursor < run.sources.len() {
1826 let source_id = run.sources[run.cursor].id;
1827 let Some(source_arc) = self.sealed_segment(source_id) else {
1828 run.cursor += 1;
1832 run.pending_ids.clear();
1833 run.pending_cursor = 0;
1834 continue;
1835 };
1836
1837 if run.sources[run.cursor].tombstones_at_open.is_none() {
1838 let source = source_arc.read();
1839 run.pending_ids = source.live_entity_ids();
1840 run.pending_cursor = 0;
1841 run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
1842 }
1843
1844 {
1845 let source = source_arc.read();
1846 while budget > 0 && run.pending_cursor < run.pending_ids.len() {
1847 let id = run.pending_ids[run.pending_cursor];
1848 run.pending_cursor += 1;
1849 budget -= 1;
1850
1851 let Some(entity) = source.get(id) else {
1854 continue;
1855 };
1856 let metadata = source.get_metadata(id);
1857 run.merged.adopt_entity(entity.clone(), metadata);
1858 run.sources[run.cursor].copied.push(id);
1859 run.entities_copied += 1;
1860 }
1861 }
1862
1863 if run.pending_cursor >= run.pending_ids.len() {
1864 run.cursor += 1;
1865 run.pending_ids.clear();
1866 run.pending_cursor = 0;
1867 }
1868 }
1869 }
1870
1871 fn finish_consolidation(&self, mut run: Consolidation) {
1877 let mut sealed = self.sealed.write();
1878
1879 let mut positions = Vec::with_capacity(run.sources.len());
1880 for source in &run.sources {
1881 if let Some(index) = sealed
1882 .iter()
1883 .position(|segment| segment.read().id() == source.id)
1884 {
1885 positions.push(index);
1886 }
1887 }
1888 if positions.is_empty() {
1889 return;
1890 }
1891 positions.sort_unstable();
1892
1893 let mut stale = Vec::new();
1897 for source in &run.sources {
1898 let Some(source_arc) = sealed
1899 .iter()
1900 .find(|segment| segment.read().id() == source.id)
1901 else {
1902 continue;
1903 };
1904 let segment = source_arc.read();
1905 if source.tombstones_at_open == Some(segment.tombstone_count()) {
1906 continue;
1907 }
1908 stale.extend(
1909 source
1910 .copied
1911 .iter()
1912 .copied()
1913 .filter(|id| !segment.contains(*id)),
1914 );
1915 }
1916 for id in stale {
1917 run.merged.evict_entity(id);
1918 }
1919
1920 let reclaimed_bytes: u64 = positions
1921 .iter()
1922 .map(|&index| sealed[index].read().resident_bytes())
1923 .sum();
1924 let reclaimed_tombstones: u64 = positions
1925 .iter()
1926 .map(|&index| sealed[index].read().tombstone_count() as u64)
1927 .sum();
1928
1929 if run.merged.seal().is_err() {
1930 return;
1931 }
1932 let merged_bytes = run.merged.resident_bytes();
1933 let merged_arc = Arc::new(RwLock::new(run.merged));
1934
1935 for &index in positions.iter().rev() {
1936 sealed.remove(index);
1937 }
1938 let insert_at = positions[0].min(sealed.len());
1939 sealed.insert(insert_at, Arc::clone(&merged_arc));
1940 drop(sealed);
1941
1942 let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
1944 {
1945 let mut entity_segment = self.entity_segment.write();
1946 for segment_id in entity_segment.values_mut() {
1947 if source_ids.contains(segment_id) {
1948 *segment_id = run.merged_id;
1949 }
1950 }
1951 }
1952
1953 {
1954 let mut stats = self.stats.write();
1955 stats.sealed_count = stats
1956 .sealed_count
1957 .saturating_sub(positions.len())
1958 .saturating_add(1);
1959 let consolidation = &mut stats.consolidation;
1960 consolidation.runs_completed += 1;
1961 consolidation.segments_merged += positions.len() as u64;
1962 consolidation.tombstones_reclaimed += reclaimed_tombstones;
1963 consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
1964 }
1965
1966 self.emit(LifecycleEvent::SegmentConsolidated {
1967 source: source_ids,
1968 target: run.merged_id,
1969 });
1970 }
1971
1972 #[cfg(test)]
1975 fn consolidation_progress(&self) -> Option<u64> {
1976 self.consolidation
1977 .read()
1978 .as_ref()
1979 .map(|run| run.entities_copied)
1980 }
1981}
1982
1983impl UnifiedSegment for Arc<RwLock<GrowingSegment>> {
1986 fn id(&self) -> SegmentId {
1987 self.read().id()
1988 }
1989
1990 fn state(&self) -> SegmentState {
1991 self.read().state()
1992 }
1993
1994 fn collection(&self) -> &str {
1995 "unknown"
1997 }
1998
1999 fn stats(&self) -> SegmentStats {
2000 self.read().stats()
2001 }
2002
2003 fn entity_count(&self) -> usize {
2004 self.read().entity_count()
2005 }
2006
2007 fn contains(&self, id: EntityId) -> bool {
2008 self.read().contains(id)
2009 }
2010
2011 fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2012 None
2015 }
2016
2017 fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2018 None
2019 }
2020
2021 fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2022 self.write().insert(entity)
2023 }
2024
2025 fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2026 self.write().update(entity)
2027 }
2028
2029 fn update_hot(
2030 &mut self,
2031 entity: UnifiedEntity,
2032 modified_columns: &[String],
2033 ) -> Result<(), SegmentError> {
2034 self.write().update_hot(entity, modified_columns)
2035 }
2036
2037 fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2038 self.write().delete(id)
2039 }
2040
2041 fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2042 self.read().get_metadata(id)
2043 }
2044
2045 fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2046 self.write().set_metadata(id, metadata)
2047 }
2048
2049 fn seal(&mut self) -> Result<(), SegmentError> {
2050 self.write().seal()
2051 }
2052
2053 fn should_seal(&self, config: &SegmentConfig) -> bool {
2054 self.read().should_seal(config)
2055 }
2056
2057 fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2058 Box::new(std::iter::empty())
2060 }
2061
2062 fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2063 Box::new(std::iter::empty())
2064 }
2065
2066 fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2067 self.read().filter_metadata(filters)
2068 }
2069}
2070
2071#[cfg(test)]
2072mod tests {
2073 use super::*;
2074 use crate::storage::schema::Value;
2075
2076 #[test]
2077 fn test_manager_basic() {
2078 let manager = SegmentManager::new("test_collection");
2079
2080 let entity = UnifiedEntity::table_row(
2081 manager.next_entity_id(),
2082 "users",
2083 1,
2084 vec![Value::text("Alice".to_string())],
2085 );
2086
2087 let id = manager.insert(entity).unwrap();
2088 assert!(manager.get(id).is_some());
2089 assert_eq!(manager.count(), 1);
2090 }
2091
2092 #[test]
2093 fn bloom_may_contain_key_probes_inserted_rid_bytes() {
2094 let manager = SegmentManager::new("test_collection");
2095
2096 let entity = UnifiedEntity::table_row(
2097 manager.next_entity_id(),
2098 "users",
2099 1,
2100 vec![Value::text("Alice".to_string())],
2101 );
2102
2103 let id = manager.insert(entity).unwrap();
2104
2105 assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
2106 assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
2107 }
2108
2109 #[test]
2110 fn test_manager_auto_seal() {
2111 let config = ManagerConfig {
2112 segment_config: SegmentConfig {
2113 max_entities: 2,
2114 ..Default::default()
2115 },
2116 ..Default::default()
2117 };
2118
2119 let manager = SegmentManager::with_config("test", config);
2120
2121 manager
2123 .insert(UnifiedEntity::vector(
2124 manager.next_entity_id(),
2125 "v",
2126 vec![0.1],
2127 ))
2128 .unwrap();
2129
2130 manager
2132 .insert(UnifiedEntity::vector(
2133 manager.next_entity_id(),
2134 "v",
2135 vec![0.2],
2136 ))
2137 .unwrap();
2138
2139 manager
2141 .insert(UnifiedEntity::vector(
2142 manager.next_entity_id(),
2143 "v",
2144 vec![0.3],
2145 ))
2146 .unwrap();
2147
2148 let stats = manager.stats();
2149 assert_eq!(stats.total_entities, 3);
2150 }
2151
2152 #[test]
2153 fn test_manager_delete() {
2154 let manager = SegmentManager::new("test");
2155
2156 let id = manager
2157 .insert(UnifiedEntity::vector(
2158 manager.next_entity_id(),
2159 "v",
2160 vec![0.1],
2161 ))
2162 .unwrap();
2163
2164 assert!(manager.get(id).is_some());
2165 assert!(manager.delete(id).unwrap());
2166 assert!(manager.get(id).is_none());
2167 }
2168
2169 #[test]
2170 fn test_manager_metadata() {
2171 let manager = SegmentManager::new("test");
2172
2173 let id = manager
2174 .insert(UnifiedEntity::table_row(
2175 manager.next_entity_id(),
2176 "hosts",
2177 1,
2178 vec![Value::text("192.168.1.1".to_string())],
2179 ))
2180 .unwrap();
2181
2182 let mut meta = Metadata::new();
2183 meta.set(
2184 "os",
2185 super::super::metadata::MetadataValue::String("linux".to_string()),
2186 );
2187
2188 manager.set_metadata(id, meta).unwrap();
2189
2190 let retrieved = manager.get_metadata(id).unwrap();
2191 assert!(retrieved.has("os"));
2192 }
2193
2194 #[test]
2195 fn test_manager_query_by_kind() {
2196 let manager = SegmentManager::new("test");
2197
2198 manager
2199 .insert(UnifiedEntity::table_row(
2200 manager.next_entity_id(),
2201 "hosts",
2202 1,
2203 vec![],
2204 ))
2205 .unwrap();
2206
2207 manager
2208 .insert(UnifiedEntity::vector(
2209 manager.next_entity_id(),
2210 "embeddings",
2211 vec![0.1],
2212 ))
2213 .unwrap();
2214
2215 manager
2216 .insert(UnifiedEntity::table_row(
2217 manager.next_entity_id(),
2218 "hosts",
2219 2,
2220 vec![],
2221 ))
2222 .unwrap();
2223
2224 let rows = manager.get_by_kind("table");
2225 assert_eq!(rows.len(), 2);
2226
2227 let vectors = manager.get_by_kind("vector");
2228 assert_eq!(vectors.len(), 1);
2229 }
2230
2231 fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
2238 SegmentManager::with_config(
2239 "test",
2240 ManagerConfig {
2241 max_sealed_segments: 1_000,
2242 consolidation_entities_per_tick: entities_per_tick,
2243 ..Default::default()
2244 },
2245 )
2246 }
2247
2248 fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
2252 let entities: Vec<UnifiedEntity> = (0..count)
2253 .map(|i| {
2254 UnifiedEntity::table_row(
2255 manager.next_entity_id(),
2256 "users",
2257 i as u64,
2258 vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
2259 )
2260 })
2261 .collect();
2262 let ids = manager.bulk_insert(entities).expect("bulk insert");
2263 manager.force_seal().expect("seal");
2264 ids
2265 }
2266
2267 fn live_ids(manager: &SegmentManager) -> Vec<u64> {
2268 let mut ids: Vec<u64> = manager
2269 .query_all(|_| true)
2270 .into_iter()
2271 .map(|entity| entity.id.raw())
2272 .collect();
2273 ids.sort_unstable();
2274 ids
2275 }
2276
2277 fn drain_maintenance(manager: &SegmentManager) -> usize {
2278 let mut ticks = 0;
2279 while manager.consolidation.read().is_some() || ticks == 0 {
2280 manager.run_maintenance().expect("maintenance");
2281 ticks += 1;
2282 assert!(ticks < 1_000, "consolidation failed to converge");
2283 }
2284 ticks
2285 }
2286
2287 #[test]
2288 fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
2289 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
2290 let ids = seal_bulk_segment(&manager, 100);
2291
2292 let bytes_before_delete = manager.resident_bytes();
2293
2294 for id in ids.iter().take(40) {
2296 assert!(manager.delete(*id).expect("delete"));
2297 }
2298
2299 let bytes_with_tombstones = manager.resident_bytes();
2302 assert!(bytes_with_tombstones > bytes_before_delete);
2303
2304 let survivors = live_ids(&manager);
2305 assert_eq!(survivors.len(), 60);
2306
2307 drain_maintenance(&manager);
2308
2309 let stats = manager.stats();
2310 assert_eq!(stats.consolidation.runs_started, 1);
2311 assert_eq!(stats.consolidation.runs_completed, 1);
2312 assert_eq!(stats.consolidation.segments_merged, 1);
2313 assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
2314 assert!(stats.consolidation.bytes_reclaimed > 0);
2315
2316 let bytes_after = manager.resident_bytes();
2318 assert!(
2319 bytes_after <= bytes_before_delete,
2320 "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
2321 );
2322 assert_eq!(
2323 bytes_with_tombstones - bytes_after,
2324 stats.consolidation.bytes_reclaimed
2325 );
2326 assert_eq!(stats.total_memory_bytes, bytes_after as usize);
2327
2328 assert_eq!(live_ids(&manager), survivors);
2329 assert_eq!(manager.sealed.read().len(), 1);
2330 assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
2331 }
2332
2333 #[test]
2334 fn query_results_are_identical_before_during_and_after_consolidation() {
2335 let manager = consolidating_manager(3);
2336 let ids = seal_bulk_segment(&manager, 20);
2337 seal_bulk_segment(&manager, 20);
2338 for id in ids.iter().take(10) {
2339 assert!(manager.delete(*id).expect("delete"));
2340 }
2341
2342 let before = live_ids(&manager);
2343 let survivor = *ids.last().expect("at least one row");
2344 let tombstoned = ids[0];
2345
2346 let mut mid_tick_observations = 0;
2349 while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
2350 manager.run_maintenance().expect("maintenance");
2351 assert_eq!(live_ids(&manager), before);
2352 assert!(manager.get(survivor).is_some());
2353 assert!(manager.get(tombstoned).is_none());
2354 assert_eq!(manager.get_by_kind("table").len(), before.len());
2355 mid_tick_observations += 1;
2356 assert!(mid_tick_observations < 1_000);
2357 }
2358
2359 assert!(
2360 mid_tick_observations > 2,
2361 "expected the run to span several ticks, saw {mid_tick_observations}"
2362 );
2363 assert_eq!(live_ids(&manager), before);
2364 assert!(manager.get(survivor).is_some());
2365 assert!(manager.get(tombstoned).is_none());
2366 assert_eq!(manager.get_by_kind("table").len(), before.len());
2367 }
2368
2369 #[test]
2370 fn get_by_kind_reaches_sealed_segments() {
2371 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
2372 let ids = seal_bulk_segment(&manager, 5);
2373 assert!(manager.delete(ids[0]).expect("delete"));
2374
2375 assert_eq!(manager.get_by_kind("table").len(), 4);
2379 assert_eq!(manager.get_by_kind("vector").len(), 0);
2380 }
2381
2382 #[test]
2383 fn consolidation_is_paced_by_the_per_tick_entity_bound() {
2384 const PER_TICK: usize = 4;
2385 let manager = consolidating_manager(PER_TICK);
2386 let ids = seal_bulk_segment(&manager, 60);
2387 for id in ids.iter().take(30) {
2388 assert!(manager.delete(*id).expect("delete"));
2389 }
2390
2391 manager.run_maintenance().expect("maintenance");
2393 assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
2394
2395 let mut previous = PER_TICK as u64;
2398 let mut ticks = 1;
2399 while let Some(copied) = manager.consolidation_progress() {
2400 assert!(
2401 copied - previous <= PER_TICK as u64,
2402 "tick copied {} entities, bound is {PER_TICK}",
2403 copied - previous
2404 );
2405 previous = copied;
2406 manager.run_maintenance().expect("maintenance");
2407 ticks += 1;
2408 assert!(ticks < 100);
2409 }
2410
2411 assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
2414 assert_eq!(manager.stats().consolidation.runs_completed, 1);
2415 assert_eq!(live_ids(&manager).len(), 30);
2416 }
2417
2418 #[test]
2419 fn sealed_segment_count_above_the_limit_triggers_a_merge() {
2420 let manager = SegmentManager::with_config(
2421 "test",
2422 ManagerConfig {
2423 max_sealed_segments: 3,
2424 ..Default::default()
2425 },
2426 );
2427 for _ in 0..4 {
2428 seal_bulk_segment(&manager, 5);
2429 }
2430 assert_eq!(manager.sealed.read().len(), 4);
2431
2432 let before = live_ids(&manager);
2433 drain_maintenance(&manager);
2434
2435 assert_eq!(manager.sealed.read().len(), 1);
2438 assert_eq!(manager.stats().consolidation.segments_merged, 4);
2439 assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
2440 assert_eq!(live_ids(&manager), before);
2441 }
2442
2443 #[test]
2444 fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
2445 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
2446 seal_bulk_segment(&manager, 10);
2447 seal_bulk_segment(&manager, 10);
2448
2449 for _ in 0..5 {
2450 manager.run_maintenance().expect("maintenance");
2451 }
2452
2453 assert_eq!(manager.stats().consolidation.runs_started, 0);
2454 assert_eq!(manager.sealed.read().len(), 2);
2455 }
2456
2457 #[test]
2458 fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
2459 use std::sync::atomic::AtomicBool;
2460
2461 let manager = Arc::new(consolidating_manager(2));
2462 let ids = seal_bulk_segment(&manager, 40);
2463 seal_bulk_segment(&manager, 40);
2464 for id in ids.iter().take(20) {
2465 assert!(manager.delete(*id).expect("delete"));
2466 }
2467
2468 let expected = live_ids(&manager);
2469 let stop = Arc::new(AtomicBool::new(false));
2470
2471 let readers: Vec<_> = (0..4)
2472 .map(|_| {
2473 let manager = Arc::clone(&manager);
2474 let stop = Arc::clone(&stop);
2475 let expected = expected.clone();
2476 std::thread::spawn(move || {
2477 let mut scans = 0u64;
2478 loop {
2481 let mut seen: Vec<u64> = manager
2482 .query_all(|_| true)
2483 .into_iter()
2484 .map(|entity| entity.id.raw())
2485 .collect();
2486 seen.sort_unstable();
2487 assert_eq!(seen, expected, "reader saw a torn sealed set");
2490 scans += 1;
2491 if stop.load(Ordering::Relaxed) {
2492 break;
2493 }
2494 }
2495 scans
2496 })
2497 })
2498 .collect();
2499
2500 drain_maintenance(&manager);
2501 stop.store(true, Ordering::Relaxed);
2502
2503 for reader in readers {
2504 let scans = reader.join().expect("reader thread");
2505 assert!(scans > 0, "reader never got to scan");
2506 }
2507
2508 assert_eq!(manager.stats().consolidation.runs_completed, 1);
2509 assert_eq!(live_ids(&manager), expected);
2510 }
2511
2512 #[test]
2513 fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
2514 let manager = consolidating_manager(4);
2515 let ids = seal_bulk_segment(&manager, 20);
2516 for id in ids.iter().take(6) {
2517 assert!(manager.delete(*id).expect("delete"));
2518 }
2519
2520 manager.run_maintenance().expect("maintenance");
2522 assert_eq!(manager.consolidation_progress(), Some(4));
2523 let doomed = ids[6];
2524 assert!(manager.delete(doomed).expect("delete"));
2525
2526 let expected = live_ids(&manager);
2527 drain_maintenance(&manager);
2528
2529 assert!(manager.get(doomed).is_none(), "deleted row came back");
2530 assert_eq!(live_ids(&manager), expected);
2531 }
2532
2533 #[test]
2534 fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
2535 let manager = consolidating_manager(3);
2536 let ids = seal_bulk_segment(&manager, 30);
2537 for id in ids.iter().take(12) {
2538 assert!(manager.delete(*id).expect("delete"));
2539 }
2540
2541 let before = live_ids(&manager);
2542 let sealed_before: Vec<SegmentId> = manager
2543 .sealed
2544 .read()
2545 .iter()
2546 .map(|segment| segment.read().id())
2547 .collect();
2548
2549 manager.run_maintenance().expect("maintenance");
2550 manager.run_maintenance().expect("maintenance");
2551 assert!(manager.consolidation.read().is_some());
2552
2553 *manager.consolidation.write() = None;
2557
2558 assert_eq!(live_ids(&manager), before);
2559 let sealed_after: Vec<SegmentId> = manager
2560 .sealed
2561 .read()
2562 .iter()
2563 .map(|segment| segment.read().id())
2564 .collect();
2565 assert_eq!(sealed_after, sealed_before);
2566 assert_eq!(manager.stats().consolidation.runs_completed, 0);
2567
2568 drain_maintenance(&manager);
2569 assert_eq!(live_ids(&manager), before);
2570 assert_eq!(manager.stats().consolidation.runs_completed, 1);
2571 }
2572
2573 #[test]
2574 fn consolidation_preserves_entity_payload_metadata_and_sequence() {
2575 let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
2576
2577 let entities: Vec<UnifiedEntity> = (0..10)
2581 .map(|i| {
2582 UnifiedEntity::table_row(
2583 manager.next_entity_id(),
2584 "users",
2585 i as u64,
2586 vec![Value::Integer(i as i64)],
2587 )
2588 })
2589 .collect();
2590 let ids = manager.bulk_insert(entities).expect("bulk insert");
2591
2592 let mut meta = Metadata::new();
2593 meta.set(
2594 "os",
2595 super::super::metadata::MetadataValue::String("linux".to_string()),
2596 );
2597 manager.set_metadata(ids[9], meta).expect("set metadata");
2598 manager.force_seal().expect("seal");
2599
2600 assert!(manager.delete(ids[0]).expect("delete"));
2601 assert!(manager.delete(ids[1]).expect("delete"));
2602 assert!(manager.delete(ids[2]).expect("delete"));
2603
2604 let before: Vec<UnifiedEntity> = {
2605 let mut entities = manager.query_all(|_| true);
2606 entities.sort_by_key(|entity| entity.id.raw());
2607 entities
2608 };
2609
2610 drain_maintenance(&manager);
2611
2612 let after: Vec<UnifiedEntity> = {
2613 let mut entities = manager.query_all(|_| true);
2614 entities.sort_by_key(|entity| entity.id.raw());
2615 entities
2616 };
2617
2618 let columns = |entity: &UnifiedEntity| match &entity.data {
2619 super::super::entity::EntityData::Row(row) => row.columns.clone(),
2620 _ => Vec::new(),
2621 };
2622
2623 assert_eq!(before.len(), after.len());
2624 for (old, new) in before.iter().zip(after.iter()) {
2625 assert_eq!(old.id, new.id);
2626 assert_eq!(columns(old), columns(new));
2627 assert_eq!(
2628 old.sequence_id, new.sequence_id,
2629 "merge must not re-issue sequence ids"
2630 );
2631 }
2632 assert!(manager
2633 .get_metadata(ids[9])
2634 .expect("metadata survives the merge")
2635 .has("os"));
2636 }
2637
2638 #[test]
2639 #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
2640 fn test_lifecycle_events() {
2641 let manager = SegmentManager::new("test");
2642
2643 manager
2644 .insert(UnifiedEntity::vector(
2645 manager.next_entity_id(),
2646 "v",
2647 vec![0.1],
2648 ))
2649 .unwrap();
2650
2651 let events = manager.drain_events();
2652
2653 assert!(events
2655 .iter()
2656 .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
2657 assert!(events
2658 .iter()
2659 .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
2660 }
2661}