Skip to main content

nodedb_lite/nodedb/core/ops/
memory_engines.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Memory governor reporting and per-engine handle accessors.
4
5use std::sync::{Arc, Mutex};
6
7use nodedb_mem::{EngineId, MemoryGovernor, ReservationToken, ScopedMemory};
8use nodedb_types::error::NodeDbResult;
9use nodedb_types::{DatabaseId, TenantId};
10
11use crate::engine::strict::StrictEngine;
12use crate::nodedb::core::types::NodeDbLite;
13use crate::nodedb::lock_ext::LockExt;
14use crate::storage::engine::StorageEngine;
15
16impl<S: StorageEngine> NodeDbLite<S> {
17    /// A scope charging `engine` against Lite's single database and tenant.
18    ///
19    /// Lite runs one database and one tenant, so every scope it builds uses
20    /// [`DatabaseId::DEFAULT`] and tenant zero.
21    pub(crate) fn memory_for(&self, engine: EngineId) -> ScopedMemory {
22        ScopedMemory::new(
23            Arc::clone(&self.governor),
24            DatabaseId::DEFAULT,
25            TenantId::new(0),
26            engine,
27        )
28    }
29
30    /// Update memory governor with current engine usage.
31    pub fn update_memory_stats(&self) {
32        if let Ok(indices) = self.vector_state.hnsw_indices.lock() {
33            let hnsw_bytes: usize = indices
34                .values()
35                .map(|idx| idx.len() * (idx.dim() * 4 + 128))
36                .sum();
37            self.charge_engine_usage(&self.vector_mem_token, EngineId::Vector, hnsw_bytes);
38        }
39        if let Ok(csr_map) = self.csr.lock() {
40            let total: usize = csr_map
41                .values()
42                .map(|idx| idx.estimated_memory_bytes())
43                .sum();
44            self.charge_engine_usage(&self.graph_mem_token, EngineId::Graph, total);
45        }
46        if let Ok(crdt) = self.crdt.lock() {
47            self.charge_engine_usage(
48                &self.crdt_mem_token,
49                EngineId::Crdt,
50                crdt.estimated_memory_bytes(),
51            );
52        }
53    }
54
55    /// Replace `slot`'s held reservation with one charging `bytes` for `engine`.
56    ///
57    /// `report_usage` reports an absolute current footprint, not a delta, so
58    /// the previous token drops before the new one is charged — otherwise the
59    /// governor would see the sum of every report instead of the current one.
60    /// Single-database, single-tenant: uses [`DatabaseId::DEFAULT`] and
61    /// `TenantId::new(0)`, matching every other Lite call site.
62    fn charge_engine_usage(
63        &self,
64        slot: &Mutex<Option<ReservationToken>>,
65        engine: EngineId,
66        bytes: usize,
67    ) {
68        let scoped = ScopedMemory::new(
69            Arc::clone(&self.governor),
70            DatabaseId::DEFAULT,
71            TenantId::new(0),
72            engine,
73        );
74        let mut token = slot.lock_or_recover();
75        token.take();
76        *token = Some(scoped.charge(bytes));
77    }
78
79    /// List currently loaded HNSW collections.
80    pub fn loaded_collections(&self) -> NodeDbResult<Vec<String>> {
81        let indices = self.vector_state.hnsw_indices.lock_or_recover();
82        Ok(indices.keys().cloned().collect())
83    }
84
85    /// Access the memory governor.
86    pub fn governor(&self) -> &Arc<MemoryGovernor> {
87        &self.governor
88    }
89
90    /// Access the sequence registry that `nextval` / `currval` / `setval`
91    /// in a SELECT list read. Register a `LiteSequenceDef` here before
92    /// querying it.
93    pub fn sequences(&self) -> &Arc<crate::sequence::LiteSequenceRegistry> {
94        self.query_engine.sequences()
95    }
96
97    /// Access the strict document engine (for direct Binary Tuple CRUD).
98    pub fn strict_engine(&self) -> &Arc<StrictEngine<S>> {
99        &self.strict
100    }
101
102    /// Access the columnar analytics engine (for direct segment operations).
103    pub fn columnar_engine(&self) -> &Arc<crate::engine::columnar::ColumnarEngine<S>> {
104        &self.columnar
105    }
106
107    /// Access the HTAP bridge (for materialized view inspection).
108    pub fn htap_bridge(&self) -> &Arc<crate::engine::htap::HtapBridge> {
109        &self.htap
110    }
111
112    /// Access the timeseries engine (continuous aggregates, ingest, flush).
113    pub fn timeseries_engine(
114        &self,
115    ) -> &Arc<Mutex<crate::engine::timeseries::engine::TimeseriesEngine>> {
116        &self.timeseries
117    }
118}