rushdb-query 3.0.3

Internal component of the RushDB server: query planner and executor. No stable API; install rushdb-server instead.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! Live aggregate retention and per-record reconstruction admission.
use std::cell::Cell;
use std::sync::Arc;

use rushdb_lmpg::reader::TryRead;
use rushdb_storage::StorageError;
use serde_json::Value;

use crate::{GroupBucket, QueryError, QuerySnapshot};

pub(super) fn check(limit: Option<usize>, bytes: usize) -> Result<(), QueryError> {
    if limit.is_some_and(|limit| bytes > limit) {
        return Err(QueryError::ResourceLimit(
            "aggregate working byte budget exceeded".into(),
        ));
    }
    Ok(())
}

pub(super) fn value_bytes(value: &Value) -> usize {
    let payload = match value {
        Value::String(value) => value.len(),
        Value::Array(values) => values
            .iter()
            .fold(0usize, |sum, v| sum.saturating_add(value_bytes(v))),
        Value::Object(values) => object_bytes(values),
        _ => 0,
    };
    std::mem::size_of::<Value>().saturating_add(payload)
}

pub(super) fn object_bytes(values: &serde_json::Map<String, Value>) -> usize {
    values.iter().fold(0usize, |sum, (key, value)| {
        sum.saturating_add(96)
            .saturating_add(key.len())
            .saturating_add(value_bytes(value))
    })
}

pub(super) fn bucket_bytes(bucket: &GroupBucket) -> usize {
    [
        &bucket.keys,
        &bucket.sums,
        &bucket.counts,
        &bucket.mins,
        &bucket.maxs,
    ]
    .into_iter()
    .fold(std::mem::size_of::<GroupBucket>(), |sum, map| {
        sum.saturating_add(object_bytes(map))
    })
    .saturating_add(
        bucket
            .derived_outputs
            .iter()
            .fold(0usize, |sum, (key, value)| {
                sum.saturating_add(96)
                    .saturating_add(key.len())
                    .saturating_add(value_bytes(value))
            }),
    )
}

/// Charge reads before JSON decoding. The factor includes decoded scalar,
/// occurrence, family-cache and tree-node overhead. Each record gets its own
/// allowance, released when its reconstruction session is dropped.
pub(super) struct RecordReader {
    view: QuerySnapshot,
    remaining: Cell<Option<usize>>,
}

impl RecordReader {
    pub(super) fn new(view: &QuerySnapshot, limit: Option<usize>) -> Self {
        Self {
            view: Arc::clone(view),
            remaining: Cell::new(limit),
        }
    }

    fn view(&self) -> QuerySnapshot {
        self.remaining.get().map_or_else(
            || Arc::clone(&self.view),
            |bytes| self.view.with_materialization_limit(bytes / 16),
        )
    }

    fn charge(&self, bytes: usize) -> Result<(), StorageError> {
        if let Some(remaining) = self.remaining.get() {
            let charge = bytes.saturating_mul(16).saturating_add(128);
            let next = remaining.checked_sub(charge).ok_or_else(|| {
                StorageError::ResourceExhausted(
                    "aggregate reconstruction exceeds byte budget".into(),
                )
            })?;
            self.remaining.set(Some(next));
        }
        Ok(())
    }
}

impl TryRead for RecordReader {
    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
        let value = self.view().get(key)?;
        if let Some(value) = &value {
            self.charge(key.len().saturating_add(value.len()))?;
        }
        Ok(value)
    }

    fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
        let rows = self.view().range_prefix(prefix)?;
        for (key, value) in &rows {
            self.charge(key.len().saturating_add(value.len()))?;
        }
        Ok(rows)
    }

    fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
        let rows = self.view().range(start, end)?;
        for (key, value) in &rows {
            self.charge(key.len().saturating_add(value.len()))?;
        }
        Ok(rows)
    }
}

pub(super) fn resource_error(error: QueryError) -> QueryError {
    match error {
        QueryError::Storage(StorageError::ResourceExhausted(message)) => {
            QueryError::ResourceLimit(message)
        }
        error => error,
    }
}