Skip to main content

writerslogic_warm/core/
query_limits.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4/// Bounds for optional reasoning and concept queries. Exhaustion returns an error, never
5/// a partial answer. Byte accounting includes retained inputs and intermediate frontiers.
6#[derive(Clone, Copy, Debug)]
7pub struct QueryLimits {
8    pub max_work: usize,
9    pub max_frontier: usize,
10    pub max_bytes: usize,
11}
12
13impl Default for QueryLimits {
14    fn default() -> Self {
15        Self {
16            max_work: 10_000_000,
17            max_frontier: 16_384,
18            max_bytes: 64 * 1024 * 1024,
19        }
20    }
21}
22
23#[derive(Debug, thiserror::Error, PartialEq, Eq)]
24#[error("query {resource} limit {limit} exceeded")]
25pub struct QueryLimitExceeded {
26    pub resource: &'static str,
27    pub limit: usize,
28}
29
30pub(crate) struct QueryBudget {
31    limits: QueryLimits,
32    work: usize,
33    bytes: usize,
34}
35
36impl QueryBudget {
37    pub(crate) fn new(limits: QueryLimits) -> Self {
38        Self {
39            limits,
40            work: 0,
41            bytes: 0,
42        }
43    }
44
45    pub(crate) fn work(&mut self, amount: usize) -> Result<(), QueryLimitExceeded> {
46        Self::charge(&mut self.work, amount, self.limits.max_work, "work")
47    }
48
49    pub(crate) fn reserve(&mut self, amount: usize) -> Result<(), QueryLimitExceeded> {
50        Self::charge(&mut self.bytes, amount, self.limits.max_bytes, "bytes")
51    }
52
53    pub(crate) fn release(&mut self, amount: usize) {
54        self.bytes = self
55            .bytes
56            .checked_sub(amount)
57            .expect("query budget released unreserved bytes");
58    }
59
60    pub(crate) fn frontier(&self, count: usize) -> Result<(), QueryLimitExceeded> {
61        if count > self.limits.max_frontier {
62            Err(QueryLimitExceeded {
63                resource: "frontier",
64                limit: self.limits.max_frontier,
65            })
66        } else {
67            Ok(())
68        }
69    }
70
71    fn charge(
72        used: &mut usize,
73        amount: usize,
74        limit: usize,
75        resource: &'static str,
76    ) -> Result<(), QueryLimitExceeded> {
77        *used = used
78            .checked_add(amount)
79            .filter(|&total| total <= limit)
80            .ok_or(QueryLimitExceeded { resource, limit })?;
81        Ok(())
82    }
83}