velesdb_core/guardrails/context.rs
1//! Query execution context with guard-rail tracking (EPIC-048).
2//!
3//! Tracks per-query resource consumption (time, depth, cardinality, memory)
4//! and enforces the configured limits.
5
6// Reason: Numeric casts in guardrails are intentional:
7// - u128->u64 for millisecond durations: durations fit within u64 (thousands of years)
8// - Used for timeout checking and logging, not precise calculations
9#![allow(clippy::cast_possible_truncation)]
10
11use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
12use std::time::{Duration, Instant};
13
14use super::limits::{GuardRailViolation, QueryLimits};
15
16/// Query execution context with guard-rail tracking (EPIC-048).
17#[derive(Debug)]
18pub struct QueryContext {
19 /// Query limits configuration.
20 pub limits: QueryLimits,
21 /// Query start time.
22 start_time: Instant,
23 /// Current traversal depth.
24 current_depth: AtomicU64,
25 /// Current cardinality (intermediate results count).
26 current_cardinality: AtomicUsize,
27 /// Estimated memory usage in bytes.
28 memory_used: AtomicUsize,
29}
30
31impl QueryContext {
32 /// Creates a new query context with the given limits.
33 #[must_use]
34 pub fn new(limits: QueryLimits) -> Self {
35 Self {
36 limits,
37 start_time: Instant::now(),
38 current_depth: AtomicU64::new(0),
39 current_cardinality: AtomicUsize::new(0),
40 memory_used: AtomicUsize::new(0),
41 }
42 }
43
44 /// Checks if the query has timed out (US-001).
45 ///
46 /// # Errors
47 ///
48 /// Returns [`GuardRailViolation::Timeout`] when elapsed time exceeds
49 /// the configured timeout.
50 pub fn check_timeout(&self) -> Result<(), GuardRailViolation> {
51 // timeout_ms == 0 means "disabled" — never fire.
52 if self.limits.timeout_ms == 0 {
53 return Ok(());
54 }
55 let elapsed_ms = self.start_time.elapsed().as_millis() as u64;
56 if elapsed_ms >= self.limits.timeout_ms {
57 return Err(GuardRailViolation::Timeout {
58 max_ms: self.limits.timeout_ms,
59 elapsed_ms,
60 });
61 }
62 Ok(())
63 }
64
65 /// Checks and updates traversal depth (US-002).
66 ///
67 /// # Errors
68 ///
69 /// Returns [`GuardRailViolation::DepthExceeded`] when `depth` is greater
70 /// than the configured maximum.
71 pub fn check_depth(&self, depth: u32) -> Result<(), GuardRailViolation> {
72 self.current_depth
73 .store(u64::from(depth), Ordering::Relaxed);
74 if depth > self.limits.max_depth {
75 return Err(GuardRailViolation::DepthExceeded {
76 max: self.limits.max_depth,
77 actual: depth,
78 });
79 }
80 Ok(())
81 }
82
83 /// Checks and updates cardinality (US-003).
84 ///
85 /// # Errors
86 ///
87 /// Returns [`GuardRailViolation::CardinalityExceeded`] when cumulative
88 /// intermediate result count exceeds the configured maximum.
89 ///
90 /// # Known Limitation
91 ///
92 /// This method is called on the final result set (post-filter, post-ORDER BY,
93 /// pre-LIMIT). It does **not** track intermediate over-fetched candidate sets
94 /// (e.g., `candidates_k = execution_limit * 10 * N` during similarity search).
95 /// Those are bounded by `MAX_LIMIT` internally and therefore do not escape.
96 /// Future work: thread `QueryContext` into ANN search to track intermediates.
97 pub fn check_cardinality(&self, count: usize) -> Result<(), GuardRailViolation> {
98 let current = self.current_cardinality.fetch_add(count, Ordering::Relaxed) + count;
99 if current > self.limits.max_cardinality {
100 return Err(GuardRailViolation::CardinalityExceeded {
101 max: self.limits.max_cardinality,
102 actual: current,
103 });
104 }
105 Ok(())
106 }
107
108 /// Checks and updates memory usage (US-004).
109 ///
110 /// # Errors
111 ///
112 /// Returns [`GuardRailViolation::MemoryExceeded`] when cumulative estimated
113 /// memory usage exceeds the configured budget.
114 pub fn check_memory(&self, bytes: usize) -> Result<(), GuardRailViolation> {
115 let current = self.memory_used.fetch_add(bytes, Ordering::Relaxed) + bytes;
116 if current > self.limits.memory_limit_bytes {
117 return Err(GuardRailViolation::MemoryExceeded {
118 max_bytes: self.limits.memory_limit_bytes,
119 used_bytes: current,
120 });
121 }
122 Ok(())
123 }
124
125 /// Returns elapsed time since query start.
126 #[must_use]
127 pub fn elapsed(&self) -> Duration {
128 self.start_time.elapsed()
129 }
130
131 /// Returns current memory usage estimate.
132 #[must_use]
133 pub fn memory_used(&self) -> usize {
134 self.memory_used.load(Ordering::Relaxed)
135 }
136}