1use super::*;
2
3impl UnifiedStore {
4 pub(crate) fn persist_entities_to_pager(
5 &self,
6 collection: &str,
7 entities: &[UnifiedEntity],
8 ) -> Result<(), StoreError> {
9 self.persist_entities_impl(collection, entities, false)
10 }
11
12 pub(crate) fn persist_entity_refs_to_pager(
16 &self,
17 collection: &str,
18 entities: &[&UnifiedEntity],
19 ) -> Result<(), StoreError> {
20 self.persist_entity_refs_impl(collection, entities, false)
21 }
22
23 pub(crate) fn persist_entity_refs_to_pager_wal_only(
26 &self,
27 collection: &str,
28 entities: &[&UnifiedEntity],
29 ) -> Result<(), StoreError> {
30 self.persist_entity_refs_impl(collection, entities, true)
31 }
32
33 fn persist_entity_refs_impl(
34 &self,
35 collection: &str,
36 entities: &[&UnifiedEntity],
37 skip_btree: bool,
38 ) -> Result<(), StoreError> {
39 if entities.is_empty() {
40 return Ok(());
41 }
42 if self.pager.is_none() && self.config.embedded_wal_path.is_none() {
43 return Ok(());
44 }
45 let fv = self.format_version();
46 let manager = self
47 .get_collection(collection)
48 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
49 let mut serialized: Vec<(Vec<u8>, Vec<u8>)> = entities
50 .iter()
51 .map(|entity| {
52 let metadata = manager.get_metadata(entity.id);
53 (
54 entity.id.raw().to_be_bytes().to_vec(),
55 Self::serialize_entity_record(entity, metadata.as_ref(), fv),
56 )
57 })
58 .collect();
59 serialized.sort_by(|a, b| a.0.cmp(&b.0));
60
61 if !skip_btree {
62 let Some(pager) = &self.pager else {
63 let records: Vec<Vec<u8>> =
64 serialized.into_iter().map(|(_id, record)| record).collect();
65 self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
66 collection: collection.to_string(),
67 records,
68 }])?;
69 return Ok(());
70 };
71 let mut btree_indices = self.btree_indices.write();
72 let btree = btree_indices
73 .entry(collection.to_string())
74 .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
75 let root_before = btree.root_page_id();
76 btree.upsert_batch_sorted(&serialized).map_err(|e| {
77 StoreError::Io(std::io::Error::other(format!("B-tree upsert error: {}", e)))
78 })?;
79 let root_after = btree.root_page_id();
80 drop(btree_indices);
81 if root_before != root_after {
82 self.mark_paged_registry_dirty();
83 }
84 }
85 let records: Vec<Vec<u8>> = serialized.into_iter().map(|(_id, record)| record).collect();
86 self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
87 collection: collection.to_string(),
88 records,
89 }])?;
90 Ok(())
91 }
92
93 pub(crate) fn persist_entities_to_pager_wal_only(
105 &self,
106 collection: &str,
107 entities: &[UnifiedEntity],
108 ) -> Result<(), StoreError> {
109 self.persist_entities_impl(collection, entities, true)
110 }
111
112 fn persist_entities_impl(
113 &self,
114 collection: &str,
115 entities: &[UnifiedEntity],
116 skip_btree: bool,
117 ) -> Result<(), StoreError> {
118 if entities.is_empty() {
119 return Ok(());
120 }
121
122 if self.pager.is_none() && self.config.embedded_wal_path.is_none() {
123 return Ok(());
124 }
125
126 let fv = self.format_version();
127 let manager = self
128 .get_collection(collection)
129 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
130 let mut serialized: Vec<(Vec<u8>, Vec<u8>)> = entities
131 .iter()
132 .map(|entity| {
133 let metadata = manager.get_metadata(entity.id);
134 (
135 entity.id.raw().to_be_bytes().to_vec(),
136 Self::serialize_entity_record(entity, metadata.as_ref(), fv),
137 )
138 })
139 .collect();
140 serialized.sort_by(|a, b| a.0.cmp(&b.0));
142
143 if !skip_btree {
144 let Some(pager) = &self.pager else {
145 let records: Vec<Vec<u8>> =
146 serialized.into_iter().map(|(_id, record)| record).collect();
147 self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
148 collection: collection.to_string(),
149 records,
150 }])?;
151 return Ok(());
152 };
153 let mut btree_indices = self.btree_indices.write();
154 let btree = btree_indices
155 .entry(collection.to_string())
156 .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
157 let root_before = btree.root_page_id();
158
159 btree.upsert_batch_sorted(&serialized).map_err(|e| {
163 StoreError::Io(std::io::Error::other(format!("B-tree upsert error: {}", e)))
164 })?;
165 let root_after = btree.root_page_id();
166 drop(btree_indices);
167 if root_before != root_after {
168 self.mark_paged_registry_dirty();
169 }
170 }
171 let records: Vec<Vec<u8>> = serialized.into_iter().map(|(_id, record)| record).collect();
181 self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
182 collection: collection.to_string(),
183 records,
184 }])?;
185
186 Ok(())
187 }
188
189 pub(crate) fn update_graph_label_index(
191 &self,
192 collection: &str,
193 label: &str,
194 entity_id: EntityId,
195 ) {
196 let key = (collection.to_string(), label.to_string());
197 let mut idx = self.graph_label_index.write();
198 idx.entry(key).or_default().push(entity_id);
199 }
200
201 pub(crate) fn remove_from_graph_label_index(&self, collection: &str, entity_id: EntityId) {
203 let mut idx = self.graph_label_index.write();
204 for ((col, _), ids) in idx.iter_mut() {
205 if col == collection {
206 ids.retain(|&id| id != entity_id);
207 }
208 }
209 idx.retain(|_, ids| !ids.is_empty());
211 }
212
213 pub(crate) fn remove_from_graph_label_index_batch(
214 &self,
215 collection: &str,
216 entity_ids: &[EntityId],
217 ) {
218 if entity_ids.is_empty() {
219 return;
220 }
221 let id_set: std::collections::HashSet<EntityId> = entity_ids.iter().copied().collect();
222 let mut idx = self.graph_label_index.write();
223 for ((col, _), ids) in idx.iter_mut() {
224 if col == collection {
225 ids.retain(|id| !id_set.contains(id));
226 }
227 }
228 idx.retain(|_, ids| !ids.is_empty());
229 }
230
231 pub fn lookup_graph_nodes_by_label(&self, label: &str) -> Vec<EntityId> {
233 let idx = self.graph_label_index.read();
234 idx.iter()
235 .filter(|((_, l), _)| l == label)
236 .flat_map(|(_, ids)| ids.iter().copied())
237 .collect()
238 }
239
240 pub fn lookup_graph_nodes_by_label_in(&self, collection: &str, label: &str) -> Vec<EntityId> {
242 let idx = self.graph_label_index.read();
243 idx.get(&(collection.to_string(), label.to_string()))
244 .cloned()
245 .unwrap_or_default()
246 }
247
248 pub fn create_collection(&self, name: impl Into<String>) -> Result<(), StoreError> {
249 let name = name.into();
250 let mut collections = self.collections.write();
251
252 if collections.contains_key(&name) {
253 return Err(StoreError::CollectionExists(name));
254 }
255
256 let manager = SegmentManager::with_config(&name, self.config.manager_config.clone());
257 collections.insert(name.clone(), Arc::new(manager));
258 drop(collections);
259
260 if let Err(err) = self.publish_operational_collection_create(&name) {
261 let mut collections = self.collections.write();
262 collections.remove(&name);
263 return Err(err);
264 }
265
266 self.mark_paged_registry_dirty();
267 self.finish_paged_write([StoreWalAction::CreateCollection { name }])?;
268
269 Ok(())
270 }
271
272 pub fn get_or_create_collection(&self, name: impl Into<String>) -> Arc<SegmentManager> {
274 let name = name.into();
275 {
277 let collections = self.collections.read();
278 if let Some(manager) = collections.get(&name) {
279 return Arc::clone(manager);
280 }
281 }
282 let mut collections = self.collections.write();
284 if let Some(manager) = collections.get(&name) {
286 return Arc::clone(manager);
287 }
288 let manager = Arc::new(SegmentManager::with_config(
289 &name,
290 self.config.manager_config.clone(),
291 ));
292 collections.insert(name, Arc::clone(&manager));
293 self.mark_paged_registry_dirty();
294 manager
295 }
296
297 pub fn get_collection(&self, name: &str) -> Option<Arc<SegmentManager>> {
299 self.collections.read().get(name).map(Arc::clone)
300 }
301
302 pub fn context_index(&self) -> &ContextIndex {
304 &self.context_index
305 }
306
307 pub fn take_replayed_turbo_inserts(&self, collection: &str) -> Option<Vec<(u64, Vec<f32>)>> {
313 let mut map = self.replayed_turbo_inserts.lock();
314 map.remove(collection)
315 }
316
317 pub(crate) fn take_replayed_probabilistic_deltas(&self) -> Vec<(u8, u8, String, Vec<Vec<u8>>)> {
318 let mut deltas = self.replayed_probabilistic_deltas.lock();
319 std::mem::take(&mut *deltas)
320 }
321
322 pub fn set_config_tree(&self, prefix: &str, json: &crate::serde_json::Value) -> usize {
325 let _ = self.get_or_create_collection("red_config");
326 let mut pairs = Vec::new();
327 flatten_config_json(prefix, json, &mut pairs);
328 let mut saved = 0;
329 for (key, value) in pairs {
330 let entity = UnifiedEntity::new(
331 EntityId::new(0),
332 EntityKind::TableRow {
333 table: Arc::from("red_config"),
334 row_id: 0,
335 },
336 EntityData::Row(RowData {
337 columns: Vec::new(),
338 named: Some(
339 [
340 ("key".to_string(), crate::storage::schema::Value::text(key)),
341 ("value".to_string(), value),
342 ]
343 .into_iter()
344 .collect(),
345 ),
346 schema: None,
347 }),
348 );
349 if self.insert_auto("red_config", entity).is_ok() {
350 saved += 1;
351 }
352 }
353 saved
354 }
355
356 pub fn get_config(&self, key: &str) -> Option<crate::storage::schema::Value> {
358 let manager = self.get_collection("red_config")?;
359 for entity in manager.query_all(|_| true) {
360 if let EntityData::Row(row) = &entity.data {
361 if let Some(named) = &row.named {
362 let key_matches = named
363 .get("key")
364 .and_then(|v| match v {
365 crate::storage::schema::Value::Text(s) => Some(s.as_ref() == key),
366 _ => None,
367 })
368 .unwrap_or(false);
369 if key_matches {
370 return named.get("value").cloned();
371 }
372 }
373 }
374 }
375 None
376 }
377
378 pub fn set_aux_metadata(&self, bytes: Vec<u8>) {
383 *self.aux_metadata.write() = bytes;
384 }
385
386 pub fn aux_metadata(&self) -> Vec<u8> {
388 self.aux_metadata.read().clone()
389 }
390
391 pub fn list_collections(&self) -> Vec<String> {
393 self.collections.read().keys().cloned().collect()
394 }
395
396 pub fn drop_collection(&self, name: &str) -> Result<(), StoreError> {
398 self.publish_operational_collection_pending_drop(name)?;
399 let manager = {
400 let mut collections = self.collections.write();
401
402 collections
403 .remove(name)
404 .ok_or_else(|| StoreError::CollectionNotFound(name.to_string()))?
405 };
406
407 let entities = manager.query_all(|_| true);
408 let entity_ids: Vec<EntityId> = entities.iter().map(|entity| entity.id).collect();
409
410 for entity_id in &entity_ids {
411 self.context_index.remove_entity(*entity_id);
412 let _ = self.unindex_cross_refs(*entity_id);
413 }
414
415 self.btree_indices.write().remove(name);
416
417 self.entity_cache.retain(|entity_id, (collection, _)| {
418 collection != name && !entity_ids.iter().any(|id| id.raw() == entity_id)
419 });
420
421 self.cross_refs.write().retain(|source_id, refs| {
422 refs.retain(|(target_id, _, target_collection)| {
423 target_collection != name && !entity_ids.iter().any(|id| id == target_id)
424 });
425 !entity_ids.iter().any(|id| id == source_id)
426 });
427
428 self.reverse_refs.write().retain(|target_id, refs| {
429 refs.retain(|(source_id, _, source_collection)| {
430 source_collection != name && !entity_ids.iter().any(|id| id == source_id)
431 });
432 !entity_ids.iter().any(|id| id == target_id)
433 });
434
435 self.mark_paged_registry_dirty();
436 self.finish_paged_write([StoreWalAction::DropCollection {
437 name: name.to_string(),
438 }])?;
439 self.publish_operational_collection_drop_finished(name)?;
440
441 Ok(())
442 }
443
444 pub fn insert(&self, collection: &str, entity: UnifiedEntity) -> Result<EntityId, StoreError> {
446 let manager = self
447 .get_collection(collection)
448 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
449
450 let mut entity = entity;
451 if entity.id.raw() == 0 {
452 entity.id = self.next_entity_id();
453 } else {
454 self.register_entity_id(entity.id);
455 }
456 if let EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
458 if *row_id == 0 {
459 *row_id = manager.next_row_id();
460 } else {
461 manager.register_row_id(*row_id);
462 }
463 }
464 entity.ensure_table_logical_id();
465 let graph_node_label: Option<String> = if let EntityKind::GraphNode(ref node) = entity.kind
467 {
468 Some(node.label.clone())
469 } else {
470 None
471 };
472
473 let id = manager.insert(entity)?;
474 self.register_entity_id(id);
475
476 if let Some(ref label) = graph_node_label {
478 self.update_graph_label_index(collection, label, id);
479 }
480
481 let mut registry_dirty = false;
483 if let Some(pager) = &self.pager {
484 if let Some(entity) = manager.get(id) {
485 let mut btree_indices = self.btree_indices.write();
486 let btree = btree_indices
487 .entry(collection.to_string())
488 .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
489 let root_before = btree.root_page_id();
490
491 let key = id.raw().to_be_bytes();
492 let metadata = manager.get_metadata(id);
493 let value = Self::serialize_entity_record(
494 &entity,
495 metadata.as_ref(),
496 self.format_version(),
497 );
498 let _ = btree.insert(&key, &value);
500 registry_dirty = root_before != btree.root_page_id();
501 }
502 }
503
504 if self.config.auto_index_refs {
506 if let Some(entity) = manager.get(id) {
507 self.index_cross_refs(&entity, collection)?;
508 }
509 }
510
511 if self.pager.is_some() || self.config.embedded_wal_path.is_some() {
515 let actions = manager
516 .get(id)
517 .map(|entity| {
518 let metadata = manager.get_metadata(id);
519 vec![StoreWalAction::upsert_entity(
520 collection,
521 &entity,
522 metadata.as_ref(),
523 self.format_version(),
524 )]
525 })
526 .unwrap_or_default();
527 if registry_dirty {
528 self.mark_paged_registry_dirty();
529 }
530 self.finish_paged_write(actions)?;
531 }
532
533 Ok(id)
534 }
535
536 pub fn bulk_insert(
541 &self,
542 collection: &str,
543 mut entities: Vec<UnifiedEntity>,
544 ) -> Result<Vec<EntityId>, StoreError> {
545 let trace = matches!(
549 std::env::var("REDDB_BULK_TIMING").ok().as_deref(),
550 Some("1") | Some("true") | Some("on")
551 );
552 let t_start = std::time::Instant::now();
553 let n = entities.len();
554 let manager = self.get_or_create_collection(collection);
555 let t_get_coll = t_start.elapsed();
556
557 let t0 = std::time::Instant::now();
566 let n_missing_entity_ids = entities.iter().filter(|e| e.id.raw() == 0).count() as u64;
567 let n_missing_row_ids = entities
568 .iter()
569 .filter(|e| matches!(e.kind, EntityKind::TableRow { row_id: 0, .. }))
570 .count() as u64;
571 let mut entity_id_range = if n_missing_entity_ids > 0 {
572 self.reserve_entity_ids(n_missing_entity_ids)
573 } else {
574 0..0
575 };
576 let mut row_id_range = if n_missing_row_ids > 0 {
577 manager.reserve_row_ids(n_missing_row_ids)
578 } else {
579 0..0
580 };
581 for entity in &mut entities {
582 if entity.id.raw() == 0 {
583 let next = entity_id_range
584 .next()
585 .expect("reserved entity-id range exhausted");
586 entity.id = EntityId::new(next);
587 } else {
588 self.register_entity_id(entity.id);
589 }
590 if let EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
591 if *row_id == 0 {
592 *row_id = row_id_range
593 .next()
594 .expect("reserved row-id range exhausted");
595 } else {
596 manager.register_row_id(*row_id);
597 }
598 }
599 entity.ensure_table_logical_id();
600 }
601 let t_assign_ids = t0.elapsed();
602
603 let graph_labels: Vec<Option<(String, EntityId)>> = entities
605 .iter()
606 .map(|e| {
607 if let EntityKind::GraphNode(ref node) = e.kind {
608 Some((node.label.clone(), e.id))
609 } else {
610 None
611 }
612 })
613 .collect();
614
615 let t0 = std::time::Instant::now();
622 let serialized: Option<Vec<(Vec<u8>, Vec<u8>)>> =
623 if self.pager.is_some() || self.config.embedded_wal_path.is_some() {
624 let fv = self.format_version();
625 let serial_map = |e: &UnifiedEntity| {
626 (
627 e.id.raw().to_be_bytes().to_vec(),
628 Self::serialize_entity_record(e, None, fv),
629 )
630 };
631 if entities.len() >= 512 {
637 use rayon::prelude::*;
638 Some(entities.par_iter().map(serial_map).collect())
639 } else {
640 Some(entities.iter().map(serial_map).collect())
641 }
642 } else {
643 None
644 };
645 let t_serialize = t0.elapsed();
646
647 let t0 = std::time::Instant::now();
649 let ids = manager.bulk_insert(entities)?;
650 let t_manager = t0.elapsed();
651 for id in &ids {
652 self.register_entity_id(*id);
653 }
654
655 for (label, entity_id) in graph_labels.iter().flatten() {
657 self.update_graph_label_index(collection, label, *entity_id);
658 }
659
660 let skip_btree_requested = matches!(
670 std::env::var("REDDB_BULK_SKIP_PERSIST_UNSAFE")
671 .ok()
672 .as_deref(),
673 Some("1") | Some("true") | Some("on")
674 );
675 let skip_btree = skip_btree_requested && self.pager.is_none();
678 if skip_btree_requested && !skip_btree {
679 static IGNORED: std::sync::OnceLock<()> = std::sync::OnceLock::new();
681 IGNORED.get_or_init(|| {
682 tracing::warn!(
683 "REDDB_BULK_SKIP_PERSIST_UNSAFE set but durable pager is \
684 active — flag ignored; bulk inserts will be persisted normally"
685 );
686 });
687 } else if skip_btree {
688 static WARNED: std::sync::OnceLock<()> = std::sync::OnceLock::new();
690 WARNED.get_or_init(|| {
691 tracing::warn!(
692 "REDDB_BULK_SKIP_PERSIST_UNSAFE set (ephemeral/no-pager mode) — \
693 bulk inserts NOT durable; data will be lost on restart"
694 );
695 });
696 }
697
698 let mut t_btree_lock = std::time::Duration::ZERO;
702 let mut t_btree_insert = std::time::Duration::ZERO;
703 let mut t_flush = std::time::Duration::ZERO;
704 if !skip_btree {
705 if let (Some(pager), Some(batch)) = (&self.pager, serialized.as_ref()) {
706 let t0 = std::time::Instant::now();
707 let mut btree_indices = self.btree_indices.write();
708 let btree = btree_indices
709 .entry(collection.to_string())
710 .or_insert_with(|| Arc::new(BTree::new(Arc::clone(pager))));
711 let root_before = btree.root_page_id();
712 t_btree_lock = t0.elapsed();
713
714 let t0 = std::time::Instant::now();
715 let _ = btree.bulk_insert_sorted(batch);
716 t_btree_insert = t0.elapsed();
717 let registry_dirty = root_before != btree.root_page_id();
718
719 let t0 = std::time::Instant::now();
720 if registry_dirty {
721 self.mark_paged_registry_dirty();
722 }
723 t_flush = t0.elapsed();
724 }
725 }
726
727 let actions = serialized
733 .map(|batch| {
734 let records: Vec<Vec<u8>> =
735 batch.into_iter().map(|(_key, record)| record).collect();
736 vec![StoreWalAction::BulkUpsertEntityRecords {
737 collection: collection.to_string(),
738 records,
739 }]
740 })
741 .unwrap_or_default();
742 self.finish_paged_write(actions)?;
743
744 if trace {
745 tracing::debug!(
746 n,
747 total = ?t_start.elapsed(),
748 get_coll = ?t_get_coll,
749 assign = ?t_assign_ids,
750 serialize = ?t_serialize,
751 manager = ?t_manager,
752 btree_lock = ?t_btree_lock,
753 btree = ?t_btree_insert,
754 flush = ?t_flush,
755 "bulk_insert timing"
756 );
757 }
758
759 Ok(ids)
760 }
761
762 pub fn insert_auto(
764 &self,
765 collection: &str,
766 entity: UnifiedEntity,
767 ) -> Result<EntityId, StoreError> {
768 let manager = self.get_or_create_collection(collection);
769 let mut entity = entity;
770 if entity.id.raw() == 0 {
771 entity.id = self.next_entity_id();
772 } else {
773 self.register_entity_id(entity.id);
774 }
775 if let EntityKind::TableRow { ref mut row_id, .. } = entity.kind {
777 if *row_id == 0 {
778 *row_id = manager.next_row_id();
779 } else {
780 manager.register_row_id(*row_id);
781 }
782 }
783 entity.ensure_table_logical_id();
784
785 let graph_node_label: Option<String> = if let EntityKind::GraphNode(ref node) = entity.kind
787 {
788 Some(node.label.clone())
789 } else {
790 None
791 };
792 self.context_index.index_entity(collection, &entity);
794
795 let id_for_serialize = entity.id;
810 let serialized_record: Option<Vec<u8>> =
811 if self.pager.is_some() || self.config.embedded_wal_path.is_some() {
812 Some(Self::serialize_entity_record(
813 &entity,
814 None,
815 self.format_version(),
816 ))
817 } else {
818 None
819 };
820 if self.config.auto_index_refs {
821 self.index_cross_refs(&entity, collection)?;
822 }
823
824 let id = manager.insert(entity)?;
825 debug_assert_eq!(id, id_for_serialize);
826 if let Some(ref label) = graph_node_label {
834 self.update_graph_label_index(collection, label, id);
835 }
836
837 let mut registry_dirty = false;
838 if let (Some(_pager), Some(record)) = (&self.pager, serialized_record.as_ref()) {
839 if let Some(btree) = self.get_or_create_btree(collection) {
840 let root_before = btree.root_page_id();
841
842 let key = id.raw().to_be_bytes();
843 btree.insert(&key, record).map_err(|e| {
844 StoreError::Io(std::io::Error::other(format!(
845 "B-tree insert error while inserting '{collection}'/{id}: {e}"
846 )))
847 })?;
848 registry_dirty = root_before != btree.root_page_id();
849 }
850 }
851
852 if self.pager.is_some() || self.config.embedded_wal_path.is_some() {
856 let actions = serialized_record
857 .map(|record| {
858 vec![StoreWalAction::UpsertEntityRecord {
859 collection: collection.to_string(),
860 record,
861 }]
862 })
863 .unwrap_or_default();
864 if registry_dirty {
865 self.mark_paged_registry_dirty();
866 }
867 self.finish_paged_write(actions)?;
868 }
869 Ok(id)
870 }
871
872 pub fn get(&self, collection: &str, id: EntityId) -> Option<UnifiedEntity> {
878 if let Some(entity) = self
880 .get_collection(collection)
881 .and_then(|manager| manager.get(id))
882 {
883 return Some(entity);
884 }
885
886 if self.pager.is_some() {
888 let btree_indices = self.btree_indices.read();
889 if let Some(btree) = btree_indices.get(collection) {
890 let key = id.raw().to_be_bytes();
891 if let Ok(Some(value)) = btree.get(&key) {
892 if let Ok((entity, _)) =
893 Self::deserialize_entity_record(&value, self.format_version())
894 {
895 return Some(entity);
896 }
897 }
898 }
899 }
900
901 None
902 }
903
904 pub fn get_table_row_by_logical_id(
910 &self,
911 collection: &str,
912 logical_id: EntityId,
913 ) -> Option<UnifiedEntity> {
914 if let Some(entity) = self.get(collection, logical_id) {
915 if matches!(entity.kind, EntityKind::TableRow { .. })
916 && entity.logical_id() == logical_id
917 && entity.xmax == 0
918 {
919 return Some(entity);
920 }
921 }
922
923 let manager = self.get_collection(collection)?;
924 let mut matches = manager.query_all(|entity| {
925 matches!(entity.kind, EntityKind::TableRow { .. }) && entity.logical_id() == logical_id
926 });
927 matches
928 .iter()
929 .find(|entity| entity.xmax == 0)
930 .cloned()
931 .or_else(|| matches.pop())
932 }
933
934 pub fn table_row_versions_by_logical_id(
935 &self,
936 collection: &str,
937 logical_id: EntityId,
938 ) -> Vec<UnifiedEntity> {
939 self.get_collection(collection)
940 .map(|manager| {
941 manager.query_all(|entity| {
942 matches!(entity.kind, EntityKind::TableRow { .. })
943 && entity.logical_id() == logical_id
944 })
945 })
946 .unwrap_or_default()
947 }
948
949 pub fn vacuum_mvcc_history(
950 &self,
951 collection: &str,
952 cutoff_xid: u64,
953 ) -> Result<MvccVacuumStats, StoreError> {
954 let manager = self
955 .get_collection(collection)
956 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
957 let entities =
958 manager.query_all(|entity| matches!(entity.kind, EntityKind::TableRow { .. }));
959 let mut logical = HashMap::<EntityId, (bool, u64)>::new();
960 for entity in &entities {
961 let entry = logical.entry(entity.logical_id()).or_insert((false, 0));
962 if entity.xmax == 0 {
963 entry.0 = true;
964 }
965 entry.1 = entry.1.max(entity.xmin);
966 }
967
968 let mut stats = MvccVacuumStats::default();
969 let mut reclaim_ids = Vec::new();
970 for entity in entities {
971 if entity.xmax == 0 {
972 continue;
973 }
974 stats.scanned_versions += 1;
975 let (has_live_version, max_xmin) = logical
976 .get(&entity.logical_id())
977 .copied()
978 .unwrap_or((false, entity.xmin));
979 let is_delete_tombstone = !has_live_version && entity.xmin == max_xmin;
980 if entity.xmax < cutoff_xid {
981 stats.reclaimed_versions += 1;
982 if is_delete_tombstone {
983 stats.reclaimed_tombstones += 1;
984 } else {
985 stats.reclaimed_history_versions += 1;
986 }
987 reclaim_ids.push(entity.id);
988 } else {
989 stats.retained_versions += 1;
990 if is_delete_tombstone {
991 stats.retained_tombstones += 1;
992 } else {
993 stats.retained_history_versions += 1;
994 }
995 }
996 }
997
998 if !reclaim_ids.is_empty() {
999 self.delete_batch(collection, &reclaim_ids)?;
1000 }
1001 Ok(stats)
1002 }
1003
1004 pub(crate) fn install_versioned_table_row_update(
1005 &self,
1006 collection: &str,
1007 old_version: UnifiedEntity,
1008 mut new_version: UnifiedEntity,
1009 metadata: Option<&Metadata>,
1010 ) -> Result<(), StoreError> {
1011 let manager = self
1012 .get_collection(collection)
1013 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1014
1015 let old_id = old_version.id;
1016 let new_id = new_version.id;
1017 let inherited_metadata = metadata.cloned().or_else(|| manager.get_metadata(old_id));
1018
1019 self.entity_cache.remove(old_id.raw());
1020 self.entity_cache.remove(new_id.raw());
1021 manager.update(old_version.clone())?;
1022 self.context_index.remove_entity(old_id);
1023
1024 self.register_entity_id(new_version.id);
1025 if let EntityKind::TableRow { ref mut row_id, .. } = new_version.kind {
1026 if *row_id == 0 {
1027 *row_id = manager.next_row_id();
1028 } else {
1029 manager.register_row_id(*row_id);
1030 }
1031 }
1032 new_version.ensure_table_logical_id();
1033 manager.insert(new_version.clone())?;
1034 if let Some(metadata) = inherited_metadata {
1035 manager.set_metadata(new_id, metadata)?;
1036 }
1037 self.context_index.index_entity(collection, &new_version);
1038 if self.config.auto_index_refs {
1039 self.index_cross_refs(&new_version, collection)?;
1040 }
1041
1042 let old_metadata = manager.get_metadata(old_id);
1043 let new_metadata = manager.get_metadata(new_id);
1044 let fv = self.format_version();
1045 let records = vec![
1046 Self::serialize_entity_record(&old_version, old_metadata.as_ref(), fv),
1047 Self::serialize_entity_record(&new_version, new_metadata.as_ref(), fv),
1048 ];
1049 self.finish_paged_write([StoreWalAction::BulkUpsertEntityRecords {
1050 collection: collection.to_string(),
1051 records,
1052 }])?;
1053
1054 Ok(())
1055 }
1056
1057 pub fn get_batch(&self, collection: &str, ids: &[EntityId]) -> Vec<Option<UnifiedEntity>> {
1063 match self.get_collection(collection) {
1064 Some(manager) => manager.get_many(ids),
1065 None => vec![None; ids.len()],
1066 }
1067 }
1068
1069 pub fn get_any(&self, id: EntityId) -> Option<(String, UnifiedEntity)> {
1071 if let Some(cached) = self.entity_cache.get(id.raw()) {
1075 return Some(cached);
1076 }
1077
1078 let collections = self.collections.read();
1080 for (name, manager) in collections.iter() {
1081 if let Some(entity) = manager.get(id) {
1082 let result = (name.clone(), entity);
1083 drop(collections);
1087 self.entity_cache.insert(id.raw(), result.clone());
1088 return Some(result);
1089 }
1090 }
1091 None
1092 }
1093
1094 pub fn entity_cache_hit_rate(&self) -> Option<f64> {
1100 self.entity_cache.hit_rate()
1101 }
1102
1103 pub fn entity_cache_stats(&self) -> super::super::entity_cache::EntityCacheStats {
1105 self.entity_cache.stats()
1106 }
1107
1108 pub fn delete(&self, collection: &str, id: EntityId) -> Result<bool, StoreError> {
1110 self.entity_cache.remove(id.raw());
1112 let manager = self
1113 .get_collection(collection)
1114 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1115
1116 let deleted = manager.delete(id)?;
1117 if !deleted {
1118 return Ok(false);
1119 }
1120
1121 let mut registry_dirty = false;
1123 if self.pager.is_some() {
1124 let btree_indices = self.btree_indices.read();
1125 if let Some(btree) = btree_indices.get(collection) {
1126 let root_before = btree.root_page_id();
1127 let key = id.raw().to_be_bytes();
1128 let _ = btree.delete(&key);
1129 registry_dirty = root_before != btree.root_page_id();
1130 }
1131 }
1132
1133 self.unindex_cross_refs(id)?;
1135
1136 self.remove_from_graph_label_index(collection, id);
1138
1139 if registry_dirty {
1140 self.mark_paged_registry_dirty();
1141 }
1142 self.finish_paged_write([StoreWalAction::DeleteEntityRecord {
1143 collection: collection.to_string(),
1144 entity_id: id.raw(),
1145 }])?;
1146
1147 Ok(true)
1148 }
1149
1150 pub fn delete_batch(
1151 &self,
1152 collection: &str,
1153 ids: &[EntityId],
1154 ) -> Result<Vec<EntityId>, StoreError> {
1155 if ids.is_empty() {
1156 return Ok(Vec::new());
1157 }
1158
1159 self.entity_cache.remove_many(ids.iter().map(|id| id.raw()));
1164
1165 let manager = self
1166 .get_collection(collection)
1167 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1168
1169 let deleted_ids = manager.delete_batch(ids)?;
1170 if deleted_ids.is_empty() {
1171 return Ok(deleted_ids);
1172 }
1173
1174 let mut registry_dirty = false;
1175 if self.pager.is_some() {
1176 let btree_indices = self.btree_indices.read();
1177 if let Some(btree) = btree_indices.get(collection) {
1178 let root_before = btree.root_page_id();
1179 for id in &deleted_ids {
1180 let key = id.raw().to_be_bytes();
1181 let _ = btree.delete(&key);
1182 }
1183 registry_dirty = root_before != btree.root_page_id();
1184 }
1185 }
1186
1187 self.unindex_cross_refs_batch(&deleted_ids)?;
1188 self.remove_from_graph_label_index_batch(collection, &deleted_ids);
1189 if registry_dirty {
1190 self.mark_paged_registry_dirty();
1191 }
1192 let actions = deleted_ids
1193 .iter()
1194 .map(|id| StoreWalAction::DeleteEntityRecord {
1195 collection: collection.to_string(),
1196 entity_id: id.raw(),
1197 })
1198 .collect::<Vec<_>>();
1199 self.finish_paged_write(actions)?;
1200
1201 Ok(deleted_ids)
1202 }
1203
1204 pub fn set_metadata(
1206 &self,
1207 collection: &str,
1208 id: EntityId,
1209 metadata: Metadata,
1210 ) -> Result<(), StoreError> {
1211 let manager = self
1212 .get_collection(collection)
1213 .ok_or_else(|| StoreError::CollectionNotFound(collection.to_string()))?;
1214
1215 manager.set_metadata(id, metadata)?;
1216 if let Some(entity) = manager.get(id) {
1217 self.persist_entities_to_pager(collection, std::slice::from_ref(&entity))?;
1218 }
1219 Ok(())
1220 }
1221
1222 pub fn get_metadata(&self, collection: &str, id: EntityId) -> Option<Metadata> {
1224 self.get_collection(collection)?.get_metadata(id)
1225 }
1226
1227 pub fn add_cross_ref(
1229 &self,
1230 source_collection: &str,
1231 source_id: EntityId,
1232 target_collection: &str,
1233 target_id: EntityId,
1234 ref_type: RefType,
1235 weight: f32,
1236 ) -> Result<(), StoreError> {
1237 let source_manager = self
1239 .get_collection(source_collection)
1240 .ok_or_else(|| StoreError::CollectionNotFound(source_collection.to_string()))?;
1241
1242 if source_manager.get(source_id).is_none() {
1243 return Err(StoreError::EntityNotFound(source_id));
1244 }
1245
1246 let target_manager = self
1248 .get_collection(target_collection)
1249 .ok_or_else(|| StoreError::CollectionNotFound(target_collection.to_string()))?;
1250
1251 if target_manager.get(target_id).is_none() {
1252 return Err(StoreError::EntityNotFound(target_id));
1253 }
1254
1255 let current_refs = self
1257 .cross_refs
1258 .read()
1259 .get(&source_id)
1260 .map_or(0, |v| v.len());
1261
1262 if current_refs >= self.config.max_cross_refs {
1263 return Err(StoreError::TooManyRefs(source_id));
1264 }
1265
1266 let mut registry_dirty = false;
1267 {
1268 let mut forward = self.cross_refs.write();
1269 let refs = forward.entry(source_id).or_default();
1270 let inserted = !refs.iter().any(|(id, kind, coll)| {
1271 *id == target_id && *kind == ref_type && coll == target_collection
1272 });
1273 if inserted {
1274 refs.push((target_id, ref_type, target_collection.to_string()));
1275 registry_dirty = true;
1276 }
1277 }
1278
1279 {
1280 let mut reverse = self.reverse_refs.write();
1281 let refs = reverse.entry(target_id).or_default();
1282 let inserted = !refs.iter().any(|(id, kind, coll)| {
1283 *id == source_id && *kind == ref_type && coll == source_collection
1284 });
1285 if inserted {
1286 refs.push((source_id, ref_type, source_collection.to_string()));
1287 registry_dirty = true;
1288 }
1289 }
1290
1291 if let Some(mut entity) = source_manager.get(source_id) {
1292 if !entity.cross_refs().iter().any(|xref| {
1293 xref.target == target_id
1294 && xref.ref_type == ref_type
1295 && xref.target_collection == target_collection
1296 }) {
1297 let cross_ref = CrossRef::with_weight(
1298 source_id,
1299 target_id,
1300 target_collection,
1301 ref_type,
1302 weight,
1303 );
1304 entity.add_cross_ref(cross_ref);
1305 let _ = source_manager.update(entity.clone());
1306 registry_dirty = true;
1307 self.persist_entities_to_pager(source_collection, std::slice::from_ref(&entity))?;
1308 }
1309 }
1310
1311 if registry_dirty {
1312 self.mark_paged_registry_dirty();
1313 if matches!(
1314 self.config.durability_mode,
1315 crate::api::DurabilityMode::Strict
1316 ) {
1317 self.flush_paged_state()?;
1318 }
1319 }
1320
1321 Ok(())
1322 }
1323
1324 pub fn get_refs_from(&self, id: EntityId) -> Vec<(EntityId, RefType, String)> {
1326 self.cross_refs.read().get(&id).cloned().unwrap_or_default()
1327 }
1328
1329 pub fn get_refs_to(&self, id: EntityId) -> Vec<(EntityId, RefType, String)> {
1331 self.reverse_refs
1332 .read()
1333 .get(&id)
1334 .cloned()
1335 .unwrap_or_default()
1336 }
1337
1338 pub fn expand_refs(
1340 &self,
1341 id: EntityId,
1342 depth: u32,
1343 ref_types: Option<&[RefType]>,
1344 ) -> Vec<(UnifiedEntity, u32, RefType)> {
1345 let mut results = Vec::new();
1346 let mut visited = std::collections::HashSet::new();
1347 visited.insert(id);
1348
1349 self.expand_refs_recursive(id, depth, ref_types, &mut visited, &mut results, 1);
1350
1351 results
1352 }
1353
1354 fn expand_refs_recursive(
1355 &self,
1356 id: EntityId,
1357 max_depth: u32,
1358 ref_types: Option<&[RefType]>,
1359 visited: &mut std::collections::HashSet<EntityId>,
1360 results: &mut Vec<(UnifiedEntity, u32, RefType)>,
1361 current_depth: u32,
1362 ) {
1363 if current_depth > max_depth {
1364 return;
1365 }
1366
1367 for (target_id, ref_type, target_collection) in self.get_refs_from(id) {
1368 if visited.contains(&target_id) {
1369 continue;
1370 }
1371
1372 if let Some(types) = ref_types {
1373 if !types.contains(&ref_type) {
1374 continue;
1375 }
1376 }
1377
1378 visited.insert(target_id);
1379
1380 if let Some(entity) = self.get(&target_collection, target_id) {
1381 results.push((entity, current_depth, ref_type));
1382
1383 self.expand_refs_recursive(
1385 target_id,
1386 max_depth,
1387 ref_types,
1388 visited,
1389 results,
1390 current_depth + 1,
1391 );
1392 }
1393 }
1394 }
1395
1396 pub(crate) fn index_cross_refs(
1398 &self,
1399 entity: &UnifiedEntity,
1400 collection: &str,
1401 ) -> Result<(), StoreError> {
1402 let mut registry_dirty = false;
1403 for cross_ref in entity.cross_refs() {
1404 if cross_ref.target_collection.is_empty() {
1405 continue;
1406 }
1407 {
1408 let mut forward = self.cross_refs.write();
1409 let refs = forward.entry(cross_ref.source).or_default();
1410 let inserted = !refs.iter().any(|(id, kind, coll)| {
1411 *id == cross_ref.target
1412 && *kind == cross_ref.ref_type
1413 && coll == &cross_ref.target_collection
1414 });
1415 if inserted {
1416 refs.push((
1417 cross_ref.target,
1418 cross_ref.ref_type,
1419 cross_ref.target_collection.clone(),
1420 ));
1421 registry_dirty = true;
1422 }
1423 }
1424
1425 {
1426 let mut reverse = self.reverse_refs.write();
1427 let refs = reverse.entry(cross_ref.target).or_default();
1428 let inserted = !refs.iter().any(|(id, kind, coll)| {
1429 *id == cross_ref.source && *kind == cross_ref.ref_type && coll == collection
1430 });
1431 if inserted {
1432 refs.push((cross_ref.source, cross_ref.ref_type, collection.to_string()));
1433 registry_dirty = true;
1434 }
1435 }
1436 }
1437
1438 if registry_dirty {
1439 self.mark_paged_registry_dirty();
1440 }
1441
1442 Ok(())
1443 }
1444
1445 pub(crate) fn unindex_cross_refs(&self, id: EntityId) -> Result<(), StoreError> {
1447 self.cross_refs.write().remove(&id);
1449
1450 let mut reverse = self.reverse_refs.write();
1452 for refs in reverse.values_mut() {
1453 refs.retain(|(source, _, _)| *source != id);
1454 }
1455 reverse.remove(&id);
1456 self.mark_paged_registry_dirty();
1457
1458 Ok(())
1459 }
1460
1461 pub(crate) fn unindex_cross_refs_batch(&self, ids: &[EntityId]) -> Result<(), StoreError> {
1462 if ids.is_empty() {
1463 return Ok(());
1464 }
1465
1466 let id_set: std::collections::HashSet<EntityId> = ids.iter().copied().collect();
1467
1468 let needs_forward_cleanup = {
1482 let forward = self.cross_refs.read();
1483 id_set.iter().any(|id| forward.contains_key(id))
1484 };
1485 let needs_reverse_cleanup = {
1486 let reverse = self.reverse_refs.read();
1487 id_set.iter().any(|id| reverse.contains_key(id))
1488 };
1489
1490 if !needs_forward_cleanup && !needs_reverse_cleanup {
1491 self.unindex_cross_refs_fast_path
1492 .fetch_add(1, Ordering::Relaxed);
1493 return Ok(());
1494 }
1495
1496 if needs_forward_cleanup {
1497 let mut forward = self.cross_refs.write();
1498 for id in &id_set {
1499 forward.remove(id);
1500 }
1501 }
1502
1503 if needs_reverse_cleanup || needs_forward_cleanup {
1504 let mut reverse = self.reverse_refs.write();
1508 if needs_forward_cleanup {
1509 for refs in reverse.values_mut() {
1510 refs.retain(|(source, _, _)| !id_set.contains(source));
1511 }
1512 }
1513 reverse.retain(|target, refs| !id_set.contains(target) && !refs.is_empty());
1514 }
1515 self.mark_paged_registry_dirty();
1516
1517 Ok(())
1518 }
1519
1520 pub fn unindex_cross_refs_fast_path_hits(&self) -> u64 {
1524 self.unindex_cross_refs_fast_path.load(Ordering::Relaxed)
1525 }
1526
1527 pub fn query_all<F>(&self, filter: F) -> Vec<(String, UnifiedEntity)>
1529 where
1530 F: Fn(&UnifiedEntity) -> bool + Clone + Send + Sync,
1531 {
1532 let pairs: Vec<(String, Arc<SegmentManager>)> = {
1543 let collections = self.collections.read();
1544 collections
1545 .iter()
1546 .map(|(name, mgr)| (name.clone(), Arc::clone(mgr)))
1547 .collect()
1548 };
1549
1550 let use_parallel = pairs.len() > 1 && crate::runtime::SystemInfo::should_parallelize();
1551 if !use_parallel {
1552 return pairs
1554 .into_iter()
1555 .flat_map(|(name, mgr)| {
1556 mgr.query_all(filter.clone())
1557 .into_iter()
1558 .map(move |e| (name.clone(), e))
1559 })
1560 .collect();
1561 }
1562
1563 let filter_ref = &filter;
1565 let collection_results: Vec<Vec<(String, UnifiedEntity)>> = std::thread::scope(|s| {
1566 pairs
1567 .iter()
1568 .map(|(name, manager)| {
1569 let name = (*name).clone();
1570 s.spawn(move || {
1571 manager
1572 .query_all(|e| filter_ref(e))
1573 .into_iter()
1574 .map(|e| (name.clone(), e))
1575 .collect::<Vec<_>>()
1576 })
1577 })
1578 .collect::<Vec<_>>()
1579 .into_iter()
1580 .map(|h| h.join().unwrap_or_default())
1581 .collect()
1582 });
1583
1584 collection_results.into_iter().flatten().collect()
1585 }
1586
1587 pub fn filter_metadata_all(
1589 &self,
1590 filters: &[(String, MetadataFilter)],
1591 ) -> Vec<(String, EntityId)> {
1592 let mut results = Vec::new();
1593 let collections = self.collections.read();
1594
1595 for (name, manager) in collections.iter() {
1596 for id in manager.filter_metadata(filters) {
1597 results.push((name.clone(), id));
1598 }
1599 }
1600
1601 results
1602 }
1603
1604 pub fn stats(&self) -> StoreStats {
1606 let pairs: Vec<(String, Arc<SegmentManager>)> = {
1612 let collections = self.collections.read();
1613 collections
1614 .iter()
1615 .map(|(name, mgr)| (name.clone(), Arc::clone(mgr)))
1616 .collect()
1617 };
1618
1619 let mut stats = StoreStats {
1620 collection_count: pairs.len(),
1621 ..Default::default()
1622 };
1623
1624 for (name, manager) in &pairs {
1625 let manager_stats = manager.stats();
1626 stats.total_entities += manager_stats.total_entities;
1627 stats.total_memory_bytes += manager_stats.total_memory_bytes;
1628 stats.collections.insert(name.clone(), manager_stats);
1629 }
1630
1631 stats
1632 }
1633
1634 pub fn run_maintenance(&self) -> Result<(), StoreError> {
1636 let collections = self.collections.read();
1637 for manager in collections.values() {
1638 manager.run_maintenance()?;
1639 }
1640 Ok(())
1641 }
1642
1643 pub fn reclaimable_segment_bytes(&self) -> u64 {
1645 let collections = self.collections.read();
1646 collections
1647 .values()
1648 .map(|manager| manager.reclaimable_bytes())
1649 .fold(0, u64::saturating_add)
1650 }
1651
1652 pub fn pressure_consolidation_tick(&self) -> bool {
1655 let managers: Vec<_> = self.collections.read().values().cloned().collect();
1656 let mut advanced = false;
1659 for manager in &managers {
1660 advanced |= manager.pressure_consolidation_tick();
1661 }
1662 advanced
1663 }
1664}
1665
1666fn flatten_config_json(
1668 prefix: &str,
1669 value: &crate::serde_json::Value,
1670 out: &mut Vec<(String, crate::storage::schema::Value)>,
1671) {
1672 use crate::storage::schema::Value;
1673 match value {
1674 crate::serde_json::Value::Object(map) => {
1675 for (k, v) in map {
1676 let key = if prefix.is_empty() {
1677 k.clone()
1678 } else {
1679 format!("{prefix}.{k}")
1680 };
1681 flatten_config_json(&key, v, out);
1682 }
1683 }
1684 crate::serde_json::Value::String(s) => {
1685 out.push((prefix.to_string(), Value::text(s.clone())));
1686 }
1687 crate::serde_json::Value::Number(n) => {
1688 if n.fract().abs() < f64::EPSILON {
1689 out.push((prefix.to_string(), Value::UnsignedInteger(*n as u64)));
1690 } else {
1691 out.push((prefix.to_string(), Value::Float(*n)));
1692 }
1693 }
1694 crate::serde_json::Value::Bool(b) => {
1695 out.push((prefix.to_string(), Value::Boolean(*b)));
1696 }
1697 crate::serde_json::Value::Null => {
1698 out.push((prefix.to_string(), Value::Null));
1699 }
1700 crate::serde_json::Value::Array(arr) => {
1701 let json_str = crate::serde_json::to_string(value).unwrap_or_default();
1702 out.push((prefix.to_string(), Value::text(json_str)));
1703 }
1704 }
1705}