Skip to main content

uqa_storage/
read_control.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Query-owned provider reads share allocation limits and cancellation with their consumers.
8
9use crate::StorageBackendResult;
10use uqa_core::memory::{MemoryBudget, MemoryError};
11pub use uqa_core::CancellationToken;
12
13pub type ValueReadVisitor<'a> = dyn FnMut(Option<&[u8]>) -> StorageBackendResult<()> + 'a;
14pub type KeyValueReadVisitor<'a> = dyn FnMut(&[u8], &[u8]) -> StorageBackendResult<()> + 'a;
15pub type KeyReadVisitor<'a> = dyn FnMut(&[u8]) -> StorageBackendResult<()> + 'a;
16
17#[derive(Clone, Debug)]
18pub struct StorageReadControl {
19    memory: MemoryBudget,
20    cancellation: CancellationToken,
21}
22impl StorageReadControl {
23    pub fn with_limit(limit: usize) -> Self {
24        Self {
25            memory: MemoryBudget::new(limit),
26            cancellation: CancellationToken::new(),
27        }
28    }
29    pub fn new(memory: &MemoryBudget, cancellation: &CancellationToken) -> Self {
30        Self {
31            memory: memory.clone(),
32            cancellation: cancellation.clone(),
33        }
34    }
35    pub fn memory(&self) -> &MemoryBudget {
36        &self.memory
37    }
38    pub fn cancellation(&self) -> &CancellationToken {
39        &self.cancellation
40    }
41    pub(crate) fn shares_context(&self, other: &Self) -> bool {
42        self.memory.shares_allowance(&other.memory)
43            && self.cancellation.shares_signal(&other.cancellation)
44    }
45    pub fn check(&self) -> StorageBackendResult<()> {
46        self.cancellation.check()?;
47        Ok(())
48    }
49
50    /// Validate an encoded value's size before materializing it. Provider workspace and returned bytes still share this control's memory allowance.
51    pub fn check_value_size(&self, bytes: usize, maximum: usize) -> StorageBackendResult<()> {
52        self.check()?;
53        if bytes > maximum {
54            return Err(MemoryError::Limit {
55                required: bytes,
56                limit: maximum,
57            }
58            .into());
59        }
60        Ok(())
61    }
62}