Skip to main content

grafeo_engine/
memory_usage.rs

1//! Hierarchical memory usage breakdown for the database.
2//!
3//! Store-level types (`StoreMemory`, `IndexMemory`, etc.) live in grafeo-common.
4//! This module defines the top-level `MemoryUsage` aggregate and engine-specific
5//! types (`CacheMemory`, `BufferManagerMemory`).
6
7pub use grafeo_common::memory::usage::{
8    IndexMemory, MvccMemory, NamedMemory, StoreMemory, StringPoolMemory,
9};
10use serde::{Deserialize, Serialize};
11
12/// Hierarchical memory usage breakdown for the entire database.
13#[derive(Debug, Clone, Default, Serialize, Deserialize)]
14pub struct MemoryUsage {
15    /// Total estimated memory usage in bytes.
16    pub total_bytes: usize,
17    /// Graph storage (nodes, edges, properties).
18    pub store: StoreMemory,
19    /// Index structures.
20    pub indexes: IndexMemory,
21    /// MVCC versioning overhead.
22    pub mvcc: MvccMemory,
23    /// Caches (query plans, etc.).
24    pub caches: CacheMemory,
25    /// String interning (ArcStr label/type registries).
26    pub string_pool: StringPoolMemory,
27    /// Buffer manager tracked allocations.
28    pub buffer_manager: BufferManagerMemory,
29}
30
31impl MemoryUsage {
32    /// Recomputes `total_bytes` from child totals.
33    pub fn compute_total(&mut self) {
34        self.total_bytes = self.store.total_bytes
35            + self.indexes.total_bytes
36            + self.mvcc.total_bytes
37            + self.caches.total_bytes
38            + self.string_pool.total_bytes
39            + self.buffer_manager.allocated_bytes;
40    }
41}
42
43/// Cache memory usage.
44#[derive(Debug, Clone, Default, Serialize, Deserialize)]
45pub struct CacheMemory {
46    /// Total cache memory.
47    pub total_bytes: usize,
48    /// Parsed plan cache.
49    pub parsed_plan_cache_bytes: usize,
50    /// Optimized plan cache.
51    pub optimized_plan_cache_bytes: usize,
52    /// Number of cached plans (parsed + optimized).
53    pub cached_plan_count: usize,
54}
55
56impl CacheMemory {
57    /// Recomputes `total_bytes` from child values.
58    pub fn compute_total(&mut self) {
59        self.total_bytes = self.parsed_plan_cache_bytes + self.optimized_plan_cache_bytes;
60    }
61}
62
63/// Buffer manager tracked allocations.
64#[derive(Debug, Clone, Default, Serialize, Deserialize)]
65pub struct BufferManagerMemory {
66    /// Budget configured for the buffer manager.
67    pub budget_bytes: usize,
68    /// Currently allocated via grants.
69    pub allocated_bytes: usize,
70    /// Graph storage region.
71    pub graph_storage_bytes: usize,
72    /// Index buffers region.
73    pub index_buffers_bytes: usize,
74    /// Execution buffers region.
75    pub execution_buffers_bytes: usize,
76    /// Spill staging region.
77    pub spill_staging_bytes: usize,
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn default_memory_usage_is_zero() {
86        let usage = MemoryUsage::default();
87        assert_eq!(usage.total_bytes, 0);
88        assert_eq!(usage.store.total_bytes, 0);
89        assert_eq!(usage.indexes.total_bytes, 0);
90        assert_eq!(usage.mvcc.total_bytes, 0);
91        assert_eq!(usage.caches.total_bytes, 0);
92        assert_eq!(usage.string_pool.total_bytes, 0);
93        assert_eq!(usage.buffer_manager.allocated_bytes, 0);
94    }
95
96    #[test]
97    fn compute_total_sums_children() {
98        let mut usage = MemoryUsage {
99            store: StoreMemory {
100                total_bytes: 100,
101                ..Default::default()
102            },
103            indexes: IndexMemory {
104                total_bytes: 200,
105                ..Default::default()
106            },
107            mvcc: MvccMemory {
108                total_bytes: 50,
109                ..Default::default()
110            },
111            caches: CacheMemory {
112                total_bytes: 30,
113                ..Default::default()
114            },
115            string_pool: StringPoolMemory {
116                total_bytes: 10,
117                ..Default::default()
118            },
119            buffer_manager: BufferManagerMemory {
120                allocated_bytes: 20,
121                ..Default::default()
122            },
123            ..Default::default()
124        };
125        usage.compute_total();
126        assert_eq!(usage.total_bytes, 410);
127    }
128
129    #[test]
130    fn serde_roundtrip() {
131        let mut usage = MemoryUsage::default();
132        usage.store.nodes_bytes = 1024;
133        usage.indexes.vector_indexes.push(NamedMemory {
134            name: "vec_idx".to_string(),
135            bytes: 512,
136            item_count: 100,
137        });
138        usage.mvcc.average_chain_depth = 1.5;
139
140        let json = serde_json::to_string(&usage).unwrap();
141        let deserialized: MemoryUsage = serde_json::from_str(&json).unwrap();
142
143        assert_eq!(deserialized.store.nodes_bytes, 1024);
144        assert_eq!(deserialized.indexes.vector_indexes.len(), 1);
145        assert_eq!(deserialized.indexes.vector_indexes[0].name, "vec_idx");
146        assert!((deserialized.mvcc.average_chain_depth - 1.5).abs() < f64::EPSILON);
147    }
148}