Skip to main content

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}