writerslogic_warm/core/
query_limits.rs1#[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}