Skip to main content

reddb_server/storage/unified/
manager.rs

1//! Segment Manager
2//!
3//! Manages the lifecycle of unified segments: creation, sealing, consolidation,
4//! and archival. Coordinates writes to growing segments and queries across
5//! all segments.
6//!
7//! # Responsibilities
8//!
9//! - Route writes to the active growing segment
10//! - Auto-seal segments when thresholds are met
11//! - Coordinate queries across multiple segments
12//! - Paced consolidation of sealed segments (ADR 0073 §5)
13//! - Archive old segments to cold storage
14
15mod graph_index_merge;
16
17use parking_lot::RwLock;
18use std::collections::HashMap;
19use std::sync::atomic::{AtomicU64, Ordering};
20use std::sync::{Arc, Weak};
21
22use super::entity::{EntityId, EntityKind, UnifiedEntity};
23use super::metadata::{Metadata, MetadataFilter};
24use super::segment::{
25    GraphEntityKind, GrowingSegment, SegmentConfig, SegmentError, SegmentId, SegmentState,
26    SegmentStats, UnifiedSegment, ZoneColPred, ZoneColPredKind,
27};
28use super::visibility_map::VisibilityMap;
29use crate::runtime::mvcc::{entity_visible_with_context, SnapshotContext};
30
31/// Fraction of a collection's sealed entities that must be tombstoned before
32/// consolidation is worth its cost.
33///
34/// Tunable implementation constant, not a contract: it is surfaced read-only in
35/// `red.stats` so an operator can see what the engine decided, and it may move
36/// between releases without notice.
37pub const CONSOLIDATION_TOMBSTONE_RATIO: f64 = 0.20;
38
39/// Fraction of a collection's sealed bytes that consolidation would return to
40/// the budget (dead entity payloads plus the tombstone sets that bury them)
41/// before consolidation is worth its cost. See [`CONSOLIDATION_TOMBSTONE_RATIO`]
42/// on tunability.
43pub const CONSOLIDATION_FRAGMENTATION_RATIO: f64 = 0.30;
44
45/// Entities copied into the merged segment per maintenance tick.
46///
47/// ADR 0038 §3: consolidation is paced. A tick copies at most this many
48/// entities and returns; the half-built merged segment is carried to the next
49/// tick. No consolidation ever runs to completion in one unbounded pass.
50pub const CONSOLIDATION_ENTITIES_PER_TICK: usize = 4_096;
51
52/// Configuration for the segment manager
53#[derive(Debug, Clone)]
54pub struct ManagerConfig {
55    /// Segment configuration
56    pub segment_config: SegmentConfig,
57    /// Maximum number of sealed segments before consolidation
58    pub max_sealed_segments: usize,
59    /// Idle time (seconds) before auto-sealing
60    pub idle_seal_secs: u64,
61    /// Enable paced sealed-segment consolidation
62    pub enable_consolidation: bool,
63    /// Entities copied per maintenance tick while a consolidation is running
64    pub consolidation_entities_per_tick: usize,
65    /// Enable background archival
66    pub enable_archival: bool,
67    /// Age threshold for archival (seconds)
68    pub archive_age_secs: u64,
69}
70
71impl Default for ManagerConfig {
72    fn default() -> Self {
73        Self {
74            segment_config: SegmentConfig::default(),
75            max_sealed_segments: 10,
76            idle_seal_secs: 300, // 5 minutes
77            enable_consolidation: true,
78            consolidation_entities_per_tick: CONSOLIDATION_ENTITIES_PER_TICK,
79            enable_archival: true,
80            archive_age_secs: 86400 * 7, // 7 days
81        }
82    }
83}
84
85/// Consolidation counters for one collection (ADR 0073 §5).
86///
87/// These replace the `compact_ops` counter, which only ever counted a
88/// do-nothing branch.
89#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
90pub struct ConsolidationStats {
91    /// Consolidations that crossed a threshold and began copying.
92    pub runs_started: u64,
93    /// Consolidations that finished the swap.
94    pub runs_completed: u64,
95    /// Source segments retired by a completed swap.
96    pub segments_merged: u64,
97    /// Tombstones garbage-collected by completed swaps.
98    pub tombstones_reclaimed: u64,
99    /// Logical source-to-merged footprint reduction; readers may retain sources.
100    pub bytes_reclaimed: u64,
101}
102
103/// Manager statistics
104#[derive(Debug, Clone, Default)]
105pub struct ManagerStats {
106    /// Total entities across all segments
107    pub total_entities: usize,
108    /// Number of growing segments
109    pub growing_count: usize,
110    /// Number of sealed segments
111    pub sealed_count: usize,
112    /// Number of archived segments
113    pub archived_count: usize,
114    /// Total memory usage
115    pub total_memory_bytes: usize,
116    /// Number of seal operations
117    pub seal_ops: u64,
118    /// Consolidation counters
119    pub consolidation: ConsolidationStats,
120}
121
122#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
123pub struct SegmentScanStats {
124    pub segments_total: u64,
125    pub segments_scanned: u64,
126    pub segments_pruned: u64,
127}
128
129/// Lifecycle events for monitoring
130#[derive(Debug, Clone)]
131pub enum LifecycleEvent {
132    SegmentCreated(SegmentId),
133    SegmentSealed(SegmentId),
134    SegmentConsolidated {
135        source: Vec<SegmentId>,
136        target: SegmentId,
137    },
138    SegmentArchived(SegmentId),
139    EntityInserted(EntityId, SegmentId),
140    EntityDeleted(EntityId, SegmentId),
141}
142
143/// A consolidation in flight, carried across maintenance ticks.
144///
145/// The merged segment is built off to the side and is invisible to readers
146/// until the swap. Nothing here is durable: sealed segments are the truth (via
147/// the normal checkpoint path), so a crash mid-consolidation simply drops this
148/// state and the next boot starts over.
149struct Consolidation {
150    /// Source segments, frozen when the run started.
151    sources: Vec<SourceProgress>,
152    /// Id reserved for the merged segment.
153    merged_id: SegmentId,
154    /// The half-built merged segment.
155    merged: GrowingSegment,
156    /// Index into `sources` of the segment currently being copied.
157    cursor: usize,
158    /// Live ids of `sources[cursor]`, snapshotted when that source was opened.
159    pending_ids: Vec<EntityId>,
160    /// How many of `pending_ids` have been consumed.
161    pending_cursor: usize,
162    /// Entities copied so far, across every tick of this run.
163    entities_copied: u64,
164}
165
166/// Per-source bookkeeping for a consolidation in flight.
167struct SourceProgress {
168    id: SegmentId,
169    /// Ids copied out of this source into the merged segment.
170    copied: Vec<EntityId>,
171    /// The source's tombstone count when the run opened it, and `None` before
172    /// that. Tombstones only ever grow, so a source whose count is unchanged at
173    /// swap time saw no deletes after we started reading it — and every row we
174    /// copied out of it is still live. Anything else needs a per-row check.
175    ///
176    /// Recording this when the source is *opened*, not when it is drained, is
177    /// what makes the check sound: a row deleted between two copy ticks would
178    /// otherwise be invisible to the swap and get resurrected.
179    tombstones_at_open: Option<usize>,
180}
181
182/// Segment manager for a collection
183pub struct SegmentManager {
184    /// Statement-scoped constraint admission. Reentrant because SQL upserts
185    /// invoke the same typed mutation APIs while holding this collection gate.
186    row_constraint_lock: Arc<parking_lot::ReentrantMutex<()>>,
187    /// Collection name
188    collection: String,
189    /// Configuration
190    config: ManagerConfig,
191    /// Next segment ID counter
192    next_segment_id: AtomicU64,
193    /// Next entity ID counter
194    next_entity_id: AtomicU64,
195    /// Per-table auto-increment row ID (1, 2, 3... per collection)
196    next_row_id: AtomicU64,
197    /// Hot-path entity counter — lock-free, updated by every insert/delete.
198    /// Replaces stats.total_entities on the write path to eliminate a lock
199    /// acquisition per row (from 4 lock ops per insert → 2).
200    total_entities_atomic: AtomicU64,
201    /// Currently active growing segment
202    growing: RwLock<Option<Arc<RwLock<GrowingSegment>>>>,
203    /// Sealed segments (immutable, queryable)
204    sealed: RwLock<Vec<Arc<RwLock<GrowingSegment>>>>,
205    /// Retired sources still owned by readers. Weak handles never prolong payload lifetime.
206    retired: RwLock<Vec<Weak<RwLock<GrowingSegment>>>>,
207    /// Archived segment IDs (stored externally)
208    archived: RwLock<Vec<SegmentId>>,
209    /// Entity to segment mapping (for fast lookups by individually-inserted entities).
210    /// Bulk-inserted entities skip this map; their segment is found by sequential scan
211    /// of growing + sealed segments in get()/update()/delete().
212    entity_segment: RwLock<HashMap<EntityId, SegmentId>>,
213    /// Shared column schema: column_name → index in Vec<Value>.
214    /// Populated on first bulk_insert. Enables columnar storage (Vec instead of HashMap per row).
215    column_schema: RwLock<Option<Arc<Vec<String>>>>,
216    /// Statistics (slow path — not updated on every insert).
217    stats: RwLock<ManagerStats>,
218    /// Consolidation in flight, if any. Only `run_maintenance` touches it, so
219    /// at most one consolidation runs per collection at a time.
220    consolidation: RwLock<Option<Consolidation>>,
221    /// Event listeners (simplified - would be channels in production)
222    events: RwLock<Vec<LifecycleEvent>>,
223    /// Visibility map: sealed segment entity ranges marked as all-visible.
224    /// Growing segment is never all-visible (writes are in-flight).
225    /// Used by index-only scan decisions.
226    visibility_map: VisibilityMap,
227}
228
229impl SegmentManager {
230    /// Create a new segment manager
231    pub fn new(collection: impl Into<String>) -> Self {
232        Self::with_config(collection, ManagerConfig::default())
233    }
234
235    /// Create with custom configuration
236    pub fn with_config(collection: impl Into<String>, config: ManagerConfig) -> Self {
237        Self {
238            row_constraint_lock: Arc::new(parking_lot::ReentrantMutex::new(())),
239            collection: collection.into(),
240            config,
241            next_segment_id: AtomicU64::new(1),
242            next_entity_id: AtomicU64::new(1),
243            next_row_id: AtomicU64::new(1),
244            total_entities_atomic: AtomicU64::new(0),
245            growing: RwLock::new(None),
246            sealed: RwLock::new(Vec::new()),
247            retired: RwLock::new(Vec::new()),
248            archived: RwLock::new(Vec::new()),
249            entity_segment: RwLock::new(HashMap::new()),
250            column_schema: RwLock::new(None),
251            stats: RwLock::new(ManagerStats::default()),
252            consolidation: RwLock::new(None),
253            events: RwLock::new(Vec::new()),
254            visibility_map: VisibilityMap::new(),
255        }
256    }
257
258    pub(crate) fn row_constraint_lock(&self) -> Arc<parking_lot::ReentrantMutex<()>> {
259        Arc::clone(&self.row_constraint_lock)
260    }
261
262    /// Get or create the shared column schema from first row's named fields.
263    pub fn get_or_init_schema(
264        &self,
265        named: &HashMap<String, reddb_types::Value>,
266    ) -> Arc<Vec<String>> {
267        {
268            let schema = self.column_schema.read();
269            if let Some(ref s) = *schema {
270                return Arc::clone(s);
271            }
272        }
273        let cols: Vec<String> = named.keys().cloned().collect();
274        let arc = Arc::new(cols);
275        *self.column_schema.write() = Some(Arc::clone(&arc));
276        arc
277    }
278
279    /// Get the column schema if it exists.
280    pub fn column_schema(&self) -> Option<Arc<Vec<String>>> {
281        self.column_schema.read().clone()
282    }
283
284    pub(crate) fn set_column_schema_if_empty(&self, columns: Vec<String>) {
285        if columns.is_empty() {
286            return;
287        }
288        let mut schema = self.column_schema.write();
289        if schema.is_none() {
290            *schema = Some(Arc::new(columns));
291        }
292    }
293
294    /// Get collection name
295    pub fn collection(&self) -> &str {
296        &self.collection
297    }
298
299    /// Get configuration
300    pub fn config(&self) -> &ManagerConfig {
301        &self.config
302    }
303
304    /// Get statistics. Entity counts cover active segments; memory also covers
305    /// retained sources and unpublished consolidation work.
306    pub fn stats(&self) -> ManagerStats {
307        let mut s = self.stats.read().clone();
308        s.total_entities = self.total_entities_atomic.load(Ordering::Relaxed) as usize;
309        s.total_memory_bytes = self.resident_bytes() as usize;
310        s
311    }
312
313    /// Approximate resident segment memory, including tombstones, retired sources,
314    /// and the in-flight merge and its bookkeeping. This feeds budget admission.
315    pub fn resident_bytes(&self) -> u64 {
316        self.memory_usage().1
317    }
318
319    /// Resident segment payload/directory bytes and consolidation bookkeeping,
320    /// including retired sources. Unlike `resident_bytes`, excludes tombstones.
321    pub fn memory_bytes(&self) -> u64 {
322        self.memory_usage().0
323    }
324
325    /// Sample payload and resident bytes under the publication lock order:
326    /// consolidation, growing, sealed, retired. No entity scan or strong handle
327    /// is retained beyond this call; expired weak entries are pruned in place.
328    fn memory_usage(&self) -> (u64, u64) {
329        let consolidation = self.consolidation.read();
330        let growing = self.growing.read();
331        let sealed = self.sealed.read();
332        let mut retired = self.retired.write();
333        let mut payload = 0u64;
334        let mut resident = 0u64;
335        {
336            let mut add = |segment: &GrowingSegment| {
337                payload = payload.saturating_add(segment.memory_bytes());
338                resident = resident.saturating_add(segment.resident_bytes());
339            };
340            for segment in growing.iter().chain(sealed.iter()) {
341                add(&segment.read());
342            }
343            retired.retain(|source| {
344                if let Some(source) = source.upgrade() {
345                    add(&source.read());
346                    true
347                } else {
348                    false
349                }
350            });
351            if let Some(run) = consolidation.as_ref() {
352                add(&run.merged);
353            }
354        }
355        let mut bookkeeping =
356            (retired.capacity() * std::mem::size_of::<Weak<RwLock<GrowingSegment>>>()) as u64;
357        if let Some(run) = consolidation.as_ref() {
358            bookkeeping = bookkeeping
359                .saturating_add(
360                    (run.sources.capacity() * std::mem::size_of::<SourceProgress>()) as u64,
361                )
362                .saturating_add(
363                    (run.pending_ids.capacity() * std::mem::size_of::<EntityId>()) as u64,
364                );
365            for source in &run.sources {
366                bookkeeping = bookkeeping.saturating_add(
367                    (source.copied.capacity() * std::mem::size_of::<EntityId>()) as u64,
368                );
369            }
370        }
371        (
372            payload.saturating_add(bookkeeping),
373            resident.saturating_add(bookkeeping),
374        )
375    }
376
377    /// Bytes a pressure-triggered consolidation could plausibly return.
378    pub fn reclaimable_bytes(&self) -> u64 {
379        self.sealed
380            .read()
381            .iter()
382            .map(|segment| segment.read().reclaimable_bytes())
383            .fold(0, u64::saturating_add)
384    }
385
386    /// Generate a new entity ID
387    pub fn next_entity_id(&self) -> EntityId {
388        EntityId::new(self.next_entity_id.fetch_add(1, Ordering::SeqCst))
389    }
390
391    /// Generate a per-table sequential row ID (1, 2, 3... per collection)
392    pub fn next_row_id(&self) -> u64 {
393        self.next_row_id.fetch_add(1, Ordering::SeqCst)
394    }
395
396    /// Reserve `n` contiguous per-table row IDs with one atomic
397    /// fetch_add. Caller assigns `row_id = start + i` per entity.
398    /// Saves N-1 atomic RMWs on bulk inserts (25k atomics → 1).
399    pub fn reserve_row_ids(&self, n: u64) -> std::ops::Range<u64> {
400        let start = self.next_row_id.fetch_add(n, Ordering::SeqCst);
401        start..start + n
402    }
403
404    /// Advance the per-table row_id counter to at least `id + 1`.
405    /// Called during load to restore the counter from existing data.
406    pub fn register_row_id(&self, id: u64) {
407        let candidate = id.saturating_add(1);
408        let mut current = self.next_row_id.load(Ordering::SeqCst);
409        while candidate > current {
410            match self.next_row_id.compare_exchange(
411                current,
412                candidate,
413                Ordering::SeqCst,
414                Ordering::SeqCst,
415            ) {
416                Ok(_) => break,
417                Err(updated) => current = updated,
418            }
419        }
420    }
421
422    /// Get or create the active growing segment.
423    ///
424    /// Fast path: read lock only — no write contention when the segment already exists.
425    /// Concurrent writers each clone the `Arc` under a shared read lock, then compete
426    /// on the segment's own write lock. This eliminates the global write-lock serialisation
427    /// that previously throttled concurrent inserts to ~238 ops/s.
428    fn get_or_create_growing(&self) -> Arc<RwLock<GrowingSegment>> {
429        // Common case: segment already exists — shared read lock, zero contention.
430        {
431            let growing = self.growing.read();
432            if let Some(segment) = growing.as_ref() {
433                return Arc::clone(segment);
434            }
435        }
436
437        // Slow path: segment missing — take exclusive write lock to create it.
438        let mut growing = self.growing.write();
439        // Double-check: another thread may have created it between the two lock acquisitions.
440        if let Some(segment) = growing.as_ref() {
441            return Arc::clone(segment);
442        }
443
444        let id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
445        let segment = GrowingSegment::new(id, &self.collection);
446        let segment_arc = Arc::new(RwLock::new(segment));
447        *growing = Some(Arc::clone(&segment_arc));
448
449        self.emit(LifecycleEvent::SegmentCreated(id));
450
451        // Update growing_count in the slow-path stats struct.
452        // This is the rare segment-creation path — locking is fine here.
453        self.stats.write().growing_count += 1;
454
455        segment_arc
456    }
457
458    /// Insert a new entity
459    pub fn insert(&self, mut entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
460        // Check if we need to seal the current segment first
461        self.maybe_seal_growing()?;
462
463        let segment_arc = self.get_or_create_growing();
464        let mut segment = segment_arc.write();
465
466        // Assign entity ID if not already set
467        if entity.id.raw() == 0 {
468            entity.id = self.next_entity_id();
469        }
470
471        let entity_id = entity.id;
472        let segment_id = segment.id();
473
474        segment.insert(entity)?;
475
476        // Lock-free counter update — eliminates the stats write lock on the hot path.
477        self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
478
479        // entity_segment map is intentionally NOT updated here.
480        // update() and update_hot() first probe the growing segment directly
481        // (growing.contains(entity.id)) before consulting this map, so entities
482        // that were just inserted are found without entity_segment. The map is
483        // only consulted for entities that may have been moved to sealed segments,
484        // which can't be updated anyway (state().is_writable() == false).
485        // Skipping this write removes one exclusive HashMap lock per insert.
486
487        self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
488
489        Ok(entity_id)
490    }
491
492    /// Restore a snapshot image without allocating a new sequence identity.
493    pub(crate) fn restore_snapshot_entity(
494        &self,
495        entity: UnifiedEntity,
496    ) -> Result<EntityId, SegmentError> {
497        self.maybe_seal_growing()?;
498        if self.get(entity.id).is_some() {
499            return Err(SegmentError::AlreadyExists(entity.id));
500        }
501        let segment_arc = self.get_or_create_growing();
502        let mut segment = segment_arc.write();
503        let entity_id = entity.id;
504        let segment_id = segment.id();
505        // Adoption rebuilds derived indexes and advances the allocator past
506        // the stored sequence, including gaps and the legitimate value zero.
507        segment.adopt_entity(entity, None);
508        self.total_entities_atomic.fetch_add(1, Ordering::Relaxed);
509        self.emit(LifecycleEvent::EntityInserted(entity_id, segment_id));
510        Ok(entity_id)
511    }
512
513    /// Insert multiple entities (batch) — sequential, one lock per item.
514    pub fn insert_batch(
515        &self,
516        entities: Vec<UnifiedEntity>,
517    ) -> Result<Vec<EntityId>, SegmentError> {
518        let mut ids = Vec::with_capacity(entities.len());
519        for entity in entities {
520            ids.push(self.insert(entity)?);
521        }
522        Ok(ids)
523    }
524
525    /// Turbo bulk insert — single lock acquisition for the entire batch.
526    /// Skips bloom filter and cross-ref indexing for maximum speed.
527    pub fn bulk_insert(
528        &self,
529        mut entities: Vec<UnifiedEntity>,
530    ) -> Result<Vec<EntityId>, SegmentError> {
531        // Assign IDs and per-table row_ids.
532        for entity in &mut entities {
533            if entity.id.raw() == 0 {
534                entity.id = self.next_entity_id();
535            }
536            if let super::entity::EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
537                if *row_id == 0 {
538                    *row_id = self.next_row_id();
539                } else {
540                    self.register_row_id(*row_id);
541                }
542            }
543        }
544
545        // Convert named HashMap → positional Vec (compact memory representation)
546        // The schema (column order) is shared across all rows in the collection.
547        if let Some(first_row) = entities.first() {
548            if let super::entity::EntityData::Row(ref row) = first_row.data {
549                if let Some(ref named) = row.named {
550                    let schema = self.get_or_init_schema(named);
551                    for entity in &mut entities {
552                        if let super::entity::EntityData::Row(ref mut row) = entity.data {
553                            if let Some(named) = row.named.take() {
554                                let mut cols = Vec::with_capacity(schema.len());
555                                for col_name in schema.iter() {
556                                    cols.push(
557                                        named
558                                            .get(col_name)
559                                            .cloned()
560                                            .unwrap_or(reddb_types::Value::Null),
561                                    );
562                                }
563                                row.columns = cols;
564                                row.schema = Some(Arc::clone(&schema));
565                            }
566                        }
567                    }
568                }
569            }
570        }
571
572        let segment_arc = self.get_or_create_growing();
573        let mut segment = segment_arc.write();
574        let segment_id = segment.id();
575
576        // Single call to GrowingSegment.bulk_insert (one lock, no bloom/cross-refs)
577        let ids = segment.bulk_insert(entities)?;
578
579        // Skip entity-segment mapping for bulk inserts (saves ~56 bytes/entity).
580        // The get() method scans growing+sealed segments directly.
581
582        // Lock-free batch counter update.
583        self.total_entities_atomic
584            .fetch_add(ids.len() as u64, Ordering::Relaxed);
585
586        Ok(ids)
587    }
588
589    /// Get an entity by ID — scans growing then sealed segments.
590    pub fn get(&self, id: EntityId) -> Option<UnifiedEntity> {
591        self.get_with(id, Clone::clone)
592    }
593
594    /// Inspect under the segment lock, for example to bound a payload clone.
595    /// The callback must not re-enter storage or runtime memory admission.
596    pub(crate) fn get_with<T>(
597        &self,
598        id: EntityId,
599        visit: impl FnOnce(&UnifiedEntity) -> T,
600    ) -> Option<T> {
601        // Growing segment first (most likely for recent inserts)
602        if let Some(growing_arc) = self.growing.read().as_ref() {
603            let growing = growing_arc.read();
604            if let Some(entity) = growing.get(id) {
605                return Some(visit(entity));
606            }
607        }
608
609        // Then sealed segments
610        let sealed = self.sealed.read();
611        for segment in sealed.iter() {
612            let seg = segment.read();
613            if let Some(entity) = seg.get(id) {
614                return Some(visit(entity));
615            }
616        }
617
618        None
619    }
620
621    /// Resolve a physical graph node reference without cloning historical payloads.
622    /// This deliberately ignores MVCC: callers must authorize the visible logical
623    /// node separately. The callback bounds every segment probe, including misses.
624    pub(crate) fn graph_node_logical_id(
625        &self,
626        id: EntityId,
627        mut before_segment: impl FnMut() -> bool,
628    ) -> Option<EntityId> {
629        if let Some(growing_arc) = self.growing.read().as_ref() {
630            if !before_segment() {
631                return None;
632            }
633            let growing = growing_arc.read();
634            if let Some(entity) = growing.get(id) {
635                return matches!(entity.kind, EntityKind::GraphNode(_))
636                    .then(|| entity.logical_id());
637            }
638        }
639        let sealed = self.sealed.read();
640        for segment in sealed.iter() {
641            if !before_segment() {
642                return None;
643            }
644            let segment = segment.read();
645            if let Some(entity) = segment.get(id) {
646                return matches!(entity.kind, EntityKind::GraphNode(_))
647                    .then(|| entity.logical_id());
648            }
649        }
650        None
651    }
652
653    /// Pin the source inventory across sealing/consolidation. Consumers hydrate
654    /// IDs through the current manager and run outside all storage locks.
655    /// Admission is also outside storage locks; a growing inventory retries only
656    /// its small handle array. Each indexed candidate buffer has a fixed size.
657    pub(crate) fn visit_graph_index_batches<E>(
658        &self,
659        kind: GraphEntityKind,
660        keys: &[EntityId],
661        mut before_work: impl FnMut() -> bool,
662        mut admit: impl FnMut(usize) -> Result<(), E>,
663        mut consume: impl FnMut(&[EntityId]) -> Result<bool, E>,
664    ) -> Result<bool, E> {
665        use super::segment::{GraphIndexCursor, SCAN_BATCH_SIZE};
666        admit(std::mem::size_of::<[EntityId; SCAN_BATCH_SIZE]>())?;
667        let mut admitted = 0;
668        let sources = loop {
669            let count = {
670                let growing = self.growing.read();
671                let sealed = self.sealed.read();
672                sealed.len() + usize::from(growing.is_some())
673            };
674            let bytes = count.saturating_mul(std::mem::size_of::<(
675                Arc<RwLock<GrowingSegment>>,
676                GraphIndexCursor,
677            )>());
678            admit(bytes.saturating_sub(admitted))?;
679            admitted = admitted.max(bytes);
680            let mut sources = Vec::with_capacity(count);
681            let growing = self.growing.read();
682            let sealed = self.sealed.read();
683            if sealed.len() + usize::from(growing.is_some()) > sources.capacity() {
684                continue;
685            }
686            for segment in growing.iter().chain(sealed.iter()) {
687                if !before_work() {
688                    return Ok(false);
689                }
690                let view = segment.read();
691                if !view.may_contain_graph_kind(kind) {
692                    continue;
693                }
694                let Some(cursor) = view.graph_index_cursor(kind, keys, &mut before_work) else {
695                    return Ok(false);
696                };
697                sources.push((Arc::clone(segment), cursor));
698            }
699            break sources;
700        };
701        let mut ids = [EntityId::new(0); SCAN_BATCH_SIZE];
702        for (source, mut cursor) in sources {
703            while !cursor.finished(keys) {
704                let mut count = 0;
705                {
706                    let segment = source.read();
707                    if !segment.may_contain_graph_kind(kind) {
708                        break;
709                    }
710                    if !segment.graph_index_batch(
711                        kind,
712                        keys,
713                        &mut cursor,
714                        &mut before_work,
715                        &mut |id| {
716                            ids[count] = id;
717                            count += 1;
718                        },
719                    ) {
720                        return Ok(false);
721                    }
722                }
723                if !before_work() || !consume(&ids[..count])? {
724                    return Ok(false);
725                }
726            }
727        }
728        Ok(true)
729    }
730
731    /// Batch-fetch multiple entities by ID in a single lock acquisition per segment.
732    ///
733    /// For indexed-scan result sets (up to ~5000 ids from range/bitmap lookup) this
734    /// is 2-3 lock acquisitions total vs N×3 with individual `get()` calls.
735    pub fn get_many(&self, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
736        let mut out: Vec<Option<UnifiedEntity>> = vec![None; ids.len()];
737        let mut remaining: Vec<usize> = (0..ids.len()).collect(); // indices still unfound
738
739        // Growing segment — one read lock for the entire batch.
740        // Non-blocking first: if a writer is active, fall back to blocking.
741        if let Some(growing_arc) = self.growing.read().as_ref() {
742            let growing = if let Some(g) = growing_arc.try_read() {
743                g
744            } else {
745                growing_arc.read()
746            };
747            remaining.retain(|&i| {
748                if let Some(entity) = growing.get(ids[i]) {
749                    out[i] = Some(entity.clone());
750                    false // remove from remaining
751                } else {
752                    true // keep — not found yet
753                }
754            });
755        }
756
757        if remaining.is_empty() {
758            return out;
759        }
760
761        // Sealed segments — one read lock per segment
762        let sealed = self.sealed.read();
763        for segment in sealed.iter() {
764            if remaining.is_empty() {
765                break;
766            }
767            let seg = segment.read();
768            remaining.retain(|&i| {
769                if let Some(entity) = seg.get(ids[i]) {
770                    out[i] = Some(entity.clone());
771                    false
772                } else {
773                    true
774                }
775            });
776        }
777
778        out
779    }
780
781    /// Visitor-pattern batch fetch. Invokes `f(&UnifiedEntity, usize_index)`
782    /// for each id that resolves, never cloning the entity.
783    ///
784    /// Used by scan hot paths (select_range, select_filtered) that
785    /// materialize each entity into an output record and don't need
786    /// an owned `UnifiedEntity`. Eliminates ~20% of scan CPU spent in
787    /// `UnifiedEntity::clone` when `get_batch` is followed by
788    /// `runtime_table_record_lean(entity)`.
789    ///
790    /// The closure runs while the segment read lock is held, so it
791    /// must be short — avoid doing I/O or taking unrelated locks in
792    /// `f`.
793    pub fn for_each_id<F>(&self, ids: &[EntityId], mut f: F)
794    where
795        F: FnMut(usize, &UnifiedEntity),
796    {
797        // Thread-local scratch buffer for the "pending" index list.
798        // Previous code allocated a fresh `Vec<usize>` of capacity
799        // N on every call — 4200 × 1000 queries / scenario on the
800        // select_range bench path. Take-and-restore pattern (vs
801        // RefCell::borrow_mut) so user closures that recurse into
802        // another `for_each_id` don't panic on a re-borrow; worst
803        // case they allocate a fresh buffer and we lose the caching
804        // win for that nested call.
805        thread_local! {
806            static REMAINING_SCRATCH: std::cell::Cell<Vec<usize>> =
807                const { std::cell::Cell::new(Vec::new()) };
808        }
809
810        let mut remaining: Vec<usize> = REMAINING_SCRATCH.with(|cell| cell.take());
811        remaining.clear();
812
813        if let Some(growing_arc) = self.growing.read().as_ref() {
814            let growing = if let Some(g) = growing_arc.try_read() {
815                g
816            } else {
817                growing_arc.read()
818            };
819            for (i, id) in ids.iter().enumerate() {
820                if let Some(entity) = growing.get(*id) {
821                    f(i, entity);
822                } else {
823                    remaining.push(i);
824                }
825            }
826        } else {
827            remaining.reserve(ids.len());
828            remaining.extend(0..ids.len());
829        }
830
831        if !remaining.is_empty() {
832            let sealed = self.sealed.read();
833            for segment in sealed.iter() {
834                if remaining.is_empty() {
835                    break;
836                }
837                let seg = segment.read();
838                remaining.retain(|&i| {
839                    if let Some(entity) = seg.get(ids[i]) {
840                        f(i, entity);
841                        false
842                    } else {
843                        true
844                    }
845                });
846            }
847        }
848
849        REMAINING_SCRATCH.with(|cell| cell.set(remaining));
850    }
851
852    /// Scan all segments for an entity
853    fn scan_for_entity(&self, id: EntityId) -> Option<UnifiedEntity> {
854        // Check growing
855        if let Some(growing_arc) = self.growing.read().as_ref() {
856            let growing = growing_arc.read();
857            if let Some(entity) = growing.get(id) {
858                return Some(entity.clone());
859            }
860        }
861
862        // Check sealed
863        let sealed = self.sealed.read();
864        for segment in sealed.iter() {
865            if let Some(entity) = segment.get(id) {
866                return Some(entity.clone());
867            }
868        }
869
870        None
871    }
872
873    fn find_sealed_segment_arc(&self, id: EntityId) -> Option<Arc<RwLock<GrowingSegment>>> {
874        let sealed = self.sealed.read();
875        sealed
876            .iter()
877            .find(|segment_arc| segment_arc.read().contains(id))
878            .map(Arc::clone)
879    }
880
881    fn rewrite_sealed_entity_into_growing(
882        &self,
883        entity: UnifiedEntity,
884        metadata: Option<&Metadata>,
885    ) -> Result<(), SegmentError> {
886        let entity_id = entity.id;
887        let sealed_arc = self
888            .find_sealed_segment_arc(entity_id)
889            .ok_or(SegmentError::NotFound(entity_id))?;
890
891        let metadata_to_apply = {
892            let mut sealed = sealed_arc.write();
893            let existing_metadata = sealed.get_metadata(entity_id);
894            if !sealed.force_delete(entity_id) {
895                return Err(SegmentError::NotFound(entity_id));
896            }
897            metadata.cloned().or(existing_metadata)
898        };
899
900        let growing_arc = self.get_or_create_growing();
901        let growing_id = {
902            let mut growing = growing_arc.write();
903            growing.insert(entity)?;
904            if let Some(metadata) = metadata_to_apply {
905                growing.set_metadata(entity_id, metadata)?;
906            }
907            growing.id()
908        };
909
910        self.entity_segment.write().insert(entity_id, growing_id);
911        Ok(())
912    }
913
914    /// Update an entity
915    pub fn update(&self, entity: UnifiedEntity) -> Result<(), SegmentError> {
916        let entity_id = entity.id;
917        let mut entity = Some(entity);
918
919        // Try growing segment directly (covers bulk-inserted entities without entity_segment map)
920        if let Some(growing_arc) = self.growing.read().as_ref() {
921            let mut growing = growing_arc.write();
922            if growing.contains(entity_id) && growing.state().is_writable() {
923                return growing.update(entity.take().expect("entity already moved"));
924            }
925        }
926
927        // Try entity_segment mapping for individually inserted entities
928        let segment_id = self.entity_segment.read().get(&entity_id).copied();
929        if let Some(seg_id) = segment_id {
930            if let Some(growing_arc) = self.growing.read().as_ref() {
931                let mut growing = growing_arc.write();
932                if growing.id() == seg_id && growing.state().is_writable() {
933                    return growing.update(entity.take().expect("entity already moved"));
934                }
935            }
936        }
937
938        if let Some(entity) = entity.take() {
939            return self.rewrite_sealed_entity_into_growing(entity, None);
940        }
941
942        Err(SegmentError::NotFound(entity_id))
943    }
944
945    /// Update an entity and, optionally, replace its metadata while holding
946    /// the segment write lock only once.
947    pub fn update_with_metadata(
948        &self,
949        entity: UnifiedEntity,
950        metadata: Option<&Metadata>,
951    ) -> Result<(), SegmentError> {
952        let entity_id = entity.id;
953        let mut entity = Some(entity);
954
955        // Try growing segment directly (covers bulk-inserted entities without entity_segment map)
956        if let Some(growing_arc) = self.growing.read().as_ref() {
957            let mut growing = growing_arc.write();
958            if growing.contains(entity_id) && growing.state().is_writable() {
959                growing.update(entity.take().expect("entity already moved"))?;
960                if let Some(metadata) = metadata {
961                    growing.set_metadata(entity_id, metadata.clone())?;
962                }
963                return Ok(());
964            }
965        }
966
967        // Try entity_segment mapping for individually inserted entities
968        let segment_id = self.entity_segment.read().get(&entity_id).copied();
969        if let Some(seg_id) = segment_id {
970            if let Some(growing_arc) = self.growing.read().as_ref() {
971                let mut growing = growing_arc.write();
972                if growing.id() == seg_id && growing.state().is_writable() {
973                    growing.update(entity.take().expect("entity already moved"))?;
974                    if let Some(metadata) = metadata {
975                        growing.set_metadata(entity_id, metadata.clone())?;
976                    }
977                    return Ok(());
978                }
979            }
980        }
981
982        if let Some(entity) = entity.take() {
983            return self.rewrite_sealed_entity_into_growing(entity, metadata);
984        }
985
986        Err(SegmentError::NotFound(entity_id))
987    }
988
989    /// HOT-update: like update but skips index work for unchanged columns.
990    /// `modified_columns` is the list of column names actually changed by the
991    /// UPDATE statement — lets us skip pk_index and cross_ref when safe.
992    pub fn update_hot(
993        &self,
994        entity: UnifiedEntity,
995        modified_columns: &[String],
996    ) -> Result<(), SegmentError> {
997        let entity_id = entity.id;
998        let mut entity = Some(entity);
999
1000        if let Some(growing_arc) = self.growing.read().as_ref() {
1001            let mut growing = growing_arc.write();
1002            if growing.contains(entity_id) && growing.state().is_writable() {
1003                return growing.update_hot(
1004                    entity.take().expect("entity already moved"),
1005                    modified_columns,
1006                );
1007            }
1008        }
1009
1010        let segment_id = self.entity_segment.read().get(&entity_id).copied();
1011        if let Some(seg_id) = segment_id {
1012            if let Some(growing_arc) = self.growing.read().as_ref() {
1013                let mut growing = growing_arc.write();
1014                if growing.id() == seg_id && growing.state().is_writable() {
1015                    return growing.update_hot(
1016                        entity.take().expect("entity already moved"),
1017                        modified_columns,
1018                    );
1019                }
1020            }
1021        }
1022
1023        if let Some(entity) = entity.take() {
1024            return self.rewrite_sealed_entity_into_growing(entity, None);
1025        }
1026
1027        Err(SegmentError::NotFound(entity_id))
1028    }
1029
1030    /// HOT-update an entity and, optionally, replace its metadata while
1031    /// holding the segment write lock only once.
1032    pub fn update_hot_with_metadata(
1033        &self,
1034        entity: UnifiedEntity,
1035        modified_columns: &[String],
1036        metadata: Option<&Metadata>,
1037    ) -> Result<(), SegmentError> {
1038        let entity_id = entity.id;
1039        let mut entity = Some(entity);
1040
1041        if let Some(growing_arc) = self.growing.read().as_ref() {
1042            let mut growing = growing_arc.write();
1043            if growing.contains(entity_id) && growing.state().is_writable() {
1044                growing.update_hot(
1045                    entity.take().expect("entity already moved"),
1046                    modified_columns,
1047                )?;
1048                if let Some(metadata) = metadata {
1049                    growing.set_metadata(entity_id, metadata.clone())?;
1050                }
1051                return Ok(());
1052            }
1053        }
1054
1055        let segment_id = self.entity_segment.read().get(&entity_id).copied();
1056        if let Some(seg_id) = segment_id {
1057            if let Some(growing_arc) = self.growing.read().as_ref() {
1058                let mut growing = growing_arc.write();
1059                if growing.id() == seg_id && growing.state().is_writable() {
1060                    growing.update_hot(
1061                        entity.take().expect("entity already moved"),
1062                        modified_columns,
1063                    )?;
1064                    if let Some(metadata) = metadata {
1065                        growing.set_metadata(entity_id, metadata.clone())?;
1066                    }
1067                    return Ok(());
1068                }
1069            }
1070        }
1071
1072        if let Some(entity) = entity.take() {
1073            return self.rewrite_sealed_entity_into_growing(entity, metadata);
1074        }
1075
1076        Err(SegmentError::NotFound(entity_id))
1077    }
1078
1079    /// Batch HOT-update multiple entities while holding the growing-segment
1080    /// write lock only once when possible.
1081    pub fn update_hot_batch_with_metadata<'a, I>(&self, items: I) -> Result<(), SegmentError>
1082    where
1083        I: IntoIterator<Item = (&'a UnifiedEntity, &'a [String], Option<&'a Metadata>)>,
1084    {
1085        let items: Vec<(&UnifiedEntity, &[String], Option<&Metadata>)> =
1086            items.into_iter().collect();
1087        if items.is_empty() {
1088            return Ok(());
1089        }
1090
1091        if let Some(growing_arc) = self.growing.read().as_ref() {
1092            let mut growing = growing_arc.write();
1093            if growing.state().is_writable() {
1094                match growing.update_hot_batch_with_metadata(items.iter().copied()) {
1095                    Ok(()) => return Ok(()),
1096                    Err(SegmentError::NotFound(_)) => {}
1097                    Err(other) => return Err(other),
1098                }
1099            }
1100        }
1101
1102        for (entity, modified_columns, metadata) in items {
1103            self.update_hot_with_metadata(entity.clone(), modified_columns, metadata)?;
1104        }
1105        Ok(())
1106    }
1107
1108    /// Delete an entity
1109    pub fn delete(&self, id: EntityId) -> Result<bool, SegmentError> {
1110        // Fast path: probe the growing segment directly — covers entities inserted via
1111        // insert() which no longer writes to entity_segment, and bulk-inserted entities.
1112        if let Some(growing_arc) = self.growing.read().as_ref() {
1113            let mut growing = growing_arc.write();
1114            if growing.contains(id) && growing.state().is_writable() {
1115                let seg_id = growing.id();
1116                let deleted = growing.delete(id)?;
1117                if deleted {
1118                    self.entity_segment.write().remove(&id);
1119                    self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1120                    self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1121                }
1122                return Ok(deleted);
1123            }
1124        }
1125
1126        // Fallback: check entity_segment map (populated for older insert() paths
1127        // or entities that were in a previous growing segment).
1128        let segment_id = self.entity_segment.read().get(&id).copied();
1129
1130        if let Some(seg_id) = segment_id {
1131            if let Some(growing_arc) = self.growing.read().as_ref() {
1132                let mut growing = growing_arc.write();
1133                if growing.id() == seg_id && growing.state().is_writable() {
1134                    let deleted = growing.delete(id)?;
1135                    if deleted {
1136                        self.entity_segment.write().remove(&id);
1137                        self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1138                        self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1139                    }
1140                    return Ok(deleted);
1141                }
1142            }
1143        }
1144
1145        // Fallback: entity is in a sealed segment (bulk-inserted, not in entity_segment map).
1146        // Single write-lock per segment to avoid TOCTOU race between contains() and force_delete().
1147        {
1148            let sealed = self.sealed.read();
1149            for segment_arc in sealed.iter() {
1150                let mut seg = segment_arc.write();
1151                let seg_id = seg.id();
1152                if seg.contains(id) {
1153                    let deleted = seg.force_delete(id);
1154                    drop(seg);
1155                    if deleted {
1156                        self.entity_segment.write().remove(&id);
1157                        self.total_entities_atomic.fetch_sub(1, Ordering::Relaxed);
1158                        self.emit(LifecycleEvent::EntityDeleted(id, seg_id));
1159                    }
1160                    return Ok(deleted);
1161                }
1162            }
1163        }
1164
1165        Ok(false)
1166    }
1167
1168    pub fn delete_batch(&self, ids: &[EntityId]) -> Result<Vec<EntityId>, SegmentError> {
1169        if ids.is_empty() {
1170            return Ok(Vec::new());
1171        }
1172
1173        let mut deleted_ids = Vec::with_capacity(ids.len());
1174
1175        if let Some(growing_arc) = self.growing.read().as_ref() {
1176            let mut growing = growing_arc.write();
1177            if growing.state().is_writable() {
1178                let seg_id = growing.id();
1179                let deleted = growing.delete_batch(ids)?;
1180                if !deleted.is_empty() {
1181                    {
1182                        let mut entity_segment = self.entity_segment.write();
1183                        for id in &deleted {
1184                            entity_segment.remove(id);
1185                        }
1186                    }
1187                    self.total_entities_atomic
1188                        .fetch_sub(deleted.len() as u64, Ordering::Relaxed);
1189                    for id in &deleted {
1190                        self.emit(LifecycleEvent::EntityDeleted(*id, seg_id));
1191                    }
1192                    deleted_ids.extend(deleted);
1193                }
1194            }
1195        }
1196
1197        if deleted_ids.len() == ids.len() {
1198            return Ok(deleted_ids);
1199        }
1200
1201        let deleted_set: std::collections::HashSet<EntityId> =
1202            deleted_ids.iter().copied().collect();
1203        for &id in ids {
1204            if deleted_set.contains(&id) {
1205                continue;
1206            }
1207            if self.delete(id)? {
1208                deleted_ids.push(id);
1209            }
1210        }
1211
1212        Ok(deleted_ids)
1213    }
1214
1215    /// Get metadata for an entity
1216    pub fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
1217        // Fast path: probe growing segment directly (no entity_segment needed).
1218        if let Some(growing_arc) = self.growing.read().as_ref() {
1219            let growing = growing_arc.read();
1220            if growing.contains(id) {
1221                return growing.get_metadata(id);
1222            }
1223        }
1224
1225        // Fallback: entity_segment map (for pre-existing or sealed entities)
1226        let segment_id = self.entity_segment.read().get(&id).copied();
1227
1228        if let Some(seg_id) = segment_id {
1229            if let Some(growing_arc) = self.growing.read().as_ref() {
1230                let growing = growing_arc.read();
1231                if growing.id() == seg_id {
1232                    return growing.get_metadata(id);
1233                }
1234            }
1235
1236            let sealed = self.sealed.read();
1237            for segment in sealed.iter() {
1238                if segment.id() == seg_id {
1239                    return segment.get_metadata(id);
1240                }
1241            }
1242        }
1243
1244        if let Some(segment_arc) = self.find_sealed_segment_arc(id) {
1245            return segment_arc.read().get_metadata(id);
1246        }
1247
1248        None
1249    }
1250
1251    /// Set metadata for an entity
1252    pub fn set_metadata(&self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
1253        // Fast path: probe growing segment directly — covers entities inserted via
1254        // insert() which no longer writes to entity_segment.
1255        if let Some(growing_arc) = self.growing.read().as_ref() {
1256            let mut growing = growing_arc.write();
1257            if growing.contains(id) && growing.state().is_writable() {
1258                return growing.set_metadata(id, metadata);
1259            }
1260        }
1261
1262        // Fallback: entity_segment map (sealed or pre-atomic-path entities)
1263        let segment_id = self.entity_segment.read().get(&id).copied();
1264
1265        if let Some(seg_id) = segment_id {
1266            if let Some(growing_arc) = self.growing.read().as_ref() {
1267                let mut growing = growing_arc.write();
1268                if growing.id() == seg_id && growing.state().is_writable() {
1269                    return growing.set_metadata(id, metadata);
1270                }
1271            }
1272        }
1273
1274        if let Some(entity) = self.get(id) {
1275            return self.rewrite_sealed_entity_into_growing(entity, Some(&metadata));
1276        }
1277
1278        Err(SegmentError::NotFound(id))
1279    }
1280
1281    /// Check if growing segment should be sealed
1282    fn maybe_seal_growing(&self) -> Result<(), SegmentError> {
1283        let should_seal = {
1284            let growing_opt = self.growing.read();
1285            if let Some(growing_arc) = growing_opt.as_ref() {
1286                let growing = growing_arc.read();
1287                growing.should_seal(&self.config.segment_config)
1288                    || growing.idle_secs() >= self.config.idle_seal_secs
1289            } else {
1290                false
1291            }
1292        };
1293
1294        if should_seal {
1295            self.seal_current()?;
1296        }
1297
1298        Ok(())
1299    }
1300
1301    /// Seal the current growing segment
1302    pub fn seal_current(&self) -> Result<SegmentId, SegmentError> {
1303        // Keep topology publication atomic for cursors that capture growing
1304        // and sealed under the same lock order: growing, then sealed.
1305        let mut growing_slot = self.growing.write();
1306        if let Some(growing_arc) = growing_slot.take() {
1307            let mut growing = growing_arc.write();
1308            let seg_id = growing.id();
1309            let entity_count = growing.stats().entity_count as u64;
1310
1311            // Seal it
1312            growing.seal()?;
1313
1314            // Move to sealed list (we need to extract it from the Arc)
1315            drop(growing); // Release write lock
1316
1317            // In a real implementation, we'd convert to SealedSegment here
1318            // For now, we keep it as-is since GrowingSegment implements UnifiedSegment
1319            self.sealed.write().push(growing_arc);
1320            drop(growing_slot);
1321
1322            // Mark sealed segment pages all-visible — they're now immutable
1323            self.mark_sealed_pages_visible(entity_count);
1324
1325            // Update stats
1326            {
1327                let mut stats = self.stats.write();
1328                stats.growing_count = stats.growing_count.saturating_sub(1);
1329                stats.sealed_count += 1;
1330                stats.seal_ops += 1;
1331            }
1332
1333            self.emit(LifecycleEvent::SegmentSealed(seg_id));
1334
1335            return Ok(seg_id);
1336        }
1337
1338        Err(SegmentError::InvalidState(SegmentState::Sealed))
1339    }
1340
1341    /// Force seal (for testing/manual control)
1342    pub fn force_seal(&self) -> Result<Option<SegmentId>, SegmentError> {
1343        let has_growing = self.growing.read().is_some();
1344        if has_growing {
1345            Ok(Some(self.seal_current()?))
1346        } else {
1347            Ok(None)
1348        }
1349    }
1350
1351    /// Fraction of "pages" in sealed segments that are marked all-visible.
1352    ///
1353    /// Sealed segments are immutable so all their rows are safe for
1354    /// index-only scans. The growing segment is never counted (writes
1355    /// may be in-flight). Uses `rows_per_page = 256` (matching 8 KB pages
1356    /// with ~32-byte rows).
1357    ///
1358    /// Returns a value in `[0.0, 1.0]`. 1.0 when all sealed rows are
1359    /// visible; 0.0 when there are no sealed segments.
1360    pub fn all_visible_fraction(&self) -> f64 {
1361        const ROWS_PER_PAGE: u32 = 256;
1362        let sealed = self.sealed.read();
1363        if sealed.is_empty() {
1364            return 0.0;
1365        }
1366        let mut total_pages: u64 = 0;
1367        for seg_arc in sealed.iter() {
1368            let seg = seg_arc.read();
1369            let entity_count = seg.stats().entity_count as u64;
1370            let pages = entity_count.div_ceil(ROWS_PER_PAGE as u64);
1371            total_pages += pages;
1372        }
1373        if total_pages == 0 {
1374            return 0.0;
1375        }
1376        let visible = self.visibility_map.all_visible_count();
1377        (visible as f64 / total_pages as f64).min(1.0)
1378    }
1379
1380    /// Mark all pages of newly sealed segments as all-visible in the
1381    /// visibility map. Called internally after `seal_current`.
1382    fn mark_sealed_pages_visible(&self, seg_entity_count: u64) {
1383        const ROWS_PER_PAGE: u32 = 256;
1384        let existing_visible = self.visibility_map.all_visible_count();
1385        // Append pages starting after the last known visible page
1386        let start_page = existing_visible as u32;
1387        let new_pages = seg_entity_count.div_ceil(ROWS_PER_PAGE as u64);
1388        let end_page = start_page + new_pages as u32;
1389        self.visibility_map.mark_range_visible(start_page, end_page);
1390    }
1391
1392    /// Probe a unique-key lookup in each segment under the same entity
1393    /// locks used by scans and writes. Lookup state is rebuilt lazily after
1394    /// recovery, consolidation, or a key-changing mutation.
1395    pub(crate) fn find_unique_key_conflict(
1396        &self,
1397        columns: &[Vec<String>],
1398        key_index: usize,
1399        signatures: &[String],
1400        exclude_id: Option<EntityId>,
1401        reserves_key: &impl Fn(&UnifiedEntity) -> bool,
1402    ) -> Option<EntityId> {
1403        if let Some(growing) = self.growing.read().as_ref() {
1404            if let Some(id) = growing.read().find_unique_key_conflict(
1405                columns,
1406                key_index,
1407                signatures,
1408                exclude_id,
1409                reserves_key,
1410            ) {
1411                return Some(id);
1412            }
1413        }
1414        for segment in self.sealed.read().iter() {
1415            if let Some(id) = segment.read().find_unique_key_conflict(
1416                columns,
1417                key_index,
1418                signatures,
1419                exclude_id,
1420                reserves_key,
1421            ) {
1422                return Some(id);
1423            }
1424        }
1425        None
1426    }
1427
1428    /// Iterate over all entities in-place without collecting into a Vec.
1429    ///
1430    /// The callback receives a reference to each entity. Return `true` to
1431    /// continue iteration, `false` to stop early (e.g. when a LIMIT is reached).
1432    /// This avoids the allocation and cloning overhead of `query_all`.
1433    pub fn for_each_entity<F>(&self, callback: F)
1434    where
1435        F: FnMut(&UnifiedEntity) -> bool,
1436    {
1437        self.for_each_entity_in_segments(|_| true, || true, callback);
1438    }
1439
1440    fn for_each_entity_in_segments<F>(
1441        &self,
1442        segment_filter: impl Fn(&GrowingSegment) -> bool,
1443        mut before_segment: impl FnMut() -> bool,
1444        mut callback: F,
1445    ) where
1446        F: FnMut(&UnifiedEntity) -> bool,
1447    {
1448        // Growing segment — direct iteration (no Box<dyn>)
1449        // Try non-blocking read first; fall back to blocking only when a writer
1450        // is actively holding the write lock (rare in read-heavy workloads).
1451        if let Some(growing_arc) = self.growing.read().as_ref() {
1452            if !before_segment() {
1453                return;
1454            }
1455            let growing = if let Some(g) = growing_arc.try_read() {
1456                g
1457            } else {
1458                growing_arc.read()
1459            };
1460            if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
1461                return;
1462            }
1463        }
1464
1465        // Sealed segments
1466        let sealed = self.sealed.read();
1467        for segment_arc in sealed.iter() {
1468            if !before_segment() {
1469                return;
1470            }
1471            let segment = segment_arc.read();
1472            if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
1473                return;
1474            }
1475        }
1476    }
1477
1478    /// Parallel fold across all entities. Each sealed segment is
1479    /// processed on its own rayon task; the growing segment stays on
1480    /// the caller thread (its read lock is briefly held).
1481    ///
1482    /// - `init` builds a fresh accumulator per thread.
1483    /// - `fold` mutates an accumulator with one entity at a time.
1484    /// - `reduce` combines two accumulators into one.
1485    ///
1486    /// The returned value is the reduction of every per-thread
1487    /// accumulator. Use this for aggregate-shape workloads (GROUP BY)
1488    /// where per-thread partial state can be merged cheaply.
1489    ///
1490    /// NOTE: when there are 0 or 1 sealed segments, the parallel path
1491    /// is skipped and the work runs sequentially to avoid rayon
1492    /// overhead on tiny tables.
1493    pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1494        &self,
1495        init: FInit,
1496        fold: FFold,
1497        reduce: FReduce,
1498    ) -> T
1499    where
1500        T: Send,
1501        FInit: Fn() -> T + Send + Sync,
1502        FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1503        FReduce: Fn(T, T) -> T + Send + Sync,
1504    {
1505        use rayon::prelude::*;
1506
1507        // Growing segment — always sequential (single writer lock,
1508        // usually small working set).
1509        let mut acc = init();
1510        if let Some(growing_arc) = self.growing.read().as_ref() {
1511            let growing = if let Some(g) = growing_arc.try_read() {
1512                g
1513            } else {
1514                growing_arc.read()
1515            };
1516            growing.for_each_fast(|entity| {
1517                acc = fold(std::mem::replace(&mut acc, init()), entity);
1518                true
1519            });
1520        }
1521
1522        // Sealed segments — snapshot the Arc list under the read lock,
1523        // then drop the lock so rayon workers can fan out without
1524        // blocking writers.
1525        let segments: Vec<_> = {
1526            let sealed = self.sealed.read();
1527            sealed.iter().cloned().collect()
1528        };
1529
1530        if segments.len() <= 1 {
1531            for seg_arc in &segments {
1532                let seg = seg_arc.read();
1533                seg.for_each_fast(|entity| {
1534                    acc = fold(std::mem::replace(&mut acc, init()), entity);
1535                    true
1536                });
1537            }
1538            return acc;
1539        }
1540
1541        let sealed_acc = segments
1542            .into_par_iter()
1543            .map(|seg_arc| {
1544                let mut local = init();
1545                let seg = seg_arc.read();
1546                seg.for_each_fast(|entity| {
1547                    local = fold(std::mem::replace(&mut local, init()), entity);
1548                    true
1549                });
1550                local
1551            })
1552            .reduce(&init, &reduce);
1553
1554        reduce(acc, sealed_acc)
1555    }
1556
1557    /// Zone-map-aware iteration across all segments.
1558    ///
1559    /// Like `for_each_entity`, but checks `zone_preds` against each segment's
1560    /// column zone maps before iterating. Segments where any predicate can
1561    /// definitively prove no rows match are skipped entirely.
1562    ///
1563    /// `zone_preds`: slice of `(column_name, ZoneColPred)` extracted from the WHERE clause.
1564    /// Empty slice → same behaviour as `for_each_entity` (no pruning).
1565    pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
1566    where
1567        F: FnMut(&UnifiedEntity) -> bool,
1568    {
1569        let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
1570    }
1571
1572    pub fn for_each_entity_zoned_with_stats<F>(
1573        &self,
1574        zone_preds: &[(&str, ZoneColPred<'_>)],
1575        mut callback: F,
1576    ) -> SegmentScanStats
1577    where
1578        F: FnMut(&UnifiedEntity) -> bool,
1579    {
1580        let mut scan_stats = SegmentScanStats::default();
1581        // Growing segment — never skip (it's receiving writes, zones are partial).
1582        // Try a non-blocking read first: if a writer is currently inserting
1583        // (holding the write lock), try_read() returns None and we fall back to
1584        // the blocking read.  In low-contention workloads (reads far outnumber
1585        // writes) the try_read() almost always succeeds and readers never stall.
1586        if let Some(growing_arc) = self.growing.read().as_ref() {
1587            scan_stats.segments_total += 1;
1588            scan_stats.segments_scanned += 1;
1589            let growing = if let Some(g) = growing_arc.try_read() {
1590                g
1591            } else {
1592                growing_arc.read()
1593            };
1594            if !growing.for_each_fast(&mut callback) {
1595                return scan_stats;
1596            }
1597        }
1598
1599        // Sealed segments — check zone maps before iterating
1600        let sealed = self.sealed.read();
1601        for segment_arc in sealed.iter() {
1602            scan_stats.segments_total += 1;
1603            let segment = segment_arc.read();
1604            if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
1605                scan_stats.segments_pruned += 1;
1606                continue; // entire segment pruned
1607            }
1608            scan_stats.segments_scanned += 1;
1609            if !segment.for_each_fast(&mut callback) {
1610                return scan_stats;
1611            }
1612        }
1613        scan_stats
1614    }
1615
1616    /// Zone-map-aware parallel query.
1617    ///
1618    /// Like `query_all` but applies `zone_preds` on the main thread to
1619    /// prune sealed segments before spawning workers — segments that
1620    /// provably contain no matching rows are skipped entirely.
1621    ///
1622    /// Zone check runs single-threaded (it reads per-segment metadata,
1623    /// not row data), so it's cheap. Surviving segments are then scanned
1624    /// in parallel using `std::thread::scope` when there are > 1 of them.
1625    pub fn query_all_zoned<F>(
1626        &self,
1627        zone_preds: &[(&str, ZoneColPred<'_>)],
1628        filter: F,
1629    ) -> Vec<UnifiedEntity>
1630    where
1631        F: Fn(&UnifiedEntity) -> bool + Sync,
1632    {
1633        self.query_all_zoned_with_stats(zone_preds, filter).0
1634    }
1635
1636    pub fn query_all_zoned_with_stats<F>(
1637        &self,
1638        zone_preds: &[(&str, ZoneColPred<'_>)],
1639        filter: F,
1640    ) -> (Vec<UnifiedEntity>, SegmentScanStats)
1641    where
1642        F: Fn(&UnifiedEntity) -> bool + Sync,
1643    {
1644        let mut results = Vec::new();
1645        let mut scan_stats = SegmentScanStats::default();
1646
1647        // Growing segment — always scan, no zone skip (zones are partial).
1648        // Non-blocking try_read() avoids stalling behind in-progress inserts.
1649        if let Some(growing_arc) = self.growing.read().as_ref() {
1650            scan_stats.segments_total += 1;
1651            scan_stats.segments_scanned += 1;
1652            let growing = if let Some(g) = growing_arc.try_read() {
1653                g
1654            } else {
1655                growing_arc.read()
1656            };
1657            results.extend(growing.iter().filter(|e| filter(e)).cloned());
1658        }
1659
1660        // Sealed segments: zone-prune on main thread, then scan in parallel.
1661        let sealed = self.sealed.read();
1662        // Collect only the segments that survive zone-predicate pruning.
1663        let surviving: Vec<_> = sealed
1664            .iter()
1665            .filter(|seg_arc| {
1666                scan_stats.segments_total += 1;
1667                if zone_preds.is_empty() {
1668                    scan_stats.segments_scanned += 1;
1669                    return true;
1670                }
1671                let seg = seg_arc.read();
1672                if seg.can_skip_zone_preds(zone_preds) {
1673                    scan_stats.segments_pruned += 1;
1674                    false
1675                } else {
1676                    scan_stats.segments_scanned += 1;
1677                    true
1678                }
1679            })
1680            .collect();
1681
1682        let use_parallel = surviving.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1683
1684        if use_parallel {
1685            let filter_ref = &filter;
1686            let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1687                surviving
1688                    .iter()
1689                    .map(|segment| {
1690                        s.spawn(move || {
1691                            segment
1692                                .read()
1693                                .iter()
1694                                .filter(|e| filter_ref(e))
1695                                .cloned()
1696                                .collect::<Vec<_>>()
1697                        })
1698                    })
1699                    .collect::<Vec<_>>()
1700                    .into_iter()
1701                    .map(|handle| handle.join().unwrap_or_default())
1702                    .collect()
1703            });
1704            for batch in segment_results {
1705                results.extend(batch);
1706            }
1707        } else {
1708            for segment_arc in surviving {
1709                let seg = segment_arc.read();
1710                results.extend(seg.iter().filter(|e| filter(e)).cloned());
1711            }
1712        }
1713
1714        (results, scan_stats)
1715    }
1716
1717    /// Visibility predicate for the `scan_*` family. `Some` delegates to the
1718    /// canonical `entity_visible_with_context`; `None` supports low-level
1719    /// callers that intentionally scan without a statement snapshot and keeps
1720    /// the documented resolver fallback — moderation gate plus hiding
1721    /// superseded physical versions — instead of showing every version.
1722    fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
1723        match snapshot {
1724            Some(ctx) => entity_visible_with_context(Some(ctx), entity),
1725            None => {
1726                !crate::runtime::ai::moderation::entity_moderation_hidden(entity)
1727                    && entity.xmax == 0
1728            }
1729        }
1730    }
1731
1732    /// Scan every segment under an explicit MVCC snapshot.
1733    ///
1734    /// Visibility is evaluated while each segment iterator is active, before
1735    /// the caller's filter. Unlike the raw iteration methods, this entry point
1736    /// does not consult thread-local snapshot state, so it preserves the same
1737    /// view when invoked from a different thread.
1738    pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
1739    where
1740        F: Fn(&UnifiedEntity) -> bool + Sync,
1741    {
1742        self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
1743    }
1744
1745    /// Collect visible graph items, pruning only under the segment read guard.
1746    pub(crate) fn scan_graph_kind(
1747        &self,
1748        snapshot: Option<&SnapshotContext>,
1749        kind: GraphEntityKind,
1750    ) -> Vec<UnifiedEntity> {
1751        self.query_all_in_segments(
1752            |segment| segment.may_contain_graph_kind(kind),
1753            |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
1754        )
1755    }
1756
1757    /// Visit visible candidates in segments that may contain the requested kind.
1758    /// Non-graph items in those segments still reach the callback so work budgets
1759    /// charge every inspected visible candidate. RLS must run outside this callback.
1760    /// `before_segment` runs even for pruned segments; false stops the entire scan.
1761    pub(crate) fn scan_graph_candidates_for_each<F>(
1762        &self,
1763        snapshot: Option<&SnapshotContext>,
1764        kind: GraphEntityKind,
1765        before_segment: impl FnMut() -> bool,
1766        mut callback: F,
1767    ) where
1768        F: FnMut(&UnifiedEntity) -> bool,
1769    {
1770        self.for_each_entity_in_segments(
1771            |segment| segment.may_contain_graph_kind(kind),
1772            before_segment,
1773            |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
1774        );
1775    }
1776
1777    /// Visit every entity visible under an explicit MVCC snapshot.
1778    ///
1779    /// Returning `false` from `callback` stops the scan early. Hidden entities
1780    /// never reach the callback and therefore cannot stop the scan.
1781    pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
1782    where
1783        F: FnMut(&UnifiedEntity) -> bool,
1784    {
1785        self.for_each_entity(|entity| {
1786            !Self::scan_entity_visible(snapshot, entity) || callback(entity)
1787        });
1788    }
1789
1790    /// Capture segment membership and append boundaries, then hydrate bounded
1791    /// batches. The consumer runs outside topology and segment locks, so scoring
1792    /// and RLS may re-enter storage and writers can progress between batches.
1793    /// Retired sources pin positions only; hydration uses current segments so
1794    /// mutations after consolidation are visible to the snapshot predicate.
1795    pub(crate) fn scan_batches(
1796        &self,
1797        snapshot: Option<&SnapshotContext>,
1798        mut before_work: impl FnMut() -> bool,
1799        mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
1800    ) -> bool {
1801        let segments = {
1802            let growing = self.growing.read();
1803            let sealed = self.sealed.read();
1804            let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
1805            for segment in growing.iter().chain(sealed.iter()) {
1806                if !before_work() {
1807                    return false;
1808                }
1809                let cursor = segment.read().scan_cursor();
1810                segments.push((Arc::clone(segment), cursor));
1811            }
1812            segments
1813        };
1814        let mut ids = Vec::new();
1815        let mut batch = Vec::new();
1816        for (segment, mut cursor) in segments {
1817            while !cursor.finished() {
1818                ids.clear();
1819                batch.clear();
1820                let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
1821                ids.reserve_exact(batch_size);
1822                batch.reserve_exact(batch_size);
1823                {
1824                    let segment = segment.read();
1825                    if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
1826                        return false;
1827                    }
1828                }
1829                // Sources can retire between batches. Consult authoritative storage
1830                // after releasing the source guard, just as the former ID-list scan did.
1831                if !before_work() {
1832                    return false;
1833                }
1834                self.for_each_id(&ids, |_, entity| {
1835                    if Self::scan_entity_visible(snapshot, entity) {
1836                        batch.push(entity.clone());
1837                    }
1838                });
1839                if !before_work() {
1840                    return false;
1841                }
1842                if !consume(&batch) {
1843                    return false;
1844                }
1845            }
1846        }
1847        true
1848    }
1849
1850    /// Fold entities visible under an explicit MVCC snapshot in parallel.
1851    pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
1852        &self,
1853        snapshot: Option<&SnapshotContext>,
1854        init: FInit,
1855        fold: FFold,
1856        reduce: FReduce,
1857    ) -> T
1858    where
1859        T: Send,
1860        FInit: Fn() -> T + Send + Sync,
1861        FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1862        FReduce: Fn(T, T) -> T + Send + Sync,
1863    {
1864        self.fold_entities_parallel(
1865            init,
1866            |acc, entity| {
1867                if Self::scan_entity_visible(snapshot, entity) {
1868                    fold(acc, entity)
1869                } else {
1870                    acc
1871                }
1872            },
1873            reduce,
1874        )
1875    }
1876
1877    /// Visit zone-map candidates visible under an explicit MVCC snapshot.
1878    pub fn scan_for_each_zoned_with_stats<F>(
1879        &self,
1880        snapshot: Option<&SnapshotContext>,
1881        zone_preds: &[(&str, ZoneColPred<'_>)],
1882        mut callback: F,
1883    ) -> SegmentScanStats
1884    where
1885        F: FnMut(&UnifiedEntity) -> bool,
1886    {
1887        self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
1888            !Self::scan_entity_visible(snapshot, entity) || callback(entity)
1889        })
1890    }
1891
1892    /// Collect zone-map candidates visible under an explicit MVCC snapshot.
1893    pub fn scan_zoned_with_stats<F>(
1894        &self,
1895        snapshot: Option<&SnapshotContext>,
1896        zone_preds: &[(&str, ZoneColPred<'_>)],
1897        filter: F,
1898    ) -> (Vec<UnifiedEntity>, SegmentScanStats)
1899    where
1900        F: Fn(&UnifiedEntity) -> bool + Sync,
1901    {
1902        self.query_all_zoned_with_stats(zone_preds, |entity| {
1903            Self::scan_entity_visible(snapshot, entity) && filter(entity)
1904        })
1905    }
1906
1907    /// Query across all segments. Uses parallel scanning for sealed segments
1908    /// when more than one sealed segment exists.
1909    pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
1910    where
1911        F: Fn(&UnifiedEntity) -> bool + Sync,
1912    {
1913        self.query_all_in_segments(|_| true, filter)
1914    }
1915
1916    fn query_all_in_segments<F>(
1917        &self,
1918        segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
1919        filter: F,
1920    ) -> Vec<UnifiedEntity>
1921    where
1922        F: Fn(&UnifiedEntity) -> bool + Sync,
1923    {
1924        let mut results = Vec::new();
1925
1926        // Query growing segment — try non-blocking read first (avoids stalling
1927        // behind an in-progress insert; falls back to blocking if writer is active).
1928        if let Some(growing_arc) = self.growing.read().as_ref() {
1929            let growing = if let Some(g) = growing_arc.try_read() {
1930                g
1931            } else {
1932                growing_arc.read()
1933            };
1934            if segment_filter(&growing) {
1935                results.extend(growing.iter().filter(|e| filter(e)).cloned());
1936            }
1937        }
1938
1939        // Query sealed segments — parallel when multiple exist AND multi-core
1940        let sealed = self.sealed.read();
1941        let use_parallel = sealed.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1942        if use_parallel {
1943            let filter_ref = &filter;
1944            let segment_filter_ref = &segment_filter;
1945            let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1946                sealed
1947                    .iter()
1948                    .map(|segment| {
1949                        s.spawn(move || {
1950                            let segment = segment.read();
1951                            if segment_filter_ref(&segment) {
1952                                segment
1953                                    .iter()
1954                                    .filter(|e| filter_ref(e))
1955                                    .cloned()
1956                                    .collect::<Vec<_>>()
1957                            } else {
1958                                Vec::new()
1959                            }
1960                        })
1961                    })
1962                    .collect::<Vec<_>>()
1963                    .into_iter()
1964                    .map(|handle| handle.join().unwrap_or_default())
1965                    .collect()
1966            });
1967            for batch in segment_results {
1968                results.extend(batch);
1969            }
1970        } else {
1971            for segment in sealed.iter() {
1972                let seg = segment.read();
1973                if segment_filter(&seg) {
1974                    results.extend(seg.iter().filter(|e| filter(e)).cloned());
1975                }
1976            }
1977        }
1978
1979        results
1980    }
1981
1982    /// Probe exact segment membership for a key without materializing matching entities.
1983    ///
1984    /// Returns `false` only when every consulted segment proves absence.
1985    /// Missing segment state is conservative and returns `true`.
1986    pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
1987        if let Some(growing_arc) = self.growing.read().as_ref() {
1988            let growing = growing_arc.read();
1989            if growing.may_contain_exact_key(key) {
1990                return true;
1991            }
1992        } else {
1993            return true;
1994        }
1995
1996        let sealed = self.sealed.read();
1997        for segment_arc in sealed.iter() {
1998            let segment = segment_arc.read();
1999            if segment.may_contain_exact_key(key) {
2000                return true;
2001            }
2002        }
2003
2004        false
2005    }
2006
2007    /// Filter by metadata across all segments
2008    pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2009        let mut results = Vec::new();
2010
2011        // Growing segment
2012        if let Some(growing_arc) = self.growing.read().as_ref() {
2013            let growing = growing_arc.read();
2014            results.extend(growing.filter_metadata(filters));
2015        }
2016
2017        // Sealed segments
2018        let sealed = self.sealed.read();
2019        for segment in sealed.iter() {
2020            results.extend(segment.filter_metadata(filters));
2021        }
2022
2023        results
2024    }
2025
2026    /// Get entities by kind
2027    pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2028        let mut results = Vec::new();
2029
2030        // Growing segment
2031        if let Some(growing_arc) = self.growing.read().as_ref() {
2032            let growing = growing_arc.read();
2033            for entity in growing.iter_kind(kind) {
2034                results.push(entity.clone());
2035            }
2036        }
2037
2038        // Sealed segments. Iterate through the read guard, not the
2039        // `Arc<RwLock<GrowingSegment>>` impl of `iter_kind` — that one cannot
2040        // hand out references across the lock, so it returns an empty iterator
2041        // and every sealed entity was silently missing from the result.
2042        let sealed = self.sealed.read();
2043        for segment_arc in sealed.iter() {
2044            let segment = segment_arc.read();
2045            for entity in segment.iter_kind(kind) {
2046                results.push(entity.clone());
2047            }
2048        }
2049
2050        results
2051    }
2052
2053    /// Count entities
2054    pub fn count(&self) -> usize {
2055        self.total_entities_atomic.load(Ordering::Relaxed) as usize
2056    }
2057
2058    /// Get all segment IDs
2059    pub fn segment_ids(&self) -> Vec<SegmentId> {
2060        let mut ids = Vec::new();
2061
2062        if let Some(growing_arc) = self.growing.read().as_ref() {
2063            ids.push(growing_arc.read().id());
2064        }
2065
2066        let sealed = self.sealed.read();
2067        for segment in sealed.iter() {
2068            ids.push(segment.id());
2069        }
2070
2071        ids.extend(self.archived.read().iter().copied());
2072
2073        ids
2074    }
2075
2076    /// Emit a lifecycle event.
2077    ///
2078    /// Perf: this used to push onto a `RwLock<Vec<LifecycleEvent>>`
2079    /// on every insert / delete / seal. Nobody consumes that vec
2080    /// today (no subscription API, `drain_events` has no callers),
2081    /// so the write lock + push was pure tax — and the vec grew
2082    /// unbounded in long-running processes.
2083    ///
2084    /// Current behaviour: no-op. If we ever want the hooks back,
2085    /// replace this with a bounded channel or an actual subscriber
2086    /// registry; the callers (`insert`, `delete`, `maybe_seal_growing`)
2087    /// already pass well-typed events.
2088    #[inline]
2089    #[allow(clippy::unused_self)]
2090    fn emit(&self, _event: LifecycleEvent) {}
2091
2092    /// Drain events. Kept for API compatibility; always returns
2093    /// empty because `emit` no longer buffers.
2094    pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2095        std::mem::take(&mut *self.events.write())
2096    }
2097
2098    /// Run maintenance (would be called periodically in production)
2099    pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2100        // Auto-seal idle segments
2101        self.maybe_seal_growing()?;
2102
2103        if self.config.enable_consolidation {
2104            self.consolidation_tick();
2105        }
2106
2107        Ok(())
2108    }
2109
2110    // ========================================================================
2111    // Consolidation — ADR 0073 §5
2112    // ========================================================================
2113
2114    /// One paced step of sealed-segment consolidation.
2115    ///
2116    /// Starts a run when a threshold has been crossed, copies at most
2117    /// `consolidation_entities_per_tick` live entities into the merged segment,
2118    /// and swaps once every source has been drained. Returns without finishing
2119    /// otherwise — the half-built merged segment waits for the next tick.
2120    fn consolidation_tick(&self) {
2121        // Drop the read guard before `start_consolidation` takes the write
2122        // guard: `parking_lot` locks are not reentrant.
2123        let running = self.consolidation.read().is_some();
2124        if !running && !self.start_consolidation() {
2125            return;
2126        }
2127
2128        // Keep the accounting guard through publication: the merged allocation
2129        // must never disappear between the in-flight and active inventories.
2130        let mut guard = self.consolidation.write();
2131        let Some(run) = guard.as_mut() else {
2132            return;
2133        };
2134        self.copy_bounded(run);
2135        if run.cursor >= run.sources.len() {
2136            if let Some(run) = guard.take() {
2137                self.finish_consolidation(run);
2138            }
2139        }
2140    }
2141
2142    /// Evaluate the trigger and, when it fires, open a consolidation run.
2143    ///
2144    /// Returns `true` when a run was started.
2145    fn start_consolidation(&self) -> bool {
2146        let sources = self.select_candidates();
2147        self.start_consolidation_from_sources(sources)
2148    }
2149
2150    fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
2151        if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
2152            return false;
2153        }
2154
2155        let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
2156        let merged = GrowingSegment::new(merged_id, &self.collection);
2157
2158        *self.consolidation.write() = Some(Consolidation {
2159            sources: sources
2160                .into_iter()
2161                .map(|id| SourceProgress {
2162                    id,
2163                    copied: Vec::new(),
2164                    tombstones_at_open: None,
2165                })
2166                .collect(),
2167            merged_id,
2168            merged,
2169            cursor: 0,
2170            pending_ids: Vec::new(),
2171            pending_cursor: 0,
2172            entities_copied: 0,
2173        });
2174
2175        self.stats.write().consolidation.runs_started += 1;
2176        true
2177    }
2178
2179    /// Run one paced consolidation tick because admission pressure needs
2180    /// reclaimable bytes, even if the background thresholds have not fired.
2181    ///
2182    /// Returns true when the call started or advanced a consolidation run.
2183    pub fn pressure_consolidation_tick(&self) -> bool {
2184        if !self.config.enable_consolidation {
2185            return false;
2186        }
2187
2188        let running = self.consolidation.read().is_some();
2189        if !running {
2190            let sources = self.dirty_sealed_segments();
2191            if !self.start_consolidation_from_sources(sources) {
2192                return false;
2193            }
2194        }
2195
2196        self.consolidation_tick();
2197        true
2198    }
2199
2200    /// Which sealed segments this collection should merge, if any.
2201    ///
2202    /// Three thresholds, any of which arms consolidation:
2203    ///
2204    /// * **sealed-segment count** above `max_sealed_segments` — every sealed
2205    ///   segment merges, which is what that knob always claimed to do.
2206    /// * **tombstone ratio** — dead entities over total entities across the
2207    ///   sealed set, past [`CONSOLIDATION_TOMBSTONE_RATIO`].
2208    /// * **fragmentation** — reclaimable bytes over held bytes, past
2209    ///   [`CONSOLIDATION_FRAGMENTATION_RATIO`].
2210    ///
2211    /// The count trigger takes every sealed segment; the two budget triggers
2212    /// take only the segments that actually carry tombstones, so a clean
2213    /// segment is never rewritten for nothing.
2214    fn select_candidates(&self) -> Vec<SegmentId> {
2215        let sealed = self.sealed.read();
2216        if sealed.is_empty() {
2217            return Vec::new();
2218        }
2219
2220        if sealed.len() > self.config.max_sealed_segments {
2221            return sealed.iter().map(|segment| segment.read().id()).collect();
2222        }
2223
2224        let mut live = 0usize;
2225        let mut tombstones = 0usize;
2226        let mut held = 0u64;
2227        let mut reclaimable = 0u64;
2228        let mut dirty = Vec::new();
2229
2230        for segment_arc in sealed.iter() {
2231            let segment = segment_arc.read();
2232            live += segment.live_entity_count();
2233            tombstones += segment.tombstone_count();
2234            held += segment.resident_bytes();
2235            reclaimable += segment.reclaimable_bytes();
2236            if segment.tombstone_count() > 0 {
2237                dirty.push(segment.id());
2238            }
2239        }
2240
2241        let total = live + tombstones;
2242        let tombstone_ratio = if total == 0 {
2243            0.0
2244        } else {
2245            tombstones as f64 / total as f64
2246        };
2247        let fragmentation = if held == 0 {
2248            0.0
2249        } else {
2250            reclaimable as f64 / held as f64
2251        };
2252
2253        if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
2254            || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
2255        {
2256            return dirty;
2257        }
2258
2259        Vec::new()
2260    }
2261
2262    fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
2263        self.sealed
2264            .read()
2265            .iter()
2266            .filter_map(|segment_arc| {
2267                let segment = segment_arc.read();
2268                (segment.reclaimable_bytes() > 0).then(|| segment.id())
2269            })
2270            .collect()
2271    }
2272
2273    /// A single candidate is still worth merging when it carries tombstones —
2274    /// the merge drops them. A single clean candidate is not.
2275    fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
2276        if sources.is_empty() {
2277            return false;
2278        }
2279        let sealed = self.sealed.read();
2280        sealed.iter().any(|segment_arc| {
2281            let segment = segment_arc.read();
2282            sources.contains(&segment.id()) && segment.tombstone_count() > 0
2283        })
2284    }
2285
2286    fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
2287        let sealed = self.sealed.read();
2288        sealed
2289            .iter()
2290            .find(|segment| segment.read().id() == id)
2291            .map(Arc::clone)
2292    }
2293
2294    /// Copy at most one tick's worth of live entities into the merged segment.
2295    ///
2296    /// The per-tick bound is the pacing contract: a tick's cost is O(bound),
2297    /// independent of how many entities the run still has to move.
2298    fn copy_bounded(&self, run: &mut Consolidation) {
2299        let mut budget = self.config.consolidation_entities_per_tick;
2300
2301        while budget > 0 && run.cursor < run.sources.len() {
2302            let source_id = run.sources[run.cursor].id;
2303            let Some(source_arc) = self.sealed_segment(source_id) else {
2304                // The source vanished. Only consolidation removes sealed
2305                // segments and it is single-threaded per collection, so this
2306                // cannot happen today — skip it rather than resurrect its rows.
2307                run.cursor += 1;
2308                run.pending_ids.clear();
2309                run.pending_cursor = 0;
2310                continue;
2311            };
2312
2313            if run.sources[run.cursor].tombstones_at_open.is_none() {
2314                let source = source_arc.read();
2315                run.pending_ids = source.live_entity_ids();
2316                run.pending_cursor = 0;
2317                run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
2318            }
2319
2320            {
2321                let source = source_arc.read();
2322                while budget > 0 && run.pending_cursor < run.pending_ids.len() {
2323                    let id = run.pending_ids[run.pending_cursor];
2324                    run.pending_cursor += 1;
2325                    budget -= 1;
2326
2327                    // Tombstoned between the snapshot and now — it must not be
2328                    // resurrected into the merged segment.
2329                    let Some(entity) = source.get(id) else {
2330                        continue;
2331                    };
2332                    let metadata = source.get_metadata(id);
2333                    run.merged.adopt_entity(entity.clone(), metadata);
2334                    run.sources[run.cursor].copied.push(id);
2335                    run.entities_copied += 1;
2336                }
2337            }
2338
2339            if run.pending_cursor >= run.pending_ids.len() {
2340                run.cursor += 1;
2341                run.pending_ids.clear();
2342                run.pending_cursor = 0;
2343            }
2344        }
2345    }
2346
2347    /// Atomically replace the source segments with the merged one.
2348    ///
2349    /// The `sealed` write lock is the boundary: a reader holds the read lock
2350    /// for the whole of its scan, so it observes either the old set or the new
2351    /// set, never a gap and never both copies of a row.
2352    fn finish_consolidation(&self, mut run: Consolidation) {
2353        let mut sealed = self.sealed.write();
2354
2355        let mut positions = Vec::with_capacity(run.sources.len());
2356        for source in &run.sources {
2357            if let Some(index) = sealed
2358                .iter()
2359                .position(|segment| segment.read().id() == source.id)
2360            {
2361                positions.push(index);
2362            }
2363        }
2364        if positions.is_empty() {
2365            return;
2366        }
2367        positions.sort_unstable();
2368
2369        // Rows deleted from a source after we copied them must not come back.
2370        // A source whose tombstone count is unchanged saw no deletes while the
2371        // merge ran, so its copied rows need no per-row check.
2372        let mut stale = Vec::new();
2373        for source in &run.sources {
2374            let Some(source_arc) = sealed
2375                .iter()
2376                .find(|segment| segment.read().id() == source.id)
2377            else {
2378                continue;
2379            };
2380            let segment = source_arc.read();
2381            if source.tombstones_at_open == Some(segment.tombstone_count()) {
2382                continue;
2383            }
2384            stale.extend(
2385                source
2386                    .copied
2387                    .iter()
2388                    .copied()
2389                    .filter(|id| !segment.contains(*id)),
2390            );
2391        }
2392        for id in stale {
2393            run.merged.evict_entity(id);
2394        }
2395
2396        let reclaimed_bytes: u64 = positions
2397            .iter()
2398            .map(|&index| sealed[index].read().resident_bytes())
2399            .sum();
2400        let reclaimed_tombstones: u64 = positions
2401            .iter()
2402            .map(|&index| sealed[index].read().tombstone_count() as u64)
2403            .sum();
2404
2405        if run.merged.seal().is_err() {
2406            return;
2407        }
2408        let merged_bytes = run.merged.resident_bytes();
2409        let merged_arc = Arc::new(RwLock::new(run.merged));
2410
2411        {
2412            let mut retired = self.retired.write();
2413            retired.retain(|source| source.strong_count() > 0);
2414            for &index in positions.iter().rev() {
2415                let source = sealed.remove(index);
2416                if Arc::strong_count(&source) > 1 {
2417                    retired.push(Arc::downgrade(&source));
2418                }
2419            }
2420        }
2421        let insert_at = positions[0].min(sealed.len());
2422        sealed.insert(insert_at, Arc::clone(&merged_arc));
2423        drop(sealed);
2424
2425        // Remap the entity→segment hints that pointed at the retired sources.
2426        let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
2427        {
2428            let mut entity_segment = self.entity_segment.write();
2429            for segment_id in entity_segment.values_mut() {
2430                if source_ids.contains(segment_id) {
2431                    *segment_id = run.merged_id;
2432                }
2433            }
2434        }
2435
2436        {
2437            let mut stats = self.stats.write();
2438            stats.sealed_count = stats
2439                .sealed_count
2440                .saturating_sub(positions.len())
2441                .saturating_add(1);
2442            let consolidation = &mut stats.consolidation;
2443            consolidation.runs_completed += 1;
2444            consolidation.segments_merged += positions.len() as u64;
2445            consolidation.tombstones_reclaimed += reclaimed_tombstones;
2446            consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
2447        }
2448
2449        self.emit(LifecycleEvent::SegmentConsolidated {
2450            source: source_ids,
2451            target: run.merged_id,
2452        });
2453    }
2454
2455    /// Entities copied by the consolidation currently in flight, or `None` when
2456    /// none is running. Exposed for the pacing tests.
2457    #[cfg(test)]
2458    fn consolidation_progress(&self) -> Option<u64> {
2459        self.consolidation
2460            .read()
2461            .as_ref()
2462            .map(|run| run.entities_copied)
2463    }
2464}
2465
2466// Implement the Arc<RwLock<GrowingSegment>> as UnifiedSegment
2467// This is needed because we store growing segments in the sealed list after sealing
2468impl UnifiedSegment for Arc<RwLock<GrowingSegment>> {
2469    fn id(&self) -> SegmentId {
2470        self.read().id()
2471    }
2472
2473    fn state(&self) -> SegmentState {
2474        self.read().state()
2475    }
2476
2477    fn collection(&self) -> &str {
2478        // This is a limitation - we'd need to store collection in the Arc wrapper
2479        "unknown"
2480    }
2481
2482    fn stats(&self) -> SegmentStats {
2483        self.read().stats()
2484    }
2485
2486    fn entity_count(&self) -> usize {
2487        self.read().entity_count()
2488    }
2489
2490    fn contains(&self, id: EntityId) -> bool {
2491        self.read().contains(id)
2492    }
2493
2494    fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2495        // This is tricky with RwLock - we can't return a reference
2496        // In production, we'd use a different approach
2497        None
2498    }
2499
2500    fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2501        None
2502    }
2503
2504    fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2505        self.write().insert(entity)
2506    }
2507
2508    fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2509        self.write().update(entity)
2510    }
2511
2512    fn update_hot(
2513        &mut self,
2514        entity: UnifiedEntity,
2515        modified_columns: &[String],
2516    ) -> Result<(), SegmentError> {
2517        self.write().update_hot(entity, modified_columns)
2518    }
2519
2520    fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2521        self.write().delete(id)
2522    }
2523
2524    fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2525        self.read().get_metadata(id)
2526    }
2527
2528    fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2529        self.write().set_metadata(id, metadata)
2530    }
2531
2532    fn seal(&mut self) -> Result<(), SegmentError> {
2533        self.write().seal()
2534    }
2535
2536    fn should_seal(&self, config: &SegmentConfig) -> bool {
2537        self.read().should_seal(config)
2538    }
2539
2540    fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2541        // Cannot return iterator with RwLock
2542        Box::new(std::iter::empty())
2543    }
2544
2545    fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2546        Box::new(std::iter::empty())
2547    }
2548
2549    fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2550        self.read().filter_metadata(filters)
2551    }
2552}
2553
2554#[cfg(test)]
2555mod tests {
2556    use std::collections::HashSet;
2557
2558    use super::*;
2559    use crate::runtime::mvcc::SnapshotContext;
2560    use crate::storage::transaction::snapshot::{Snapshot, SnapshotManager};
2561    use crate::storage::unified::entity::{EntityData, EntityKind, RowData};
2562    use reddb_types::Value;
2563
2564    #[test]
2565    fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
2566        let manager = SegmentManager::with_config(
2567            "vectors",
2568            ManagerConfig {
2569                max_sealed_segments: 1,
2570                consolidation_entities_per_tick: 10_000,
2571                ..Default::default()
2572            },
2573        );
2574        for part in 0..3 {
2575            manager
2576                .bulk_insert(
2577                    (1..=300)
2578                        .map(|offset| {
2579                            UnifiedEntity::vector(
2580                                EntityId::new(part * 300 + offset),
2581                                "vectors",
2582                                vec![1.0, 0.0],
2583                            )
2584                        })
2585                        .collect(),
2586                )
2587                .expect("batch");
2588            if part < 2 {
2589                manager.force_seal().expect("seal");
2590            }
2591        }
2592        let original_growing = manager.growing.read().as_ref().expect("growing").clone();
2593        let mut changed = false;
2594        let mut seen = Vec::new();
2595        assert!(manager.scan_batches(
2596            None,
2597            || true,
2598            |batch| {
2599                assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
2600                assert!(
2601                    manager.growing.try_write().is_some(),
2602                    "consumer holds no topology guard"
2603                );
2604                assert!(
2605                    original_growing.try_write().is_some(),
2606                    "consumer holds no segment guard"
2607                );
2608                for entity in batch {
2609                    assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
2610                    seen.push(entity.id.raw());
2611                }
2612                if !changed {
2613                    changed = true;
2614                    manager.force_seal().expect("seal during scan");
2615                    manager
2616                        .insert(UnifiedEntity::vector(
2617                            EntityId::new(5000),
2618                            "vectors",
2619                            vec![1.0, 0.0],
2620                        ))
2621                        .expect("new write");
2622                    drain_maintenance(&manager);
2623                    assert!(manager.stats().consolidation.runs_completed > 0);
2624                    assert!(manager
2625                        .delete(EntityId::new(1))
2626                        .expect("delete after source retirement"));
2627                    assert!(manager.get(EntityId::new(1)).is_none());
2628                }
2629                true
2630            }
2631        ));
2632        seen.sort_unstable();
2633        assert_eq!(
2634            seen,
2635            (2..=900).collect::<Vec<_>>(),
2636            "no duplicate, missing or late candidates"
2637        );
2638        let mut after = Vec::new();
2639        assert!(manager.scan_batches(
2640            None,
2641            || true,
2642            |batch| {
2643                after.extend(batch.iter().map(|entity| entity.id.raw()));
2644                true
2645            }
2646        ));
2647        after.sort_unstable();
2648        let mut expected: Vec<_> = (2..=900).collect();
2649        expected.push(5000);
2650        assert_eq!(after, expected, "adopted entities have scan positions");
2651    }
2652
2653    #[test]
2654    fn bounded_scan_applies_snapshot_to_relocated_versions() {
2655        for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
2656            let manager = SegmentManager::with_config(
2657                "vectors",
2658                ManagerConfig {
2659                    max_sealed_segments: 1,
2660                    ..Default::default()
2661                },
2662            );
2663            for ids in [[1, 2], [3, 4]] {
2664                let entities = ids.map(|id| {
2665                    let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
2666                    entity.set_xmin(3);
2667                    entity
2668                });
2669                if bulk {
2670                    manager
2671                        .bulk_insert(entities.into())
2672                        .expect("initial versions");
2673                } else {
2674                    for entity in entities {
2675                        manager.insert(entity).expect("initial version");
2676                    }
2677                }
2678                manager.force_seal().expect("seal source");
2679            }
2680            let snapshot = SnapshotContext {
2681                snapshot: Snapshot {
2682                    xid: 10,
2683                    in_progress: HashSet::new(),
2684                },
2685                manager: Arc::new(SnapshotManager::new()),
2686                own_xids: HashSet::new(),
2687                requires_index_fallback: false,
2688                serializable_reader: None,
2689            };
2690            let mut changed = false;
2691            let mut seen = Vec::new();
2692            assert!(manager.scan_batches(
2693                Some(&snapshot),
2694                || true,
2695                |batch| {
2696                    seen.extend(batch.iter().map(|entity| entity.id.raw()));
2697                    if !changed {
2698                        changed = true;
2699                        if consolidate {
2700                            drain_maintenance(&manager);
2701                            assert!(manager.stats().consolidation.runs_completed > 0);
2702                        }
2703                        for (id, xmax) in [(3, 8), (4, 12)] {
2704                            let mut entity =
2705                                manager.get(EntityId::new(id)).expect("retained version");
2706                            entity.set_xmax(xmax);
2707                            manager
2708                                .update(entity)
2709                                .expect("update after source retirement");
2710                        }
2711                    }
2712                    true
2713                }
2714            ));
2715            seen.sort_unstable();
2716            assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
2717            let mut again = Vec::new();
2718            assert!(manager.scan_batches(
2719                Some(&snapshot),
2720                || true,
2721                |batch| {
2722                    again.extend(batch.iter().map(|entity| entity.id.raw()));
2723                    true
2724                }
2725            ));
2726            again.sort_unstable();
2727            assert_eq!(
2728                again,
2729                vec![1, 2, 4],
2730                "old slots do not duplicate relocated versions"
2731            );
2732            let mut current = Vec::new();
2733            assert!(manager.scan_batches(
2734                None,
2735                || true,
2736                |batch| {
2737                    current.extend(batch.iter().map(|entity| entity.id.raw()));
2738                    true
2739                }
2740            ));
2741            current.sort_unstable();
2742            assert_eq!(
2743                current,
2744                vec![1, 2],
2745                "current view excludes both deleted versions"
2746            );
2747        }
2748    }
2749
2750    #[test]
2751    fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
2752        let manager = SegmentManager::new("mixed");
2753        let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
2754        manager
2755            .bulk_insert((1..=300).map(vector).collect())
2756            .expect("flat");
2757        manager.insert(vector(500)).expect("hash insert");
2758        manager
2759            .bulk_insert(vec![vector(502), vector(505)])
2760            .expect("gapped bulk");
2761        manager.delete(EntityId::new(2)).expect("flat delete");
2762        manager.delete(EntityId::new(505)).expect("hash delete");
2763        let mut changed = false;
2764        let mut seen = Vec::new();
2765        assert!(manager.scan_batches(
2766            None,
2767            || true,
2768            |batch| {
2769                seen.extend(batch.iter().map(|entity| entity.id.raw()));
2770                if !changed {
2771                    changed = true;
2772                    manager
2773                        .insert(vector(700))
2774                        .expect("append after cursor boundary");
2775                    manager
2776                        .update(UnifiedEntity::graph_node(
2777                            EntityId::new(500),
2778                            "mixed",
2779                            "Node",
2780                            HashMap::new(),
2781                        ))
2782                        .expect("structural update");
2783                }
2784                true
2785            }
2786        ));
2787        seen.sort_unstable();
2788        let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
2789        expected.extend([500, 502]);
2790        assert_eq!(seen, expected);
2791    }
2792
2793    #[test]
2794    fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
2795        let manager = SegmentManager::new("vectors");
2796        let mut entities = Vec::new();
2797        for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
2798            let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
2799            entity.set_xmin(xmin);
2800            entity.set_xmax(xmax);
2801            entities.push(entity);
2802        }
2803        manager.bulk_insert(entities).expect("versions");
2804        manager.force_seal().expect("seal");
2805        let snapshot = SnapshotContext {
2806            snapshot: Snapshot {
2807                xid: 10,
2808                in_progress: HashSet::from([4]),
2809            },
2810            manager: Arc::new(SnapshotManager::new()),
2811            own_xids: HashSet::new(),
2812            requires_index_fallback: false,
2813            serializable_reader: None,
2814        };
2815        let mut work = 0;
2816        let mut callbacks = 0;
2817        assert!(!manager.scan_batches(
2818            Some(&snapshot),
2819            || {
2820                work += 1;
2821                work <= 2
2822            },
2823            |_| {
2824                callbacks += 1;
2825                true
2826            }
2827        ));
2828        assert_eq!(callbacks, 0, "a stopped batch is not consumed");
2829        for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
2830            let mut seen = Vec::new();
2831            assert!(manager.scan_batches(
2832                view,
2833                || true,
2834                |batch| {
2835                    seen.extend(batch.iter().map(|entity| entity.id.raw()));
2836                    true
2837                }
2838            ));
2839            seen.sort_unstable();
2840            assert_eq!(seen, expected);
2841        }
2842    }
2843
2844    #[test]
2845    fn test_manager_basic() {
2846        let manager = SegmentManager::new("test_collection");
2847
2848        let entity = UnifiedEntity::table_row(
2849            manager.next_entity_id(),
2850            "users",
2851            1,
2852            vec![Value::text("Alice".to_string())],
2853        );
2854
2855        let id = manager.insert(entity).unwrap();
2856        assert!(manager.get(id).is_some());
2857        assert_eq!(manager.count(), 1);
2858    }
2859
2860    #[test]
2861    fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
2862        let manager = Arc::new(SegmentManager::new("mixed"));
2863        manager
2864            .bulk_insert(
2865                (1..=128)
2866                    .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
2867                    .collect(),
2868            )
2869            .expect("vector-only segment");
2870        manager.force_seal().expect("seal vectors");
2871        let mut old =
2872            UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
2873        old.set_xmin(3);
2874        old.set_xmax(12);
2875        let mut new =
2876            UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
2877        new.set_xmin(12);
2878        let edge = UnifiedEntity::graph_edge(
2879            EntityId::new(202),
2880            "mixed",
2881            "a",
2882            "b",
2883            1.0,
2884            Default::default(),
2885        );
2886        manager
2887            .bulk_insert(vec![old, new, edge])
2888            .expect("physical graph versions");
2889        manager.force_seal().expect("seal graph");
2890        manager
2891            .insert(UnifiedEntity::graph_node(
2892                EntityId::new(203),
2893                "mixed",
2894                "person",
2895                Default::default(),
2896            ))
2897            .expect("growing node");
2898        let snapshot = SnapshotContext {
2899            snapshot: Snapshot {
2900                xid: 10,
2901                in_progress: HashSet::new(),
2902            },
2903            manager: Arc::new(SnapshotManager::new()),
2904            own_xids: HashSet::new(),
2905            requires_index_fallback: false,
2906            serializable_reader: None,
2907        };
2908        std::thread::spawn(move || {
2909            for view in [None, Some(&snapshot)] {
2910                for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
2911                    let expected: HashSet<_> = manager
2912                        .scan(view, |item| kind.matches(&item.kind))
2913                        .into_iter()
2914                        .map(|item| item.id)
2915                        .collect();
2916                    let actual: HashSet<_> = manager
2917                        .scan_graph_kind(view, kind)
2918                        .into_iter()
2919                        .map(|item| item.id)
2920                        .collect();
2921                    assert_eq!(actual, expected);
2922                    let mut visited = 0;
2923                    let mut candidates = HashSet::new();
2924                    manager.scan_graph_candidates_for_each(
2925                        view,
2926                        kind,
2927                        || true,
2928                        |item| {
2929                            visited += 1;
2930                            if kind.matches(&item.kind) {
2931                                candidates.insert(item.id);
2932                            }
2933                            true
2934                        },
2935                    );
2936                    assert_eq!(candidates, expected);
2937                    assert_eq!(
2938                        visited,
2939                        if matches!(kind, GraphEntityKind::Node) {
2940                            3
2941                        } else {
2942                            2
2943                        }
2944                    );
2945                }
2946            }
2947            let ids: HashSet<_> = manager
2948                .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
2949                .into_iter()
2950                .map(|item| item.id.raw())
2951                .collect();
2952            assert_eq!(ids, HashSet::from([200, 203]));
2953            let ids: HashSet<_> = manager
2954                .scan_graph_kind(None, GraphEntityKind::Node)
2955                .into_iter()
2956                .map(|item| item.id.raw())
2957                .collect();
2958            assert_eq!(ids, HashSet::from([201, 203]));
2959            let mut visits = 0;
2960            manager.scan_graph_candidates_for_each(
2961                None,
2962                GraphEntityKind::Node,
2963                || true,
2964                |_| {
2965                    visits += 1;
2966                    false
2967                },
2968            );
2969            assert_eq!(
2970                visits, 1,
2971                "early termination must also stop subsequent segments"
2972            );
2973        })
2974        .join()
2975        .expect("snapshot worker");
2976    }
2977
2978    #[test]
2979    fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
2980        let manager = Arc::new(SegmentManager::new("accounts"));
2981        let cases = [
2982            ("legacy row", 0, 0, true),
2983            ("committed creator", 3, 0, true),
2984            ("future creator", 11, 0, false),
2985            ("in-progress creator", 4, 0, false),
2986            ("committed deleter", 3, 7, false),
2987            ("future deleter", 3, 11, true),
2988            ("in-progress deleter", 3, 4, true),
2989        ];
2990
2991        for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
2992            let row_id = index as u64 + 1;
2993            let mut entity = UnifiedEntity::table_row(
2994                EntityId::new(row_id),
2995                "accounts",
2996                row_id,
2997                vec![Value::Integer(row_id as i64)],
2998            );
2999            entity.set_xmin(*xmin);
3000            entity.set_xmax(*xmax);
3001            manager.insert(entity).unwrap();
3002        }
3003
3004        let snapshot = SnapshotContext {
3005            snapshot: Snapshot {
3006                xid: 10,
3007                in_progress: HashSet::from([4]),
3008            },
3009            manager: Arc::new(SnapshotManager::new()),
3010            own_xids: HashSet::new(),
3011            requires_index_fallback: false,
3012            serializable_reader: None,
3013        };
3014
3015        let raw_visible_row_ids: HashSet<u64> = manager
3016            .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
3017            .into_iter()
3018            .filter_map(|entity| match entity.kind {
3019                EntityKind::TableRow { row_id, .. } => Some(row_id),
3020                _ => None,
3021            })
3022            .collect();
3023
3024        let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
3025            .join()
3026            .unwrap();
3027        let visible_row_ids: HashSet<u64> = scanned
3028            .into_iter()
3029            .filter_map(|entity| match entity.kind {
3030                EntityKind::TableRow { row_id, .. } => Some(row_id),
3031                _ => None,
3032            })
3033            .collect();
3034        assert_eq!(visible_row_ids, raw_visible_row_ids);
3035
3036        for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
3037            let row_id = index as u64 + 1;
3038            assert_eq!(
3039                visible_row_ids.contains(&row_id),
3040                *expected_visible,
3041                "{name} visibility mismatch"
3042            );
3043        }
3044    }
3045
3046    #[test]
3047    fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
3048        // A `None` snapshot is the autocommit / prepared-fast-path read: it
3049        // must keep the resolver fallback (xmax == 0), never show every
3050        // version, and never panic (issue #2136 review).
3051        let manager = SegmentManager::new("accounts");
3052        let mut live =
3053            UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
3054        live.set_xmin(2);
3055        live.set_xmax(0);
3056        manager.insert(live).unwrap();
3057
3058        let mut superseded =
3059            UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
3060        superseded.set_xmin(2);
3061        superseded.set_xmax(9);
3062        manager.insert(superseded).unwrap();
3063
3064        let rows: Vec<u64> = manager
3065            .scan(None, |_| true)
3066            .into_iter()
3067            .filter_map(|entity| match entity.kind {
3068                EntityKind::TableRow { row_id, .. } => Some(row_id),
3069                _ => None,
3070            })
3071            .collect();
3072        assert_eq!(
3073            rows,
3074            vec![1],
3075            "frameless scan must hide superseded versions"
3076        );
3077    }
3078
3079    #[test]
3080    fn query_scan_callers_use_snapshot_api_for_visibility() {
3081        let table = include_str!("../../runtime/query_exec/table.rs");
3082        let aggregate = include_str!("../../runtime/query_exec/aggregate.rs");
3083
3084        assert_eq!(
3085            table.matches(".for_each_entity(|entity|").count(),
3086            1,
3087            "only the visibility-independent schema inspection may use the raw table iterator"
3088        );
3089        assert!(!table.contains(".query_all_zoned_with_stats("));
3090        assert!(!table.contains(".for_each_entity_zoned_with_stats("));
3091        assert!(!table.contains("TableRowMvccReadResolver::captured"));
3092
3093        assert!(!aggregate.contains("TableRowMvccReadResolver"));
3094        assert!(!aggregate.contains(".for_each_entity("));
3095        assert!(!aggregate.contains(".fold_entities_parallel("));
3096    }
3097
3098    #[test]
3099    fn remaining_visibility_scan_callers_use_snapshot_api() {
3100        let aggregate_pushdown =
3101            include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs");
3102        let ask = include_str!("../../runtime/ask_pipeline.rs");
3103        let dml = include_str!("../../runtime/impl_dml.rs");
3104        let dml_targets = include_str!("../../runtime/dml_target_scan.rs");
3105        let events = include_str!("../../runtime/impl_events.rs");
3106        let graph_commands = include_str!("../../runtime/impl_graph_commands.rs");
3107        let graph_dsl = include_str!("../../runtime/graph_dsl.rs");
3108        let graph_tvf = include_str!("../../runtime/graph_tvf.rs");
3109        let kv = include_str!("../../runtime/impl_kv.rs");
3110        let queue = include_str!("../../runtime/impl_queue.rs");
3111        let primary_queue = include_str!("../../runtime/primary_queue_store.rs");
3112        let record_search = include_str!("../../runtime/record_search.rs");
3113        let search = include_str!("../../runtime/impl_search.rs");
3114        let vector = include_str!("../../runtime/query_exec/vector.rs");
3115
3116        assert!(!aggregate_pushdown.contains("TableRowMvccReadResolver"));
3117        assert!(!aggregate_pushdown.contains(".for_each_entity("));
3118        assert!(!ask.contains("for entity in manager.query_all(|_| true)"));
3119        assert!(!dml.contains("let entities = manager.query_all(|_| true);"));
3120        // The DML target scan keeps raw iteration on purpose: its
3121        // visibility is kind-conditional (live table rows at the tip
3122        // version, everything else under the snapshot), which the uniform
3123        // scan API cannot express. `candidate_visible` is the single gate.
3124        assert_eq!(
3125            dml_targets
3126                .matches("self.candidate_visible(snapshot.as_ref(), entity)")
3127                .count(),
3128            2,
3129            "both raw DML iterations must route through candidate_visible"
3130        );
3131        // The event-backfill dedup set is deliberately visibility-free:
3132        // it must include consumed messages or backfill re-enqueues them.
3133        assert_eq!(
3134            events.matches("manager.query_all(").count(),
3135            1,
3136            "only the visibility-free backfill dedup may use query_all"
3137        );
3138        assert!(!graph_commands.contains("TableRowMvccReadResolver"));
3139        assert!(!graph_dsl.contains("entity_visible_with_context(snap_ctx.as_ref(), e)"));
3140        assert!(!graph_tvf.contains("entity_visible_with_context(snap_ctx.as_ref(), &entity)"));
3141        assert!(!kv.contains("TableRowMvccReadResolver"));
3142        assert!(!queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3143        assert!(!primary_queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3144        assert!(!record_search.contains("TableRowMvccReadResolver"));
3145        assert!(!search.contains("for entity in manager.query_all(|_| true)"));
3146        assert!(!vector.contains("manager.query_all("));
3147    }
3148
3149    /// Strip `#[cfg(test)] mod … { … }` blocks so the raw-iteration gate
3150    /// only sees production code. Blocks are matched at column zero and
3151    /// closed by brace balance, which is why a file may carry a test
3152    /// module in the middle and still keep production code after it.
3153    fn production_source(src: &str) -> String {
3154        let mut kept: Vec<&str> = Vec::new();
3155        let lines: Vec<&str> = src.lines().collect();
3156        let mut idx = 0;
3157        while idx < lines.len() {
3158            let starts_test_mod = lines[idx].starts_with("#[cfg(test)]")
3159                && lines
3160                    .get(idx + 1)
3161                    .is_some_and(|next| next.starts_with("mod "));
3162            if !starts_test_mod {
3163                kept.push(lines[idx]);
3164                idx += 1;
3165                continue;
3166            }
3167            idx += 1;
3168            let mut depth: isize = 0;
3169            while idx < lines.len() {
3170                depth += lines[idx].matches('{').count() as isize;
3171                depth -= lines[idx].matches('}').count() as isize;
3172                idx += 1;
3173                if depth <= 0 {
3174                    break;
3175                }
3176            }
3177        }
3178        kept.join("\n")
3179    }
3180
3181    /// Count raw (visibility-bypassing) iteration entry points.
3182    fn raw_iteration_sites(src: &str) -> usize {
3183        src.matches("query_all(").count()
3184            + src.matches(".for_each_entity").count()
3185            + src.matches("fold_entities_parallel(").count()
3186    }
3187
3188    #[test]
3189    fn raw_iteration_confined_to_maintenance_allowlist() {
3190        // The query paths migrated by Spec #2107 batches 1 and 2. Inside
3191        // this perimeter raw iteration is maintenance-only: it bypasses
3192        // MVCC visibility entirely, so every remaining site is either an
3193        // internal catalog/bookkeeping read that must observe all physical
3194        // versions, or one of the three exceptions the batch reviews
3195        // documented. A new raw-iteration site changes the count and fails
3196        // this gate — route it through `SegmentManager::scan*` instead.
3197        let allowlist: &[(&str, &str, usize, &str)] = &[
3198            (
3199                "runtime/query_exec/table.rs",
3200                include_str!("../../runtime/query_exec/table.rs"),
3201                1,
3202                "documented exception: schema inspection (does the collection ever \
3203                 observe these fields?) is visibility-independent",
3204            ),
3205            (
3206                "runtime/query_exec/aggregate.rs",
3207                include_str!("../../runtime/query_exec/aggregate.rs"),
3208                0,
3209                "fully migrated by batch 1",
3210            ),
3211            (
3212                "runtime/query_exec/aggregate_pushdown_dispatch.rs",
3213                include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs"),
3214                0,
3215                "fully migrated by batch 1",
3216            ),
3217            (
3218                "runtime/query_exec/vector.rs",
3219                include_str!("../../runtime/query_exec/vector.rs"),
3220                0,
3221                "fully migrated by batch 2",
3222            ),
3223            (
3224                "runtime/ask_pipeline.rs",
3225                include_str!("../../runtime/ask_pipeline.rs"),
3226                0,
3227                "fully migrated by batch 2",
3228            ),
3229            (
3230                "runtime/impl_dml.rs",
3231                include_str!("../../runtime/impl_dml.rs"),
3232                1,
3233                "documented exception: ON CONFLICT resolves a unique key against \
3234                 physical versions on purpose. A snapshot scan cannot see a row an \
3235                 uncommitted concurrent writer already holds the key for, which is \
3236                 exactly the case that must raise a serialization conflict rather \
3237                 than insert a duplicate; the same pass reaps index entries left by \
3238                 aborted writers. \
3239                 TODO(#2263 follow-up): narrow the candidates through the unique \
3240                 index instead of walking the collection",
3241            ),
3242            (
3243                "runtime/dml_target_scan.rs",
3244                include_str!("../../runtime/dml_target_scan.rs"),
3245                2,
3246                "documented exception: kind-conditional visibility the uniform scan \
3247                 API cannot express; both iterations gate on candidate_visible",
3248            ),
3249            (
3250                "runtime/impl_events.rs",
3251                include_str!("../../runtime/impl_events.rs"),
3252                1,
3253                "documented exception: backfill dedup must include consumed messages, \
3254                 so it is deliberately visibility-free",
3255            ),
3256            (
3257                "runtime/impl_graph_commands.rs",
3258                include_str!("../../runtime/impl_graph_commands.rs"),
3259                0,
3260                "native node properties resolve one snapshot-visible, RLS-admitted version",
3261            ),
3262            (
3263                "runtime/graph_dsl.rs",
3264                include_str!("../../runtime/graph_dsl.rs"),
3265                1,
3266                "remaining legacy debt: unused edge-property helper; search expansion \
3267                 now uses snapshot visitors and RLS-gated hydration",
3268            ),
3269            (
3270                "runtime/graph_tvf.rs",
3271                include_str!("../../runtime/graph_tvf.rs"),
3272                0,
3273                "fully migrated by batch 2",
3274            ),
3275            (
3276                "runtime/impl_kv.rs",
3277                include_str!("../../runtime/impl_kv.rs"),
3278                4,
3279                "two non-versioned-collection fallbacks (the versioned arms scan under \
3280                 the captured snapshot). Vault reads use snapshot scans. Offline vault \
3281                 export preserves all \
3282                 physical versions across tenants; restore's emptiness check must reject \
3283                 every physical row, including versions invisible to a snapshot",
3284            ),
3285            (
3286                "runtime/impl_queue.rs",
3287                include_str!("../../runtime/impl_queue.rs"),
3288                10,
3289                "maintenance: red_queue_meta bookkeeping (config, groups, acks, \
3290                 pending, dedup TTL reaper, gauge counts) — queue message bodies are \
3291                 read through scan. \
3292                 TODO(#2138 follow-up): QUEUE PENDING projects meta rows to clients",
3293            ),
3294            (
3295                "runtime/primary_queue_store.rs",
3296                include_str!("../../runtime/primary_queue_store.rs"),
3297                3,
3298                "maintenance: generic red_queue_meta read/delete helpers and the \
3299                 next-position allocator, which must not collide with rows invisible \
3300                 to the writer's snapshot",
3301            ),
3302            (
3303                "runtime/record_search.rs",
3304                include_str!("../../runtime/record_search.rs"),
3305                0,
3306                "fully migrated by batch 2",
3307            ),
3308            (
3309                "runtime/search_graph.rs",
3310                include_str!("../../runtime/search_graph.rs"),
3311                0,
3312                "search graph preparation uses snapshot visitors and gated hydration",
3313            ),
3314            (
3315                "runtime/impl_search.rs",
3316                include_str!("../../runtime/impl_search.rs"),
3317                1,
3318                "maintenance: ask-audit retention sweep must see every physical row",
3319            ),
3320        ];
3321
3322        for (path, src, allowed, reason) in allowlist {
3323            let found = raw_iteration_sites(&production_source(src));
3324            assert_eq!(
3325                found, *allowed,
3326                "{path}: {found} raw-iteration sites, allowlist permits {allowed} \
3327                 ({reason}). Raw iteration bypasses MVCC visibility — use \
3328                 SegmentManager::scan / scan_for_each / scan_fold_parallel / \
3329                 scan_zoned_with_stats, or extend the allowlist with a maintenance \
3330                 justification."
3331            );
3332        }
3333    }
3334
3335    #[test]
3336    fn explicit_snapshot_scan_ignores_ambient_thread_local() {
3337        let manager = Arc::new(SegmentManager::new("accounts"));
3338        // (xmin, xmax) — mix of committed, in-progress and future writers.
3339        let versions = [(0, 0), (3, 0), (11, 0), (4, 0), (3, 7), (3, 11), (3, 4)];
3340        for (index, (xmin, xmax)) in versions.iter().enumerate() {
3341            let row_id = index as u64 + 1;
3342            let mut entity = UnifiedEntity::table_row(
3343                EntityId::new(row_id),
3344                "accounts",
3345                row_id,
3346                vec![Value::Integer(row_id as i64)],
3347            );
3348            entity.set_xmin(*xmin);
3349            entity.set_xmax(*xmax);
3350            manager.insert(entity).unwrap();
3351        }
3352
3353        let explicit = SnapshotContext {
3354            snapshot: Snapshot {
3355                xid: 10,
3356                in_progress: HashSet::from([4]),
3357            },
3358            manager: Arc::new(SnapshotManager::new()),
3359            own_xids: HashSet::new(),
3360            requires_index_fallback: false,
3361            serializable_reader: None,
3362        };
3363        // A decoy statement frame that would hide almost everything if the
3364        // scan ever consulted ambient state instead of its argument.
3365        let decoy = SnapshotContext {
3366            snapshot: Snapshot {
3367                xid: 1,
3368                in_progress: HashSet::new(),
3369            },
3370            manager: Arc::new(SnapshotManager::new()),
3371            own_xids: HashSet::new(),
3372            requires_index_fallback: false,
3373            serializable_reader: None,
3374        };
3375
3376        let row_ids = |entities: Vec<UnifiedEntity>| -> HashSet<u64> {
3377            entities
3378                .into_iter()
3379                .filter_map(|entity| match entity.kind {
3380                    EntityKind::TableRow { row_id, .. } => Some(row_id),
3381                    _ => None,
3382                })
3383                .collect()
3384        };
3385
3386        let guard = crate::runtime::execution_context::CurrentSnapshotGuard::install(decoy);
3387        let on_installing_thread = row_ids(manager.scan(Some(&explicit), |_| true));
3388        drop(guard);
3389
3390        let spawned_manager = Arc::clone(&manager);
3391        let spawned_snapshot = explicit.clone();
3392        let on_fresh_thread = std::thread::spawn(move || {
3393            row_ids(spawned_manager.scan(Some(&spawned_snapshot), |_| true))
3394        })
3395        .join()
3396        .unwrap();
3397
3398        assert_eq!(
3399            on_installing_thread, on_fresh_thread,
3400            "an explicit snapshot must produce the same visible set on any thread"
3401        );
3402        assert_eq!(on_installing_thread, HashSet::from([1, 2, 6, 7]));
3403
3404        // The raw iterators are the maintenance contract: no visibility
3405        // filtering at all, on any thread, regardless of installed frames.
3406        let raw_manager = Arc::clone(&manager);
3407        let raw_count = std::thread::spawn(move || raw_manager.query_all(|_| true).len())
3408            .join()
3409            .unwrap();
3410        assert_eq!(
3411            raw_count,
3412            versions.len(),
3413            "raw iteration must expose every physical version"
3414        );
3415    }
3416
3417    #[test]
3418    fn bloom_may_contain_key_probes_inserted_rid_bytes() {
3419        let manager = SegmentManager::new("test_collection");
3420
3421        let entity = UnifiedEntity::table_row(
3422            manager.next_entity_id(),
3423            "users",
3424            1,
3425            vec![Value::text("Alice".to_string())],
3426        );
3427
3428        let id = manager.insert(entity).unwrap();
3429
3430        assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3431        assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
3432        assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
3433
3434        assert!(manager.delete(id).unwrap());
3435        assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3436    }
3437
3438    #[test]
3439    fn zoned_scan_reports_pruned_segments_and_preserves_results() {
3440        let config = ManagerConfig {
3441            segment_config: SegmentConfig {
3442                max_entities: 2,
3443                ..Default::default()
3444            },
3445            ..Default::default()
3446        };
3447        let manager = SegmentManager::with_config("events", config);
3448
3449        for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
3450            let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
3451            let entity = UnifiedEntity::new(
3452                manager.next_entity_id(),
3453                EntityKind::TableRow {
3454                    table: "events".into(),
3455                    row_id: (idx + 1) as u64,
3456                },
3457                EntityData::Row(row),
3458            );
3459            manager.insert(entity).unwrap();
3460        }
3461
3462        let full = manager.query_all(|entity| match &entity.data {
3463            EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
3464            _ => false,
3465        });
3466        let probe = Value::Integer(100);
3467        let (pruned, stats) =
3468            manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
3469                match &entity.data {
3470                    EntityData::Row(row) => {
3471                        matches!(row.get_field("ts"), Some(Value::Integer(100)))
3472                    }
3473                    _ => false,
3474                }
3475            });
3476
3477        assert_eq!(pruned.len(), full.len());
3478        assert_eq!(pruned[0].id, full[0].id);
3479        assert_eq!(stats.segments_total, 3);
3480        assert_eq!(stats.segments_scanned, 2);
3481        assert_eq!(stats.segments_pruned, 1);
3482    }
3483
3484    #[test]
3485    fn test_manager_auto_seal() {
3486        let config = ManagerConfig {
3487            segment_config: SegmentConfig {
3488                max_entities: 2,
3489                ..Default::default()
3490            },
3491            ..Default::default()
3492        };
3493
3494        let manager = SegmentManager::with_config("test", config);
3495
3496        // Insert first entity
3497        manager
3498            .insert(UnifiedEntity::vector(
3499                manager.next_entity_id(),
3500                "v",
3501                vec![0.1],
3502            ))
3503            .unwrap();
3504
3505        // Insert second entity (triggers seal check)
3506        manager
3507            .insert(UnifiedEntity::vector(
3508                manager.next_entity_id(),
3509                "v",
3510                vec![0.2],
3511            ))
3512            .unwrap();
3513
3514        // Insert third entity (should trigger auto-seal of first segment)
3515        manager
3516            .insert(UnifiedEntity::vector(
3517                manager.next_entity_id(),
3518                "v",
3519                vec![0.3],
3520            ))
3521            .unwrap();
3522
3523        let stats = manager.stats();
3524        assert_eq!(stats.total_entities, 3);
3525    }
3526
3527    #[test]
3528    fn test_manager_delete() {
3529        let manager = SegmentManager::new("test");
3530
3531        let id = manager
3532            .insert(UnifiedEntity::vector(
3533                manager.next_entity_id(),
3534                "v",
3535                vec![0.1],
3536            ))
3537            .unwrap();
3538
3539        assert!(manager.get(id).is_some());
3540        assert!(manager.delete(id).unwrap());
3541        assert!(manager.get(id).is_none());
3542    }
3543
3544    #[test]
3545    fn test_manager_metadata() {
3546        let manager = SegmentManager::new("test");
3547
3548        let id = manager
3549            .insert(UnifiedEntity::table_row(
3550                manager.next_entity_id(),
3551                "hosts",
3552                1,
3553                vec![Value::text("192.168.1.1".to_string())],
3554            ))
3555            .unwrap();
3556
3557        let mut meta = Metadata::new();
3558        meta.set(
3559            "os",
3560            super::super::metadata::MetadataValue::String("linux".to_string()),
3561        );
3562
3563        manager.set_metadata(id, meta).unwrap();
3564
3565        let retrieved = manager.get_metadata(id).unwrap();
3566        assert!(retrieved.has("os"));
3567    }
3568
3569    #[test]
3570    fn test_manager_query_by_kind() {
3571        let manager = SegmentManager::new("test");
3572
3573        manager
3574            .insert(UnifiedEntity::table_row(
3575                manager.next_entity_id(),
3576                "hosts",
3577                1,
3578                vec![],
3579            ))
3580            .unwrap();
3581
3582        manager
3583            .insert(UnifiedEntity::vector(
3584                manager.next_entity_id(),
3585                "embeddings",
3586                vec![0.1],
3587            ))
3588            .unwrap();
3589
3590        manager
3591            .insert(UnifiedEntity::table_row(
3592                manager.next_entity_id(),
3593                "hosts",
3594                2,
3595                vec![],
3596            ))
3597            .unwrap();
3598
3599        let rows = manager.get_by_kind("table");
3600        assert_eq!(rows.len(), 2);
3601
3602        let vectors = manager.get_by_kind("vector");
3603        assert_eq!(vectors.len(), 1);
3604    }
3605
3606    // ========================================================================
3607    // Consolidation — ADR 0073 §5
3608    // ========================================================================
3609
3610    /// A manager whose sealed set never trips the count trigger, so tests can
3611    /// exercise the tombstone/fragmentation triggers in isolation.
3612    fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
3613        SegmentManager::with_config(
3614            "test",
3615            ManagerConfig {
3616                max_sealed_segments: 1_000,
3617                consolidation_entities_per_tick: entities_per_tick,
3618                ..Default::default()
3619            },
3620        )
3621    }
3622
3623    /// Bulk-load `count` rows into a fresh segment and seal it. Bulk-inserted
3624    /// entities land in flat storage, where a tombstone keeps its payload
3625    /// resident — which is the memory consolidation reclaims.
3626    fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
3627        let entities: Vec<UnifiedEntity> = (0..count)
3628            .map(|i| {
3629                UnifiedEntity::table_row(
3630                    manager.next_entity_id(),
3631                    "users",
3632                    i as u64,
3633                    vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
3634                )
3635            })
3636            .collect();
3637        let ids = manager.bulk_insert(entities).expect("bulk insert");
3638        manager.force_seal().expect("seal");
3639        ids
3640    }
3641
3642    fn live_ids(manager: &SegmentManager) -> Vec<u64> {
3643        let mut ids: Vec<u64> = manager
3644            .query_all(|_| true)
3645            .into_iter()
3646            .map(|entity| entity.id.raw())
3647            .collect();
3648        ids.sort_unstable();
3649        ids
3650    }
3651
3652    fn drain_maintenance(manager: &SegmentManager) -> usize {
3653        let mut ticks = 0;
3654        while manager.consolidation.read().is_some() || ticks == 0 {
3655            manager.run_maintenance().expect("maintenance");
3656            ticks += 1;
3657            assert!(ticks < 1_000, "consolidation failed to converge");
3658        }
3659        ticks
3660    }
3661
3662    #[test]
3663    fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
3664        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3665        let ids = manager
3666            .bulk_insert(
3667                (1..=1024)
3668                    .map(|id| {
3669                        UnifiedEntity::graph_edge(
3670                            EntityId::new(id),
3671                            "step",
3672                            "2000",
3673                            "3000",
3674                            1.0,
3675                            HashMap::new(),
3676                        )
3677                    })
3678                    .collect(),
3679            )
3680            .expect("edges");
3681        manager.force_seal().expect("seal");
3682        for id in ids.iter().take(300) {
3683            manager.delete(*id).expect("fragment source");
3684        }
3685        let source = Arc::downgrade(&manager.sealed.read()[0]);
3686        let mut seen = Vec::new();
3687        let mut changed = false;
3688        let mut metadata_bytes = 0;
3689        assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
3690            || true,
3691            |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
3692            |batch| {
3693                assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3694                assert!(manager.growing.try_write().is_some(), "no growing topology guard");
3695                assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
3696                assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
3697                if !changed {
3698                    changed = true;
3699                    drain_maintenance(&manager);
3700                    assert!(manager.stats().consolidation.runs_completed > 0);
3701                    assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
3702                    let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
3703                    if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
3704                    manager.update(updated).expect("update authoritative merged copy");
3705                }
3706                for id in batch {
3707                    let entity = manager.get(*id).expect("hydrate current manager");
3708                    if id.raw() == 1024 {
3709                        assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
3710                    }
3711                    seen.push(id.raw());
3712                }
3713                Ok(true)
3714            }).expect("cursor"));
3715        assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
3716        assert!(
3717            source.upgrade().is_none(),
3718            "finished cursor releases retired source"
3719        );
3720        assert!(
3721            metadata_bytes < 4096,
3722            "fixed ID buffer plus source handles: {metadata_bytes}"
3723        );
3724    }
3725
3726    #[test]
3727    fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
3728        let manager = SegmentManager::new("graph");
3729        manager
3730            .bulk_insert(
3731                (1..=1024)
3732                    .map(|id| {
3733                        UnifiedEntity::graph_edge(
3734                            EntityId::new(id),
3735                            "step",
3736                            "2000",
3737                            "3000",
3738                            1.0,
3739                            HashMap::new(),
3740                        )
3741                    })
3742                    .collect(),
3743            )
3744            .expect("edges");
3745        let mut delivered = 0;
3746        let denied = manager.visit_graph_index_batches(
3747            GraphEntityKind::Edge,
3748            &[EntityId::new(2000)],
3749            || true,
3750            |_| Err("out of memory"),
3751            |ids| {
3752                delivered += ids.len();
3753                Ok(true)
3754            },
3755        );
3756        assert_eq!(denied, Err("out of memory"));
3757        assert_eq!(delivered, 0);
3758        let mut work = 0;
3759        assert!(!manager
3760            .visit_graph_index_batches(
3761                GraphEntityKind::Edge,
3762                &[EntityId::new(2000)],
3763                || {
3764                    work += 1;
3765                    work < 100
3766                },
3767                |_| Ok::<_, ()>(()),
3768                |ids| {
3769                    delivered += ids.len();
3770                    Ok(true)
3771                }
3772            )
3773            .expect("cancelled cursor"));
3774        assert_eq!(
3775            delivered, 0,
3776            "cancelled partial batch cannot reach consumer"
3777        );
3778        assert!(!manager
3779            .visit_graph_index_batches(
3780                GraphEntityKind::Edge,
3781                &[EntityId::new(2000)],
3782                || true,
3783                |_| Ok::<_, ()>(()),
3784                |ids| {
3785                    delivered += ids.len();
3786                    Ok(false)
3787                }
3788            )
3789            .expect("consumer stop"));
3790        assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
3791    }
3792
3793    #[test]
3794    fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
3795        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3796        let ids = seal_bulk_segment(&manager, 100);
3797        for id in ids.iter().take(40) {
3798            manager.delete(*id).expect("delete before scan");
3799        }
3800        let source_bytes = manager.resident_bytes();
3801        let source = Arc::downgrade(&manager.sealed.read()[0]);
3802        std::thread::scope(|scope| {
3803            let (ready, waiting) = std::sync::mpsc::channel();
3804            let (release_first, first_wait) = std::sync::mpsc::channel();
3805            let (release_last, last_wait) = std::sync::mpsc::channel();
3806            let reader = |resume: std::sync::mpsc::Receiver<()>,
3807                          ready: std::sync::mpsc::Sender<()>| {
3808                assert!(manager.scan_batches(
3809                    None,
3810                    || true,
3811                    |_| {
3812                        ready.send(()).expect("signal captured cursor");
3813                        let _ = resume.recv();
3814                        true
3815                    }
3816                ));
3817            };
3818            let ready_first = ready.clone();
3819            let first = scope.spawn(move || reader(first_wait, ready_first));
3820            let last = scope.spawn(move || reader(last_wait, ready));
3821            for _ in 0..2 {
3822                waiting
3823                    .recv_timeout(std::time::Duration::from_secs(10))
3824                    .expect("reader ready");
3825            }
3826            drain_maintenance(&manager);
3827            assert!(
3828                source.upgrade().is_some(),
3829                "both cursors pin the retired source"
3830            );
3831            let active_bytes = manager.sealed.read()[0].read().resident_bytes();
3832            let held = manager.resident_bytes();
3833            assert!(
3834                held >= active_bytes + source_bytes,
3835                "retired source remains resident"
3836            );
3837            release_first.send(()).expect("release first reader");
3838            first.join().expect("first reader");
3839            assert_eq!(
3840                manager.resident_bytes(),
3841                held,
3842                "one reader still pins the source"
3843            );
3844            release_last.send(()).expect("release last reader");
3845            last.join().expect("last reader");
3846            assert!(
3847                source.upgrade().is_none(),
3848                "accounting must not retain payloads"
3849            );
3850            assert_eq!(held - manager.resident_bytes(), source_bytes);
3851            assert_eq!(
3852                manager.stats().total_memory_bytes as u64,
3853                manager.resident_bytes()
3854            );
3855        });
3856    }
3857
3858    #[test]
3859    fn retained_memory_gates_runtime_growth_until_the_cursor_is_cancelled() {
3860        use crate::storage::memory_pools::MemoryPool;
3861        use crate::{RedDBOptions, RedDBRuntime};
3862
3863        let budget = 128 * 1024;
3864        let runtime =
3865            RedDBRuntime::with_options(RedDBOptions::in_memory().with_memory_budget(budget))
3866                .expect("runtime");
3867        runtime
3868            .db()
3869            .store()
3870            .create_collection("test")
3871            .expect("collection");
3872        let manager = runtime
3873            .db()
3874            .store()
3875            .get_collection("test")
3876            .expect("manager");
3877        let ids = seal_bulk_segment(&manager, 100);
3878        for id in ids.iter().take(40) {
3879            manager.delete(*id).expect("delete before scan");
3880        }
3881        let source_bytes = manager.resident_bytes();
3882        let mut growth = 0;
3883        let mut held_used = 0;
3884        assert!(!manager.scan_batches(
3885            None,
3886            || true,
3887            |_| {
3888                drain_maintenance(&manager);
3889                runtime.refresh_memory_accounting();
3890                // Other collections (including the runtime catalog) share this
3891                // arena, so derive headroom from the complete runtime sample.
3892                held_used = runtime.memory_accounting().total_used_bytes();
3893                growth = budget
3894                    .checked_sub(held_used)
3895                    .expect("fixture fits while the cursor is open")
3896                    + source_bytes / 2;
3897                assert!(growth > 0);
3898                assert!(
3899                    runtime
3900                        .admit_non_evictable_growth(
3901                            MemoryPool::SegmentArena,
3902                            "growth with pinned source",
3903                            growth,
3904                        )
3905                        .is_err(),
3906                    "retained source must participate in runtime admission"
3907                );
3908                false
3909            }
3910        ));
3911        runtime.refresh_memory_accounting();
3912        assert_eq!(
3913            held_used - runtime.memory_accounting().total_used_bytes(),
3914            source_bytes
3915        );
3916        let _reservation = runtime
3917            .admit_non_evictable_growth(
3918                MemoryPool::SegmentArena,
3919                "growth after cursor cancellation",
3920                growth,
3921            )
3922            .expect("the same growth fits once the cursor releases its source");
3923    }
3924
3925    #[test]
3926    fn retained_memory_includes_an_inflight_consolidation_copy() {
3927        let manager = consolidating_manager(1);
3928        let ids = seal_bulk_segment(&manager, 100);
3929        for id in ids.iter().take(40) {
3930            manager.delete(*id).expect("delete");
3931        }
3932        let before = manager.resident_bytes();
3933        assert!(manager.pressure_consolidation_tick());
3934        let copied = manager
3935            .consolidation
3936            .read()
3937            .as_ref()
3938            .expect("paced merge")
3939            .merged
3940            .resident_bytes();
3941        assert!(copied > 0);
3942        assert!(
3943            manager.resident_bytes() >= before + copied,
3944            "unpublished copy consumes memory"
3945        );
3946        drain_maintenance(&manager);
3947        assert!(
3948            manager.resident_bytes() < before,
3949            "completed merge releases the old storage"
3950        );
3951    }
3952
3953    #[test]
3954    fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
3955        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3956        let ids = seal_bulk_segment(&manager, 100);
3957
3958        let bytes_before_delete = manager.resident_bytes();
3959
3960        // 40% tombstones — past CONSOLIDATION_TOMBSTONE_RATIO.
3961        for id in ids.iter().take(40) {
3962            assert!(manager.delete(*id).expect("delete"));
3963        }
3964
3965        // Tombstones cost memory: the payload stays resident in flat storage
3966        // and the tombstone set grows.
3967        let bytes_with_tombstones = manager.resident_bytes();
3968        assert!(bytes_with_tombstones > bytes_before_delete);
3969
3970        let survivors = live_ids(&manager);
3971        assert_eq!(survivors.len(), 60);
3972
3973        drain_maintenance(&manager);
3974
3975        let stats = manager.stats();
3976        assert_eq!(stats.consolidation.runs_started, 1);
3977        assert_eq!(stats.consolidation.runs_completed, 1);
3978        assert_eq!(stats.consolidation.segments_merged, 1);
3979        assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
3980        assert!(stats.consolidation.bytes_reclaimed > 0);
3981
3982        // The reclaimed memory is at least the tombstoned entities' share.
3983        let bytes_after = manager.resident_bytes();
3984        assert!(
3985            bytes_after <= bytes_before_delete,
3986            "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
3987        );
3988        assert_eq!(
3989            bytes_with_tombstones - bytes_after,
3990            stats.consolidation.bytes_reclaimed
3991        );
3992        assert_eq!(stats.total_memory_bytes, bytes_after as usize);
3993
3994        assert_eq!(live_ids(&manager), survivors);
3995        assert_eq!(manager.sealed.read().len(), 1);
3996        assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
3997    }
3998
3999    #[test]
4000    fn query_results_are_identical_before_during_and_after_consolidation() {
4001        let manager = consolidating_manager(3);
4002        let ids = seal_bulk_segment(&manager, 20);
4003        seal_bulk_segment(&manager, 20);
4004        for id in ids.iter().take(10) {
4005            assert!(manager.delete(*id).expect("delete"));
4006        }
4007
4008        let before = live_ids(&manager);
4009        let survivor = *ids.last().expect("at least one row");
4010        let tombstoned = ids[0];
4011
4012        // Every mid-tick observation matches the pre-consolidation snapshot:
4013        // point reads, full scans, and the kind index (rebuilt by the merge).
4014        let mut mid_tick_observations = 0;
4015        while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
4016            manager.run_maintenance().expect("maintenance");
4017            assert_eq!(live_ids(&manager), before);
4018            assert!(manager.get(survivor).is_some());
4019            assert!(manager.get(tombstoned).is_none());
4020            assert_eq!(manager.get_by_kind("table").len(), before.len());
4021            mid_tick_observations += 1;
4022            assert!(mid_tick_observations < 1_000);
4023        }
4024
4025        assert!(
4026            mid_tick_observations > 2,
4027            "expected the run to span several ticks, saw {mid_tick_observations}"
4028        );
4029        assert_eq!(live_ids(&manager), before);
4030        assert!(manager.get(survivor).is_some());
4031        assert!(manager.get(tombstoned).is_none());
4032        assert_eq!(manager.get_by_kind("table").len(), before.len());
4033    }
4034
4035    #[test]
4036    fn get_by_kind_reaches_sealed_segments() {
4037        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4038        let ids = seal_bulk_segment(&manager, 5);
4039        assert!(manager.delete(ids[0]).expect("delete"));
4040
4041        // The sealed list is walked through the segment read guard; the
4042        // `Arc<RwLock<GrowingSegment>>` impl of `iter_kind` cannot hand out
4043        // references across the lock and yields nothing.
4044        assert_eq!(manager.get_by_kind("table").len(), 4);
4045        assert_eq!(manager.get_by_kind("vector").len(), 0);
4046    }
4047
4048    #[test]
4049    fn consolidation_is_paced_by_the_per_tick_entity_bound() {
4050        const PER_TICK: usize = 4;
4051        let manager = consolidating_manager(PER_TICK);
4052        let ids = seal_bulk_segment(&manager, 60);
4053        for id in ids.iter().take(30) {
4054            assert!(manager.delete(*id).expect("delete"));
4055        }
4056
4057        // First tick opens the run and copies exactly one bound's worth.
4058        manager.run_maintenance().expect("maintenance");
4059        assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
4060
4061        // Each subsequent tick advances by no more than the bound, and the run
4062        // is still in flight — it never completes in one unbounded pass.
4063        let mut previous = PER_TICK as u64;
4064        let mut ticks = 1;
4065        while let Some(copied) = manager.consolidation_progress() {
4066            assert!(
4067                copied - previous <= PER_TICK as u64,
4068                "tick copied {} entities, bound is {PER_TICK}",
4069                copied - previous
4070            );
4071            previous = copied;
4072            manager.run_maintenance().expect("maintenance");
4073            ticks += 1;
4074            assert!(ticks < 100);
4075        }
4076
4077        // 30 survivors at 4 per tick: the run cannot have finished before the
4078        // eighth tick.
4079        assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
4080        assert_eq!(manager.stats().consolidation.runs_completed, 1);
4081        assert_eq!(live_ids(&manager).len(), 30);
4082    }
4083
4084    #[test]
4085    fn sealed_segment_count_above_the_limit_triggers_a_merge() {
4086        let manager = SegmentManager::with_config(
4087            "test",
4088            ManagerConfig {
4089                max_sealed_segments: 3,
4090                ..Default::default()
4091            },
4092        );
4093        for _ in 0..4 {
4094            seal_bulk_segment(&manager, 5);
4095        }
4096        assert_eq!(manager.sealed.read().len(), 4);
4097
4098        let before = live_ids(&manager);
4099        drain_maintenance(&manager);
4100
4101        // The knob finally does what it claimed: the sealed set collapses even
4102        // though not a single row was deleted.
4103        assert_eq!(manager.sealed.read().len(), 1);
4104        assert_eq!(manager.stats().consolidation.segments_merged, 4);
4105        assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
4106        assert_eq!(live_ids(&manager), before);
4107    }
4108
4109    #[test]
4110    fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
4111        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4112        seal_bulk_segment(&manager, 10);
4113        seal_bulk_segment(&manager, 10);
4114
4115        for _ in 0..5 {
4116            manager.run_maintenance().expect("maintenance");
4117        }
4118
4119        assert_eq!(manager.stats().consolidation.runs_started, 0);
4120        assert_eq!(manager.sealed.read().len(), 2);
4121    }
4122
4123    #[test]
4124    fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
4125        use std::sync::atomic::AtomicBool;
4126
4127        let manager = Arc::new(consolidating_manager(2));
4128        let ids = seal_bulk_segment(&manager, 40);
4129        seal_bulk_segment(&manager, 40);
4130        for id in ids.iter().take(20) {
4131            assert!(manager.delete(*id).expect("delete"));
4132        }
4133
4134        let expected = live_ids(&manager);
4135        let stop = Arc::new(AtomicBool::new(false));
4136
4137        let readers: Vec<_> = (0..4)
4138            .map(|_| {
4139                let manager = Arc::clone(&manager);
4140                let stop = Arc::clone(&stop);
4141                let expected = expected.clone();
4142                std::thread::spawn(move || {
4143                    let mut scans = 0u64;
4144                    // Do-while: every reader scans at least once even when the
4145                    // drain finishes before this thread is first scheduled.
4146                    loop {
4147                        let mut seen: Vec<u64> = manager
4148                            .query_all(|_| true)
4149                            .into_iter()
4150                            .map(|entity| entity.id.raw())
4151                            .collect();
4152                        seen.sort_unstable();
4153                        // Sorting is not enough: a duplicated row would keep
4154                        // the set equal but change the length.
4155                        assert_eq!(seen, expected, "reader saw a torn sealed set");
4156                        scans += 1;
4157                        if stop.load(Ordering::Relaxed) {
4158                            break;
4159                        }
4160                    }
4161                    scans
4162                })
4163            })
4164            .collect();
4165
4166        drain_maintenance(&manager);
4167        stop.store(true, Ordering::Relaxed);
4168
4169        for reader in readers {
4170            let scans = reader.join().expect("reader thread");
4171            assert!(scans > 0, "reader never got to scan");
4172        }
4173
4174        assert_eq!(manager.stats().consolidation.runs_completed, 1);
4175        assert_eq!(live_ids(&manager), expected);
4176    }
4177
4178    #[test]
4179    fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
4180        let manager = consolidating_manager(4);
4181        let ids = seal_bulk_segment(&manager, 20);
4182        for id in ids.iter().take(6) {
4183            assert!(manager.delete(*id).expect("delete"));
4184        }
4185
4186        // Copy the first bound's worth, then delete a row we just copied.
4187        manager.run_maintenance().expect("maintenance");
4188        assert_eq!(manager.consolidation_progress(), Some(4));
4189        let doomed = ids[6];
4190        assert!(manager.delete(doomed).expect("delete"));
4191
4192        let expected = live_ids(&manager);
4193        drain_maintenance(&manager);
4194
4195        assert!(manager.get(doomed).is_none(), "deleted row came back");
4196        assert_eq!(live_ids(&manager), expected);
4197    }
4198
4199    #[test]
4200    fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
4201        let manager = consolidating_manager(3);
4202        let ids = seal_bulk_segment(&manager, 30);
4203        for id in ids.iter().take(12) {
4204            assert!(manager.delete(*id).expect("delete"));
4205        }
4206
4207        let before = live_ids(&manager);
4208        let sealed_before: Vec<SegmentId> = manager
4209            .sealed
4210            .read()
4211            .iter()
4212            .map(|segment| segment.read().id())
4213            .collect();
4214
4215        manager.run_maintenance().expect("maintenance");
4216        manager.run_maintenance().expect("maintenance");
4217        assert!(manager.consolidation.read().is_some());
4218
4219        // A crash between ticks drops the half-built merged segment. Sealed
4220        // segments are the durable truth, so nothing was lost — and the run is
4221        // simply started over.
4222        *manager.consolidation.write() = None;
4223
4224        assert_eq!(live_ids(&manager), before);
4225        let sealed_after: Vec<SegmentId> = manager
4226            .sealed
4227            .read()
4228            .iter()
4229            .map(|segment| segment.read().id())
4230            .collect();
4231        assert_eq!(sealed_after, sealed_before);
4232        assert_eq!(manager.stats().consolidation.runs_completed, 0);
4233
4234        drain_maintenance(&manager);
4235        assert_eq!(live_ids(&manager), before);
4236        assert_eq!(manager.stats().consolidation.runs_completed, 1);
4237    }
4238
4239    #[test]
4240    fn consolidation_preserves_entity_payload_metadata_and_sequence() {
4241        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4242
4243        // Metadata must be attached before the seal: `set_metadata` on a sealed
4244        // row rewrites it into the growing segment, which would take it out of
4245        // the segment under test.
4246        let entities: Vec<UnifiedEntity> = (0..10)
4247            .map(|i| {
4248                UnifiedEntity::table_row(
4249                    manager.next_entity_id(),
4250                    "users",
4251                    i as u64,
4252                    vec![Value::Integer(i as i64)],
4253                )
4254            })
4255            .collect();
4256        let ids = manager.bulk_insert(entities).expect("bulk insert");
4257
4258        let mut meta = Metadata::new();
4259        meta.set(
4260            "os",
4261            super::super::metadata::MetadataValue::String("linux".to_string()),
4262        );
4263        manager.set_metadata(ids[9], meta).expect("set metadata");
4264        manager.force_seal().expect("seal");
4265
4266        assert!(manager.delete(ids[0]).expect("delete"));
4267        assert!(manager.delete(ids[1]).expect("delete"));
4268        assert!(manager.delete(ids[2]).expect("delete"));
4269
4270        let before: Vec<UnifiedEntity> = {
4271            let mut entities = manager.query_all(|_| true);
4272            entities.sort_by_key(|entity| entity.id.raw());
4273            entities
4274        };
4275
4276        drain_maintenance(&manager);
4277
4278        let after: Vec<UnifiedEntity> = {
4279            let mut entities = manager.query_all(|_| true);
4280            entities.sort_by_key(|entity| entity.id.raw());
4281            entities
4282        };
4283
4284        let columns = |entity: &UnifiedEntity| match &entity.data {
4285            super::super::entity::EntityData::Row(row) => row.columns.clone(),
4286            _ => Vec::new(),
4287        };
4288
4289        assert_eq!(before.len(), after.len());
4290        for (old, new) in before.iter().zip(after.iter()) {
4291            assert_eq!(old.id, new.id);
4292            assert_eq!(columns(old), columns(new));
4293            assert_eq!(
4294                old.sequence_id, new.sequence_id,
4295                "merge must not re-issue sequence ids"
4296            );
4297        }
4298        assert!(manager
4299            .get_metadata(ids[9])
4300            .expect("metadata survives the merge")
4301            .has("os"));
4302    }
4303
4304    #[test]
4305    #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
4306    fn test_lifecycle_events() {
4307        let manager = SegmentManager::new("test");
4308
4309        manager
4310            .insert(UnifiedEntity::vector(
4311                manager.next_entity_id(),
4312                "v",
4313                vec![0.1],
4314            ))
4315            .unwrap();
4316
4317        let events = manager.drain_events();
4318
4319        // Should have: SegmentCreated, EntityInserted
4320        assert!(events
4321            .iter()
4322            .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
4323        assert!(events
4324            .iter()
4325            .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
4326    }
4327}