reddb_engine/query/unified/
memory.rs

1//! Retained result capacity, including spare slots and nested model payloads.
2//! Shared values are charged per result owner, independently of Arc counts.
3use super::*;
4use std::alloc::Layout;
5use std::mem::size_of;
6
7impl UnifiedResult {
8    /// Conservative retained heap estimate without allocating scratch. This is
9    /// admission for a completed result, not accounting for executor workspaces.
10    pub fn retained_memory_bytes(&self) -> usize {
11        let mut bytes = vector_slots(&self.columns)
12            .saturating_add(
13                self.columns
14                    .iter()
15                    .map(String::capacity)
16                    .fold(0, usize::saturating_add),
17            )
18            .saturating_add(vector_slots(&self.records))
19            .saturating_add(
20                self.pre_serialized_json
21                    .as_ref()
22                    .map_or(0, String::capacity),
23            );
24        let first_schema = self.records.first().map(|record| &record.schema);
25        for (index, record) in self.records.iter().enumerate() {
26            bytes = bytes.saturating_add(record_heap_bytes(record));
27            // The common shared schema is counted once, with no HashSet or
28            // per-row scratch. Other schemas are conservatively charged per row.
29            if index == 0 || first_schema.is_none_or(|schema| !Arc::ptr_eq(schema, &record.schema))
30            {
31                bytes = bytes.saturating_add(schema_bytes(&record.schema));
32            }
33        }
34        if let Some(stats) = &self.stats.vector {
35            bytes = bytes
36                .saturating_add(vector_slots(&stats.operators))
37                .saturating_add(
38                    stats
39                        .operators
40                        .iter()
41                        .map(|operator| operator.operator.capacity())
42                        .fold(0, usize::saturating_add),
43                )
44                .saturating_add(stats.mode_requested.capacity())
45                .saturating_add(stats.mode_executed.capacity())
46                .saturating_add(stats.access_path.capacity())
47                .saturating_add(stats.fallback_reason.as_ref().map_or(0, String::capacity));
48        }
49        bytes
50    }
51}
52
53fn vector_slots<T>(values: &Vec<T>) -> usize {
54    values.capacity().saturating_mul(size_of::<T>())
55}
56
57fn arc_text_bytes(value: &Arc<str>) -> usize {
58    Layout::array::<u8>(value.len())
59        .map(arc_allocation_bytes)
60        .unwrap_or(usize::MAX)
61}
62
63fn arc_allocation_bytes(payload: Layout) -> usize {
64    // Arc's strong/weak counters precede its payload. Include the trailing
65    // alignment: e.g. one UTF-8 byte requests 24 bytes on a 64-bit host, not 17.
66    Layout::new::<[usize; 2]>()
67        .extend(payload)
68        .map(|(layout, _)| layout.pad_to_align().size())
69        .unwrap_or(usize::MAX)
70}
71
72fn schema_bytes(schema: &Arc<Vec<Arc<str>>>) -> usize {
73    arc_allocation_bytes(Layout::new::<Vec<Arc<str>>>())
74        .saturating_add(vector_slots(schema))
75        .saturating_add(
76            schema
77                .iter()
78                .map(arc_text_bytes)
79                .fold(0, usize::saturating_add),
80        )
81}
82
83fn map_bytes<K, V>(map: &HashMap<K, V>, mut payload: impl FnMut(&K, &V) -> usize) -> usize {
84    // Capacity is the usable entry count, not the raw bucket count. Twice the
85    // usable slots plus control/alignment space conservatively covers sparse
86    // SwissTable allocations, including the smallest nonempty table.
87    let slots = if map.capacity() == 0 {
88        0
89    } else {
90        map.capacity()
91            .saturating_mul(2)
92            .saturating_mul(size_of::<(K, V)>().saturating_add(1))
93            .saturating_add(64)
94    };
95    map.iter().fold(slots, |bytes, (key, value)| {
96        bytes.saturating_add(payload(key, value))
97    })
98}
99
100fn properties_bytes(properties: &HashMap<String, Value>) -> usize {
101    map_bytes(properties, |key, value| {
102        key.capacity().saturating_add(value_heap_bytes(value))
103    })
104}
105
106fn edge_heap_bytes(edge: &MatchedEdge) -> usize {
107    edge.from
108        .capacity()
109        .saturating_add(edge.to.capacity())
110        .saturating_add(edge.edge_label.capacity())
111        .saturating_add(properties_bytes(&edge.properties))
112}
113
114fn record_heap_bytes(record: &UnifiedRecord) -> usize {
115    let mut bytes = vector_slots(&record.values).saturating_add(
116        record
117            .values
118            .iter()
119            .map(value_heap_bytes)
120            .fold(0, usize::saturating_add),
121    );
122    if let Some(overflow) = &record.overflow {
123        bytes = bytes.saturating_add(map_bytes(overflow, |key, value| {
124            arc_text_bytes(key).saturating_add(value_heap_bytes(value))
125        }));
126    }
127    bytes = bytes
128        .saturating_add(map_bytes(&record.nodes, |key, node| {
129            key.capacity()
130                .saturating_add(node.id.capacity())
131                .saturating_add(node.label.capacity())
132                .saturating_add(node.node_label.capacity())
133                .saturating_add(properties_bytes(&node.properties))
134        }))
135        .saturating_add(map_bytes(&record.edges, |key, edge| {
136            key.capacity().saturating_add(edge_heap_bytes(edge))
137        }))
138        .saturating_add(vector_slots(&record.paths))
139        .saturating_add(vector_slots(&record.vector_results));
140    for path in &record.paths {
141        bytes = bytes
142            .saturating_add(vector_slots(&path.nodes))
143            .saturating_add(
144                path.nodes
145                    .iter()
146                    .map(String::capacity)
147                    .fold(0, usize::saturating_add),
148            )
149            .saturating_add(vector_slots(&path.edges))
150            .saturating_add(
151                path.edges
152                    .iter()
153                    .map(edge_heap_bytes)
154                    .fold(0, usize::saturating_add),
155            );
156    }
157    for vector in &record.vector_results {
158        bytes = bytes
159            .saturating_add(vector.collection.capacity())
160            .saturating_add(vector.vector.as_ref().map_or(0, vector_slots))
161            .saturating_add(vector.metadata.as_ref().map_or(0, properties_bytes))
162            .saturating_add(vector.linked_node.as_ref().map_or(0, String::capacity))
163            .saturating_add(
164                vector
165                    .linked_row
166                    .as_ref()
167                    .map_or(0, |(table, _)| table.capacity()),
168            );
169    }
170    bytes
171}
172
173fn value_heap_bytes(value: &Value) -> usize {
174    // Bounded explicit stack: API-supplied values deeper than the parser's
175    // maximum cannot receive a representable retained-memory allowance.
176    const LEVELS: usize = reddb_rql::limits::JSON_LITERAL_MAX_DEPTH as usize + 1;
177    let mut stack: [std::slice::Iter<'_, Value>; LEVELS] = std::array::from_fn(|_| [].iter());
178    let mut depth = 0;
179    let mut current = Some(value);
180    let mut bytes = 0usize;
181    loop {
182        if let Some(value) = current.take() {
183            if let Value::Array(values) = value {
184                if depth == LEVELS {
185                    return usize::MAX;
186                }
187                bytes = bytes.saturating_add(vector_slots(values));
188                stack[depth] = values.iter();
189                depth += 1;
190            } else {
191                bytes = bytes.saturating_add(scalar_heap_bytes(value));
192            }
193        }
194        if depth == 0 {
195            return bytes;
196        }
197        current = stack[depth - 1].next();
198        if current.is_none() {
199            depth -= 1;
200        }
201    }
202}
203
204fn scalar_heap_bytes(value: &Value) -> usize {
205    match value {
206        Value::Text(value) => arc_text_bytes(value),
207        Value::Blob(value) | Value::Json(value) | Value::Secret(value) => value.capacity(),
208        Value::Vector(value) => vector_slots(value),
209        Value::NodeRef(value)
210        | Value::EdgeRef(value)
211        | Value::VectorRef(value, _)
212        | Value::RowRef(value, _)
213        | Value::DocRef(value, _)
214        | Value::TableRef(value)
215        | Value::Email(value)
216        | Value::Url(value)
217        | Value::Password(value)
218        | Value::DecimalText(value)
219        | Value::AssetCode(value) => value.capacity(),
220        Value::KeyRef(collection, key) => collection.capacity().saturating_add(key.capacity()),
221        Value::Money { asset_code, .. } => asset_code.capacity(),
222        Value::Array(_) => unreachable!("arrays use the bounded stack"),
223        Value::Null
224        | Value::Integer(_)
225        | Value::UnsignedInteger(_)
226        | Value::Float(_)
227        | Value::Boolean(_)
228        | Value::Timestamp(_)
229        | Value::Duration(_)
230        | Value::IpAddr(_)
231        | Value::MacAddr(_)
232        | Value::Uuid(_)
233        | Value::Color(_)
234        | Value::Phone(_)
235        | Value::Semver(_)
236        | Value::Cidr(_, _)
237        | Value::Date(_)
238        | Value::Time(_)
239        | Value::Decimal(_)
240        | Value::EnumValue(_)
241        | Value::TimestampMs(_)
242        | Value::Ipv4(_)
243        | Value::Ipv6(_)
244        | Value::Subnet(_, _)
245        | Value::Port(_)
246        | Value::Latitude(_)
247        | Value::Longitude(_)
248        | Value::GeoPoint(_, _)
249        | Value::Country2(_)
250        | Value::Country3(_)
251        | Value::Lang2(_)
252        | Value::Lang5(_)
253        | Value::Currency(_)
254        | Value::ColorAlpha(_)
255        | Value::BigInt(_)
256        | Value::PageRef(_) => 0,
257    }
258}
reddb_engine/query/unified/memory.rs

reddb_engine/query/unified/
memory.rs
