Skip to main content

reddb_server/storage/unified/store/
impl_entities.rs

1use super::*;
2
3impl UnifiedStore {
4    /// Vault versions are WAL-backed in memory until a committed checkpoint.
5    /// Preparing a commit must not dirty pager pages: eviction or an auth
6    /// flush could otherwise persist a version before its WAL batch exists.
7    pub(crate) fn insert_vault_version(
8        &self,
9        collection: &str,
10        entity: UnifiedEntity,
11        metadata: Metadata,
12    ) -> Result<(), StoreError> {
13        assert!(matches!(entity.data, EntityData::Row(_)));
14        let manager = self
15            .get_collection(collection)
16            .ok_or_else(|| StoreError::CollectionNotFound(collection.into()))?;
17        self.pending_vault_versions.write().insert(entity.id);
18        let id = manager.insert(entity)?;
19        self.register_entity_id(id);
20        if !metadata.fields.is_empty() {
21            manager.set_metadata(id, metadata)?;
22        }
23        let entity = manager
24            .get(id)
25            .expect("inserted vault version remains in memory");
26        self.persist_entities_to_pager_wal_only(collection, &[entity])
27    }
28
29    pub(crate) fn persist_entities_to_pager(
30        &self,
31        collection: &str,
32        entities: &[UnifiedEntity],
33    ) -> Result<(), StoreError> {
34        self.persist_entities_impl(collection, entities, /* skip_btree= */ false)
35    }
36
37    /// Same as `persist_entities_to_pager` but takes `&[&UnifiedEntity]` so
38    /// callers that already hold references (SQL UPDATE fast path) don't
39    /// have to clone a `Vec<UnifiedEntity>` just to satisfy the signature.
40    pub(crate) fn persist_entity_refs_to_pager(
41        &self,
42        collection: &str,
43        entities: &[&UnifiedEntity],
44    ) -> Result<(), StoreError> {
45        self.persist_entity_refs_impl(collection, entities, /* skip_btree= */ false)
46    }
47
48    /// Reference-taking WAL-only variant — see
49    /// `persist_entities_to_pager_wal_only` for semantics.
50    pub(crate) fn persist_entity_refs_to_pager_wal_only(
51        &self,
52        collection: &str,
53        entities: &[&UnifiedEntity],
54    ) -> Result<(), StoreError> {
55        self.persist_entity_refs_impl(collection, entities, /* skip_btree= */ true)
56    }
57
58    fn persist_entity_refs_impl(
59        &self,
60        collection: &str,
61        entities: &[&UnifiedEntity],
62        skip_btree: bool,
63    ) -> Result<(), StoreError> {
64        if entities.is_empty() {
65            return Ok(());
66        }
67        if self.pager.is_none() && self.config.embedded_wal_path.is_none() {
68            return Ok(());
69        }
70        let fv = self.format_version();
71        let manager = self
72            .get_collection(collection)
73            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
74        let mut serialized: Vec<(Vec<u8>, Vec<u8>)> = entities
75            .iter()
76            .map(|entity| {
77                let metadata = manager.get_metadata(entity.id);
78                (
79                    entity.id.raw().to_be_bytes().to_vec(),
80                    Self::serialize_entity_record(entity, metadata.as_ref(), fv),
81                )
82            })
83            .collect();
84        serialized.sort_by(|a, b| a.0.cmp(&b.0));
85
86        if !skip_btree {
87            let Some(pager) = &self.pager else {
88                let records: Vec<Vec<u8>> =
89                    serialized.into_iter().map(|(_id, record)| record).collect();
90                self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
91                    collection: collection.to_string(),
92                    records,
93                }])?;
94                return Ok(());
95            };
96            let mut btree_indices = self.btree_indices.write();
97            let btree = btree_indices
98                .entry(collection.to_string())
99                .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
100            let root_before = btree.root_page_id();
101            btree.upsert_batch_sorted(&serialized).map_err(|e| {
102                StoreError::Io(std::io::Error::other(format!("B-tree upsert error: {}", e)))
103            })?;
104            let root_after = btree.root_page_id();
105            drop(btree_indices);
106            if root_before != root_after {
107                self.mark_paged_registry_dirty();
108            }
109        }
110        let records: Vec<Vec<u8>> = serialized.into_iter().map(|(_id, record)| record).collect();
111        self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
112            collection: collection.to_string(),
113            records,
114        }])?;
115        Ok(())
116    }
117
118    /// PG-HOT-like UPDATE persist: emit the WAL record so the update
119    /// survives a crash, but skip the in-line B-tree upsert. Reads
120    /// continue to see the new entity because `manager.get()` prefers
121    /// the segment over the B-tree; WAL replay re-applies the upsert
122    /// to the B-tree on recovery so the two views converge before
123    /// any stale B-tree page can be surfaced.
124    ///
125    /// Safe to use only when the caller has already updated the
126    /// segment in-place (so live reads see the new value). Cuts out
127    /// the dominant cost of single-row UPDATEs: per-entity serialize
128    /// + leaf descent + page decompress / compress / write_page.
129    pub(crate) fn persist_entities_to_pager_wal_only(
130        &self,
131        collection: &str,
132        entities: &[UnifiedEntity],
133    ) -> Result<(), StoreError> {
134        self.persist_entities_impl(collection, entities, /* skip_btree= */ true)
135    }
136
137    fn persist_entities_impl(
138        &self,
139        collection: &str,
140        entities: &[UnifiedEntity],
141        skip_btree: bool,
142    ) -> Result<(), StoreError> {
143        if entities.is_empty() {
144            return Ok(());
145        }
146
147        if self.pager.is_none() && self.config.embedded_wal_path.is_none() {
148            return Ok(());
149        }
150
151        let fv = self.format_version();
152        let manager = self
153            .get_collection(collection)
154            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
155        let mut serialized: Vec<(Vec<u8>, Vec<u8>)> = entities
156            .iter()
157            .map(|entity| {
158                let metadata = manager.get_metadata(entity.id);
159                (
160                    entity.id.raw().to_be_bytes().to_vec(),
161                    Self::serialize_entity_record(entity, metadata.as_ref(), fv),
162                )
163            })
164            .collect();
165        // u64 IDs encoded as big-endian — lex order = numeric order.
166        serialized.sort_by(|a, b| a.0.cmp(&b.0));
167
168        if !skip_btree {
169            let Some(pager) = &self.pager else {
170                let records: Vec<Vec<u8>> =
171                    serialized.into_iter().map(|(_id, record)| record).collect();
172                self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
173                    collection: collection.to_string(),
174                    records,
175                }])?;
176                return Ok(());
177            };
178            let mut btree_indices = self.btree_indices.write();
179            let btree = btree_indices
180                .entry(collection.to_string())
181                .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
182            let root_before = btree.root_page_id();
183
184            // Walks each distinct leaf once, applies all in-place overwrites
185            // that belong there under one read+write. Keys that miss or grow
186            // fall back to the per-key `upsert` path internally.
187            btree.upsert_batch_sorted(&serialized).map_err(|e| {
188                StoreError::Io(std::io::Error::other(format!("B-tree upsert error: {}", e)))
189            })?;
190            let root_after = btree.root_page_id();
191            drop(btree_indices);
192            if root_before != root_after {
193                self.mark_paged_registry_dirty();
194            }
195        }
196        // One WAL action covering the whole bulk. Previously each entity
197        // produced its own `UpsertEntityRecord` action, which then became
198        // its own `WalRecord::PageWrite` triple inside
199        // `StoreCommitCoordinator::append_actions` — N WAL records per
200        // bulk. The batched variant is one WAL record per bulk.
201        //
202        // WAL is emitted even in the `skip_btree` variant — recovery
203        // replay reapplies the upsert to the B-tree, so durability
204        // is unchanged.
205        let records: Vec<Vec<u8>> = serialized.into_iter().map(|(_id, record)| record).collect();
206        self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
207            collection: collection.to_string(),
208            records,
209        }])?;
210
211        Ok(())
212    }
213
214    /// Insert a label→entity_id mapping into the graph label index.
215    pub(crate) fn update_graph_label_index(
216        &self,
217        collection: &str,
218        label: &str,
219        entity_id: EntityId,
220    ) {
221        let key = (collection.to_string(), label.to_string());
222        let mut idx = self.graph_label_index.write();
223        idx.entry(key).or_default().push(entity_id);
224    }
225
226    /// Remove a specific entity_id from the graph label index (called on delete).
227    pub(crate) fn remove_from_graph_label_index(&self, collection: &str, entity_id: EntityId) {
228        let mut idx = self.graph_label_index.write();
229        for ((col, _), ids) in idx.iter_mut() {
230            if col == collection {
231                ids.retain(|&id| id != entity_id);
232            }
233        }
234        // Prune empty entries to keep the index compact
235        idx.retain(|_, ids| !ids.is_empty());
236    }
237
238    pub(crate) fn remove_from_graph_label_index_batch(
239        &self,
240        collection: &str,
241        entity_ids: &[EntityId],
242    ) {
243        if entity_ids.is_empty() {
244            return;
245        }
246        let id_set: std::collections::HashSet<EntityId> = entity_ids.iter().copied().collect();
247        let mut idx = self.graph_label_index.write();
248        for ((col, _), ids) in idx.iter_mut() {
249            if col == collection {
250                ids.retain(|id| !id_set.contains(id));
251            }
252        }
253        idx.retain(|_, ids| !ids.is_empty());
254    }
255
256    /// Look up entity IDs for a graph node label across all collections.
257    pub fn lookup_graph_nodes_by_label(&self, label: &str) -> Vec<EntityId> {
258        let idx = self.graph_label_index.read();
259        idx.iter()
260            .filter(|((_, l), _)| l == label)
261            .flat_map(|(_, ids)| ids.iter().copied())
262            .collect()
263    }
264
265    /// Look up entity IDs for a graph node label scoped to a single collection.
266    pub fn lookup_graph_nodes_by_label_in(&self, collection: &str, label: &str) -> Vec<EntityId> {
267        let idx = self.graph_label_index.read();
268        idx.get(&(collection.to_string(), label.to_string()))
269            .cloned()
270            .unwrap_or_default()
271    }
272
273    pub fn create_collection(&self, name: impl Into<String>) -> Result<(), StoreError> {
274        let name = name.into();
275        let mut collections = self.collections.write();
276
277        if collections.contains_key(&name) {
278            return Err(StoreError::CollectionExists(name));
279        }
280
281        let manager = SegmentManager::with_config(&name, self.config.manager_config.clone());
282        collections.insert(name.clone(), Arc::new(manager));
283        drop(collections);
284
285        let create = match self.prepare_operational_collection_create(&name) {
286            Ok(create) => create,
287            Err(err) => {
288                let mut collections = self.collections.write();
289                collections.remove(&name);
290                return Err(err);
291            }
292        };
293        if let Some(create) = create {
294            if let Err(err) = create.publish().map_err(StoreError::Io) {
295                let mut collections = self.collections.write();
296                collections.remove(&name);
297                return Err(err);
298            }
299        }
300
301        self.mark_paged_registry_dirty();
302        self.finish_paged_write([StoreWalAction::CreateCollection { name }])?;
303
304        Ok(())
305    }
306
307    /// Get or create a collection
308    pub fn get_or_create_collection(&self, name: impl Into<String>) -> Arc<SegmentManager> {
309        let name = name.into();
310        // Fast path: shared read lock — zero contention for existing collections
311        {
312            let collections = self.collections.read();
313            if let Some(manager) = collections.get(&name) {
314                return Arc::clone(manager);
315            }
316        }
317        // Slow path: exclusive write lock — only when collection is missing
318        let mut collections = self.collections.write();
319        // Double-check after acquiring write lock (another thread may have created it)
320        if let Some(manager) = collections.get(&name) {
321            return Arc::clone(manager);
322        }
323        let manager = Arc::new(SegmentManager::with_config(
324            &name,
325            self.config.manager_config.clone(),
326        ));
327        collections.insert(name, Arc::clone(&manager));
328        self.mark_paged_registry_dirty();
329        manager
330    }
331
332    /// Get a collection
333    pub fn get_collection(&self, name: &str) -> Option<Arc<SegmentManager>> {
334        self.collections.read().get(name).map(Arc::clone)
335    }
336
337    /// Get the context index for cross-structure search.
338    pub fn context_index(&self) -> &ContextIndex {
339        &self.context_index
340    }
341
342    /// Drain WAL-replayed `VectorInsert` records for `collection`
343    /// (issue #694). Returns `None` when nothing was captured at open
344    /// time (in-memory mode, fresh database, or non-turbo collection).
345    /// One-shot: the caller owns the records after this call, the
346    /// store's internal buffer is cleared for that collection.
347    pub fn take_replayed_turbo_inserts(&self, collection: &str) -> Option<Vec<(u64, Vec<f32>)>> {
348        let mut map = self.replayed_turbo_inserts.lock();
349        map.remove(collection)
350    }
351
352    pub(crate) fn take_replayed_probabilistic_deltas(&self) -> Vec<(u8, u8, String, Vec<Vec<u8>>)> {
353        let mut deltas = self.replayed_probabilistic_deltas.lock();
354        std::mem::take(&mut *deltas)
355    }
356
357    /// Set multiple config KV pairs at once from a JSON tree.
358    /// Keys are flattened with dot-notation: `{"a":{"b":1}}` → `a.b = 1`.
359    ///
360    /// `red_config` is last-writer-wins (ADR 0014): each write replaces the
361    /// key's earlier rows. Appending instead grew the store on every boot,
362    /// since boot re-seeds keys such as `red.system.*`.
363    pub fn set_config_tree(&self, prefix: &str, json: &crate::serde_json::Value) -> usize {
364        let manager = self.get_or_create_collection("red_config");
365        let mut pairs = Vec::new();
366        flatten_config_json(prefix, json, &mut pairs);
367        // Existing rows per key, with the value each one holds.
368        let mut previous: HashMap<String, Vec<(EntityId, Option<reddb_types::Value>)>> = pairs
369            .iter()
370            .map(|(key, _)| (key.clone(), Vec::new()))
371            .collect();
372        for entity in manager.query_all(|_| true) {
373            if let EntityData::Row(row) = &entity.data {
374                let Some(named) = row.named.as_ref() else {
375                    continue;
376                };
377                if let Some(reddb_types::Value::Text(key)) = named.get("key") {
378                    if let Some(rows) = previous.get_mut(key.as_ref()) {
379                        rows.push((entity.id, named.get("value").cloned()));
380                    }
381                }
382            }
383        }
384        let mut saved = 0;
385        for (key, value) in pairs {
386            let stale = previous.remove(&key).unwrap_or_default();
387            // Rewriting a key with the value it already holds is a no-op.
388            // Boot re-seeds the same config every time, and each write is a
389            // durable WAL append with its own fsync.
390            if let [(_, Some(current))] = stale.as_slice() {
391                if *current == value {
392                    saved += 1;
393                    continue;
394                }
395            }
396            let stale: Vec<EntityId> = stale.into_iter().map(|(id, _)| id).collect();
397            // Insert before removing the old rows: a failure in between
398            // leaves a duplicate that readers resolve to the newest row,
399            // never a missing key.
400            let entity = UnifiedEntity::new(
401                EntityId::new(0),
402                EntityKind::TableRow {
403                    table: Arc::from("red_config"),
404                    row_id: 0,
405                },
406                EntityData::Row(RowData {
407                    columns: Vec::new(),
408                    named: Some(
409                        [
410                            ("key".to_string(), reddb_types::Value::text(key)),
411                            ("value".to_string(), value),
412                        ]
413                        .into_iter()
414                        .collect(),
415                    ),
416                    schema: None,
417                }),
418            );
419            if self.insert_auto("red_config", entity).is_ok() {
420                saved += 1;
421                if !stale.is_empty() {
422                    let _ = self.delete_batch("red_config", &stale);
423                }
424            }
425        }
426        saved
427    }
428
429    /// Read a single config value from `red_config` by dot-notation key.
430    /// When a key has several rows the newest one wins, like every other
431    /// `red_config` reader.
432    pub fn get_config(&self, key: &str) -> Option<reddb_types::Value> {
433        let manager = self.get_collection("red_config")?;
434        let mut latest: Option<(u64, reddb_types::Value)> = None;
435        for entity in manager.query_all(|_| true) {
436            if let EntityData::Row(row) = &entity.data {
437                if let Some(named) = &row.named {
438                    let key_matches = named
439                        .get("key")
440                        .and_then(|v| match v {
441                            reddb_types::Value::Text(s) => Some(s.as_ref() == key),
442                            _ => None,
443                        })
444                        .unwrap_or(false);
445                    if key_matches && latest.as_ref().is_none_or(|(id, _)| entity.id.raw() >= *id) {
446                        if let Some(value) = named.get("value") {
447                            latest = Some((entity.id.raw(), value.clone()));
448                        }
449                    }
450                }
451            }
452        }
453        latest.map(|(_, value)| value)
454    }
455
456    /// Replace the opaque store-level auxiliary metadata blob. Persisted
457    /// inside the binary dump (store format V10+); see the `aux_metadata`
458    /// field. RedDB uses it to carry collection contracts through the
459    /// single-file artifact.
460    pub fn set_aux_metadata(&self, bytes: Vec<u8>) {
461        *self.aux_metadata.write() = bytes;
462    }
463
464    /// Read the opaque store-level auxiliary metadata blob. Empty when unset.
465    pub fn aux_metadata(&self) -> Vec<u8> {
466        self.aux_metadata.read().clone()
467    }
468
469    /// List all collections
470    pub fn list_collections(&self) -> Vec<String> {
471        self.collections.read().keys().cloned().collect()
472    }
473
474    /// Drop a collection
475    pub fn drop_collection(&self, name: &str) -> Result<(), StoreError> {
476        let _vault_guard = self.vault_publication_lock.lock();
477        let drop = self.prepare_operational_collection_drop(name)?;
478        let manager = {
479            let mut collections = self.collections.write();
480
481            collections
482                .remove(name)
483                .ok_or_else(|| StoreError::CollectionNotFound(name.to_string()))?
484        };
485
486        let entities = manager.query_all(|_| true);
487        let entity_ids: Vec<EntityId> = entities.iter().map(|entity| entity.id).collect();
488
489        for entity_id in &entity_ids {
490            self.context_index.remove_entity(*entity_id);
491            let _ = self.unindex_cross_refs(*entity_id);
492        }
493
494        if let Some(tree) = self.btree_indices.write().remove(name) {
495            self.retire_btree(tree);
496        }
497
498        self.entity_cache.retain(|entity_id, (collection, _)| {
499            collection != name && !entity_ids.iter().any(|id| id.raw() == entity_id)
500        });
501
502        self.cross_refs.write().retain(|source_id, refs| {
503            refs.retain(|(target_id, _, target_collection)| {
504                target_collection != name && !entity_ids.iter().any(|id| id == target_id)
505            });
506            !entity_ids.iter().any(|id| id == source_id)
507        });
508
509        self.reverse_refs.write().retain(|target_id, refs| {
510            refs.retain(|(source_id, _, source_collection)| {
511                source_collection != name && !entity_ids.iter().any(|id| id == source_id)
512            });
513            !entity_ids.iter().any(|id| id == target_id)
514        });
515
516        self.mark_paged_registry_dirty();
517        self.finish_paged_write([StoreWalAction::DropCollection {
518            name: name.to_string(),
519        }])?;
520        if let Some(drop) = drop {
521            drop.finish().map_err(StoreError::Io)?;
522        }
523
524        Ok(())
525    }
526
527    /// Insert an entity into a collection
528    pub fn insert(&self, collection: &str, entity: UnifiedEntity) -> Result<EntityId, StoreError> {
529        let manager = self
530            .get_collection(collection)
531            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
532
533        let mut entity = entity;
534        if entity.id.raw() == 0 {
535            entity.id = self.next_entity_id();
536        } else {
537            self.register_entity_id(entity.id);
538        }
539        // Assign per-table sequential row_id if not set
540        if let EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
541            if *row_id == 0 {
542                *row_id = manager.next_row_id();
543            } else {
544                manager.register_row_id(*row_id);
545            }
546        }
547        entity.ensure_table_logical_id();
548        // Capture graph node label before entity is moved into the manager
549        let graph_node_label: Option<String> = if let EntityKind::GraphNode(ref node) = entity.kind
550        {
551            Some(node.label.clone())
552        } else {
553            None
554        };
555
556        let id = manager.insert(entity)?;
557        self.register_entity_id(id);
558
559        // Update graph label index for GraphNode entities
560        if let Some(ref label) = graph_node_label {
561            self.update_graph_label_index(collection, label, id);
562        }
563
564        // Also insert into B-tree index if pager is active
565        let mut registry_dirty = false;
566        if let Some(pager) = &self.pager {
567            if let Some(entity) = manager.get(id) {
568                let mut btree_indices = self.btree_indices.write();
569                let btree = btree_indices
570                    .entry(collection.to_string())
571                    .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
572                let root_before = btree.root_page_id();
573
574                let key = id.raw().to_be_bytes();
575                let metadata = manager.get_metadata(id);
576                let value = Self::serialize_entity_record(
577                    &entity,
578                    metadata.as_ref(),
579                    self.format_version(),
580                );
581                // Ignore duplicate key errors (update scenario)
582                let _ = btree.insert(&key, &value);
583                registry_dirty = root_before != btree.root_page_id();
584            }
585        }
586
587        // Index cross-references if enabled
588        if self.config.auto_index_refs {
589            if let Some(entity) = manager.get(id) {
590                self.index_cross_refs(&entity, collection)?;
591            }
592        }
593
594        // Perf: skip WAL-action construction when the store is
595        // pagerless. For in-memory benchmarks this saved another
596        // `manager.get(id)` + `serialize_entity_record` per call.
597        if self.pager.is_some() || self.config.embedded_wal_path.is_some() {
598            let actions = manager
599                .get(id)
600                .map(|entity| {
601                    let metadata = manager.get_metadata(id);
602                    vec![StoreWalAction::upsert_entity(
603                        collection,
604                        &entity,
605                        metadata.as_ref(),
606                        self.format_version(),
607                    )]
608                })
609                .unwrap_or_default();
610            if registry_dirty {
611                self.mark_paged_registry_dirty();
612            }
613            self.finish_paged_write(actions)?;
614        }
615
616        Ok(id)
617    }
618
619    /// Turbo bulk insert — optimized fast path.
620    ///
621    /// Single lock for the entire batch. Skips bloom filter,
622    /// context index, and cross-ref indexing. B-tree writes are batched.
623    pub fn bulk_insert(
624        &self,
625        collection: &str,
626        mut entities: Vec<UnifiedEntity>,
627    ) -> Result<Vec<EntityId>, StoreError> {
628        // REDDB_BULK_TIMING=1 prints a per-call breakdown of the bulk
629        // insert path to stderr. Off by default — used by the reddb
630        // benchmark harness to locate ingest bottlenecks.
631        let trace = matches!(
632            std::env::var("REDDB_BULK_TIMING").ok().as_deref(),
633            Some("1") | Some("true") | Some("on")
634        );
635        let t_start = std::time::Instant::now();
636        let n = entities.len();
637        let manager = self.get_or_create_collection(collection);
638        let t_get_coll = t_start.elapsed();
639
640        // Assign IDs and per-table row_ids before serialization. Bulk
641        // insert must follow the same global ID semantics as
642        // insert()/insert_auto(). Reserve ID ranges in single fetch_add
643        // calls instead of one-per-entity. The overwhelming common case
644        // is "every entity comes in with id==0 and kind==TableRow with
645        // row_id==0" (the wire bulk insert path). Any entity that
646        // already carries an id or row_id falls back to the individual
647        // register_* path to preserve those semantics exactly.
648        let t0 = std::time::Instant::now();
649        let n_missing_entity_ids = entities.iter().filter(|e| e.id.raw() == 0).count() as u64;
650        let n_missing_row_ids = entities
651            .iter()
652            .filter(|e| matches!(e.kind, EntityKind::TableRow { row_id: 0, .. }))
653            .count() as u64;
654        let mut entity_id_range = if n_missing_entity_ids > 0 {
655            self.reserve_entity_ids(n_missing_entity_ids)
656        } else {
657            0..0
658        };
659        let mut row_id_range = if n_missing_row_ids > 0 {
660            manager.reserve_row_ids(n_missing_row_ids)
661        } else {
662            0..0
663        };
664        for entity in &mut entities {
665            if entity.id.raw() == 0 {
666                let next = entity_id_range
667                    .next()
668                    .expect("reserved entity-id range exhausted");
669                entity.id = EntityId::new(next);
670            } else {
671                self.register_entity_id(entity.id);
672            }
673            if let EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
674                if *row_id == 0 {
675                    *row_id = row_id_range
676                        .next()
677                        .expect("reserved row-id range exhausted");
678                } else {
679                    manager.register_row_id(*row_id);
680                }
681            }
682            entity.ensure_table_logical_id();
683        }
684        let t_assign_ids = t0.elapsed();
685
686        // Capture graph node labels before entities are moved into the segment manager
687        let graph_labels: Vec<Option<(String, EntityId)>> = entities
688            .iter()
689            .map(|e| {
690                if let EntityKind::GraphNode(ref node) = e.kind {
691                    Some((node.label.clone(), e.id))
692                } else {
693                    None
694                }
695            })
696            .collect();
697
698        // Pre-serialize for B-tree while we still have references.
699        // `serialize_entity_record` is pure; with `n >= 1024` the
700        // rayon dispatch cost is amortised and the 16-core serialize
701        // fan-out becomes the biggest single win on insert_bulk. For
702        // smaller batches, stay serial — micro-batches pay more in
703        // work-stealing overhead than they save.
704        let t0 = std::time::Instant::now();
705        let serialized: Option<Vec<(Vec<u8>, Vec<u8>)>> =
706            if self.pager.is_some() || self.config.embedded_wal_path.is_some() {
707                let fv = self.format_version();
708                let serial_map = |e: &UnifiedEntity| {
709                    (
710                        e.id.raw().to_be_bytes().to_vec(),
711                        Self::serialize_entity_record(e, None, fv),
712                    )
713                };
714                // Gate chosen to match the bench's `BULK_BATCH_SIZE=1000`.
715                // Rayon's dispatch overhead is ~30µs/batch — on a 15-col
716                // row serializing for ~40µs the break-even is around
717                // 200-300 rows on a 16-core host. Keep a safety margin
718                // at 512.
719                if entities.len() >= 512 {
720                    use rayon::prelude::*;
721                    Some(entities.par_iter().map(serial_map).collect())
722                } else {
723                    Some(entities.iter().map(serial_map).collect())
724                }
725            } else {
726                None
727            };
728        let t_serialize = t0.elapsed();
729
730        // Move entities into segment
731        let t0 = std::time::Instant::now();
732        let ids = manager.bulk_insert(entities)?;
733        let t_manager = t0.elapsed();
734        for id in &ids {
735            self.register_entity_id(*id);
736        }
737
738        // Update graph label index for bulk-inserted GraphNode entities
739        for (label, entity_id) in graph_labels.iter().flatten() {
740            self.update_graph_label_index(collection, label, *entity_id);
741        }
742
743        // REDDB_BULK_SKIP_PERSIST_UNSAFE=1 skips the persistent B-tree index
744        // during bulk ingest.
745        //
746        // UNSAFE: for ephemeral benchmark containers ONLY.
747        // This flag is silently ignored when a pager (durable storage) is active.
748        // In persistent mode, bulk inserts ALWAYS write to the B-tree so the data
749        // survives a cold restart without any manual rebuild step.
750        //
751        // The flag is only honoured when self.pager is None (in-memory / ephemeral).
752        let skip_btree_requested = matches!(
753            std::env::var("REDDB_BULK_SKIP_PERSIST_UNSAFE")
754                .ok()
755                .as_deref(),
756            Some("1") | Some("true") | Some("on")
757        );
758        // Honour the flag only when there is no durable pager.
759        // If a pager exists we are in persistent mode → always persist.
760        let skip_btree = skip_btree_requested && self.pager.is_none();
761        if skip_btree_requested && !skip_btree {
762            // Flag was set but we ignored it because we have a real pager.
763            static IGNORED: std::sync::OnceLock<()> = std::sync::OnceLock::new();
764            IGNORED.get_or_init(|| {
765                tracing::warn!(
766                    "REDDB_BULK_SKIP_PERSIST_UNSAFE set but durable pager is \
767                     active — flag ignored; bulk inserts will be persisted normally"
768                );
769            });
770        } else if skip_btree {
771            // Ephemeral mode and flag is active — warn once.
772            static WARNED: std::sync::OnceLock<()> = std::sync::OnceLock::new();
773            WARNED.get_or_init(|| {
774                tracing::warn!(
775                    "REDDB_BULK_SKIP_PERSIST_UNSAFE set (ephemeral/no-pager mode) — \
776                     bulk inserts NOT durable; data will be lost on restart"
777                );
778            });
779        }
780
781        // Batch B-tree write from pre-serialized data.
782        // Uses sorted bulk insert: walks to a leaf once, appends many entries,
783        // writes each leaf exactly once per batch — O(N) instead of O(N²).
784        let mut t_btree_lock = std::time::Duration::ZERO;
785        let mut t_btree_insert = std::time::Duration::ZERO;
786        let mut t_flush = std::time::Duration::ZERO;
787        if !skip_btree {
788            if let (Some(pager), Some(batch)) = (&self.pager, serialized.as_ref()) {
789                let t0 = std::time::Instant::now();
790                let mut btree_indices = self.btree_indices.write();
791                let btree = btree_indices
792                    .entry(collection.to_string())
793                    .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
794                let root_before = btree.root_page_id();
795                t_btree_lock = t0.elapsed();
796
797                let t0 = std::time::Instant::now();
798                let _ = btree.bulk_insert_sorted(batch);
799                t_btree_insert = t0.elapsed();
800                let registry_dirty = root_before != btree.root_page_id();
801
802                let t0 = std::time::Instant::now();
803                if registry_dirty {
804                    self.mark_paged_registry_dirty();
805                }
806                t_flush = t0.elapsed();
807            }
808        }
809
810        // Fold 25k per-entity `UpsertEntityRecord` actions into one
811        // `BulkUpsertEntityRecords` — same on-disk semantics (replay
812        // applies each record by id), one `WalRecord::PageWrite`
813        // instead of N, and `record.clone()` drops out because we
814        // consume `serialized`.
815        let actions = serialized
816            .map(|batch| {
817                let records: Vec<Vec<u8>> =
818                    batch.into_iter().map(|(_key, record)| record).collect();
819                vec![StoreWalAction::BulkUpsertEntityRecords {
820                    collection: collection.to_string(),
821                    records,
822                }]
823            })
824            .unwrap_or_default();
825        self.finish_paged_write(actions)?;
826
827        if trace {
828            tracing::debug!(
829                n,
830                total = ?t_start.elapsed(),
831                get_coll = ?t_get_coll,
832                assign = ?t_assign_ids,
833                serialize = ?t_serialize,
834                manager = ?t_manager,
835                btree_lock = ?t_btree_lock,
836                btree = ?t_btree_insert,
837                flush = ?t_flush,
838                "bulk_insert timing"
839            );
840        }
841
842        Ok(ids)
843    }
844
845    /// Insert an entity, creating collection if needed
846    pub fn insert_auto(
847        &self,
848        collection: &str,
849        entity: UnifiedEntity,
850    ) -> Result<EntityId, StoreError> {
851        // Include the WAL append in the checkpoint publication fence.
852        let _publication = self.vault_publication_lock.lock();
853        let manager = self.get_or_create_collection(collection);
854        let mut entity = entity;
855        if entity.id.raw() == 0 {
856            entity.id = self.next_entity_id();
857        } else {
858            self.register_entity_id(entity.id);
859        }
860        // Assign per-table sequential row_id if not set
861        if let EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
862            if *row_id == 0 {
863                *row_id = manager.next_row_id();
864            } else {
865                manager.register_row_id(*row_id);
866            }
867        }
868        entity.ensure_table_logical_id();
869
870        // Capture graph node label before entity is moved into the manager
871        let graph_node_label: Option<String> = if let EntityKind::GraphNode(ref node) = entity.kind
872        {
873            Some(node.label.clone())
874        } else {
875            None
876        };
877        // Index into context index before consuming the entity
878        self.context_index.index_entity(collection, &entity);
879
880        // Serialize the canonical segment image after insertion: the segment
881        // assigns its sequence identity. Borrow it instead of cloning payloads.
882        let persist_record = self.pager.is_some() || self.config.embedded_wal_path.is_some();
883        if self.config.auto_index_refs {
884            self.index_cross_refs(&entity, collection)?;
885        }
886
887        // #2376: `persist()` rebuilds every collection B-tree from the
888        // manager under `btree_indices.write()`. Hold the read side from the
889        // manager insert through the B-tree insert so a rebuild observes
890        // both or neither. Otherwise the rebuild re-inserts this row from the
891        // manager and one of the two B-tree inserts fails on a duplicate key,
892        // or lands mid-rebuild and tears the tree.
893        let paged_btree = if persist_record {
894            self.get_or_create_btree(collection)
895        } else {
896            None
897        };
898        let persist_fence = paged_btree
899            .as_ref()
900            .map(|_| self.btree_indices.read_recursive());
901        let before = paged_btree.as_ref().map(|_| manager.checkpoint_revisions());
902        let unversioned = entity.xmin == 0 && entity.xmax == 0;
903
904        let id = manager.insert(entity)?;
905        let serialized_record = if persist_record {
906            Some(
907                manager
908                    .get_with(id, |entity| {
909                        Self::serialize_entity_record(entity, None, self.format_version())
910                    })
911                    .ok_or(StoreError::EntityNotFound(id))?,
912            )
913        } else {
914            None
915        };
916        // `register_entity_id` already advances the atomic counter on
917        // the allocation path above (`self.next_entity_id()`), so the
918        // second call here is a no-op CAS loop on the hot path. Only
919        // needed for the caller-supplied-id branch which happens via
920        // the `register_entity_id` call on line 573.
921
922        // Update graph label index for GraphNode entities
923        if let Some(ref label) = graph_node_label {
924            self.update_graph_label_index(collection, label, id);
925        }
926
927        let mut registry_dirty = false;
928        if let (Some(btree), Some(record)) = (paged_btree.as_ref(), serialized_record.as_ref()) {
929            let root_before = btree.root_page_id();
930
931            let key = id.raw().to_be_bytes();
932            btree.insert(&key, record).map_err(|e| {
933                StoreError::Io(std::io::Error::other(format!(
934                    "B-tree insert error while inserting '{collection}'/{id}: {e}"
935                )))
936            })?;
937            registry_dirty = root_before != btree.root_page_id();
938        }
939        drop(persist_fence);
940
941        // Perf: pagerless → skip WAL-action construction (saves a
942        // third manager.get + entity serialize per insert). For
943        // in-memory runtimes finish_paged_write is a no-op.
944        if self.pager.is_some() || self.config.embedded_wal_path.is_some() {
945            let actions = serialized_record
946                .map(|record| {
947                    vec![StoreWalAction::UpsertEntityRecord {
948                        collection: collection.to_string(),
949                        record,
950                    }]
951                })
952                .unwrap_or_default();
953            if registry_dirty {
954                self.mark_paged_registry_dirty();
955            }
956            self.finish_paged_write(actions)?;
957        }
958        if unversioned && !Self::deferred_store_wal_capture_active() {
959            if let Some(before) = before {
960                self.note_paged_insert(collection, &manager, &before, id);
961            }
962        }
963        Ok(id)
964    }
965
966    /// Get an entity from a collection
967    ///
968    /// Prefers the live SegmentManager view so reads after update/delete observe
969    /// the current in-memory state even when the paged B-tree image has not been
970    /// refreshed yet. Falls back to the B-tree image for recovery-oriented reads.
971    pub fn get(&self, collection: &str, id: EntityId) -> Option<UnifiedEntity> {
972        // Prefer the live manager state to avoid stale reads after manager.update().
973        if let Some(entity) = self
974            .get_collection(collection)
975            .and_then(|manager| manager.get(id))
976        {
977            return Some(entity);
978        }
979
980        // Fall back to the paged B-tree image if the manager does not currently hold the row.
981        if self.pager.is_some() {
982            let btree_indices = self.btree_indices.read();
983            if let Some(btree) = btree_indices.get(collection) {
984                let key = id.raw().to_be_bytes();
985                if let Ok(Some(value)) = btree.get(&key) {
986                    if let Ok((entity, _)) =
987                        Self::deserialize_entity_record(&value, self.format_version())
988                    {
989                        return Some(entity);
990                    }
991                }
992            }
993        }
994
995        None
996    }
997
998    /// Resolve a table row by stable logical identity.
999    ///
1000    /// Legacy rows use `logical_id == id`, so the physical-id lookup remains
1001    /// the fast path. Versioned rows may have a different physical id and fall
1002    /// back to a scan until the history-store/index slices add a dedicated map.
1003    pub fn get_table_row_by_logical_id(
1004        &self,
1005        collection: &str,
1006        logical_id: EntityId,
1007    ) -> Option<UnifiedEntity> {
1008        if let Some(entity) = self.get(collection, logical_id) {
1009            if matches!(entity.kind, EntityKind::TableRow { .. })
1010                && entity.logical_id() == logical_id
1011                && entity.xmax == 0
1012            {
1013                return Some(entity);
1014            }
1015        }
1016
1017        let manager = self.get_collection(collection)?;
1018        let mut matches = manager.query_all(|entity| {
1019            matches!(entity.kind, EntityKind::TableRow { .. }) && entity.logical_id() == logical_id
1020        });
1021        matches
1022            .iter()
1023            .find(|entity| entity.xmax == 0)
1024            .cloned()
1025            .or_else(|| matches.pop())
1026    }
1027
1028    pub fn table_row_versions_by_logical_id(
1029        &self,
1030        collection: &str,
1031        logical_id: EntityId,
1032    ) -> Vec<UnifiedEntity> {
1033        self.get_collection(collection)
1034            .map(|manager| {
1035                manager.query_all(|entity| {
1036                    matches!(entity.kind, EntityKind::TableRow { .. })
1037                        && entity.logical_id() == logical_id
1038                })
1039            })
1040            .unwrap_or_default()
1041    }
1042
1043    pub fn vacuum_mvcc_history(
1044        &self,
1045        collection: &str,
1046        cutoff_xid: u64,
1047    ) -> Result<MvccVacuumStats, StoreError> {
1048        let manager = self
1049            .get_collection(collection)
1050            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1051        let entities =
1052            manager.query_all(|entity| matches!(entity.kind, EntityKind::TableRow { .. }));
1053        let mut logical = HashMap::<EntityId, (bool, u64)>::new();
1054        for entity in &entities {
1055            let entry = logical.entry(entity.logical_id()).or_insert((false, 0));
1056            if entity.xmax == 0 {
1057                entry.0 = true;
1058            }
1059            entry.1 = entry.1.max(entity.xmin);
1060        }
1061
1062        let mut stats = MvccVacuumStats::default();
1063        let mut reclaim_ids = Vec::new();
1064        for entity in entities {
1065            if entity.xmax == 0 {
1066                continue;
1067            }
1068            stats.scanned_versions += 1;
1069            let (has_live_version, max_xmin) = logical
1070                .get(&entity.logical_id())
1071                .copied()
1072                .unwrap_or((false, entity.xmin));
1073            let is_delete_tombstone = !has_live_version && entity.xmin == max_xmin;
1074            if entity.xmax < cutoff_xid {
1075                stats.reclaimed_versions += 1;
1076                if is_delete_tombstone {
1077                    stats.reclaimed_tombstones += 1;
1078                } else {
1079                    stats.reclaimed_history_versions += 1;
1080                }
1081                reclaim_ids.push(entity.id);
1082            } else {
1083                stats.retained_versions += 1;
1084                if is_delete_tombstone {
1085                    stats.retained_tombstones += 1;
1086                } else {
1087                    stats.retained_history_versions += 1;
1088                }
1089            }
1090        }
1091
1092        if !reclaim_ids.is_empty() {
1093            self.delete_batch(collection, &reclaim_ids)?;
1094        }
1095        Ok(stats)
1096    }
1097
1098    pub(crate) fn install_versioned_table_row_update(
1099        &self,
1100        collection: &str,
1101        old_version: UnifiedEntity,
1102        mut new_version: UnifiedEntity,
1103        metadata: Option<&Metadata>,
1104    ) -> Result<(), StoreError> {
1105        let manager = self
1106            .get_collection(collection)
1107            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1108
1109        let old_id = old_version.id;
1110        let new_id = new_version.id;
1111        let inherited_metadata = metadata.cloned().or_else(|| manager.get_metadata(old_id));
1112
1113        self.entity_cache.remove(old_id.raw());
1114        self.entity_cache.remove(new_id.raw());
1115        manager.update(old_version.clone())?;
1116        self.context_index.remove_entity(old_id);
1117
1118        self.register_entity_id(new_version.id);
1119        if let EntityKind::TableRow { ref mut row_id, .. } = new_version.kind {
1120            if *row_id == 0 {
1121                *row_id = manager.next_row_id();
1122            } else {
1123                manager.register_row_id(*row_id);
1124            }
1125        }
1126        new_version.ensure_table_logical_id();
1127        manager.insert(new_version.clone())?;
1128        if let Some(metadata) = inherited_metadata {
1129            manager.set_metadata(new_id, metadata)?;
1130        }
1131        self.context_index.index_entity(collection, &new_version);
1132        if self.config.auto_index_refs {
1133            self.index_cross_refs(&new_version, collection)?;
1134        }
1135
1136        let old_metadata = manager.get_metadata(old_id);
1137        let new_metadata = manager.get_metadata(new_id);
1138        let fv = self.format_version();
1139        let records = vec![
1140            Self::serialize_entity_record(&old_version, old_metadata.as_ref(), fv),
1141            Self::serialize_entity_record(&new_version, new_metadata.as_ref(), fv),
1142        ];
1143        self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
1144            collection: collection.to_string(),
1145            records,
1146        }])?;
1147
1148        Ok(())
1149    }
1150
1151    /// Batch-fetch multiple entities from the same collection in minimal lock acquisitions.
1152    ///
1153    /// Preferred over N individual `get()` calls in indexed-scan loops (sorted index,
1154    /// bitmap, hash). Reduces lock acquisitions from N×3 to 2-3 total.
1155    /// Preserves input order: `result[i]` corresponds to `ids[i]`.
1156    pub fn get_batch(&self, collection: &str, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
1157        match self.get_collection(collection) {
1158            Some(manager) => manager.get_many(ids),
1159            None => vec![None; ids.len()],
1160        }
1161    }
1162
1163    /// Get an entity from any collection
1164    pub fn get_any(&self, id: EntityId) -> Option<(String, UnifiedEntity)> {
1165        // Sharded LRU probe: hits / misses are counted by `EntityCache::get`,
1166        // so external observers can read `entity_cache_hit_rate()` without
1167        // any extra wiring on the call site.
1168        if let Some(cached) = self.entity_cache.get(id.raw()) {
1169            return Some(cached);
1170        }
1171
1172        // Cache miss — fall back to the full collection scan, then memoise.
1173        let collections = self.collections.read();
1174        for (name, manager) in collections.iter() {
1175            if let Some(entity) = manager.get(id) {
1176                let result = (name.clone(), entity);
1177                // Drop the collections read guard before taking any cache
1178                // lock — the cache shards are independent, but releasing
1179                // here keeps `collections` contention low.
1180                drop(collections);
1181                self.entity_cache.insert(id.raw(), result.clone());
1182                return Some(result);
1183            }
1184        }
1185        None
1186    }
1187
1188    /// Hit rate of the store's entity cache (`get_any` lookups).
1189    ///
1190    /// Returns `None` until the cache has served at least one lookup.
1191    /// Exposed for observability — e.g. dashboards distinguishing graph
1192    /// workloads (high hit rate) from OLTP DML (≈ 0 % hit rate).
1193    pub fn entity_cache_hit_rate(&self) -> Option<f64> {
1194        self.entity_cache.hit_rate()
1195    }
1196
1197    /// Snapshot of cache hit / miss / eviction counters and current size.
1198    pub fn entity_cache_stats(&self) -> super::super::entity_cache::EntityCacheStats {
1199        self.entity_cache.stats()
1200    }
1201
1202    /// Delete an entity
1203    pub fn delete(&self, collection: &str, id: EntityId) -> Result<bool, StoreError> {
1204        // Invalidate entity cache (read-lock probe; no write lock when cold).
1205        self.entity_cache.remove(id.raw());
1206        let manager = self
1207            .get_collection(collection)
1208            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1209
1210        let deleted = manager.delete(id)?;
1211        if !deleted {
1212            return Ok(false);
1213        }
1214
1215        // Remove from B-tree index if active
1216        let mut registry_dirty = false;
1217        if self.pager.is_some() {
1218            let btree_indices = self.btree_indices.read();
1219            if let Some(btree) = btree_indices.get(collection) {
1220                let root_before = btree.root_page_id();
1221                let key = id.raw().to_be_bytes();
1222                let _ = btree.delete(&key);
1223                registry_dirty = root_before != btree.root_page_id();
1224            }
1225        }
1226
1227        // Remove cross-references
1228        self.unindex_cross_refs(id)?;
1229
1230        // Remove from graph label index
1231        self.remove_from_graph_label_index(collection, id);
1232
1233        if registry_dirty {
1234            self.mark_paged_registry_dirty();
1235        }
1236        self.finish_paged_write([StoreWalAction::DeleteEntityRecord {
1237            collection: collection.to_string(),
1238            entity_id: id.raw(),
1239        }])?;
1240
1241        Ok(true)
1242    }
1243
1244    pub fn delete_batch(
1245        &self,
1246        collection: &str,
1247        ids: &[EntityId],
1248    ) -> Result<Vec<EntityId>, StoreError> {
1249        if ids.is_empty() {
1250            return Ok(Vec::new());
1251        }
1252
1253        // Sharded invalidation. The bounded LRU's shard write lock is
1254        // only acquired when the shard actually carries one of these ids,
1255        // so the bench-typical zero-hit `delete_sequential` workload never
1256        // takes a write lock here at all.
1257        self.entity_cache.remove_many(ids.iter().map(|id| id.raw()));
1258
1259        let manager = self
1260            .get_collection(collection)
1261            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1262
1263        let deleted_ids = manager.delete_batch(ids)?;
1264        if deleted_ids.is_empty() {
1265            return Ok(deleted_ids);
1266        }
1267
1268        let mut registry_dirty = false;
1269        if self.pager.is_some() {
1270            let btree_indices = self.btree_indices.read();
1271            if let Some(btree) = btree_indices.get(collection) {
1272                let root_before = btree.root_page_id();
1273                for id in &deleted_ids {
1274                    let key = id.raw().to_be_bytes();
1275                    let _ = btree.delete(&key);
1276                }
1277                registry_dirty = root_before != btree.root_page_id();
1278            }
1279        }
1280
1281        self.unindex_cross_refs_batch(&deleted_ids)?;
1282        self.remove_from_graph_label_index_batch(collection, &deleted_ids);
1283        if registry_dirty {
1284            self.mark_paged_registry_dirty();
1285        }
1286        let actions = deleted_ids
1287            .iter()
1288            .map(|id| StoreWalAction::DeleteEntityRecord {
1289                collection: collection.to_string(),
1290                entity_id: id.raw(),
1291            })
1292            .collect::<Vec<_>>();
1293        self.finish_paged_write(actions)?;
1294
1295        Ok(deleted_ids)
1296    }
1297
1298    /// Set metadata for an entity
1299    pub fn set_metadata(
1300        &self,
1301        collection: &str,
1302        id: EntityId,
1303        metadata: Metadata,
1304    ) -> Result<(), StoreError> {
1305        let manager = self
1306            .get_collection(collection)
1307            .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1308
1309        manager.set_metadata(id, metadata)?;
1310        if let Some(entity) = manager.get(id) {
1311            self.persist_entities_to_pager(collection, std::slice::from_ref(&entity))?;
1312        }
1313        Ok(())
1314    }
1315
1316    /// Get metadata for an entity
1317    pub fn get_metadata(&self, collection: &str, id: EntityId) -> Option<Metadata> {
1318        self.get_collection(collection)?.get_metadata(id)
1319    }
1320
1321    /// Add a cross-reference between entities
1322    pub fn add_cross_ref(
1323        &self,
1324        source_collection: &str,
1325        source_id: EntityId,
1326        target_collection: &str,
1327        target_id: EntityId,
1328        ref_type: RefType,
1329        weight: f32,
1330    ) -> Result<(), StoreError> {
1331        // Check source exists
1332        let source_manager = self
1333            .get_collection(source_collection)
1334            .ok_or_else(|| StoreError::CollectionNotFound(source_collection.to_string()))?;
1335
1336        if source_manager.get(source_id).is_none() {
1337            return Err(StoreError::EntityNotFound(source_id));
1338        }
1339
1340        // Check target exists
1341        let target_manager = self
1342            .get_collection(target_collection)
1343            .ok_or_else(|| StoreError::CollectionNotFound(target_collection.to_string()))?;
1344
1345        if target_manager.get(target_id).is_none() {
1346            return Err(StoreError::EntityNotFound(target_id));
1347        }
1348
1349        // Check limits
1350        let current_refs = self
1351            .cross_refs
1352            .read()
1353            .get(&source_id)
1354            .map_or(0, |v| v.len());
1355
1356        if current_refs >= self.config.max_cross_refs {
1357            return Err(StoreError::TooManyRefs(source_id));
1358        }
1359
1360        let mut registry_dirty = false;
1361        {
1362            let mut forward = self.cross_refs.write();
1363            let refs = forward.entry(source_id).or_default();
1364            let inserted = !refs.iter().any(|(id, kind, coll)| {
1365                *id == target_id && *kind == ref_type && coll == target_collection
1366            });
1367            if inserted {
1368                refs.push((target_id, ref_type, target_collection.to_string()));
1369                registry_dirty = true;
1370            }
1371        }
1372
1373        {
1374            let mut reverse = self.reverse_refs.write();
1375            let refs = reverse.entry(target_id).or_default();
1376            let inserted = !refs.iter().any(|(id, kind, coll)| {
1377                *id == source_id && *kind == ref_type && coll == source_collection
1378            });
1379            if inserted {
1380                refs.push((source_id, ref_type, source_collection.to_string()));
1381                registry_dirty = true;
1382            }
1383        }
1384
1385        if let Some(mut entity) = source_manager.get(source_id) {
1386            if !entity.cross_refs().iter().any(|xref| {
1387                xref.target == target_id
1388                    && xref.ref_type == ref_type
1389                    && xref.target_collection == target_collection
1390            }) {
1391                let cross_ref = CrossRef::with_weight(
1392                    source_id,
1393                    target_id,
1394                    target_collection,
1395                    ref_type,
1396                    weight,
1397                );
1398                entity.add_cross_ref(cross_ref);
1399                let _ = source_manager.update(entity.clone());
1400                registry_dirty = true;
1401                self.persist_entities_to_pager(source_collection, std::slice::from_ref(&entity))?;
1402            }
1403        }
1404
1405        if registry_dirty {
1406            self.mark_paged_registry_dirty();
1407            if matches!(
1408                self.config.durability_mode,
1409                crate::api::DurabilityMode::Strict
1410            ) {
1411                self.flush_paged_state()?;
1412            }
1413        }
1414
1415        Ok(())
1416    }
1417
1418    /// Get cross-references from an entity
1419    pub fn get_refs_from(&self, id: EntityId) -> Vec<(EntityId, RefType, String)> {
1420        self.cross_refs.read().get(&id).cloned().unwrap_or_default()
1421    }
1422
1423    /// Get cross-references to an entity
1424    pub fn get_refs_to(&self, id: EntityId) -> Vec<(EntityId, RefType, String)> {
1425        self.reverse_refs
1426            .read()
1427            .get(&id)
1428            .cloned()
1429            .unwrap_or_default()
1430    }
1431
1432    /// Expand cross-references to get related entities
1433    pub fn expand_refs(
1434        &self,
1435        id: EntityId,
1436        depth: u32,
1437        ref_types: Option<&[RefType]>,
1438    ) -> Vec<(UnifiedEntity, u32, RefType)> {
1439        let mut results = Vec::new();
1440        let mut visited = std::collections::HashSet::new();
1441        visited.insert(id);
1442
1443        self.expand_refs_recursive(id, depth, ref_types, &mut visited, &mut results, 1);
1444
1445        results
1446    }
1447
1448    fn expand_refs_recursive(
1449        &self,
1450        id: EntityId,
1451        max_depth: u32,
1452        ref_types: Option<&[RefType]>,
1453        visited: &mut std::collections::HashSet<EntityId>,
1454        results: &mut Vec<(UnifiedEntity, u32, RefType)>,
1455        current_depth: u32,
1456    ) {
1457        if current_depth > max_depth {
1458            return;
1459        }
1460
1461        for (target_id, ref_type, target_collection) in self.get_refs_from(id) {
1462            if visited.contains(&target_id) {
1463                continue;
1464            }
1465
1466            if let Some(types) = ref_types {
1467                if !types.contains(&ref_type) {
1468                    continue;
1469                }
1470            }
1471
1472            visited.insert(target_id);
1473
1474            if let Some(entity) = self.get(&target_collection, target_id) {
1475                results.push((entity, current_depth, ref_type));
1476
1477                // Recurse
1478                self.expand_refs_recursive(
1479                    target_id,
1480                    max_depth,
1481                    ref_types,
1482                    visited,
1483                    results,
1484                    current_depth + 1,
1485                );
1486            }
1487        }
1488    }
1489
1490    /// Index cross-references from an entity
1491    pub(crate) fn index_cross_refs(
1492        &self,
1493        entity: &UnifiedEntity,
1494        collection: &str,
1495    ) -> Result<(), StoreError> {
1496        let mut registry_dirty = false;
1497        for cross_ref in entity.cross_refs() {
1498            if cross_ref.target_collection.is_empty() {
1499                continue;
1500            }
1501            {
1502                let mut forward = self.cross_refs.write();
1503                let refs = forward.entry(cross_ref.source).or_default();
1504                let inserted = !refs.iter().any(|(id, kind, coll)| {
1505                    *id == cross_ref.target
1506                        && *kind == cross_ref.ref_type
1507                        && coll == &cross_ref.target_collection
1508                });
1509                if inserted {
1510                    refs.push((
1511                        cross_ref.target,
1512                        cross_ref.ref_type,
1513                        cross_ref.target_collection.clone(),
1514                    ));
1515                    registry_dirty = true;
1516                }
1517            }
1518
1519            {
1520                let mut reverse = self.reverse_refs.write();
1521                let refs = reverse.entry(cross_ref.target).or_default();
1522                let inserted = !refs.iter().any(|(id, kind, coll)| {
1523                    *id == cross_ref.source && *kind == cross_ref.ref_type && coll == collection
1524                });
1525                if inserted {
1526                    refs.push((cross_ref.source, cross_ref.ref_type, collection.to_string()));
1527                    registry_dirty = true;
1528                }
1529            }
1530        }
1531
1532        if registry_dirty {
1533            self.mark_paged_registry_dirty();
1534        }
1535
1536        Ok(())
1537    }
1538
1539    /// Remove cross-references for an entity
1540    pub(crate) fn unindex_cross_refs(&self, id: EntityId) -> Result<(), StoreError> {
1541        // Remove forward refs
1542        self.cross_refs.write().remove(&id);
1543
1544        // Remove from reverse refs (scan all)
1545        let mut reverse = self.reverse_refs.write();
1546        for refs in reverse.values_mut() {
1547            refs.retain(|(source, _, _)| *source != id);
1548        }
1549        reverse.remove(&id);
1550        self.mark_paged_registry_dirty();
1551
1552        Ok(())
1553    }
1554
1555    pub(crate) fn unindex_cross_refs_batch(&self, ids: &[EntityId]) -> Result<(), StoreError> {
1556        if ids.is_empty() {
1557            return Ok(());
1558        }
1559
1560        let id_set: std::collections::HashSet<EntityId> = ids.iter().copied().collect();
1561
1562        // Read-only probe: under steady-state delete workloads most rows have
1563        // no cross-refs at all, so both maps are entirely cold for the batch.
1564        // Acquiring the write lock to scan `reverse_refs.values_mut()` over an
1565        // unrelated dictionary was the dominant cost in #85's `delete_sequential`
1566        // profile. The read-lock probe below is O(|ids|) hashmap lookups against
1567        // each map, vs. O(|reverse_refs|) for the write-path scan it replaces.
1568        //
1569        // Decision: skip the write path iff
1570        //   - no deleted id is a key in `cross_refs`  (no outbound refs to drop,
1571        //     and so no `(deleted_id, _, _)` source entry exists anywhere in
1572        //     `reverse_refs.values()` either — the two maps are kept in sync by
1573        //     `add_cross_ref`), AND
1574        //   - no deleted id is a key in `reverse_refs` (no inbound refs to drop).
1575        let needs_forward_cleanup = {
1576            let forward = self.cross_refs.read();
1577            id_set.iter().any(|id| forward.contains_key(id))
1578        };
1579        let needs_reverse_cleanup = {
1580            let reverse = self.reverse_refs.read();
1581            id_set.iter().any(|id| reverse.contains_key(id))
1582        };
1583
1584        if !needs_forward_cleanup && !needs_reverse_cleanup {
1585            self.unindex_cross_refs_fast_path
1586                .fetch_add(1, Ordering::Relaxed);
1587            return Ok(());
1588        }
1589
1590        if needs_forward_cleanup {
1591            let mut forward = self.cross_refs.write();
1592            for id in &id_set {
1593                forward.remove(id);
1594            }
1595        }
1596
1597        if needs_reverse_cleanup || needs_forward_cleanup {
1598            // The inner `values_mut()` scan is only needed when at least one
1599            // deleted id was a *source* somewhere — which, by the invariant
1600            // above, can only happen when it had outbound forward refs.
1601            let mut reverse = self.reverse_refs.write();
1602            if needs_forward_cleanup {
1603                for refs in reverse.values_mut() {
1604                    refs.retain(|(source, _, _)| !id_set.contains(source));
1605                }
1606            }
1607            reverse.retain(|target, refs| !id_set.contains(target) && !refs.is_empty());
1608        }
1609        self.mark_paged_registry_dirty();
1610
1611        Ok(())
1612    }
1613
1614    /// Test/observability hook: number of times `unindex_cross_refs_batch`
1615    /// took the read-only fast path. Used to pin the early-exit in tests
1616    /// and as a cheap signal in delete-heavy benchmarks.
1617    pub fn unindex_cross_refs_fast_path_hits(&self) -> u64 {
1618        self.unindex_cross_refs_fast_path.load(Ordering::Relaxed)
1619    }
1620
1621    /// Query across all collections with a filter
1622    pub fn query_all<F>(&self, filter: F) -> Vec<(String, UnifiedEntity)>
1623    where
1624        F: Fn(&UnifiedEntity) -> bool + Clone + Send + Sync,
1625    {
1626        // Issue #815 — snapshot the (name, manager) handles under a brief
1627        // read lock, then release the `collections` lock *before* the
1628        // potentially long per-collection scan. Holding `collections.read()`
1629        // across a full-store scan let a large `/catalog` snapshot block (and
1630        // be blocked by) the replica apply loop's collection-creation writes:
1631        // under parking_lot a writer waiting behind the long read parks every
1632        // subsequent reader, so `/catalog` and readiness wedged during a big
1633        // WAL re-apply. The managers are `Arc`-shared with their own internal
1634        // locks, so scanning them off the map lock is safe and keeps the HTTP
1635        // surface responsive.
1636        let pairs: Vec<(String, Arc<SegmentManager>)> = {
1637            let collections = self.collections.read();
1638            collections
1639                .iter()
1640                .map(|(name, mgr)| (name.clone(), Arc::clone(mgr)))
1641                .collect()
1642        };
1643
1644        let use_parallel = pairs.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1645        if !use_parallel {
1646            // Single collection — no parallelism overhead
1647            return pairs
1648                .into_iter()
1649                .flat_map(|(name, mgr)| {
1650                    mgr.query_all(filter.clone())
1651                        .into_iter()
1652                        .map(move |e| (name.clone(), e))
1653                })
1654                .collect();
1655        }
1656
1657        // Multiple collections — scan in parallel
1658        let filter_ref = &filter;
1659        let collection_results: Vec<Vec<(String, UnifiedEntity)>> = std::thread::scope(|s| {
1660            pairs
1661                .iter()
1662                .map(|(name, manager)| {
1663                    let name = (*name).clone();
1664                    s.spawn(move || {
1665                        manager
1666                            .query_all(|e| filter_ref(e))
1667                            .into_iter()
1668                            .map(|e| (name.clone(), e))
1669                            .collect::<Vec<_>>()
1670                    })
1671                })
1672                .collect::<Vec<_>>()
1673                .into_iter()
1674                .map(|h| h.join().unwrap_or_default())
1675                .collect()
1676        });
1677
1678        collection_results.into_iter().flatten().collect()
1679    }
1680
1681    /// Scan every collection under an explicit MVCC snapshot.
1682    pub fn scan<F>(
1683        &self,
1684        snapshot: Option<&crate::runtime::impl_core::SnapshotContext>,
1685        filter: F,
1686    ) -> Vec<(String, UnifiedEntity)>
1687    where
1688        F: Fn(&UnifiedEntity) -> bool + Clone + Send + Sync,
1689    {
1690        let pairs: Vec<(String, Arc<SegmentManager>)> = {
1691            let collections = self.collections.read();
1692            collections
1693                .iter()
1694                .map(|(name, manager)| (name.clone(), Arc::clone(manager)))
1695                .collect()
1696        };
1697
1698        let use_parallel = pairs.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1699        if !use_parallel {
1700            return pairs
1701                .into_iter()
1702                .flat_map(|(name, manager)| {
1703                    manager
1704                        .scan(snapshot, filter.clone())
1705                        .into_iter()
1706                        .map(move |entity| (name.clone(), entity))
1707                })
1708                .collect();
1709        }
1710
1711        let filter_ref = &filter;
1712        let collection_results: Vec<Vec<(String, UnifiedEntity)>> = std::thread::scope(|scope| {
1713            pairs
1714                .iter()
1715                .map(|(name, manager)| {
1716                    let name = name.clone();
1717                    scope.spawn(move || {
1718                        manager
1719                            .scan(snapshot, |entity| filter_ref(entity))
1720                            .into_iter()
1721                            .map(|entity| (name.clone(), entity))
1722                            .collect::<Vec<_>>()
1723                    })
1724                })
1725                .collect::<Vec<_>>()
1726                .into_iter()
1727                .map(|handle| handle.join().unwrap_or_default())
1728                .collect()
1729        });
1730
1731        collection_results.into_iter().flatten().collect()
1732    }
1733
1734    /// Filter by metadata across all collections
1735    pub fn filter_metadata_all(
1736        &self,
1737        filters: &[(String, MetadataFilter)],
1738    ) -> Vec<(String, EntityId)> {
1739        let mut results = Vec::new();
1740        let collections = self.collections.read();
1741
1742        for (name, manager) in collections.iter() {
1743            for id in manager.filter_metadata(filters) {
1744                results.push((name.clone(), id));
1745            }
1746        }
1747
1748        results
1749    }
1750
1751    /// Get statistics
1752    pub fn stats(&self) -> StoreStats {
1753        // Issue #815 — snapshot the manager handles, drop the `collections`
1754        // read lock, then gather per-collection stats off the map lock. The
1755        // readiness probe (via `health()`) walks `stats()`, so holding the
1756        // store lock across a full-store stats scan was part of the same
1757        // wedge that froze `/catalog` during a large replica re-apply.
1758        let pairs: Vec<(String, Arc<SegmentManager>)> = {
1759            let collections = self.collections.read();
1760            collections
1761                .iter()
1762                .map(|(name, mgr)| (name.clone(), Arc::clone(mgr)))
1763                .collect()
1764        };
1765
1766        let mut stats = StoreStats {
1767            collection_count: pairs.len(),
1768            ..Default::default()
1769        };
1770
1771        for (name, manager) in &pairs {
1772            let manager_stats = manager.stats();
1773            stats.total_entities += manager_stats.total_entities;
1774            stats.total_memory_bytes += manager_stats.total_memory_bytes;
1775            stats.collections.insert(name.clone(), manager_stats);
1776        }
1777
1778        stats
1779    }
1780
1781    /// Run maintenance on all collections
1782    pub fn run_maintenance(&self) -> Result<(), StoreError> {
1783        let collections = self.collections.read();
1784        for manager in collections.values() {
1785            manager.run_maintenance()?;
1786        }
1787        Ok(())
1788    }
1789
1790    /// Bytes pressure consolidation could plausibly return across collections.
1791    pub fn reclaimable_segment_bytes(&self) -> u64 {
1792        let collections = self.collections.read();
1793        collections
1794            .values()
1795            .map(|manager| manager.reclaimable_bytes())
1796            .fold(0, u64::saturating_add)
1797    }
1798
1799    /// Advance pressure-triggered consolidation by one paced tick per
1800    /// collection. Returns true if any collection started or advanced a run.
1801    pub fn pressure_consolidation_tick(&self) -> bool {
1802        let managers: Vec<_> = self.collections.read().values().cloned().collect();
1803        // A plain loop, not `any()`: every collection must receive its tick —
1804        // short-circuiting on the first advanced run would starve the rest.
1805        let mut advanced = false;
1806        for manager in &managers {
1807            advanced |= manager.pressure_consolidation_tick();
1808        }
1809        advanced
1810    }
1811}
1812
1813/// Flatten a JSON value into dot-notation key-value pairs for red_config.
1814fn flatten_config_json(
1815    prefix: &str,
1816    value: &crate::serde_json::Value,
1817    out: &mut Vec<(String, reddb_types::Value)>,
1818) {
1819    use reddb_types::Value;
1820    match value {
1821        crate::serde_json::Value::Object(map) => {
1822            for (k, v) in map {
1823                let key = if prefix.is_empty() {
1824                    k.clone()
1825                } else {
1826                    format!("{prefix}.{k}")
1827                };
1828                flatten_config_json(&key, v, out);
1829            }
1830        }
1831        crate::serde_json::Value::String(s) => {
1832            out.push((prefix.to_string(), Value::text(s.clone())));
1833        }
1834        crate::serde_json::Value::Integer(n) => {
1835            out.push((prefix.to_string(), Value::Integer(*n)));
1836        }
1837        crate::serde_json::Value::Number(n) => {
1838            if n.fract().abs() < f64::EPSILON {
1839                out.push((prefix.to_string(), Value::UnsignedInteger(*n as u64)));
1840            } else {
1841                out.push((prefix.to_string(), Value::Float(*n)));
1842            }
1843        }
1844        crate::serde_json::Value::Decimal(n) => {
1845            out.push((prefix.to_string(), Value::DecimalText(n.clone())));
1846        }
1847        crate::serde_json::Value::Bool(b) => {
1848            out.push((prefix.to_string(), Value::Boolean(*b)));
1849        }
1850        crate::serde_json::Value::Null => {
1851            out.push((prefix.to_string(), Value::Null));
1852        }
1853        crate::serde_json::Value::Array(arr) => {
1854            let json_str = crate::serde_json::to_string(value).unwrap_or_default();
1855            out.push((prefix.to_string(), Value::text(json_str)));
1856        }
1857    }
1858}