questdb/egress/
symbol_dict.rs

1/*******************************************************************************
2 *     ___                  _   ____  ____
3 *    / _ \ _   _  ___  ___| |_|  _ \| __ )
4 *   | | | | | | |/ _ \/ __| __| | | |  _ \
5 *   | |_| | |_| |  __/\__ \ |_| |_| | |_) |
6 *    \__\_\\__,_|\___||___/\__|____/|____/
7 *
8 *  Copyright (c) 2014-2019 Appsicle
9 *  Copyright (c) 2019-2025 QuestDB
10 *
11 *  Licensed under the Apache License, Version 2.0 (the "License");
12 *  you may not use this file except in compliance with the License.
13 *  You may obtain a copy of the License at
14 *
15 *  http://www.apache.org/licenses/LICENSE-2.0
16 *
17 *  Unless required by applicable law or agreed to in writing, software
18 *  distributed under the License is distributed on an "AS IS" BASIS,
19 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 *  See the License for the specific language governing permissions and
21 *  limitations under the License.
22 *
23 ******************************************************************************/
24
25//! Connection-scoped symbol dictionary.
26//!
27//! Each `RESULT_BATCH` carrying `FLAG_DELTA_SYMBOL_DICT` appends entries
28//! to the dictionary; `SYMBOL` columns transmit only integer codes that
29//! index into it. The dictionary persists across queries on the same
30//! connection until a `CACHE_RESET` with the dict bit clears it.
31//!
32//! Wire format of the delta section (when `FLAG_DELTA_SYMBOL_DICT` is set):
33//!
34//! ```text
35//! delta_start: varint     first conn-id assigned in this batch
36//! delta_count: varint     number of new entries
37//! repeat delta_count times:
38//!   entry_len: varint
39//!   entry:     bytes      UTF-8 symbol string
40//! ```
41//!
42//! `delta_start` MUST equal the dictionary's current length; after a
43//! reset, the next delta MUST start at 0.
44
45use crate::egress::wire::varint;
46use crate::error::{Result, fmt};
47
48/// Hard cap on the connection-scoped SYMBOL dict's UTF-8 heap size in
49/// bytes. Mirrors `MAX_CONN_DICT_HEAP_BYTES` in the Java reference
50/// client. Well-behaved servers approaching this cap are expected to
51/// emit `CACHE_RESET(RESET_MASK_DICT)`; crossing it without a reset is
52/// a protocol violation and we error rather than grow without bound.
53pub(crate) const MAX_CONN_DICT_HEAP_BYTES: usize = 256 * 1024 * 1024;
54
55/// Hard cap on the connection-scoped SYMBOL dict entry count. Mirrors
56/// `MAX_CONN_DICT_SIZE` in the Java reference client.
57pub(crate) const MAX_CONN_DICT_SIZE: usize = 8_388_608;
58
59/// Byte range of one symbol string within [`SymbolDict::arena`].
60#[derive(Debug, Clone, Copy)]
61#[repr(C)]
62pub struct SymbolEntry {
63    pub offset: u32,
64    pub len: u32,
65}
66
67/// Connection-scoped symbol dictionary.
68#[derive(Debug, Default, Clone)]
69pub struct SymbolDict {
70    arena: Vec<u8>,
71    entries: Vec<SymbolEntry>,
72}
73
74impl SymbolDict {
75    pub fn new() -> Self {
76        Self::default()
77    }
78
79    /// Number of entries currently stored. Also the next conn-id to assign.
80    pub fn len(&self) -> usize {
81        self.entries.len()
82    }
83
84    pub fn is_empty(&self) -> bool {
85        self.entries.is_empty()
86    }
87
88    /// UTF-8 bytes currently held in the arena.
89    pub fn heap_bytes(&self) -> usize {
90        self.arena.len()
91    }
92
93    /// Resolve a connection-scoped symbol id to its UTF-8 string.
94    pub fn get(&self, id: u32) -> Option<&str> {
95        let entry = self.entries.get(id as usize)?;
96        let start = entry.offset as usize;
97        let end = start + entry.len as usize;
98        debug_assert!(
99            end <= self.arena.len(),
100            "entry {id} offset+len={end} exceeds arena len {}",
101            self.arena.len()
102        );
103        // Safety: every byte slice that reaches the arena was UTF-8 validated
104        // by `apply_delta` before being copied in.
105        Some(unsafe { std::str::from_utf8_unchecked(&self.arena[start..end]) })
106    }
107
108    /// Raw UTF-8 heap holding every entry's bytes back-to-back.
109    pub fn arena(&self) -> &[u8] {
110        &self.arena
111    }
112
113    /// Entry table: index `i` addresses the conn-id-`i` symbol's bytes
114    /// within [`arena`](Self::arena).
115    pub fn entries(&self) -> &[SymbolEntry] {
116        &self.entries
117    }
118
119    /// Clear all state. Triggered by a `CACHE_RESET` with the dict bit.
120    /// Shrinks the backing allocations so a previously-saturated dict
121    /// doesn't keep its full heap reserved after the reset.
122    pub fn reset(&mut self) {
123        self.entries.clear();
124        self.arena.clear();
125        self.entries.shrink_to(1024);
126        self.arena.shrink_to(64 * 1024);
127    }
128
129    /// Apply a delta whose first new id is `delta_start` and whose entries
130    /// are produced in order. Validates UTF-8 and the sequencing invariant.
131    pub fn apply_delta<'a, I>(&mut self, delta_start: u64, entries: I) -> Result<()>
132    where
133        I: IntoIterator<Item = &'a [u8]>,
134    {
135        let expected = self.entries.len() as u64;
136        if delta_start != expected {
137            return Err(fmt!(
138                ProtocolError,
139                "symbol dict delta_start={} but registry len={}",
140                delta_start,
141                expected
142            ));
143        }
144        for bytes in entries {
145            self.push_one(bytes)?;
146        }
147        Ok(())
148    }
149
150    /// Decode + apply a delta directly from the wire bytes. Returns the
151    /// number of bytes consumed.
152    ///
153    /// All-or-nothing: if any entry in the delta is malformed, the dict
154    /// is rolled back to its pre-call state. Without this, a partial
155    /// failure would leave `self.entries.len()` between the old and new
156    /// expected values, and every subsequent delta would mismatch the
157    /// `delta_start` check above and break the connection until reset.
158    pub fn apply_delta_from_bytes(&mut self, bytes: &[u8]) -> Result<usize> {
159        let mut cursor = 0usize;
160        let (delta_start, n) = varint::decode_u64(&bytes[cursor..])?;
161        cursor += n;
162        let (delta_count, n) = varint::decode_u64(&bytes[cursor..])?;
163        cursor += n;
164
165        let expected = self.entries.len() as u64;
166        if delta_start != expected {
167            return Err(fmt!(
168                ProtocolError,
169                "symbol dict delta_start={} but registry len={}",
170                delta_start,
171                expected
172            ));
173        }
174
175        // Upfront cap on delta_count: a corrupt batch with delta_count
176        // = u64::MAX would otherwise iterate up to MAX_CONN_DICT_SIZE
177        // (8M) times — burning real CPU on per-entry varint decode +
178        // UTF-8 validation + heap-size checks — before push_one finally
179        // refuses to grow past the soft cap. Reject the malformed
180        // count up front against the headroom remaining in the dict.
181        //
182        // `saturating_sub` is defensive: `push_one` rejects every path
183        // that would let `entries.len()` exceed `MAX_CONN_DICT_SIZE`,
184        // so this subtraction can't actually underflow today. But a
185        // future refactor introducing a new write path that misses the
186        // cap would otherwise silently underflow in release mode
187        // (wrapping to a huge `headroom`) and disable this very guard.
188        // Saturating to 0 keeps the guard correct even under that bug:
189        // any positive `delta_count` would then be rejected.
190        let headroom = MAX_CONN_DICT_SIZE.saturating_sub(self.entries.len()) as u64;
191        if delta_count > headroom {
192            return Err(fmt!(
193                ProtocolError,
194                "symbol dict delta_count={} exceeds remaining capacity {} \
195                 (current entries={}, max={})",
196                delta_count,
197                headroom,
198                self.entries.len(),
199                MAX_CONN_DICT_SIZE
200            ));
201        }
202
203        let snapshot_entries = self.entries.len();
204        let snapshot_arena = self.arena.len();
205        let result: Result<usize> = (|| {
206            for i in 0..delta_count {
207                let (entry_len, n) = varint::decode_usize(&bytes[cursor..])?;
208                cursor += n;
209                let end = cursor.checked_add(entry_len).ok_or_else(|| {
210                    fmt!(
211                        ProtocolError,
212                        "symbol dict entry length overflow at i={}",
213                        i
214                    )
215                })?;
216                if end > bytes.len() {
217                    return Err(fmt!(
218                        ProtocolError,
219                        "symbol dict truncated at entry {}: need {} bytes, have {}",
220                        i,
221                        entry_len,
222                        bytes.len() - cursor
223                    ));
224                }
225                self.push_one(&bytes[cursor..end])?;
226                cursor = end;
227            }
228            Ok(cursor)
229        })();
230        if result.is_err() {
231            self.entries.truncate(snapshot_entries);
232            self.arena.truncate(snapshot_arena);
233        }
234        result
235    }
236
237    fn push_one(&mut self, bytes: &[u8]) -> Result<()> {
238        let s = std::str::from_utf8(bytes).map_err(|e| {
239            fmt!(
240                InvalidUtf8,
241                "symbol dict entry {} is not valid UTF-8: {}",
242                self.entries.len(),
243                e
244            )
245        })?;
246        if self.entries.len() >= MAX_CONN_DICT_SIZE {
247            return Err(fmt!(
248                ProtocolError,
249                "symbol dict full: {} entries (max {}); server must emit \
250                 CACHE_RESET(dict) before adding more",
251                self.entries.len(),
252                MAX_CONN_DICT_SIZE
253            ));
254        }
255        let new_heap = self
256            .arena
257            .len()
258            .checked_add(s.len())
259            .ok_or_else(|| fmt!(ProtocolError, "symbol dict heap overflow"))?;
260        if new_heap > MAX_CONN_DICT_HEAP_BYTES {
261            return Err(fmt!(
262                ProtocolError,
263                "symbol dict heap would reach {} bytes (max {}); server \
264                 must emit CACHE_RESET(dict) before adding more",
265                new_heap,
266                MAX_CONN_DICT_HEAP_BYTES
267            ));
268        }
269        let offset = u32::try_from(self.arena.len())
270            .map_err(|_| fmt!(ProtocolError, "symbol dict arena exceeds u32"))?;
271        let len = u32::try_from(s.len())
272            .map_err(|_| fmt!(ProtocolError, "symbol dict entry exceeds u32 length"))?;
273        self.arena.extend_from_slice(s.as_bytes());
274        self.entries.push(SymbolEntry { offset, len });
275        Ok(())
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::egress::wire::varint::encode_u64;
283    use crate::error::ErrorCode;
284
285    fn build_delta(start: u64, entries: &[&str]) -> Vec<u8> {
286        let mut out = Vec::new();
287        encode_u64(start, &mut out);
288        encode_u64(entries.len() as u64, &mut out);
289        for e in entries {
290            encode_u64(e.len() as u64, &mut out);
291            out.extend_from_slice(e.as_bytes());
292        }
293        out
294    }
295
296    #[test]
297    fn empty_dict() {
298        let d = SymbolDict::new();
299        assert_eq!(d.len(), 0);
300        assert!(d.is_empty());
301        assert_eq!(d.heap_bytes(), 0);
302        assert!(d.get(0).is_none());
303    }
304
305    #[test]
306    fn apply_first_delta_via_iter() {
307        let mut d = SymbolDict::new();
308        let entries: Vec<&[u8]> = vec![b"AAPL", b"MSFT", b"GOOG"];
309        d.apply_delta(0, entries).unwrap();
310        assert_eq!(d.len(), 3);
311        assert_eq!(d.get(0), Some("AAPL"));
312        assert_eq!(d.get(1), Some("MSFT"));
313        assert_eq!(d.get(2), Some("GOOG"));
314        assert_eq!(d.get(3), None);
315        assert_eq!(d.heap_bytes(), 4 + 4 + 4);
316    }
317
318    #[test]
319    fn second_delta_appends() {
320        let mut d = SymbolDict::new();
321        d.apply_delta(0, [b"a".as_slice()]).unwrap();
322        d.apply_delta(1, [b"bb".as_slice(), b"ccc".as_slice()])
323            .unwrap();
324        assert_eq!(d.len(), 3);
325        assert_eq!(d.get(2), Some("ccc"));
326    }
327
328    #[test]
329    fn delta_start_mismatch_rejected() {
330        let mut d = SymbolDict::new();
331        d.apply_delta(0, [b"x".as_slice()]).unwrap();
332        // Server claims new entries start at 5, but we have only 1.
333        let err = d.apply_delta(5, [b"y".as_slice()]).unwrap_err();
334        assert_eq!(err.code(), ErrorCode::ProtocolError);
335    }
336
337    #[test]
338    fn from_bytes_roundtrip() {
339        let mut d = SymbolDict::new();
340        let bytes = build_delta(0, &["AAPL", "MSFT"]);
341        let consumed = d.apply_delta_from_bytes(&bytes).unwrap();
342        assert_eq!(consumed, bytes.len());
343        assert_eq!(d.get(0), Some("AAPL"));
344        assert_eq!(d.get(1), Some("MSFT"));
345
346        let bytes2 = build_delta(2, &["GOOG"]);
347        d.apply_delta_from_bytes(&bytes2).unwrap();
348        assert_eq!(d.get(2), Some("GOOG"));
349    }
350
351    #[test]
352    fn from_bytes_partial_failure_rolls_back() {
353        // Build a delta where the first entry is fine and the second is
354        // truncated. Without rollback, the dict would commit the first
355        // entry and the delta_start of every subsequent batch would
356        // mismatch.
357        let mut d = SymbolDict::new();
358        d.apply_delta(0, [b"first".as_slice()]).unwrap();
359        let snapshot_len = d.len();
360        let snapshot_heap = d.heap_bytes();
361
362        let mut bytes = Vec::new();
363        encode_u64(snapshot_len as u64, &mut bytes); // delta_start
364        encode_u64(2, &mut bytes); // delta_count
365        encode_u64(2, &mut bytes); // entry 0 len
366        bytes.extend_from_slice(b"ok");
367        encode_u64(10, &mut bytes); // entry 1 claims 10 bytes
368        bytes.extend_from_slice(b"abc"); // only 3 follow → truncated
369
370        let err = d.apply_delta_from_bytes(&bytes).unwrap_err();
371        assert_eq!(err.code(), ErrorCode::ProtocolError);
372        // Dict reverted to snapshot: subsequent delta_start check uses
373        // the original length.
374        assert_eq!(d.len(), snapshot_len);
375        assert_eq!(d.heap_bytes(), snapshot_heap);
376        let next = build_delta(snapshot_len as u64, &["recovered"]);
377        d.apply_delta_from_bytes(&next).unwrap();
378        assert_eq!(d.get(snapshot_len as u32), Some("recovered"));
379    }
380
381    #[test]
382    fn from_bytes_truncated_entry_rejected() {
383        let mut d = SymbolDict::new();
384        let mut bytes = build_delta(0, &["hello"]);
385        bytes.truncate(bytes.len() - 1); // chop one byte off the entry
386        let err = d.apply_delta_from_bytes(&bytes).unwrap_err();
387        assert_eq!(err.code(), ErrorCode::ProtocolError);
388    }
389
390    #[test]
391    fn from_bytes_invalid_utf8_rejected() {
392        let mut bytes = Vec::new();
393        encode_u64(0, &mut bytes);
394        encode_u64(1, &mut bytes);
395        encode_u64(2, &mut bytes);
396        bytes.extend_from_slice(&[0xFF, 0xFE]); // invalid UTF-8
397        let mut d = SymbolDict::new();
398        let err = d.apply_delta_from_bytes(&bytes).unwrap_err();
399        assert_eq!(err.code(), ErrorCode::InvalidUtf8);
400    }
401
402    #[test]
403    fn reset_clears_state() {
404        let mut d = SymbolDict::new();
405        d.apply_delta(0, [b"x".as_slice(), b"yy".as_slice()])
406            .unwrap();
407        assert_eq!(d.len(), 2);
408        d.reset();
409        assert_eq!(d.len(), 0);
410        assert_eq!(d.heap_bytes(), 0);
411        // After reset, next delta must start at 0.
412        d.apply_delta(0, [b"new".as_slice()]).unwrap();
413        assert_eq!(d.get(0), Some("new"));
414    }
415
416    #[test]
417    fn delta_with_zero_entries_is_noop() {
418        let mut d = SymbolDict::new();
419        d.apply_delta(0, std::iter::empty::<&[u8]>()).unwrap();
420        let bytes = build_delta(0, &[]);
421        let consumed = d.apply_delta_from_bytes(&bytes).unwrap();
422        assert_eq!(consumed, bytes.len());
423        assert_eq!(d.len(), 0);
424    }
425
426    #[test]
427    fn delta_count_exceeding_capacity_rejected_upfront() {
428        // A corrupt batch with `delta_count = u64::MAX` must fail fast,
429        // not iterate up to MAX_CONN_DICT_SIZE times burning CPU on
430        // per-entry varint decode + UTF-8 + heap-size checks.
431        let mut d = SymbolDict::new();
432        let mut bytes = Vec::new();
433        encode_u64(0, &mut bytes); // delta_start
434        encode_u64(u64::MAX, &mut bytes); // delta_count
435        // No entries follow: if the cap weren't enforced upfront, the
436        // first iteration would error on truncated entry-length varint
437        // — which is also a ProtocolError but only after the loop has
438        // started. We can't directly observe iteration count, but we
439        // can pin the error message: the upfront cap surfaces
440        // "exceeds remaining capacity", the per-entry path surfaces
441        // "truncated".
442        let err = d.apply_delta_from_bytes(&bytes).unwrap_err();
443        assert_eq!(err.code(), ErrorCode::ProtocolError);
444        assert!(
445            err.msg().contains("exceeds remaining capacity"),
446            "expected upfront-cap rejection, got: {}",
447            err.msg()
448        );
449        assert_eq!(d.len(), 0);
450    }
451
452    #[test]
453    fn unicode_entries_preserved() {
454        let mut d = SymbolDict::new();
455        let bytes = build_delta(0, &["café", "日本語"]);
456        d.apply_delta_from_bytes(&bytes).unwrap();
457        assert_eq!(d.get(0), Some("café"));
458        assert_eq!(d.get(1), Some("日本語"));
459    }
460}
questdb/egress/symbol_dict.rs

questdb/egress/
symbol_dict.rs
