Skip to main content

reddb_server/runtime/
owned_query_result.rs

1//! Shared ownership for completed materialized request results. Executor
2//! workspaces and serialization allocations require separate admission.
3use super::{memory_admission::QueryBufferReservation, RedDBRuntime, RuntimeQueryResult};
4use crate::{RedDBError, RedDBResult};
5use std::{borrow::Cow, fmt, ops::Deref, sync::Arc};
6
7pub(crate) const SHARED_RESULT_EXECUTOR: &str = "shared_query_result";
8
9/// Immutable result returned by the request seam. Clones share all records,
10/// payloads and metadata; no infallible Clone can bypass shared admission.
11#[derive(Clone)]
12pub struct OwnedQueryResult {
13    inner: Arc<ResultOwner>,
14}
15
16struct ResultOwner {
17    // Declaration order frees all data before returning the reservation.
18    value: RuntimeQueryResult,
19    _reservation: Option<QueryBufferReservation>,
20}
21
22impl Deref for OwnedQueryResult {
23    type Target = RuntimeQueryResult;
24    fn deref(&self) -> &Self::Target {
25        &self.inner.value
26    }
27}
28
29impl AsRef<RuntimeQueryResult> for OwnedQueryResult {
30    fn as_ref(&self) -> &RuntimeQueryResult {
31        self
32    }
33}
34
35impl fmt::Debug for OwnedQueryResult {
36    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
37        self.inner.value.fmt(formatter)
38    }
39}
40
41impl OwnedQueryResult {
42    /// Admitted results stay borrowed so their allocations cannot escape their
43    /// charge by moving into an unaccounted output. A unique, unleased command
44    /// acknowledgement keeps its existing move fast path. New rendered output
45    /// still requires separate admission.
46    pub(crate) fn render<T>(self, render: impl FnOnce(Cow<'_, RuntimeQueryResult>) -> T) -> T {
47        match Arc::try_unwrap(self.inner) {
48            Ok(owner) if owner._reservation.is_none() => render(Cow::Owned(owner.value)),
49            // Keep the owner intact so fields also drop in the required order
50            // if rendering unwinds.
51            Ok(owner) => render(Cow::Borrowed(&owner.value)),
52            Err(owner) => render(Cow::Borrowed(&owner.value)),
53        }
54    }
55}
56
57impl RedDBRuntime {
58    pub(crate) fn own_query_result(
59        &self,
60        value: RuntimeQueryResult,
61    ) -> RedDBResult<OwnedQueryResult> {
62        // These are the same pure materialized plan families as query_budget.
63        // Command results (e.g. queue consumption / ASK / sequence functions)
64        // must not receive a late refusal after an acknowledged side effect.
65        let is_read = value.is_materialized_read();
66        let reservation = if is_read {
67            let owner_bytes = std::alloc::Layout::new::<[usize; 2]>()
68                .extend(std::alloc::Layout::new::<ResultOwner>())
69                .expect("fixed result owner layout is representable")
70                .0
71                .pad_to_align()
72                .size();
73            let bytes = value
74                .result
75                .retained_memory_bytes()
76                .saturating_add(value.query.capacity())
77                .saturating_add(value.bookmark.as_ref().map_or(0, String::capacity))
78                .saturating_add(value.notice.as_ref().map_or(0, String::capacity))
79                .saturating_add(owner_bytes);
80            Some(self.try_reserve_query_buffer(bytes).ok_or_else(|| {
81                self.memory_accounting().record_admission_denied();
82                RedDBError::MaterializationLimitExceeded {
83                    executor: SHARED_RESULT_EXECUTOR,
84                    limit: usize::try_from(self.memory_budget().resolved_bytes)
85                        .unwrap_or(usize::MAX),
86                    current: usize::try_from(self.memory_accounting().total_used_bytes())
87                        .unwrap_or(usize::MAX)
88                        .saturating_add(bytes),
89                }
90            })?)
91        } else {
92            None
93        };
94        Ok(OwnedQueryResult {
95            inner: Arc::new(ResultOwner {
96                value,
97                _reservation: reservation,
98            }),
99        })
100    }
101}