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    /// Visit borrowed rows with metadata under the same segment read guard.
1441    /// Callbacks must not re-enter this manager while the guard is held.
1442    pub(crate) fn for_each_entity_with_metadata(
1443        &self,
1444        mut callback: impl FnMut(&UnifiedEntity, Option<Metadata>) -> bool,
1445    ) {
1446        if let Some(growing_arc) = self.growing.read().as_ref() {
1447            let growing = growing_arc.read();
1448            if !growing.for_each_fast(|entity| callback(entity, growing.get_metadata(entity.id))) {
1449                return;
1450            }
1451        }
1452        let sealed = self.sealed.read();
1453        for segment_arc in sealed.iter() {
1454            let segment = segment_arc.read();
1455            if !segment.for_each_fast(|entity| callback(entity, segment.get_metadata(entity.id))) {
1456                return;
1457            }
1458        }
1459    }
1460
1461    fn for_each_entity_in_segments<F>(
1462        &self,
1463        segment_filter: impl Fn(&GrowingSegment) -> bool,
1464        mut before_segment: impl FnMut() -> bool,
1465        mut callback: F,
1466    ) where
1467        F: FnMut(&UnifiedEntity) -> bool,
1468    {
1469        // Growing segment — direct iteration (no Box<dyn>)
1470        // Try non-blocking read first; fall back to blocking only when a writer
1471        // is actively holding the write lock (rare in read-heavy workloads).
1472        if let Some(growing_arc) = self.growing.read().as_ref() {
1473            if !before_segment() {
1474                return;
1475            }
1476            let growing = if let Some(g) = growing_arc.try_read() {
1477                g
1478            } else {
1479                growing_arc.read()
1480            };
1481            if segment_filter(&growing) && !growing.for_each_fast(&mut callback) {
1482                return;
1483            }
1484        }
1485
1486        // Sealed segments
1487        let sealed = self.sealed.read();
1488        for segment_arc in sealed.iter() {
1489            if !before_segment() {
1490                return;
1491            }
1492            let segment = segment_arc.read();
1493            if segment_filter(&segment) && !segment.for_each_fast(&mut callback) {
1494                return;
1495            }
1496        }
1497    }
1498
1499    /// Parallel fold across all entities. Each sealed segment is
1500    /// processed on its own rayon task; the growing segment stays on
1501    /// the caller thread (its read lock is briefly held).
1502    ///
1503    /// - `init` builds a fresh accumulator per thread.
1504    /// - `fold` mutates an accumulator with one entity at a time.
1505    /// - `reduce` combines two accumulators into one.
1506    ///
1507    /// The returned value is the reduction of every per-thread
1508    /// accumulator. Use this for aggregate-shape workloads (GROUP BY)
1509    /// where per-thread partial state can be merged cheaply.
1510    ///
1511    /// NOTE: when there are 0 or 1 sealed segments, the parallel path
1512    /// is skipped and the work runs sequentially to avoid rayon
1513    /// overhead on tiny tables.
1514    pub fn fold_entities_parallel<T, FInit, FFold, FReduce>(
1515        &self,
1516        init: FInit,
1517        fold: FFold,
1518        reduce: FReduce,
1519    ) -> T
1520    where
1521        T: Send,
1522        FInit: Fn() -> T + Send + Sync,
1523        FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1524        FReduce: Fn(T, T) -> T + Send + Sync,
1525    {
1526        use rayon::prelude::*;
1527
1528        // Growing segment — always sequential (single writer lock,
1529        // usually small working set).
1530        let mut acc = init();
1531        if let Some(growing_arc) = self.growing.read().as_ref() {
1532            let growing = if let Some(g) = growing_arc.try_read() {
1533                g
1534            } else {
1535                growing_arc.read()
1536            };
1537            growing.for_each_fast(|entity| {
1538                acc = fold(std::mem::replace(&mut acc, init()), entity);
1539                true
1540            });
1541        }
1542
1543        // Sealed segments — snapshot the Arc list under the read lock,
1544        // then drop the lock so rayon workers can fan out without
1545        // blocking writers.
1546        let segments: Vec<_> = {
1547            let sealed = self.sealed.read();
1548            sealed.iter().cloned().collect()
1549        };
1550
1551        if segments.len() <= 1 {
1552            for seg_arc in &segments {
1553                let seg = seg_arc.read();
1554                seg.for_each_fast(|entity| {
1555                    acc = fold(std::mem::replace(&mut acc, init()), entity);
1556                    true
1557                });
1558            }
1559            return acc;
1560        }
1561
1562        let sealed_acc = segments
1563            .into_par_iter()
1564            .map(|seg_arc| {
1565                let mut local = init();
1566                let seg = seg_arc.read();
1567                seg.for_each_fast(|entity| {
1568                    local = fold(std::mem::replace(&mut local, init()), entity);
1569                    true
1570                });
1571                local
1572            })
1573            .reduce(&init, &reduce);
1574
1575        reduce(acc, sealed_acc)
1576    }
1577
1578    /// Zone-map-aware iteration across all segments.
1579    ///
1580    /// Like `for_each_entity`, but checks `zone_preds` against each segment's
1581    /// column zone maps before iterating. Segments where any predicate can
1582    /// definitively prove no rows match are skipped entirely.
1583    ///
1584    /// `zone_preds`: slice of `(column_name, ZoneColPred)` extracted from the WHERE clause.
1585    /// Empty slice → same behaviour as `for_each_entity` (no pruning).
1586    pub fn for_each_entity_zoned<F>(&self, zone_preds: &[(&str, ZoneColPred<'_>)], mut callback: F)
1587    where
1588        F: FnMut(&UnifiedEntity) -> bool,
1589    {
1590        let _ = self.for_each_entity_zoned_with_stats(zone_preds, &mut callback);
1591    }
1592
1593    pub fn for_each_entity_zoned_with_stats<F>(
1594        &self,
1595        zone_preds: &[(&str, ZoneColPred<'_>)],
1596        mut callback: F,
1597    ) -> SegmentScanStats
1598    where
1599        F: FnMut(&UnifiedEntity) -> bool,
1600    {
1601        let mut scan_stats = SegmentScanStats::default();
1602        // Growing segment — never skip (it's receiving writes, zones are partial).
1603        // Try a non-blocking read first: if a writer is currently inserting
1604        // (holding the write lock), try_read() returns None and we fall back to
1605        // the blocking read.  In low-contention workloads (reads far outnumber
1606        // writes) the try_read() almost always succeeds and readers never stall.
1607        if let Some(growing_arc) = self.growing.read().as_ref() {
1608            scan_stats.segments_total += 1;
1609            scan_stats.segments_scanned += 1;
1610            let growing = if let Some(g) = growing_arc.try_read() {
1611                g
1612            } else {
1613                growing_arc.read()
1614            };
1615            if !growing.for_each_fast(&mut callback) {
1616                return scan_stats;
1617            }
1618        }
1619
1620        // Sealed segments — check zone maps before iterating
1621        let sealed = self.sealed.read();
1622        for segment_arc in sealed.iter() {
1623            scan_stats.segments_total += 1;
1624            let segment = segment_arc.read();
1625            if !zone_preds.is_empty() && segment.can_skip_zone_preds(zone_preds) {
1626                scan_stats.segments_pruned += 1;
1627                continue; // entire segment pruned
1628            }
1629            scan_stats.segments_scanned += 1;
1630            if !segment.for_each_fast(&mut callback) {
1631                return scan_stats;
1632            }
1633        }
1634        scan_stats
1635    }
1636
1637    /// Zone-map-aware parallel query.
1638    ///
1639    /// Like `query_all` but applies `zone_preds` on the main thread to
1640    /// prune sealed segments before spawning workers — segments that
1641    /// provably contain no matching rows are skipped entirely.
1642    ///
1643    /// Zone check runs single-threaded (it reads per-segment metadata,
1644    /// not row data), so it's cheap. Surviving segments are then scanned
1645    /// in parallel using `std::thread::scope` when there are > 1 of them.
1646    pub fn query_all_zoned<F>(
1647        &self,
1648        zone_preds: &[(&str, ZoneColPred<'_>)],
1649        filter: F,
1650    ) -> Vec<UnifiedEntity>
1651    where
1652        F: Fn(&UnifiedEntity) -> bool + Sync,
1653    {
1654        self.query_all_zoned_with_stats(zone_preds, filter).0
1655    }
1656
1657    pub fn query_all_zoned_with_stats<F>(
1658        &self,
1659        zone_preds: &[(&str, ZoneColPred<'_>)],
1660        filter: F,
1661    ) -> (Vec<UnifiedEntity>, SegmentScanStats)
1662    where
1663        F: Fn(&UnifiedEntity) -> bool + Sync,
1664    {
1665        let mut results = Vec::new();
1666        let mut scan_stats = SegmentScanStats::default();
1667
1668        // Growing segment — always scan, no zone skip (zones are partial).
1669        // Non-blocking try_read() avoids stalling behind in-progress inserts.
1670        if let Some(growing_arc) = self.growing.read().as_ref() {
1671            scan_stats.segments_total += 1;
1672            scan_stats.segments_scanned += 1;
1673            let growing = if let Some(g) = growing_arc.try_read() {
1674                g
1675            } else {
1676                growing_arc.read()
1677            };
1678            results.extend(growing.iter().filter(|e| filter(e)).cloned());
1679        }
1680
1681        // Sealed segments: zone-prune on main thread, then scan in parallel.
1682        let sealed = self.sealed.read();
1683        // Collect only the segments that survive zone-predicate pruning.
1684        let surviving: Vec<_> = sealed
1685            .iter()
1686            .filter(|seg_arc| {
1687                scan_stats.segments_total += 1;
1688                if zone_preds.is_empty() {
1689                    scan_stats.segments_scanned += 1;
1690                    return true;
1691                }
1692                let seg = seg_arc.read();
1693                if seg.can_skip_zone_preds(zone_preds) {
1694                    scan_stats.segments_pruned += 1;
1695                    false
1696                } else {
1697                    scan_stats.segments_scanned += 1;
1698                    true
1699                }
1700            })
1701            .collect();
1702
1703        let use_parallel = surviving.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1704
1705        if use_parallel {
1706            let filter_ref = &filter;
1707            let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1708                surviving
1709                    .iter()
1710                    .map(|segment| {
1711                        s.spawn(move || {
1712                            segment
1713                                .read()
1714                                .iter()
1715                                .filter(|e| filter_ref(e))
1716                                .cloned()
1717                                .collect::<Vec<_>>()
1718                        })
1719                    })
1720                    .collect::<Vec<_>>()
1721                    .into_iter()
1722                    .map(|handle| handle.join().unwrap_or_default())
1723                    .collect()
1724            });
1725            for batch in segment_results {
1726                results.extend(batch);
1727            }
1728        } else {
1729            for segment_arc in surviving {
1730                let seg = segment_arc.read();
1731                results.extend(seg.iter().filter(|e| filter(e)).cloned());
1732            }
1733        }
1734
1735        (results, scan_stats)
1736    }
1737
1738    /// Visibility predicate for the `scan_*` family. `Some` delegates to the
1739    /// canonical `entity_visible_with_context`; `None` supports low-level
1740    /// callers that intentionally scan without a statement snapshot and keeps
1741    /// the documented resolver fallback — moderation gate plus hiding
1742    /// superseded physical versions — instead of showing every version.
1743    fn scan_entity_visible(snapshot: Option<&SnapshotContext>, entity: &UnifiedEntity) -> bool {
1744        match snapshot {
1745            Some(ctx) => entity_visible_with_context(Some(ctx), entity),
1746            None => {
1747                !crate::runtime::ai::moderation::entity_moderation_hidden(entity)
1748                    && entity.xmax == 0
1749            }
1750        }
1751    }
1752
1753    /// Scan every segment under an explicit MVCC snapshot.
1754    ///
1755    /// Visibility is evaluated while each segment iterator is active, before
1756    /// the caller's filter. Unlike the raw iteration methods, this entry point
1757    /// does not consult thread-local snapshot state, so it preserves the same
1758    /// view when invoked from a different thread.
1759    pub fn scan<F>(&self, snapshot: Option<&SnapshotContext>, filter: F) -> Vec<UnifiedEntity>
1760    where
1761        F: Fn(&UnifiedEntity) -> bool + Sync,
1762    {
1763        self.query_all(|entity| Self::scan_entity_visible(snapshot, entity) && filter(entity))
1764    }
1765
1766    /// Collect visible graph items, pruning only under the segment read guard.
1767    pub(crate) fn scan_graph_kind(
1768        &self,
1769        snapshot: Option<&SnapshotContext>,
1770        kind: GraphEntityKind,
1771    ) -> Vec<UnifiedEntity> {
1772        self.query_all_in_segments(
1773            |segment| segment.may_contain_graph_kind(kind),
1774            |entity| Self::scan_entity_visible(snapshot, entity) && kind.matches(&entity.kind),
1775        )
1776    }
1777
1778    /// Visit visible candidates in segments that may contain the requested kind.
1779    /// Non-graph items in those segments still reach the callback so work budgets
1780    /// charge every inspected visible candidate. RLS must run outside this callback.
1781    /// `before_segment` runs even for pruned segments; false stops the entire scan.
1782    pub(crate) fn scan_graph_candidates_for_each<F>(
1783        &self,
1784        snapshot: Option<&SnapshotContext>,
1785        kind: GraphEntityKind,
1786        before_segment: impl FnMut() -> bool,
1787        mut callback: F,
1788    ) where
1789        F: FnMut(&UnifiedEntity) -> bool,
1790    {
1791        self.for_each_entity_in_segments(
1792            |segment| segment.may_contain_graph_kind(kind),
1793            before_segment,
1794            |entity| !Self::scan_entity_visible(snapshot, entity) || callback(entity),
1795        );
1796    }
1797
1798    /// Visit every entity visible under an explicit MVCC snapshot.
1799    ///
1800    /// Returning `false` from `callback` stops the scan early. Hidden entities
1801    /// never reach the callback and therefore cannot stop the scan.
1802    pub fn scan_for_each<F>(&self, snapshot: Option<&SnapshotContext>, mut callback: F)
1803    where
1804        F: FnMut(&UnifiedEntity) -> bool,
1805    {
1806        self.for_each_entity(|entity| {
1807            !Self::scan_entity_visible(snapshot, entity) || callback(entity)
1808        });
1809    }
1810
1811    /// Capture segment membership and append boundaries, then hydrate bounded
1812    /// batches. The consumer runs outside topology and segment locks, so scoring
1813    /// and RLS may re-enter storage and writers can progress between batches.
1814    /// Retired sources pin positions only; hydration uses current segments so
1815    /// mutations after consolidation are visible to the snapshot predicate.
1816    pub(crate) fn scan_batches(
1817        &self,
1818        snapshot: Option<&SnapshotContext>,
1819        mut before_work: impl FnMut() -> bool,
1820        mut consume: impl FnMut(&[UnifiedEntity]) -> bool,
1821    ) -> bool {
1822        let segments = {
1823            let growing = self.growing.read();
1824            let sealed = self.sealed.read();
1825            let mut segments = Vec::with_capacity(sealed.len() + usize::from(growing.is_some()));
1826            for segment in growing.iter().chain(sealed.iter()) {
1827                if !before_work() {
1828                    return false;
1829                }
1830                let cursor = segment.read().scan_cursor();
1831                segments.push((Arc::clone(segment), cursor));
1832            }
1833            segments
1834        };
1835        let mut ids = Vec::new();
1836        let mut batch = Vec::new();
1837        for (segment, mut cursor) in segments {
1838            while !cursor.finished() {
1839                ids.clear();
1840                batch.clear();
1841                let batch_size = cursor.remaining().min(super::segment::SCAN_BATCH_SIZE);
1842                ids.reserve_exact(batch_size);
1843                batch.reserve_exact(batch_size);
1844                {
1845                    let segment = segment.read();
1846                    if !segment.scan_batch(&mut cursor, &mut before_work, &mut |id| ids.push(id)) {
1847                        return false;
1848                    }
1849                }
1850                // Sources can retire between batches. Consult authoritative storage
1851                // after releasing the source guard, just as the former ID-list scan did.
1852                if !before_work() {
1853                    return false;
1854                }
1855                self.for_each_id(&ids, |_, entity| {
1856                    if Self::scan_entity_visible(snapshot, entity) {
1857                        batch.push(entity.clone());
1858                    }
1859                });
1860                if !before_work() {
1861                    return false;
1862                }
1863                if !consume(&batch) {
1864                    return false;
1865                }
1866            }
1867        }
1868        true
1869    }
1870
1871    /// Fold entities visible under an explicit MVCC snapshot in parallel.
1872    pub fn scan_fold_parallel<T, FInit, FFold, FReduce>(
1873        &self,
1874        snapshot: Option<&SnapshotContext>,
1875        init: FInit,
1876        fold: FFold,
1877        reduce: FReduce,
1878    ) -> T
1879    where
1880        T: Send,
1881        FInit: Fn() -> T + Send + Sync,
1882        FFold: Fn(T, &UnifiedEntity) -> T + Send + Sync,
1883        FReduce: Fn(T, T) -> T + Send + Sync,
1884    {
1885        self.fold_entities_parallel(
1886            init,
1887            |acc, entity| {
1888                if Self::scan_entity_visible(snapshot, entity) {
1889                    fold(acc, entity)
1890                } else {
1891                    acc
1892                }
1893            },
1894            reduce,
1895        )
1896    }
1897
1898    /// Visit zone-map candidates visible under an explicit MVCC snapshot.
1899    pub fn scan_for_each_zoned_with_stats<F>(
1900        &self,
1901        snapshot: Option<&SnapshotContext>,
1902        zone_preds: &[(&str, ZoneColPred<'_>)],
1903        mut callback: F,
1904    ) -> SegmentScanStats
1905    where
1906        F: FnMut(&UnifiedEntity) -> bool,
1907    {
1908        self.for_each_entity_zoned_with_stats(zone_preds, |entity| {
1909            !Self::scan_entity_visible(snapshot, entity) || callback(entity)
1910        })
1911    }
1912
1913    /// Collect zone-map candidates visible under an explicit MVCC snapshot.
1914    pub fn scan_zoned_with_stats<F>(
1915        &self,
1916        snapshot: Option<&SnapshotContext>,
1917        zone_preds: &[(&str, ZoneColPred<'_>)],
1918        filter: F,
1919    ) -> (Vec<UnifiedEntity>, SegmentScanStats)
1920    where
1921        F: Fn(&UnifiedEntity) -> bool + Sync,
1922    {
1923        self.query_all_zoned_with_stats(zone_preds, |entity| {
1924            Self::scan_entity_visible(snapshot, entity) && filter(entity)
1925        })
1926    }
1927
1928    /// Query across all segments. Uses parallel scanning for sealed segments
1929    /// when more than one sealed segment exists.
1930    pub fn query_all<F>(&self, filter: F) -> Vec<UnifiedEntity>
1931    where
1932        F: Fn(&UnifiedEntity) -> bool + Sync,
1933    {
1934        self.query_all_in_segments(|_| true, filter)
1935    }
1936
1937    fn query_all_in_segments<F>(
1938        &self,
1939        segment_filter: impl Fn(&GrowingSegment) -> bool + Sync,
1940        filter: F,
1941    ) -> Vec<UnifiedEntity>
1942    where
1943        F: Fn(&UnifiedEntity) -> bool + Sync,
1944    {
1945        let mut results = Vec::new();
1946
1947        // Query growing segment — try non-blocking read first (avoids stalling
1948        // behind an in-progress insert; falls back to blocking if writer is active).
1949        if let Some(growing_arc) = self.growing.read().as_ref() {
1950            let growing = if let Some(g) = growing_arc.try_read() {
1951                g
1952            } else {
1953                growing_arc.read()
1954            };
1955            if segment_filter(&growing) {
1956                results.extend(growing.iter().filter(|e| filter(e)).cloned());
1957            }
1958        }
1959
1960        // Query sealed segments — parallel when multiple exist AND multi-core
1961        let sealed = self.sealed.read();
1962        let use_parallel = sealed.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1963        if use_parallel {
1964            let filter_ref = &filter;
1965            let segment_filter_ref = &segment_filter;
1966            let segment_results: Vec<Vec<UnifiedEntity>> = std::thread::scope(|s| {
1967                sealed
1968                    .iter()
1969                    .map(|segment| {
1970                        s.spawn(move || {
1971                            let segment = segment.read();
1972                            if segment_filter_ref(&segment) {
1973                                segment
1974                                    .iter()
1975                                    .filter(|e| filter_ref(e))
1976                                    .cloned()
1977                                    .collect::<Vec<_>>()
1978                            } else {
1979                                Vec::new()
1980                            }
1981                        })
1982                    })
1983                    .collect::<Vec<_>>()
1984                    .into_iter()
1985                    .map(|handle| handle.join().unwrap_or_default())
1986                    .collect()
1987            });
1988            for batch in segment_results {
1989                results.extend(batch);
1990            }
1991        } else {
1992            for segment in sealed.iter() {
1993                let seg = segment.read();
1994                if segment_filter(&seg) {
1995                    results.extend(seg.iter().filter(|e| filter(e)).cloned());
1996                }
1997            }
1998        }
1999
2000        results
2001    }
2002
2003    /// Probe exact segment membership for a key without materializing matching entities.
2004    ///
2005    /// Returns `false` only when every consulted segment proves absence.
2006    /// Missing segment state is conservative and returns `true`.
2007    pub fn bloom_may_contain_key(&self, key: &[u8]) -> bool {
2008        if let Some(growing_arc) = self.growing.read().as_ref() {
2009            let growing = growing_arc.read();
2010            if growing.may_contain_exact_key(key) {
2011                return true;
2012            }
2013        } else {
2014            return true;
2015        }
2016
2017        let sealed = self.sealed.read();
2018        for segment_arc in sealed.iter() {
2019            let segment = segment_arc.read();
2020            if segment.may_contain_exact_key(key) {
2021                return true;
2022            }
2023        }
2024
2025        false
2026    }
2027
2028    /// Filter by metadata across all segments
2029    pub fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2030        let mut results = Vec::new();
2031
2032        // Growing segment
2033        if let Some(growing_arc) = self.growing.read().as_ref() {
2034            let growing = growing_arc.read();
2035            results.extend(growing.filter_metadata(filters));
2036        }
2037
2038        // Sealed segments
2039        let sealed = self.sealed.read();
2040        for segment in sealed.iter() {
2041            results.extend(segment.filter_metadata(filters));
2042        }
2043
2044        results
2045    }
2046
2047    /// Get entities by kind
2048    pub fn get_by_kind(&self, kind: &str) -> Vec<UnifiedEntity> {
2049        let mut results = Vec::new();
2050
2051        // Growing segment
2052        if let Some(growing_arc) = self.growing.read().as_ref() {
2053            let growing = growing_arc.read();
2054            for entity in growing.iter_kind(kind) {
2055                results.push(entity.clone());
2056            }
2057        }
2058
2059        // Sealed segments. Iterate through the read guard, not the
2060        // `Arc<RwLock<GrowingSegment>>` impl of `iter_kind` — that one cannot
2061        // hand out references across the lock, so it returns an empty iterator
2062        // and every sealed entity was silently missing from the result.
2063        let sealed = self.sealed.read();
2064        for segment_arc in sealed.iter() {
2065            let segment = segment_arc.read();
2066            for entity in segment.iter_kind(kind) {
2067                results.push(entity.clone());
2068            }
2069        }
2070
2071        results
2072    }
2073
2074    /// Count entities
2075    pub fn count(&self) -> usize {
2076        self.total_entities_atomic.load(Ordering::Relaxed) as usize
2077    }
2078
2079    /// Get all segment IDs
2080    pub fn segment_ids(&self) -> Vec<SegmentId> {
2081        let mut ids = Vec::new();
2082
2083        if let Some(growing_arc) = self.growing.read().as_ref() {
2084            ids.push(growing_arc.read().id());
2085        }
2086
2087        let sealed = self.sealed.read();
2088        for segment in sealed.iter() {
2089            ids.push(segment.id());
2090        }
2091
2092        ids.extend(self.archived.read().iter().copied());
2093
2094        ids
2095    }
2096
2097    /// Emit a lifecycle event.
2098    ///
2099    /// Perf: this used to push onto a `RwLock<Vec<LifecycleEvent>>`
2100    /// on every insert / delete / seal. Nobody consumes that vec
2101    /// today (no subscription API, `drain_events` has no callers),
2102    /// so the write lock + push was pure tax — and the vec grew
2103    /// unbounded in long-running processes.
2104    ///
2105    /// Current behaviour: no-op. If we ever want the hooks back,
2106    /// replace this with a bounded channel or an actual subscriber
2107    /// registry; the callers (`insert`, `delete`, `maybe_seal_growing`)
2108    /// already pass well-typed events.
2109    #[inline]
2110    #[allow(clippy::unused_self)]
2111    fn emit(&self, _event: LifecycleEvent) {}
2112
2113    /// Drain events. Kept for API compatibility; always returns
2114    /// empty because `emit` no longer buffers.
2115    pub fn drain_events(&self) -> Vec<LifecycleEvent> {
2116        std::mem::take(&mut *self.events.write())
2117    }
2118
2119    /// Run maintenance (would be called periodically in production)
2120    pub fn run_maintenance(&self) -> Result<(), SegmentError> {
2121        // Auto-seal idle segments
2122        self.maybe_seal_growing()?;
2123
2124        if self.config.enable_consolidation {
2125            self.consolidation_tick();
2126        }
2127
2128        Ok(())
2129    }
2130
2131    // ========================================================================
2132    // Consolidation — ADR 0073 §5
2133    // ========================================================================
2134
2135    /// One paced step of sealed-segment consolidation.
2136    ///
2137    /// Starts a run when a threshold has been crossed, copies at most
2138    /// `consolidation_entities_per_tick` live entities into the merged segment,
2139    /// and swaps once every source has been drained. Returns without finishing
2140    /// otherwise — the half-built merged segment waits for the next tick.
2141    fn consolidation_tick(&self) {
2142        // Drop the read guard before `start_consolidation` takes the write
2143        // guard: `parking_lot` locks are not reentrant.
2144        let running = self.consolidation.read().is_some();
2145        if !running && !self.start_consolidation() {
2146            return;
2147        }
2148
2149        // Keep the accounting guard through publication: the merged allocation
2150        // must never disappear between the in-flight and active inventories.
2151        let mut guard = self.consolidation.write();
2152        let Some(run) = guard.as_mut() else {
2153            return;
2154        };
2155        self.copy_bounded(run);
2156        if run.cursor >= run.sources.len() {
2157            if let Some(run) = guard.take() {
2158                self.finish_consolidation(run);
2159            }
2160        }
2161    }
2162
2163    /// Evaluate the trigger and, when it fires, open a consolidation run.
2164    ///
2165    /// Returns `true` when a run was started.
2166    fn start_consolidation(&self) -> bool {
2167        let sources = self.select_candidates();
2168        self.start_consolidation_from_sources(sources)
2169    }
2170
2171    fn start_consolidation_from_sources(&self, sources: Vec<SegmentId>) -> bool {
2172        if sources.len() < 2 && !self.candidates_carry_tombstones(&sources) {
2173            return false;
2174        }
2175
2176        let merged_id = self.next_segment_id.fetch_add(1, Ordering::SeqCst);
2177        let merged = GrowingSegment::new(merged_id, &self.collection);
2178
2179        *self.consolidation.write() = Some(Consolidation {
2180            sources: sources
2181                .into_iter()
2182                .map(|id| SourceProgress {
2183                    id,
2184                    copied: Vec::new(),
2185                    tombstones_at_open: None,
2186                })
2187                .collect(),
2188            merged_id,
2189            merged,
2190            cursor: 0,
2191            pending_ids: Vec::new(),
2192            pending_cursor: 0,
2193            entities_copied: 0,
2194        });
2195
2196        self.stats.write().consolidation.runs_started += 1;
2197        true
2198    }
2199
2200    /// Run one paced consolidation tick because admission pressure needs
2201    /// reclaimable bytes, even if the background thresholds have not fired.
2202    ///
2203    /// Returns true when the call started or advanced a consolidation run.
2204    pub fn pressure_consolidation_tick(&self) -> bool {
2205        if !self.config.enable_consolidation {
2206            return false;
2207        }
2208
2209        let running = self.consolidation.read().is_some();
2210        if !running {
2211            let sources = self.dirty_sealed_segments();
2212            if !self.start_consolidation_from_sources(sources) {
2213                return false;
2214            }
2215        }
2216
2217        self.consolidation_tick();
2218        true
2219    }
2220
2221    /// Which sealed segments this collection should merge, if any.
2222    ///
2223    /// Three thresholds, any of which arms consolidation:
2224    ///
2225    /// * **sealed-segment count** above `max_sealed_segments` — every sealed
2226    ///   segment merges, which is what that knob always claimed to do.
2227    /// * **tombstone ratio** — dead entities over total entities across the
2228    ///   sealed set, past [`CONSOLIDATION_TOMBSTONE_RATIO`].
2229    /// * **fragmentation** — reclaimable bytes over held bytes, past
2230    ///   [`CONSOLIDATION_FRAGMENTATION_RATIO`].
2231    ///
2232    /// The count trigger takes every sealed segment; the two budget triggers
2233    /// take only the segments that actually carry tombstones, so a clean
2234    /// segment is never rewritten for nothing.
2235    fn select_candidates(&self) -> Vec<SegmentId> {
2236        let sealed = self.sealed.read();
2237        if sealed.is_empty() {
2238            return Vec::new();
2239        }
2240
2241        if sealed.len() > self.config.max_sealed_segments {
2242            return sealed.iter().map(|segment| segment.read().id()).collect();
2243        }
2244
2245        let mut live = 0usize;
2246        let mut tombstones = 0usize;
2247        let mut held = 0u64;
2248        let mut reclaimable = 0u64;
2249        let mut dirty = Vec::new();
2250
2251        for segment_arc in sealed.iter() {
2252            let segment = segment_arc.read();
2253            live += segment.live_entity_count();
2254            tombstones += segment.tombstone_count();
2255            held += segment.resident_bytes();
2256            reclaimable += segment.reclaimable_bytes();
2257            if segment.tombstone_count() > 0 {
2258                dirty.push(segment.id());
2259            }
2260        }
2261
2262        let total = live + tombstones;
2263        let tombstone_ratio = if total == 0 {
2264            0.0
2265        } else {
2266            tombstones as f64 / total as f64
2267        };
2268        let fragmentation = if held == 0 {
2269            0.0
2270        } else {
2271            reclaimable as f64 / held as f64
2272        };
2273
2274        if tombstone_ratio >= CONSOLIDATION_TOMBSTONE_RATIO
2275            || fragmentation >= CONSOLIDATION_FRAGMENTATION_RATIO
2276        {
2277            return dirty;
2278        }
2279
2280        Vec::new()
2281    }
2282
2283    fn dirty_sealed_segments(&self) -> Vec<SegmentId> {
2284        self.sealed
2285            .read()
2286            .iter()
2287            .filter_map(|segment_arc| {
2288                let segment = segment_arc.read();
2289                (segment.reclaimable_bytes() > 0).then(|| segment.id())
2290            })
2291            .collect()
2292    }
2293
2294    /// A single candidate is still worth merging when it carries tombstones —
2295    /// the merge drops them. A single clean candidate is not.
2296    fn candidates_carry_tombstones(&self, sources: &[SegmentId]) -> bool {
2297        if sources.is_empty() {
2298            return false;
2299        }
2300        let sealed = self.sealed.read();
2301        sealed.iter().any(|segment_arc| {
2302            let segment = segment_arc.read();
2303            sources.contains(&segment.id()) && segment.tombstone_count() > 0
2304        })
2305    }
2306
2307    fn sealed_segment(&self, id: SegmentId) -> Option<Arc<RwLock<GrowingSegment>>> {
2308        let sealed = self.sealed.read();
2309        sealed
2310            .iter()
2311            .find(|segment| segment.read().id() == id)
2312            .map(Arc::clone)
2313    }
2314
2315    /// Copy at most one tick's worth of live entities into the merged segment.
2316    ///
2317    /// The per-tick bound is the pacing contract: a tick's cost is O(bound),
2318    /// independent of how many entities the run still has to move.
2319    fn copy_bounded(&self, run: &mut Consolidation) {
2320        let mut budget = self.config.consolidation_entities_per_tick;
2321
2322        while budget > 0 && run.cursor < run.sources.len() {
2323            let source_id = run.sources[run.cursor].id;
2324            let Some(source_arc) = self.sealed_segment(source_id) else {
2325                // The source vanished. Only consolidation removes sealed
2326                // segments and it is single-threaded per collection, so this
2327                // cannot happen today — skip it rather than resurrect its rows.
2328                run.cursor += 1;
2329                run.pending_ids.clear();
2330                run.pending_cursor = 0;
2331                continue;
2332            };
2333
2334            if run.sources[run.cursor].tombstones_at_open.is_none() {
2335                let source = source_arc.read();
2336                run.pending_ids = source.live_entity_ids();
2337                run.pending_cursor = 0;
2338                run.sources[run.cursor].tombstones_at_open = Some(source.tombstone_count());
2339            }
2340
2341            {
2342                let source = source_arc.read();
2343                while budget > 0 && run.pending_cursor < run.pending_ids.len() {
2344                    let id = run.pending_ids[run.pending_cursor];
2345                    run.pending_cursor += 1;
2346                    budget -= 1;
2347
2348                    // Tombstoned between the snapshot and now — it must not be
2349                    // resurrected into the merged segment.
2350                    let Some(entity) = source.get(id) else {
2351                        continue;
2352                    };
2353                    let metadata = source.get_metadata(id);
2354                    run.merged.adopt_entity(entity.clone(), metadata);
2355                    run.sources[run.cursor].copied.push(id);
2356                    run.entities_copied += 1;
2357                }
2358            }
2359
2360            if run.pending_cursor >= run.pending_ids.len() {
2361                run.cursor += 1;
2362                run.pending_ids.clear();
2363                run.pending_cursor = 0;
2364            }
2365        }
2366    }
2367
2368    /// Atomically replace the source segments with the merged one.
2369    ///
2370    /// The `sealed` write lock is the boundary: a reader holds the read lock
2371    /// for the whole of its scan, so it observes either the old set or the new
2372    /// set, never a gap and never both copies of a row.
2373    fn finish_consolidation(&self, mut run: Consolidation) {
2374        let mut sealed = self.sealed.write();
2375
2376        let mut positions = Vec::with_capacity(run.sources.len());
2377        for source in &run.sources {
2378            if let Some(index) = sealed
2379                .iter()
2380                .position(|segment| segment.read().id() == source.id)
2381            {
2382                positions.push(index);
2383            }
2384        }
2385        if positions.is_empty() {
2386            return;
2387        }
2388        positions.sort_unstable();
2389
2390        // Rows deleted from a source after we copied them must not come back.
2391        // A source whose tombstone count is unchanged saw no deletes while the
2392        // merge ran, so its copied rows need no per-row check.
2393        let mut stale = Vec::new();
2394        for source in &run.sources {
2395            let Some(source_arc) = sealed
2396                .iter()
2397                .find(|segment| segment.read().id() == source.id)
2398            else {
2399                continue;
2400            };
2401            let segment = source_arc.read();
2402            if source.tombstones_at_open == Some(segment.tombstone_count()) {
2403                continue;
2404            }
2405            stale.extend(
2406                source
2407                    .copied
2408                    .iter()
2409                    .copied()
2410                    .filter(|id| !segment.contains(*id)),
2411            );
2412        }
2413        for id in stale {
2414            run.merged.evict_entity(id);
2415        }
2416
2417        let reclaimed_bytes: u64 = positions
2418            .iter()
2419            .map(|&index| sealed[index].read().resident_bytes())
2420            .sum();
2421        let reclaimed_tombstones: u64 = positions
2422            .iter()
2423            .map(|&index| sealed[index].read().tombstone_count() as u64)
2424            .sum();
2425
2426        if run.merged.seal().is_err() {
2427            return;
2428        }
2429        let merged_bytes = run.merged.resident_bytes();
2430        let merged_arc = Arc::new(RwLock::new(run.merged));
2431
2432        {
2433            let mut retired = self.retired.write();
2434            retired.retain(|source| source.strong_count() > 0);
2435            for &index in positions.iter().rev() {
2436                let source = sealed.remove(index);
2437                if Arc::strong_count(&source) > 1 {
2438                    retired.push(Arc::downgrade(&source));
2439                }
2440            }
2441        }
2442        let insert_at = positions[0].min(sealed.len());
2443        sealed.insert(insert_at, Arc::clone(&merged_arc));
2444        drop(sealed);
2445
2446        // Remap the entity→segment hints that pointed at the retired sources.
2447        let source_ids: Vec<SegmentId> = run.sources.iter().map(|source| source.id).collect();
2448        {
2449            let mut entity_segment = self.entity_segment.write();
2450            for segment_id in entity_segment.values_mut() {
2451                if source_ids.contains(segment_id) {
2452                    *segment_id = run.merged_id;
2453                }
2454            }
2455        }
2456
2457        {
2458            let mut stats = self.stats.write();
2459            stats.sealed_count = stats
2460                .sealed_count
2461                .saturating_sub(positions.len())
2462                .saturating_add(1);
2463            let consolidation = &mut stats.consolidation;
2464            consolidation.runs_completed += 1;
2465            consolidation.segments_merged += positions.len() as u64;
2466            consolidation.tombstones_reclaimed += reclaimed_tombstones;
2467            consolidation.bytes_reclaimed += reclaimed_bytes.saturating_sub(merged_bytes);
2468        }
2469
2470        self.emit(LifecycleEvent::SegmentConsolidated {
2471            source: source_ids,
2472            target: run.merged_id,
2473        });
2474    }
2475
2476    /// Entities copied by the consolidation currently in flight, or `None` when
2477    /// none is running. Exposed for the pacing tests.
2478    #[cfg(test)]
2479    fn consolidation_progress(&self) -> Option<u64> {
2480        self.consolidation
2481            .read()
2482            .as_ref()
2483            .map(|run| run.entities_copied)
2484    }
2485}
2486
2487// Implement the Arc<RwLock<GrowingSegment>> as UnifiedSegment
2488// This is needed because we store growing segments in the sealed list after sealing
2489impl UnifiedSegment for Arc<RwLock<GrowingSegment>> {
2490    fn id(&self) -> SegmentId {
2491        self.read().id()
2492    }
2493
2494    fn state(&self) -> SegmentState {
2495        self.read().state()
2496    }
2497
2498    fn collection(&self) -> &str {
2499        // This is a limitation - we'd need to store collection in the Arc wrapper
2500        "unknown"
2501    }
2502
2503    fn stats(&self) -> SegmentStats {
2504        self.read().stats()
2505    }
2506
2507    fn entity_count(&self) -> usize {
2508        self.read().entity_count()
2509    }
2510
2511    fn contains(&self, id: EntityId) -> bool {
2512        self.read().contains(id)
2513    }
2514
2515    fn get(&self, id: EntityId) -> Option<&UnifiedEntity> {
2516        // This is tricky with RwLock - we can't return a reference
2517        // In production, we'd use a different approach
2518        None
2519    }
2520
2521    fn get_mut(&mut self, _id: EntityId) -> Option<&mut UnifiedEntity> {
2522        None
2523    }
2524
2525    fn insert(&mut self, entity: UnifiedEntity) -> Result<EntityId, SegmentError> {
2526        self.write().insert(entity)
2527    }
2528
2529    fn update(&mut self, entity: UnifiedEntity) -> Result<(), SegmentError> {
2530        self.write().update(entity)
2531    }
2532
2533    fn update_hot(
2534        &mut self,
2535        entity: UnifiedEntity,
2536        modified_columns: &[String],
2537    ) -> Result<(), SegmentError> {
2538        self.write().update_hot(entity, modified_columns)
2539    }
2540
2541    fn delete(&mut self, id: EntityId) -> Result<bool, SegmentError> {
2542        self.write().delete(id)
2543    }
2544
2545    fn get_metadata(&self, id: EntityId) -> Option<Metadata> {
2546        self.read().get_metadata(id)
2547    }
2548
2549    fn set_metadata(&mut self, id: EntityId, metadata: Metadata) -> Result<(), SegmentError> {
2550        self.write().set_metadata(id, metadata)
2551    }
2552
2553    fn seal(&mut self) -> Result<(), SegmentError> {
2554        self.write().seal()
2555    }
2556
2557    fn should_seal(&self, config: &SegmentConfig) -> bool {
2558        self.read().should_seal(config)
2559    }
2560
2561    fn iter(&self) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2562        // Cannot return iterator with RwLock
2563        Box::new(std::iter::empty())
2564    }
2565
2566    fn iter_kind(&self, _kind_filter: &str) -> Box<dyn Iterator<Item = &UnifiedEntity> + '_> {
2567        Box::new(std::iter::empty())
2568    }
2569
2570    fn filter_metadata(&self, filters: &[(String, MetadataFilter)]) -> Vec<EntityId> {
2571        self.read().filter_metadata(filters)
2572    }
2573}
2574
2575#[cfg(test)]
2576mod tests {
2577    use std::collections::HashSet;
2578
2579    use super::*;
2580    use crate::runtime::mvcc::SnapshotContext;
2581    use crate::storage::transaction::snapshot::{Snapshot, SnapshotManager};
2582    use crate::storage::unified::entity::{EntityData, EntityKind, RowData};
2583    use reddb_types::Value;
2584
2585    #[test]
2586    fn bounded_scan_survives_sealing_consolidation_and_new_writes() {
2587        let manager = SegmentManager::with_config(
2588            "vectors",
2589            ManagerConfig {
2590                max_sealed_segments: 1,
2591                consolidation_entities_per_tick: 10_000,
2592                ..Default::default()
2593            },
2594        );
2595        for part in 0..3 {
2596            manager
2597                .bulk_insert(
2598                    (1..=300)
2599                        .map(|offset| {
2600                            UnifiedEntity::vector(
2601                                EntityId::new(part * 300 + offset),
2602                                "vectors",
2603                                vec![1.0, 0.0],
2604                            )
2605                        })
2606                        .collect(),
2607                )
2608                .expect("batch");
2609            if part < 2 {
2610                manager.force_seal().expect("seal");
2611            }
2612        }
2613        let original_growing = manager.growing.read().as_ref().expect("growing").clone();
2614        let mut changed = false;
2615        let mut seen = Vec::new();
2616        assert!(manager.scan_batches(
2617            None,
2618            || true,
2619            |batch| {
2620                assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
2621                assert!(
2622                    manager.growing.try_write().is_some(),
2623                    "consumer holds no topology guard"
2624                );
2625                assert!(
2626                    original_growing.try_write().is_some(),
2627                    "consumer holds no segment guard"
2628                );
2629                for entity in batch {
2630                    assert!(manager.get(entity.id).is_some(), "reentrant storage lookup");
2631                    seen.push(entity.id.raw());
2632                }
2633                if !changed {
2634                    changed = true;
2635                    manager.force_seal().expect("seal during scan");
2636                    manager
2637                        .insert(UnifiedEntity::vector(
2638                            EntityId::new(5000),
2639                            "vectors",
2640                            vec![1.0, 0.0],
2641                        ))
2642                        .expect("new write");
2643                    drain_maintenance(&manager);
2644                    assert!(manager.stats().consolidation.runs_completed > 0);
2645                    assert!(manager
2646                        .delete(EntityId::new(1))
2647                        .expect("delete after source retirement"));
2648                    assert!(manager.get(EntityId::new(1)).is_none());
2649                }
2650                true
2651            }
2652        ));
2653        seen.sort_unstable();
2654        assert_eq!(
2655            seen,
2656            (2..=900).collect::<Vec<_>>(),
2657            "no duplicate, missing or late candidates"
2658        );
2659        let mut after = Vec::new();
2660        assert!(manager.scan_batches(
2661            None,
2662            || true,
2663            |batch| {
2664                after.extend(batch.iter().map(|entity| entity.id.raw()));
2665                true
2666            }
2667        ));
2668        after.sort_unstable();
2669        let mut expected: Vec<_> = (2..=900).collect();
2670        expected.push(5000);
2671        assert_eq!(after, expected, "adopted entities have scan positions");
2672    }
2673
2674    #[test]
2675    fn bounded_scan_applies_snapshot_to_relocated_versions() {
2676        for (consolidate, bulk) in [(false, false), (false, true), (true, false), (true, true)] {
2677            let manager = SegmentManager::with_config(
2678                "vectors",
2679                ManagerConfig {
2680                    max_sealed_segments: 1,
2681                    ..Default::default()
2682                },
2683            );
2684            for ids in [[1, 2], [3, 4]] {
2685                let entities = ids.map(|id| {
2686                    let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
2687                    entity.set_xmin(3);
2688                    entity
2689                });
2690                if bulk {
2691                    manager
2692                        .bulk_insert(entities.into())
2693                        .expect("initial versions");
2694                } else {
2695                    for entity in entities {
2696                        manager.insert(entity).expect("initial version");
2697                    }
2698                }
2699                manager.force_seal().expect("seal source");
2700            }
2701            let snapshot = SnapshotContext {
2702                snapshot: Snapshot {
2703                    xid: 10,
2704                    in_progress: HashSet::new(),
2705                },
2706                manager: Arc::new(SnapshotManager::new()),
2707                own_xids: HashSet::new(),
2708                requires_index_fallback: false,
2709                serializable_reader: None,
2710            };
2711            let mut changed = false;
2712            let mut seen = Vec::new();
2713            assert!(manager.scan_batches(
2714                Some(&snapshot),
2715                || true,
2716                |batch| {
2717                    seen.extend(batch.iter().map(|entity| entity.id.raw()));
2718                    if !changed {
2719                        changed = true;
2720                        if consolidate {
2721                            drain_maintenance(&manager);
2722                            assert!(manager.stats().consolidation.runs_completed > 0);
2723                        }
2724                        for (id, xmax) in [(3, 8), (4, 12)] {
2725                            let mut entity =
2726                                manager.get(EntityId::new(id)).expect("retained version");
2727                            entity.set_xmax(xmax);
2728                            manager
2729                                .update(entity)
2730                                .expect("update after source retirement");
2731                        }
2732                    }
2733                    true
2734                }
2735            ));
2736            seen.sort_unstable();
2737            assert_eq!(seen, vec![1, 2, 4], "snapshot uses authoritative xmax");
2738            let mut again = Vec::new();
2739            assert!(manager.scan_batches(
2740                Some(&snapshot),
2741                || true,
2742                |batch| {
2743                    again.extend(batch.iter().map(|entity| entity.id.raw()));
2744                    true
2745                }
2746            ));
2747            again.sort_unstable();
2748            assert_eq!(
2749                again,
2750                vec![1, 2, 4],
2751                "old slots do not duplicate relocated versions"
2752            );
2753            let mut current = Vec::new();
2754            assert!(manager.scan_batches(
2755                None,
2756                || true,
2757                |batch| {
2758                    current.extend(batch.iter().map(|entity| entity.id.raw()));
2759                    true
2760                }
2761            ));
2762            current.sort_unstable();
2763            assert_eq!(
2764                current,
2765                vec![1, 2],
2766                "current view excludes both deleted versions"
2767            );
2768        }
2769    }
2770
2771    #[test]
2772    fn bounded_scan_keeps_flat_and_hash_positions_through_mutation() {
2773        let manager = SegmentManager::new("mixed");
2774        let vector = |id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![1.0]);
2775        manager
2776            .bulk_insert((1..=300).map(vector).collect())
2777            .expect("flat");
2778        manager.insert(vector(500)).expect("hash insert");
2779        manager
2780            .bulk_insert(vec![vector(502), vector(505)])
2781            .expect("gapped bulk");
2782        manager.delete(EntityId::new(2)).expect("flat delete");
2783        manager.delete(EntityId::new(505)).expect("hash delete");
2784        let mut changed = false;
2785        let mut seen = Vec::new();
2786        assert!(manager.scan_batches(
2787            None,
2788            || true,
2789            |batch| {
2790                seen.extend(batch.iter().map(|entity| entity.id.raw()));
2791                if !changed {
2792                    changed = true;
2793                    manager
2794                        .insert(vector(700))
2795                        .expect("append after cursor boundary");
2796                    manager
2797                        .update(UnifiedEntity::graph_node(
2798                            EntityId::new(500),
2799                            "mixed",
2800                            "Node",
2801                            HashMap::new(),
2802                        ))
2803                        .expect("structural update");
2804                }
2805                true
2806            }
2807        ));
2808        seen.sort_unstable();
2809        let mut expected: Vec<_> = (1..=300).filter(|id| *id != 2).collect();
2810        expected.extend([500, 502]);
2811        assert_eq!(seen, expected);
2812    }
2813
2814    #[test]
2815    fn bounded_scan_preserves_snapshot_and_discards_interrupted_batches() {
2816        let manager = SegmentManager::new("vectors");
2817        let mut entities = Vec::new();
2818        for (id, xmin, xmax) in [(1, 3, 12), (2, 12, 0), (3, 3, 0), (4, 3, 4)] {
2819            let mut entity = UnifiedEntity::vector(EntityId::new(id), "vectors", vec![1.0]);
2820            entity.set_xmin(xmin);
2821            entity.set_xmax(xmax);
2822            entities.push(entity);
2823        }
2824        manager.bulk_insert(entities).expect("versions");
2825        manager.force_seal().expect("seal");
2826        let snapshot = SnapshotContext {
2827            snapshot: Snapshot {
2828                xid: 10,
2829                in_progress: HashSet::from([4]),
2830            },
2831            manager: Arc::new(SnapshotManager::new()),
2832            own_xids: HashSet::new(),
2833            requires_index_fallback: false,
2834            serializable_reader: None,
2835        };
2836        let mut work = 0;
2837        let mut callbacks = 0;
2838        assert!(!manager.scan_batches(
2839            Some(&snapshot),
2840            || {
2841                work += 1;
2842                work <= 2
2843            },
2844            |_| {
2845                callbacks += 1;
2846                true
2847            }
2848        ));
2849        assert_eq!(callbacks, 0, "a stopped batch is not consumed");
2850        for (view, expected) in [(Some(&snapshot), vec![1, 3, 4]), (None, vec![2, 3])] {
2851            let mut seen = Vec::new();
2852            assert!(manager.scan_batches(
2853                view,
2854                || true,
2855                |batch| {
2856                    seen.extend(batch.iter().map(|entity| entity.id.raw()));
2857                    true
2858                }
2859            ));
2860            seen.sort_unstable();
2861            assert_eq!(seen, expected);
2862        }
2863    }
2864
2865    #[test]
2866    fn test_manager_basic() {
2867        let manager = SegmentManager::new("test_collection");
2868
2869        let entity = UnifiedEntity::table_row(
2870            manager.next_entity_id(),
2871            "users",
2872            1,
2873            vec![Value::text("Alice".to_string())],
2874        );
2875
2876        let id = manager.insert(entity).unwrap();
2877        assert!(manager.get(id).is_some());
2878        assert_eq!(manager.count(), 1);
2879    }
2880
2881    #[test]
2882    fn graph_pruning_preserves_physical_versions_and_charges_only_selected_segments() {
2883        let manager = Arc::new(SegmentManager::new("mixed"));
2884        manager
2885            .bulk_insert(
2886                (1..=128)
2887                    .map(|id| UnifiedEntity::vector(EntityId::new(id), "mixed", vec![]))
2888                    .collect(),
2889            )
2890            .expect("vector-only segment");
2891        manager.force_seal().expect("seal vectors");
2892        let mut old =
2893            UnifiedEntity::graph_node(EntityId::new(200), "mixed", "person", Default::default());
2894        old.set_xmin(3);
2895        old.set_xmax(12);
2896        let mut new =
2897            UnifiedEntity::graph_node(EntityId::new(201), "mixed", "person", Default::default());
2898        new.set_xmin(12);
2899        let edge = UnifiedEntity::graph_edge(
2900            EntityId::new(202),
2901            "mixed",
2902            "a",
2903            "b",
2904            1.0,
2905            Default::default(),
2906        );
2907        manager
2908            .bulk_insert(vec![old, new, edge])
2909            .expect("physical graph versions");
2910        manager.force_seal().expect("seal graph");
2911        manager
2912            .insert(UnifiedEntity::graph_node(
2913                EntityId::new(203),
2914                "mixed",
2915                "person",
2916                Default::default(),
2917            ))
2918            .expect("growing node");
2919        let snapshot = SnapshotContext {
2920            snapshot: Snapshot {
2921                xid: 10,
2922                in_progress: HashSet::new(),
2923            },
2924            manager: Arc::new(SnapshotManager::new()),
2925            own_xids: HashSet::new(),
2926            requires_index_fallback: false,
2927            serializable_reader: None,
2928        };
2929        std::thread::spawn(move || {
2930            for view in [None, Some(&snapshot)] {
2931                for kind in [GraphEntityKind::Node, GraphEntityKind::Edge] {
2932                    let expected: HashSet<_> = manager
2933                        .scan(view, |item| kind.matches(&item.kind))
2934                        .into_iter()
2935                        .map(|item| item.id)
2936                        .collect();
2937                    let actual: HashSet<_> = manager
2938                        .scan_graph_kind(view, kind)
2939                        .into_iter()
2940                        .map(|item| item.id)
2941                        .collect();
2942                    assert_eq!(actual, expected);
2943                    let mut visited = 0;
2944                    let mut candidates = HashSet::new();
2945                    manager.scan_graph_candidates_for_each(
2946                        view,
2947                        kind,
2948                        || true,
2949                        |item| {
2950                            visited += 1;
2951                            if kind.matches(&item.kind) {
2952                                candidates.insert(item.id);
2953                            }
2954                            true
2955                        },
2956                    );
2957                    assert_eq!(candidates, expected);
2958                    assert_eq!(
2959                        visited,
2960                        if matches!(kind, GraphEntityKind::Node) {
2961                            3
2962                        } else {
2963                            2
2964                        }
2965                    );
2966                }
2967            }
2968            let ids: HashSet<_> = manager
2969                .scan_graph_kind(Some(&snapshot), GraphEntityKind::Node)
2970                .into_iter()
2971                .map(|item| item.id.raw())
2972                .collect();
2973            assert_eq!(ids, HashSet::from([200, 203]));
2974            let ids: HashSet<_> = manager
2975                .scan_graph_kind(None, GraphEntityKind::Node)
2976                .into_iter()
2977                .map(|item| item.id.raw())
2978                .collect();
2979            assert_eq!(ids, HashSet::from([201, 203]));
2980            let mut visits = 0;
2981            manager.scan_graph_candidates_for_each(
2982                None,
2983                GraphEntityKind::Node,
2984                || true,
2985                |_| {
2986                    visits += 1;
2987                    false
2988                },
2989            );
2990            assert_eq!(
2991                visits, 1,
2992                "early termination must also stop subsequent segments"
2993            );
2994        })
2995        .join()
2996        .expect("snapshot worker");
2997    }
2998
2999    #[test]
3000    fn snapshot_scan_applies_table_driven_visibility_on_fresh_thread() {
3001        let manager = Arc::new(SegmentManager::new("accounts"));
3002        let cases = [
3003            ("legacy row", 0, 0, true),
3004            ("committed creator", 3, 0, true),
3005            ("future creator", 11, 0, false),
3006            ("in-progress creator", 4, 0, false),
3007            ("committed deleter", 3, 7, false),
3008            ("future deleter", 3, 11, true),
3009            ("in-progress deleter", 3, 4, true),
3010        ];
3011
3012        for (index, (_, xmin, xmax, _)) in cases.iter().enumerate() {
3013            let row_id = index as u64 + 1;
3014            let mut entity = UnifiedEntity::table_row(
3015                EntityId::new(row_id),
3016                "accounts",
3017                row_id,
3018                vec![Value::Integer(row_id as i64)],
3019            );
3020            entity.set_xmin(*xmin);
3021            entity.set_xmax(*xmax);
3022            manager.insert(entity).unwrap();
3023        }
3024
3025        let snapshot = SnapshotContext {
3026            snapshot: Snapshot {
3027                xid: 10,
3028                in_progress: HashSet::from([4]),
3029            },
3030            manager: Arc::new(SnapshotManager::new()),
3031            own_xids: HashSet::new(),
3032            requires_index_fallback: false,
3033            serializable_reader: None,
3034        };
3035
3036        let raw_visible_row_ids: HashSet<u64> = manager
3037            .query_all(|entity| entity_visible_with_context(Some(&snapshot), entity))
3038            .into_iter()
3039            .filter_map(|entity| match entity.kind {
3040                EntityKind::TableRow { row_id, .. } => Some(row_id),
3041                _ => None,
3042            })
3043            .collect();
3044
3045        let scanned = std::thread::spawn(move || manager.scan(Some(&snapshot), |_| true))
3046            .join()
3047            .unwrap();
3048        let visible_row_ids: HashSet<u64> = scanned
3049            .into_iter()
3050            .filter_map(|entity| match entity.kind {
3051                EntityKind::TableRow { row_id, .. } => Some(row_id),
3052                _ => None,
3053            })
3054            .collect();
3055        assert_eq!(visible_row_ids, raw_visible_row_ids);
3056
3057        for (index, (name, _, _, expected_visible)) in cases.iter().enumerate() {
3058            let row_id = index as u64 + 1;
3059            assert_eq!(
3060                visible_row_ids.contains(&row_id),
3061                *expected_visible,
3062                "{name} visibility mismatch"
3063            );
3064        }
3065    }
3066
3067    #[test]
3068    fn frameless_scan_hides_superseded_versions_and_moderated_rows() {
3069        // A `None` snapshot is the autocommit / prepared-fast-path read: it
3070        // must keep the resolver fallback (xmax == 0), never show every
3071        // version, and never panic (issue #2136 review).
3072        let manager = SegmentManager::new("accounts");
3073        let mut live =
3074            UnifiedEntity::table_row(EntityId::new(1), "accounts", 1, vec![Value::Integer(1)]);
3075        live.set_xmin(2);
3076        live.set_xmax(0);
3077        manager.insert(live).unwrap();
3078
3079        let mut superseded =
3080            UnifiedEntity::table_row(EntityId::new(2), "accounts", 2, vec![Value::Integer(2)]);
3081        superseded.set_xmin(2);
3082        superseded.set_xmax(9);
3083        manager.insert(superseded).unwrap();
3084
3085        let rows: Vec<u64> = manager
3086            .scan(None, |_| true)
3087            .into_iter()
3088            .filter_map(|entity| match entity.kind {
3089                EntityKind::TableRow { row_id, .. } => Some(row_id),
3090                _ => None,
3091            })
3092            .collect();
3093        assert_eq!(
3094            rows,
3095            vec![1],
3096            "frameless scan must hide superseded versions"
3097        );
3098    }
3099
3100    #[test]
3101    fn query_scan_callers_use_snapshot_api_for_visibility() {
3102        let table = include_str!("../../runtime/query_exec/table.rs");
3103        let aggregate = include_str!("../../runtime/query_exec/aggregate.rs");
3104
3105        assert_eq!(
3106            table.matches(".for_each_entity(|entity|").count(),
3107            1,
3108            "only the visibility-independent schema inspection may use the raw table iterator"
3109        );
3110        assert!(!table.contains(".query_all_zoned_with_stats("));
3111        assert!(!table.contains(".for_each_entity_zoned_with_stats("));
3112        assert!(!table.contains("TableRowMvccReadResolver::captured"));
3113
3114        assert!(!aggregate.contains("TableRowMvccReadResolver"));
3115        assert!(!aggregate.contains(".for_each_entity("));
3116        assert!(!aggregate.contains(".fold_entities_parallel("));
3117    }
3118
3119    #[test]
3120    fn remaining_visibility_scan_callers_use_snapshot_api() {
3121        let aggregate_pushdown =
3122            include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs");
3123        let ask = include_str!("../../runtime/ask_pipeline.rs");
3124        let dml = include_str!("../../runtime/impl_dml.rs");
3125        let dml_targets = include_str!("../../runtime/dml_target_scan.rs");
3126        let events = include_str!("../../runtime/impl_events.rs");
3127        let graph_commands = include_str!("../../runtime/impl_graph_commands.rs");
3128        let graph_dsl = include_str!("../../runtime/graph_dsl.rs");
3129        let graph_tvf = include_str!("../../runtime/graph_tvf.rs");
3130        let kv = include_str!("../../runtime/impl_kv.rs");
3131        let queue = include_str!("../../runtime/impl_queue.rs");
3132        let primary_queue = include_str!("../../runtime/primary_queue_store.rs");
3133        let record_search = include_str!("../../runtime/record_search.rs");
3134        let search = include_str!("../../runtime/impl_search.rs");
3135        let vector = include_str!("../../runtime/query_exec/vector.rs");
3136
3137        assert!(!aggregate_pushdown.contains("TableRowMvccReadResolver"));
3138        assert!(!aggregate_pushdown.contains(".for_each_entity("));
3139        assert!(!ask.contains("for entity in manager.query_all(|_| true)"));
3140        assert!(!dml.contains("let entities = manager.query_all(|_| true);"));
3141        // The DML target scan keeps raw iteration on purpose: its
3142        // visibility is kind-conditional (live table rows at the tip
3143        // version, everything else under the snapshot), which the uniform
3144        // scan API cannot express. `candidate_visible` is the single gate.
3145        assert_eq!(
3146            dml_targets
3147                .matches("self.candidate_visible(snapshot.as_ref(), entity)")
3148                .count(),
3149            2,
3150            "both raw DML iterations must route through candidate_visible"
3151        );
3152        // The event-backfill dedup set is deliberately visibility-free:
3153        // it must include consumed messages or backfill re-enqueues them.
3154        assert_eq!(
3155            events.matches("manager.query_all(").count(),
3156            1,
3157            "only the visibility-free backfill dedup may use query_all"
3158        );
3159        assert!(!graph_commands.contains("TableRowMvccReadResolver"));
3160        assert!(!graph_dsl.contains("entity_visible_with_context(snap_ctx.as_ref(), e)"));
3161        assert!(!graph_tvf.contains("entity_visible_with_context(snap_ctx.as_ref(), &entity)"));
3162        assert!(!kv.contains("TableRowMvccReadResolver"));
3163        assert!(!queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3164        assert!(!primary_queue.contains("entity_visible_with_context(snap_ctx.as_ref(), entity)"));
3165        assert!(!record_search.contains("TableRowMvccReadResolver"));
3166        assert!(!search.contains("for entity in manager.query_all(|_| true)"));
3167        assert!(!vector.contains("manager.query_all("));
3168    }
3169
3170    /// Strip `#[cfg(test)] mod … { … }` blocks so the raw-iteration gate
3171    /// only sees production code. Blocks are matched at column zero and
3172    /// closed by brace balance, which is why a file may carry a test
3173    /// module in the middle and still keep production code after it.
3174    fn production_source(src: &str) -> String {
3175        let mut kept: Vec<&str> = Vec::new();
3176        let lines: Vec<&str> = src.lines().collect();
3177        let mut idx = 0;
3178        while idx < lines.len() {
3179            let starts_test_mod = lines[idx].starts_with("#[cfg(test)]")
3180                && lines
3181                    .get(idx + 1)
3182                    .is_some_and(|next| next.starts_with("mod "));
3183            if !starts_test_mod {
3184                kept.push(lines[idx]);
3185                idx += 1;
3186                continue;
3187            }
3188            idx += 1;
3189            let mut depth: isize = 0;
3190            while idx < lines.len() {
3191                depth += lines[idx].matches('{').count() as isize;
3192                depth -= lines[idx].matches('}').count() as isize;
3193                idx += 1;
3194                if depth <= 0 {
3195                    break;
3196                }
3197            }
3198        }
3199        kept.join("\n")
3200    }
3201
3202    /// Count raw (visibility-bypassing) iteration entry points.
3203    fn raw_iteration_sites(src: &str) -> usize {
3204        src.matches("query_all(").count()
3205            + src.matches(".for_each_entity").count()
3206            + src.matches("fold_entities_parallel(").count()
3207    }
3208
3209    #[test]
3210    fn raw_iteration_confined_to_maintenance_allowlist() {
3211        // The query paths migrated by Spec #2107 batches 1 and 2. Inside
3212        // this perimeter raw iteration is maintenance-only: it bypasses
3213        // MVCC visibility entirely, so every remaining site is either an
3214        // internal catalog/bookkeeping read that must observe all physical
3215        // versions, or one of the three exceptions the batch reviews
3216        // documented. A new raw-iteration site changes the count and fails
3217        // this gate — route it through `SegmentManager::scan*` instead.
3218        let allowlist: &[(&str, &str, usize, &str)] = &[
3219            (
3220                "runtime/query_exec/table.rs",
3221                include_str!("../../runtime/query_exec/table.rs"),
3222                1,
3223                "documented exception: schema inspection (does the collection ever \
3224                 observe these fields?) is visibility-independent",
3225            ),
3226            (
3227                "runtime/query_exec/aggregate.rs",
3228                include_str!("../../runtime/query_exec/aggregate.rs"),
3229                0,
3230                "fully migrated by batch 1",
3231            ),
3232            (
3233                "runtime/query_exec/aggregate_pushdown_dispatch.rs",
3234                include_str!("../../runtime/query_exec/aggregate_pushdown_dispatch.rs"),
3235                0,
3236                "fully migrated by batch 1",
3237            ),
3238            (
3239                "runtime/query_exec/vector.rs",
3240                include_str!("../../runtime/query_exec/vector.rs"),
3241                0,
3242                "fully migrated by batch 2",
3243            ),
3244            (
3245                "runtime/ask_pipeline.rs",
3246                include_str!("../../runtime/ask_pipeline.rs"),
3247                0,
3248                "fully migrated by batch 2",
3249            ),
3250            (
3251                "runtime/impl_dml.rs",
3252                include_str!("../../runtime/impl_dml.rs"),
3253                1,
3254                "documented exception: ON CONFLICT resolves a unique key against \
3255                 physical versions on purpose. A snapshot scan cannot see a row an \
3256                 uncommitted concurrent writer already holds the key for, which is \
3257                 exactly the case that must raise a serialization conflict rather \
3258                 than insert a duplicate; the same pass reaps index entries left by \
3259                 aborted writers. \
3260                 TODO(#2263 follow-up): narrow the candidates through the unique \
3261                 index instead of walking the collection",
3262            ),
3263            (
3264                "runtime/dml_target_scan.rs",
3265                include_str!("../../runtime/dml_target_scan.rs"),
3266                2,
3267                "documented exception: kind-conditional visibility the uniform scan \
3268                 API cannot express; both iterations gate on candidate_visible",
3269            ),
3270            (
3271                "runtime/impl_events.rs",
3272                include_str!("../../runtime/impl_events.rs"),
3273                1,
3274                "documented exception: backfill dedup must include consumed messages, \
3275                 so it is deliberately visibility-free",
3276            ),
3277            (
3278                "runtime/impl_graph_commands.rs",
3279                include_str!("../../runtime/impl_graph_commands.rs"),
3280                0,
3281                "native node properties resolve one snapshot-visible, RLS-admitted version",
3282            ),
3283            (
3284                "runtime/graph_dsl.rs",
3285                include_str!("../../runtime/graph_dsl.rs"),
3286                1,
3287                "remaining legacy debt: unused edge-property helper; search expansion \
3288                 now uses snapshot visitors and RLS-gated hydration",
3289            ),
3290            (
3291                "runtime/graph_tvf.rs",
3292                include_str!("../../runtime/graph_tvf.rs"),
3293                0,
3294                "fully migrated by batch 2",
3295            ),
3296            (
3297                "runtime/impl_kv.rs",
3298                include_str!("../../runtime/impl_kv.rs"),
3299                4,
3300                "two non-versioned-collection fallbacks (the versioned arms scan under \
3301                 the captured snapshot). Vault reads use snapshot scans. Offline vault \
3302                 export preserves all \
3303                 physical versions across tenants; restore's emptiness check must reject \
3304                 every physical row, including versions invisible to a snapshot",
3305            ),
3306            (
3307                "runtime/impl_queue.rs",
3308                include_str!("../../runtime/impl_queue.rs"),
3309                10,
3310                "maintenance: red_queue_meta bookkeeping (config, groups, acks, \
3311                 pending, dedup TTL reaper, gauge counts) — queue message bodies are \
3312                 read through scan. \
3313                 TODO(#2138 follow-up): QUEUE PENDING projects meta rows to clients",
3314            ),
3315            (
3316                "runtime/primary_queue_store.rs",
3317                include_str!("../../runtime/primary_queue_store.rs"),
3318                3,
3319                "maintenance: generic red_queue_meta read/delete helpers and the \
3320                 next-position allocator, which must not collide with rows invisible \
3321                 to the writer's snapshot",
3322            ),
3323            (
3324                "runtime/record_search.rs",
3325                include_str!("../../runtime/record_search.rs"),
3326                0,
3327                "fully migrated by batch 2",
3328            ),
3329            (
3330                "runtime/search_graph.rs",
3331                include_str!("../../runtime/search_graph.rs"),
3332                0,
3333                "search graph preparation uses snapshot visitors and gated hydration",
3334            ),
3335            (
3336                "runtime/impl_search.rs",
3337                include_str!("../../runtime/impl_search.rs"),
3338                1,
3339                "maintenance: ask-audit retention sweep must see every physical row",
3340            ),
3341        ];
3342
3343        for (path, src, allowed, reason) in allowlist {
3344            let found = raw_iteration_sites(&production_source(src));
3345            assert_eq!(
3346                found, *allowed,
3347                "{path}: {found} raw-iteration sites, allowlist permits {allowed} \
3348                 ({reason}). Raw iteration bypasses MVCC visibility — use \
3349                 SegmentManager::scan / scan_for_each / scan_fold_parallel / \
3350                 scan_zoned_with_stats, or extend the allowlist with a maintenance \
3351                 justification."
3352            );
3353        }
3354    }
3355
3356    #[test]
3357    fn explicit_snapshot_scan_ignores_ambient_thread_local() {
3358        let manager = Arc::new(SegmentManager::new("accounts"));
3359        // (xmin, xmax) — mix of committed, in-progress and future writers.
3360        let versions = [(0, 0), (3, 0), (11, 0), (4, 0), (3, 7), (3, 11), (3, 4)];
3361        for (index, (xmin, xmax)) in versions.iter().enumerate() {
3362            let row_id = index as u64 + 1;
3363            let mut entity = UnifiedEntity::table_row(
3364                EntityId::new(row_id),
3365                "accounts",
3366                row_id,
3367                vec![Value::Integer(row_id as i64)],
3368            );
3369            entity.set_xmin(*xmin);
3370            entity.set_xmax(*xmax);
3371            manager.insert(entity).unwrap();
3372        }
3373
3374        let explicit = SnapshotContext {
3375            snapshot: Snapshot {
3376                xid: 10,
3377                in_progress: HashSet::from([4]),
3378            },
3379            manager: Arc::new(SnapshotManager::new()),
3380            own_xids: HashSet::new(),
3381            requires_index_fallback: false,
3382            serializable_reader: None,
3383        };
3384        // A decoy statement frame that would hide almost everything if the
3385        // scan ever consulted ambient state instead of its argument.
3386        let decoy = SnapshotContext {
3387            snapshot: Snapshot {
3388                xid: 1,
3389                in_progress: HashSet::new(),
3390            },
3391            manager: Arc::new(SnapshotManager::new()),
3392            own_xids: HashSet::new(),
3393            requires_index_fallback: false,
3394            serializable_reader: None,
3395        };
3396
3397        let row_ids = |entities: Vec<UnifiedEntity>| -> HashSet<u64> {
3398            entities
3399                .into_iter()
3400                .filter_map(|entity| match entity.kind {
3401                    EntityKind::TableRow { row_id, .. } => Some(row_id),
3402                    _ => None,
3403                })
3404                .collect()
3405        };
3406
3407        let guard = crate::runtime::execution_context::CurrentSnapshotGuard::install(decoy);
3408        let on_installing_thread = row_ids(manager.scan(Some(&explicit), |_| true));
3409        drop(guard);
3410
3411        let spawned_manager = Arc::clone(&manager);
3412        let spawned_snapshot = explicit.clone();
3413        let on_fresh_thread = std::thread::spawn(move || {
3414            row_ids(spawned_manager.scan(Some(&spawned_snapshot), |_| true))
3415        })
3416        .join()
3417        .unwrap();
3418
3419        assert_eq!(
3420            on_installing_thread, on_fresh_thread,
3421            "an explicit snapshot must produce the same visible set on any thread"
3422        );
3423        assert_eq!(on_installing_thread, HashSet::from([1, 2, 6, 7]));
3424
3425        // The raw iterators are the maintenance contract: no visibility
3426        // filtering at all, on any thread, regardless of installed frames.
3427        let raw_manager = Arc::clone(&manager);
3428        let raw_count = std::thread::spawn(move || raw_manager.query_all(|_| true).len())
3429            .join()
3430            .unwrap();
3431        assert_eq!(
3432            raw_count,
3433            versions.len(),
3434            "raw iteration must expose every physical version"
3435        );
3436    }
3437
3438    #[test]
3439    fn bloom_may_contain_key_probes_inserted_rid_bytes() {
3440        let manager = SegmentManager::new("test_collection");
3441
3442        let entity = UnifiedEntity::table_row(
3443            manager.next_entity_id(),
3444            "users",
3445            1,
3446            vec![Value::text("Alice".to_string())],
3447        );
3448
3449        let id = manager.insert(entity).unwrap();
3450
3451        assert!(manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3452        assert!(manager.bloom_may_contain_key(id.raw().to_string().as_bytes()));
3453        assert!(!manager.bloom_may_contain_key(&u64::MAX.to_le_bytes()));
3454
3455        assert!(manager.delete(id).unwrap());
3456        assert!(!manager.bloom_may_contain_key(&id.raw().to_le_bytes()));
3457    }
3458
3459    #[test]
3460    fn zoned_scan_reports_pruned_segments_and_preserves_results() {
3461        let config = ManagerConfig {
3462            segment_config: SegmentConfig {
3463                max_entities: 2,
3464                ..Default::default()
3465            },
3466            ..Default::default()
3467        };
3468        let manager = SegmentManager::with_config("events", config);
3469
3470        for (idx, value) in [10, 20, 100, 110, 200].into_iter().enumerate() {
3471            let row = RowData::with_names(vec![Value::Integer(value)], vec!["ts".to_string()]);
3472            let entity = UnifiedEntity::new(
3473                manager.next_entity_id(),
3474                EntityKind::TableRow {
3475                    table: "events".into(),
3476                    row_id: (idx + 1) as u64,
3477                },
3478                EntityData::Row(row),
3479            );
3480            manager.insert(entity).unwrap();
3481        }
3482
3483        let full = manager.query_all(|entity| match &entity.data {
3484            EntityData::Row(row) => matches!(row.get_field("ts"), Some(Value::Integer(100))),
3485            _ => false,
3486        });
3487        let probe = Value::Integer(100);
3488        let (pruned, stats) =
3489            manager.query_all_zoned_with_stats(&[("ts", ZoneColPred::Eq(&probe))], |entity| {
3490                match &entity.data {
3491                    EntityData::Row(row) => {
3492                        matches!(row.get_field("ts"), Some(Value::Integer(100)))
3493                    }
3494                    _ => false,
3495                }
3496            });
3497
3498        assert_eq!(pruned.len(), full.len());
3499        assert_eq!(pruned[0].id, full[0].id);
3500        assert_eq!(stats.segments_total, 3);
3501        assert_eq!(stats.segments_scanned, 2);
3502        assert_eq!(stats.segments_pruned, 1);
3503    }
3504
3505    #[test]
3506    fn test_manager_auto_seal() {
3507        let config = ManagerConfig {
3508            segment_config: SegmentConfig {
3509                max_entities: 2,
3510                ..Default::default()
3511            },
3512            ..Default::default()
3513        };
3514
3515        let manager = SegmentManager::with_config("test", config);
3516
3517        // Insert first entity
3518        manager
3519            .insert(UnifiedEntity::vector(
3520                manager.next_entity_id(),
3521                "v",
3522                vec![0.1],
3523            ))
3524            .unwrap();
3525
3526        // Insert second entity (triggers seal check)
3527        manager
3528            .insert(UnifiedEntity::vector(
3529                manager.next_entity_id(),
3530                "v",
3531                vec![0.2],
3532            ))
3533            .unwrap();
3534
3535        // Insert third entity (should trigger auto-seal of first segment)
3536        manager
3537            .insert(UnifiedEntity::vector(
3538                manager.next_entity_id(),
3539                "v",
3540                vec![0.3],
3541            ))
3542            .unwrap();
3543
3544        let stats = manager.stats();
3545        assert_eq!(stats.total_entities, 3);
3546    }
3547
3548    #[test]
3549    fn test_manager_delete() {
3550        let manager = SegmentManager::new("test");
3551
3552        let id = manager
3553            .insert(UnifiedEntity::vector(
3554                manager.next_entity_id(),
3555                "v",
3556                vec![0.1],
3557            ))
3558            .unwrap();
3559
3560        assert!(manager.get(id).is_some());
3561        assert!(manager.delete(id).unwrap());
3562        assert!(manager.get(id).is_none());
3563    }
3564
3565    #[test]
3566    fn test_manager_metadata() {
3567        let manager = SegmentManager::new("test");
3568
3569        let id = manager
3570            .insert(UnifiedEntity::table_row(
3571                manager.next_entity_id(),
3572                "hosts",
3573                1,
3574                vec![Value::text("192.168.1.1".to_string())],
3575            ))
3576            .unwrap();
3577
3578        let mut meta = Metadata::new();
3579        meta.set(
3580            "os",
3581            super::super::metadata::MetadataValue::String("linux".to_string()),
3582        );
3583
3584        manager.set_metadata(id, meta).unwrap();
3585
3586        let retrieved = manager.get_metadata(id).unwrap();
3587        assert!(retrieved.has("os"));
3588    }
3589
3590    #[test]
3591    fn test_manager_query_by_kind() {
3592        let manager = SegmentManager::new("test");
3593
3594        manager
3595            .insert(UnifiedEntity::table_row(
3596                manager.next_entity_id(),
3597                "hosts",
3598                1,
3599                vec![],
3600            ))
3601            .unwrap();
3602
3603        manager
3604            .insert(UnifiedEntity::vector(
3605                manager.next_entity_id(),
3606                "embeddings",
3607                vec![0.1],
3608            ))
3609            .unwrap();
3610
3611        manager
3612            .insert(UnifiedEntity::table_row(
3613                manager.next_entity_id(),
3614                "hosts",
3615                2,
3616                vec![],
3617            ))
3618            .unwrap();
3619
3620        let rows = manager.get_by_kind("table");
3621        assert_eq!(rows.len(), 2);
3622
3623        let vectors = manager.get_by_kind("vector");
3624        assert_eq!(vectors.len(), 1);
3625    }
3626
3627    // ========================================================================
3628    // Consolidation — ADR 0073 §5
3629    // ========================================================================
3630
3631    /// A manager whose sealed set never trips the count trigger, so tests can
3632    /// exercise the tombstone/fragmentation triggers in isolation.
3633    fn consolidating_manager(entities_per_tick: usize) -> SegmentManager {
3634        SegmentManager::with_config(
3635            "test",
3636            ManagerConfig {
3637                max_sealed_segments: 1_000,
3638                consolidation_entities_per_tick: entities_per_tick,
3639                ..Default::default()
3640            },
3641        )
3642    }
3643
3644    /// Bulk-load `count` rows into a fresh segment and seal it. Bulk-inserted
3645    /// entities land in flat storage, where a tombstone keeps its payload
3646    /// resident — which is the memory consolidation reclaims.
3647    fn seal_bulk_segment(manager: &SegmentManager, count: usize) -> Vec<EntityId> {
3648        let entities: Vec<UnifiedEntity> = (0..count)
3649            .map(|i| {
3650                UnifiedEntity::table_row(
3651                    manager.next_entity_id(),
3652                    "users",
3653                    i as u64,
3654                    vec![Value::Integer(i as i64), Value::text(format!("row-{i}"))],
3655                )
3656            })
3657            .collect();
3658        let ids = manager.bulk_insert(entities).expect("bulk insert");
3659        manager.force_seal().expect("seal");
3660        ids
3661    }
3662
3663    fn live_ids(manager: &SegmentManager) -> Vec<u64> {
3664        let mut ids: Vec<u64> = manager
3665            .query_all(|_| true)
3666            .into_iter()
3667            .map(|entity| entity.id.raw())
3668            .collect();
3669        ids.sort_unstable();
3670        ids
3671    }
3672
3673    fn drain_maintenance(manager: &SegmentManager) -> usize {
3674        let mut ticks = 0;
3675        while manager.consolidation.read().is_some() || ticks == 0 {
3676            manager.run_maintenance().expect("maintenance");
3677            ticks += 1;
3678            assert!(ticks < 1_000, "consolidation failed to converge");
3679        }
3680        ticks
3681    }
3682
3683    #[test]
3684    fn graph_index_batches_pin_sources_and_hydrate_current_after_consolidation() {
3685        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3686        let ids = manager
3687            .bulk_insert(
3688                (1..=1024)
3689                    .map(|id| {
3690                        UnifiedEntity::graph_edge(
3691                            EntityId::new(id),
3692                            "step",
3693                            "2000",
3694                            "3000",
3695                            1.0,
3696                            HashMap::new(),
3697                        )
3698                    })
3699                    .collect(),
3700            )
3701            .expect("edges");
3702        manager.force_seal().expect("seal");
3703        for id in ids.iter().take(300) {
3704            manager.delete(*id).expect("fragment source");
3705        }
3706        let source = Arc::downgrade(&manager.sealed.read()[0]);
3707        let mut seen = Vec::new();
3708        let mut changed = false;
3709        let mut metadata_bytes = 0;
3710        assert!(manager.visit_graph_index_batches(GraphEntityKind::Edge, &[EntityId::new(2000)],
3711            || true,
3712            |bytes| { metadata_bytes += bytes; Ok::<_, ()>(()) },
3713            |batch| {
3714                assert!(batch.len() <= super::super::segment::SCAN_BATCH_SIZE);
3715                assert!(manager.growing.try_write().is_some(), "no growing topology guard");
3716                assert!(manager.sealed.try_write().is_some(), "no sealed topology guard");
3717                assert!(source.upgrade().expect("source pinned").try_write().is_some(), "no source guard");
3718                if !changed {
3719                    changed = true;
3720                    drain_maintenance(&manager);
3721                    assert!(manager.stats().consolidation.runs_completed > 0);
3722                    assert!(source.upgrade().is_some(), "retired cursor source remains pinned");
3723                    let mut updated = manager.get(EntityId::new(1024)).expect("future edge");
3724                    if let EntityKind::GraphEdge(edge) = &mut updated.kind { edge.weight = 2000; }
3725                    manager.update(updated).expect("update authoritative merged copy");
3726                }
3727                for id in batch {
3728                    let entity = manager.get(*id).expect("hydrate current manager");
3729                    if id.raw() == 1024 {
3730                        assert!(matches!(&entity.kind, EntityKind::GraphEdge(edge) if edge.weight == 2000));
3731                    }
3732                    seen.push(id.raw());
3733                }
3734                Ok(true)
3735            }).expect("cursor"));
3736        assert_eq!(seen, (301..=1024).collect::<Vec<_>>());
3737        assert!(
3738            source.upgrade().is_none(),
3739            "finished cursor releases retired source"
3740        );
3741        assert!(
3742            metadata_bytes < 4096,
3743            "fixed ID buffer plus source handles: {metadata_bytes}"
3744        );
3745    }
3746
3747    #[test]
3748    fn graph_index_batches_propagate_admission_and_cancel_before_partial_delivery() {
3749        let manager = SegmentManager::new("graph");
3750        manager
3751            .bulk_insert(
3752                (1..=1024)
3753                    .map(|id| {
3754                        UnifiedEntity::graph_edge(
3755                            EntityId::new(id),
3756                            "step",
3757                            "2000",
3758                            "3000",
3759                            1.0,
3760                            HashMap::new(),
3761                        )
3762                    })
3763                    .collect(),
3764            )
3765            .expect("edges");
3766        let mut delivered = 0;
3767        let denied = manager.visit_graph_index_batches(
3768            GraphEntityKind::Edge,
3769            &[EntityId::new(2000)],
3770            || true,
3771            |_| Err("out of memory"),
3772            |ids| {
3773                delivered += ids.len();
3774                Ok(true)
3775            },
3776        );
3777        assert_eq!(denied, Err("out of memory"));
3778        assert_eq!(delivered, 0);
3779        let mut work = 0;
3780        assert!(!manager
3781            .visit_graph_index_batches(
3782                GraphEntityKind::Edge,
3783                &[EntityId::new(2000)],
3784                || {
3785                    work += 1;
3786                    work < 100
3787                },
3788                |_| Ok::<_, ()>(()),
3789                |ids| {
3790                    delivered += ids.len();
3791                    Ok(true)
3792                }
3793            )
3794            .expect("cancelled cursor"));
3795        assert_eq!(
3796            delivered, 0,
3797            "cancelled partial batch cannot reach consumer"
3798        );
3799        assert!(!manager
3800            .visit_graph_index_batches(
3801                GraphEntityKind::Edge,
3802                &[EntityId::new(2000)],
3803                || true,
3804                |_| Ok::<_, ()>(()),
3805                |ids| {
3806                    delivered += ids.len();
3807                    Ok(false)
3808                }
3809            )
3810            .expect("consumer stop"));
3811        assert_eq!(delivered, super::super::segment::SCAN_BATCH_SIZE);
3812    }
3813
3814    #[test]
3815    fn retained_memory_tracks_two_readers_until_the_last_cursor_finishes() {
3816        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3817        let ids = seal_bulk_segment(&manager, 100);
3818        for id in ids.iter().take(40) {
3819            manager.delete(*id).expect("delete before scan");
3820        }
3821        let source_bytes = manager.resident_bytes();
3822        let source = Arc::downgrade(&manager.sealed.read()[0]);
3823        std::thread::scope(|scope| {
3824            let (ready, waiting) = std::sync::mpsc::channel();
3825            let (release_first, first_wait) = std::sync::mpsc::channel();
3826            let (release_last, last_wait) = std::sync::mpsc::channel();
3827            let reader = |resume: std::sync::mpsc::Receiver<()>,
3828                          ready: std::sync::mpsc::Sender<()>| {
3829                assert!(manager.scan_batches(
3830                    None,
3831                    || true,
3832                    |_| {
3833                        ready.send(()).expect("signal captured cursor");
3834                        let _ = resume.recv();
3835                        true
3836                    }
3837                ));
3838            };
3839            let ready_first = ready.clone();
3840            let first = scope.spawn(move || reader(first_wait, ready_first));
3841            let last = scope.spawn(move || reader(last_wait, ready));
3842            for _ in 0..2 {
3843                waiting
3844                    .recv_timeout(std::time::Duration::from_secs(10))
3845                    .expect("reader ready");
3846            }
3847            drain_maintenance(&manager);
3848            assert!(
3849                source.upgrade().is_some(),
3850                "both cursors pin the retired source"
3851            );
3852            let active_bytes = manager.sealed.read()[0].read().resident_bytes();
3853            let held = manager.resident_bytes();
3854            assert!(
3855                held >= active_bytes + source_bytes,
3856                "retired source remains resident"
3857            );
3858            release_first.send(()).expect("release first reader");
3859            first.join().expect("first reader");
3860            assert_eq!(
3861                manager.resident_bytes(),
3862                held,
3863                "one reader still pins the source"
3864            );
3865            release_last.send(()).expect("release last reader");
3866            last.join().expect("last reader");
3867            assert!(
3868                source.upgrade().is_none(),
3869                "accounting must not retain payloads"
3870            );
3871            assert_eq!(held - manager.resident_bytes(), source_bytes);
3872            assert_eq!(
3873                manager.stats().total_memory_bytes as u64,
3874                manager.resident_bytes()
3875            );
3876        });
3877    }
3878
3879    #[test]
3880    fn retained_memory_gates_runtime_growth_until_the_cursor_is_cancelled() {
3881        use crate::storage::memory_pools::MemoryPool;
3882        use crate::{RedDBOptions, RedDBRuntime};
3883
3884        let budget = 128 * 1024;
3885        let runtime =
3886            RedDBRuntime::with_options(RedDBOptions::in_memory().with_memory_budget(budget))
3887                .expect("runtime");
3888        runtime
3889            .db()
3890            .store()
3891            .create_collection("test")
3892            .expect("collection");
3893        let manager = runtime
3894            .db()
3895            .store()
3896            .get_collection("test")
3897            .expect("manager");
3898        let ids = seal_bulk_segment(&manager, 100);
3899        for id in ids.iter().take(40) {
3900            manager.delete(*id).expect("delete before scan");
3901        }
3902        let source_bytes = manager.resident_bytes();
3903        let mut growth = 0;
3904        let mut held_used = 0;
3905        assert!(!manager.scan_batches(
3906            None,
3907            || true,
3908            |_| {
3909                drain_maintenance(&manager);
3910                runtime.refresh_memory_accounting();
3911                // Other collections (including the runtime catalog) share this
3912                // arena, so derive headroom from the complete runtime sample.
3913                held_used = runtime.memory_accounting().total_used_bytes();
3914                growth = budget
3915                    .checked_sub(held_used)
3916                    .expect("fixture fits while the cursor is open")
3917                    + source_bytes / 2;
3918                assert!(growth > 0);
3919                assert!(
3920                    runtime
3921                        .admit_non_evictable_growth(
3922                            MemoryPool::SegmentArena,
3923                            "growth with pinned source",
3924                            growth,
3925                        )
3926                        .is_err(),
3927                    "retained source must participate in runtime admission"
3928                );
3929                false
3930            }
3931        ));
3932        runtime.refresh_memory_accounting();
3933        assert_eq!(
3934            held_used - runtime.memory_accounting().total_used_bytes(),
3935            source_bytes
3936        );
3937        let _reservation = runtime
3938            .admit_non_evictable_growth(
3939                MemoryPool::SegmentArena,
3940                "growth after cursor cancellation",
3941                growth,
3942            )
3943            .expect("the same growth fits once the cursor releases its source");
3944    }
3945
3946    #[test]
3947    fn retained_memory_includes_an_inflight_consolidation_copy() {
3948        let manager = consolidating_manager(1);
3949        let ids = seal_bulk_segment(&manager, 100);
3950        for id in ids.iter().take(40) {
3951            manager.delete(*id).expect("delete");
3952        }
3953        let before = manager.resident_bytes();
3954        assert!(manager.pressure_consolidation_tick());
3955        let copied = manager
3956            .consolidation
3957            .read()
3958            .as_ref()
3959            .expect("paced merge")
3960            .merged
3961            .resident_bytes();
3962        assert!(copied > 0);
3963        assert!(
3964            manager.resident_bytes() >= before + copied,
3965            "unpublished copy consumes memory"
3966        );
3967        drain_maintenance(&manager);
3968        assert!(
3969            manager.resident_bytes() < before,
3970            "completed merge releases the old storage"
3971        );
3972    }
3973
3974    #[test]
3975    fn delete_heavy_workload_reclaims_tombstoned_memory_and_counts_it() {
3976        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
3977        let ids = seal_bulk_segment(&manager, 100);
3978
3979        let bytes_before_delete = manager.resident_bytes();
3980
3981        // 40% tombstones — past CONSOLIDATION_TOMBSTONE_RATIO.
3982        for id in ids.iter().take(40) {
3983            assert!(manager.delete(*id).expect("delete"));
3984        }
3985
3986        // Tombstones cost memory: the payload stays resident in flat storage
3987        // and the tombstone set grows.
3988        let bytes_with_tombstones = manager.resident_bytes();
3989        assert!(bytes_with_tombstones > bytes_before_delete);
3990
3991        let survivors = live_ids(&manager);
3992        assert_eq!(survivors.len(), 60);
3993
3994        drain_maintenance(&manager);
3995
3996        let stats = manager.stats();
3997        assert_eq!(stats.consolidation.runs_started, 1);
3998        assert_eq!(stats.consolidation.runs_completed, 1);
3999        assert_eq!(stats.consolidation.segments_merged, 1);
4000        assert_eq!(stats.consolidation.tombstones_reclaimed, 40);
4001        assert!(stats.consolidation.bytes_reclaimed > 0);
4002
4003        // The reclaimed memory is at least the tombstoned entities' share.
4004        let bytes_after = manager.resident_bytes();
4005        assert!(
4006            bytes_after <= bytes_before_delete,
4007            "consolidation must return the tombstones' memory: {bytes_after} > {bytes_before_delete}"
4008        );
4009        assert_eq!(
4010            bytes_with_tombstones - bytes_after,
4011            stats.consolidation.bytes_reclaimed
4012        );
4013        assert_eq!(stats.total_memory_bytes, bytes_after as usize);
4014
4015        assert_eq!(live_ids(&manager), survivors);
4016        assert_eq!(manager.sealed.read().len(), 1);
4017        assert_eq!(manager.sealed.read()[0].read().tombstone_count(), 0);
4018    }
4019
4020    #[test]
4021    fn query_results_are_identical_before_during_and_after_consolidation() {
4022        let manager = consolidating_manager(3);
4023        let ids = seal_bulk_segment(&manager, 20);
4024        seal_bulk_segment(&manager, 20);
4025        for id in ids.iter().take(10) {
4026            assert!(manager.delete(*id).expect("delete"));
4027        }
4028
4029        let before = live_ids(&manager);
4030        let survivor = *ids.last().expect("at least one row");
4031        let tombstoned = ids[0];
4032
4033        // Every mid-tick observation matches the pre-consolidation snapshot:
4034        // point reads, full scans, and the kind index (rebuilt by the merge).
4035        let mut mid_tick_observations = 0;
4036        while manager.consolidation.read().is_some() || mid_tick_observations == 0 {
4037            manager.run_maintenance().expect("maintenance");
4038            assert_eq!(live_ids(&manager), before);
4039            assert!(manager.get(survivor).is_some());
4040            assert!(manager.get(tombstoned).is_none());
4041            assert_eq!(manager.get_by_kind("table").len(), before.len());
4042            mid_tick_observations += 1;
4043            assert!(mid_tick_observations < 1_000);
4044        }
4045
4046        assert!(
4047            mid_tick_observations > 2,
4048            "expected the run to span several ticks, saw {mid_tick_observations}"
4049        );
4050        assert_eq!(live_ids(&manager), before);
4051        assert!(manager.get(survivor).is_some());
4052        assert!(manager.get(tombstoned).is_none());
4053        assert_eq!(manager.get_by_kind("table").len(), before.len());
4054    }
4055
4056    #[test]
4057    fn get_by_kind_reaches_sealed_segments() {
4058        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4059        let ids = seal_bulk_segment(&manager, 5);
4060        assert!(manager.delete(ids[0]).expect("delete"));
4061
4062        // The sealed list is walked through the segment read guard; the
4063        // `Arc<RwLock<GrowingSegment>>` impl of `iter_kind` cannot hand out
4064        // references across the lock and yields nothing.
4065        assert_eq!(manager.get_by_kind("table").len(), 4);
4066        assert_eq!(manager.get_by_kind("vector").len(), 0);
4067    }
4068
4069    #[test]
4070    fn consolidation_is_paced_by_the_per_tick_entity_bound() {
4071        const PER_TICK: usize = 4;
4072        let manager = consolidating_manager(PER_TICK);
4073        let ids = seal_bulk_segment(&manager, 60);
4074        for id in ids.iter().take(30) {
4075            assert!(manager.delete(*id).expect("delete"));
4076        }
4077
4078        // First tick opens the run and copies exactly one bound's worth.
4079        manager.run_maintenance().expect("maintenance");
4080        assert_eq!(manager.consolidation_progress(), Some(PER_TICK as u64));
4081
4082        // Each subsequent tick advances by no more than the bound, and the run
4083        // is still in flight — it never completes in one unbounded pass.
4084        let mut previous = PER_TICK as u64;
4085        let mut ticks = 1;
4086        while let Some(copied) = manager.consolidation_progress() {
4087            assert!(
4088                copied - previous <= PER_TICK as u64,
4089                "tick copied {} entities, bound is {PER_TICK}",
4090                copied - previous
4091            );
4092            previous = copied;
4093            manager.run_maintenance().expect("maintenance");
4094            ticks += 1;
4095            assert!(ticks < 100);
4096        }
4097
4098        // 30 survivors at 4 per tick: the run cannot have finished before the
4099        // eighth tick.
4100        assert!(ticks >= 8, "consolidation completed in only {ticks} ticks");
4101        assert_eq!(manager.stats().consolidation.runs_completed, 1);
4102        assert_eq!(live_ids(&manager).len(), 30);
4103    }
4104
4105    #[test]
4106    fn sealed_segment_count_above_the_limit_triggers_a_merge() {
4107        let manager = SegmentManager::with_config(
4108            "test",
4109            ManagerConfig {
4110                max_sealed_segments: 3,
4111                ..Default::default()
4112            },
4113        );
4114        for _ in 0..4 {
4115            seal_bulk_segment(&manager, 5);
4116        }
4117        assert_eq!(manager.sealed.read().len(), 4);
4118
4119        let before = live_ids(&manager);
4120        drain_maintenance(&manager);
4121
4122        // The knob finally does what it claimed: the sealed set collapses even
4123        // though not a single row was deleted.
4124        assert_eq!(manager.sealed.read().len(), 1);
4125        assert_eq!(manager.stats().consolidation.segments_merged, 4);
4126        assert_eq!(manager.stats().consolidation.tombstones_reclaimed, 0);
4127        assert_eq!(live_ids(&manager), before);
4128    }
4129
4130    #[test]
4131    fn a_clean_sealed_set_below_the_limit_is_never_consolidated() {
4132        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4133        seal_bulk_segment(&manager, 10);
4134        seal_bulk_segment(&manager, 10);
4135
4136        for _ in 0..5 {
4137            manager.run_maintenance().expect("maintenance");
4138        }
4139
4140        assert_eq!(manager.stats().consolidation.runs_started, 0);
4141        assert_eq!(manager.sealed.read().len(), 2);
4142    }
4143
4144    #[test]
4145    fn concurrent_readers_never_observe_a_missing_or_duplicated_entity() {
4146        use std::sync::atomic::AtomicBool;
4147
4148        let manager = Arc::new(consolidating_manager(2));
4149        let ids = seal_bulk_segment(&manager, 40);
4150        seal_bulk_segment(&manager, 40);
4151        for id in ids.iter().take(20) {
4152            assert!(manager.delete(*id).expect("delete"));
4153        }
4154
4155        let expected = live_ids(&manager);
4156        let stop = Arc::new(AtomicBool::new(false));
4157
4158        let readers: Vec<_> = (0..4)
4159            .map(|_| {
4160                let manager = Arc::clone(&manager);
4161                let stop = Arc::clone(&stop);
4162                let expected = expected.clone();
4163                std::thread::spawn(move || {
4164                    let mut scans = 0u64;
4165                    // Do-while: every reader scans at least once even when the
4166                    // drain finishes before this thread is first scheduled.
4167                    loop {
4168                        let mut seen: Vec<u64> = manager
4169                            .query_all(|_| true)
4170                            .into_iter()
4171                            .map(|entity| entity.id.raw())
4172                            .collect();
4173                        seen.sort_unstable();
4174                        // Sorting is not enough: a duplicated row would keep
4175                        // the set equal but change the length.
4176                        assert_eq!(seen, expected, "reader saw a torn sealed set");
4177                        scans += 1;
4178                        if stop.load(Ordering::Relaxed) {
4179                            break;
4180                        }
4181                    }
4182                    scans
4183                })
4184            })
4185            .collect();
4186
4187        drain_maintenance(&manager);
4188        stop.store(true, Ordering::Relaxed);
4189
4190        for reader in readers {
4191            let scans = reader.join().expect("reader thread");
4192            assert!(scans > 0, "reader never got to scan");
4193        }
4194
4195        assert_eq!(manager.stats().consolidation.runs_completed, 1);
4196        assert_eq!(live_ids(&manager), expected);
4197    }
4198
4199    #[test]
4200    fn a_row_deleted_mid_consolidation_is_not_resurrected_by_the_swap() {
4201        let manager = consolidating_manager(4);
4202        let ids = seal_bulk_segment(&manager, 20);
4203        for id in ids.iter().take(6) {
4204            assert!(manager.delete(*id).expect("delete"));
4205        }
4206
4207        // Copy the first bound's worth, then delete a row we just copied.
4208        manager.run_maintenance().expect("maintenance");
4209        assert_eq!(manager.consolidation_progress(), Some(4));
4210        let doomed = ids[6];
4211        assert!(manager.delete(doomed).expect("delete"));
4212
4213        let expected = live_ids(&manager);
4214        drain_maintenance(&manager);
4215
4216        assert!(manager.get(doomed).is_none(), "deleted row came back");
4217        assert_eq!(live_ids(&manager), expected);
4218    }
4219
4220    #[test]
4221    fn discarding_an_in_flight_consolidation_leaves_the_data_untouched() {
4222        let manager = consolidating_manager(3);
4223        let ids = seal_bulk_segment(&manager, 30);
4224        for id in ids.iter().take(12) {
4225            assert!(manager.delete(*id).expect("delete"));
4226        }
4227
4228        let before = live_ids(&manager);
4229        let sealed_before: Vec<SegmentId> = manager
4230            .sealed
4231            .read()
4232            .iter()
4233            .map(|segment| segment.read().id())
4234            .collect();
4235
4236        manager.run_maintenance().expect("maintenance");
4237        manager.run_maintenance().expect("maintenance");
4238        assert!(manager.consolidation.read().is_some());
4239
4240        // A crash between ticks drops the half-built merged segment. Sealed
4241        // segments are the durable truth, so nothing was lost — and the run is
4242        // simply started over.
4243        *manager.consolidation.write() = None;
4244
4245        assert_eq!(live_ids(&manager), before);
4246        let sealed_after: Vec<SegmentId> = manager
4247            .sealed
4248            .read()
4249            .iter()
4250            .map(|segment| segment.read().id())
4251            .collect();
4252        assert_eq!(sealed_after, sealed_before);
4253        assert_eq!(manager.stats().consolidation.runs_completed, 0);
4254
4255        drain_maintenance(&manager);
4256        assert_eq!(live_ids(&manager), before);
4257        assert_eq!(manager.stats().consolidation.runs_completed, 1);
4258    }
4259
4260    #[test]
4261    fn consolidation_preserves_entity_payload_metadata_and_sequence() {
4262        let manager = consolidating_manager(CONSOLIDATION_ENTITIES_PER_TICK);
4263
4264        // Metadata must be attached before the seal: `set_metadata` on a sealed
4265        // row rewrites it into the growing segment, which would take it out of
4266        // the segment under test.
4267        let entities: Vec<UnifiedEntity> = (0..10)
4268            .map(|i| {
4269                UnifiedEntity::table_row(
4270                    manager.next_entity_id(),
4271                    "users",
4272                    i as u64,
4273                    vec![Value::Integer(i as i64)],
4274                )
4275            })
4276            .collect();
4277        let ids = manager.bulk_insert(entities).expect("bulk insert");
4278
4279        let mut meta = Metadata::new();
4280        meta.set(
4281            "os",
4282            super::super::metadata::MetadataValue::String("linux".to_string()),
4283        );
4284        manager.set_metadata(ids[9], meta).expect("set metadata");
4285        manager.force_seal().expect("seal");
4286
4287        assert!(manager.delete(ids[0]).expect("delete"));
4288        assert!(manager.delete(ids[1]).expect("delete"));
4289        assert!(manager.delete(ids[2]).expect("delete"));
4290
4291        let before: Vec<UnifiedEntity> = {
4292            let mut entities = manager.query_all(|_| true);
4293            entities.sort_by_key(|entity| entity.id.raw());
4294            entities
4295        };
4296
4297        drain_maintenance(&manager);
4298
4299        let after: Vec<UnifiedEntity> = {
4300            let mut entities = manager.query_all(|_| true);
4301            entities.sort_by_key(|entity| entity.id.raw());
4302            entities
4303        };
4304
4305        let columns = |entity: &UnifiedEntity| match &entity.data {
4306            super::super::entity::EntityData::Row(row) => row.columns.clone(),
4307            _ => Vec::new(),
4308        };
4309
4310        assert_eq!(before.len(), after.len());
4311        for (old, new) in before.iter().zip(after.iter()) {
4312            assert_eq!(old.id, new.id);
4313            assert_eq!(columns(old), columns(new));
4314            assert_eq!(
4315                old.sequence_id, new.sequence_id,
4316                "merge must not re-issue sequence ids"
4317            );
4318        }
4319        assert!(manager
4320            .get_metadata(ids[9])
4321            .expect("metadata survives the merge")
4322            .has("os"));
4323    }
4324
4325    #[test]
4326    #[ignore = "lifecycle events intentionally no-op since the emit-channel refactor; drain_events returns empty — see SegmentManager::emit"]
4327    fn test_lifecycle_events() {
4328        let manager = SegmentManager::new("test");
4329
4330        manager
4331            .insert(UnifiedEntity::vector(
4332                manager.next_entity_id(),
4333                "v",
4334                vec![0.1],
4335            ))
4336            .unwrap();
4337
4338        let events = manager.drain_events();
4339
4340        // Should have: SegmentCreated, EntityInserted
4341        assert!(events
4342            .iter()
4343            .any(|e| matches!(e, LifecycleEvent::SegmentCreated(_))));
4344        assert!(events
4345            .iter()
4346            .any(|e| matches!(e, LifecycleEvent::EntityInserted(_, _))));
4347    }
4348}