rushdb_query/
page.rs

1//! Budgeted page reconstruction for the structured query read path.
2
3use super::{
4    attach_collects_with_budget, reconstruct, Lmpg, QueryError, QuerySnapshot, SearchRequest,
5};
6
7/// Resolve the effective working-set budget: the caller's remaining
8/// request budget, else the configured per-query cap.
9///
10/// `Some(0)` from a caller is a *hard zero cap* (an exhausted request budget),
11/// distinct from a configured `query_memory_max_bytes == 0`, which disables
12/// the cap. Conflating the two let an exhausted request run with no limit.
13#[must_use]
14pub fn effective_budget(work_budget: Option<usize>, config: u64) -> Option<u64> {
15    match work_budget {
16        Some(bytes) => Some(u64::try_from(bytes).unwrap_or(u64::MAX)),
17        None => (config != 0).then_some(config),
18    }
19}
20
21/// Charge one record's serialized size against `budget`, accumulating into
22/// `charged`. Errors once the running total exceeds the limit.
23///
24/// `reservation`, when given, is grown to the real running total:
25/// best-effort attribution against the shared graph-phase budget, not a new
26/// rejection path — only `budget` above can fail this call, so a tight
27/// shared pool just means the gauge under-reports.
28pub(super) fn charge_record(
29    budget: Option<u64>,
30    charged: &mut u64,
31    reservation: Option<&mut rushdb_storage::resources::ByteReservation>,
32    record: &serde_json::Value,
33) -> Result<(), QueryError> {
34    if let Some(limit) = budget {
35        let bytes =
36            serde_json::to_vec(record).map_or(0, |v| u64::try_from(v.len()).unwrap_or(u64::MAX));
37        *charged = charged.saturating_add(bytes);
38        if let Some(reservation) = reservation {
39            let _ = reservation.resize(usize::try_from(*charged).unwrap_or(usize::MAX));
40        }
41        if *charged > limit {
42            return Err(QueryError::ResourceLimit(format!(
43                "query working set exceeded {limit} bytes"
44            )));
45        }
46    }
47    Ok(())
48}
49
50/// Reconstruct the page's records in order, charging each reconstructed record
51/// and each `$collect`/traversal expansion cumulatively against the budget, so
52/// N page records cannot each retain the full remaining allowance.
53pub(super) fn reconstruct_page_with_budget(
54    lmpg: &Lmpg,
55    req: &SearchRequest,
56    is_cancelled: &dyn Fn() -> bool,
57    snap: &QuerySnapshot,
58    page_roots: Vec<u64>,
59    budget: Option<u64>,
60) -> Result<Vec<serde_json::Value>, QueryError> {
61    let mut charged: u64 = 0;
62    // A named child of the shared phase budget, so graph's own
63    // (already-exact) working-set bytes are individually observable.
64    let mut reservation = lmpg
65        .store
66        .graph_phase_budget()
67        .and_then(|phase_budget| phase_budget.try_reserve(0).ok());
68    let mut session = reconstruct::ReconstructSession::new(snap, lmpg.tenant);
69    let mut records = Vec::with_capacity(page_roots.len());
70    for ord in page_roots {
71        if is_cancelled() {
72            return Err(QueryError::Cancelled);
73        }
74        let Some(json) = session.reconstruct_record(ord)? else {
75            continue;
76        };
77        charge_record(budget, &mut charged, reservation.as_mut(), &json)?;
78        // Charge `$collect`/traversal expansion against the remaining budget so
79        // a collect cannot pull an unbounded number of related records into
80        // the response.
81        let remaining = budget
82            .map(|limit| usize::try_from(limit.saturating_sub(charged)).unwrap_or(usize::MAX));
83        let record = attach_collects_with_budget(
84            lmpg,
85            snap,
86            ord,
87            json,
88            &req.collects,
89            remaining,
90            is_cancelled,
91        )?;
92        // Cumulative: the expanded record (with collected children) counts
93        // against the same budget.
94        charge_record(budget, &mut charged, reservation.as_mut(), &record)?;
95        records.push(record);
96    }
97    Ok(records)
98}
rushdb_query/page.rs

rushdb_query/
page.rs
