mongreldb_query/
query_registry.rs

1use crate::{MongrelQueryError, Result, SqlTestHook};
2use mongreldb_core::{CancellationReason, ExecutionControl};
3use sha2::{Digest, Sha256};
4use std::collections::{HashMap, VecDeque};
5use std::fmt;
6use std::str::FromStr;
7use std::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};
8use std::sync::{Arc, Mutex, Weak};
9use std::time::{Duration, Instant};
10
11const DEFAULT_MAX_ACTIVE: usize = 1_024;
12const DEFAULT_MAX_FINISHED: usize = 2_048;
13const DEFAULT_MAX_FINISHED_BYTES: usize = 2 * 1024 * 1024;
14const DEFAULT_FINISHED_TTL: Duration = Duration::from_secs(60);
15const MAX_METADATA_BYTES: usize = 256;
16
17#[derive(Clone, Copy, PartialEq, Eq, Hash)]
18pub struct QueryId([u8; 16]);
19
20impl QueryId {
21    pub fn random() -> Result<Self> {
22        let mut bytes = [0; 16];
23        getrandom::getrandom(&mut bytes).map_err(|error| {
24            MongrelQueryError::Core(mongreldb_core::MongrelError::Other(format!(
25                "query id randomness failed: {error}"
26            )))
27        })?;
28        Ok(Self(bytes))
29    }
30
31    pub fn as_bytes(&self) -> &[u8; 16] {
32        &self.0
33    }
34}
35
36impl fmt::Display for QueryId {
37    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
38        for byte in self.0 {
39            write!(formatter, "{byte:02x}")?;
40        }
41        Ok(())
42    }
43}
44
45impl fmt::Debug for QueryId {
46    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
47        fmt::Display::fmt(self, formatter)
48    }
49}
50
51impl FromStr for QueryId {
52    type Err = MongrelQueryError;
53
54    fn from_str(value: &str) -> Result<Self> {
55        if value.len() != 32 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
56            return Err(MongrelQueryError::Core(
57                mongreldb_core::MongrelError::InvalidArgument(
58                    "query id must be exactly 32 hexadecimal characters".into(),
59                ),
60            ));
61        }
62        let mut bytes = [0; 16];
63        for (index, chunk) in value.as_bytes().chunks_exact(2).enumerate() {
64            let text = std::str::from_utf8(chunk).map_err(|_| {
65                MongrelQueryError::Core(mongreldb_core::MongrelError::InvalidArgument(
66                    "query id is not valid UTF-8".into(),
67                ))
68            })?;
69            bytes[index] = u8::from_str_radix(text, 16).map_err(|_| {
70                MongrelQueryError::Core(mongreldb_core::MongrelError::InvalidArgument(
71                    "query id contains invalid hexadecimal".into(),
72                ))
73            })?;
74        }
75        Ok(Self(bytes))
76    }
77}
78
79impl serde::Serialize for QueryId {
80    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
81    where
82        S: serde::Serializer,
83    {
84        serializer.serialize_str(&self.to_string())
85    }
86}
87
88impl<'de> serde::Deserialize<'de> for QueryId {
89    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
90    where
91        D: serde::Deserializer<'de>,
92    {
93        let value = <String as serde::Deserialize>::deserialize(deserializer)?;
94        value.parse().map_err(serde::de::Error::custom)
95    }
96}
97
98#[derive(Debug, Clone, Default)]
99pub struct SqlQueryOptions {
100    pub query_id: Option<QueryId>,
101    pub timeout: Option<Duration>,
102    pub owner: Option<String>,
103    pub session_id: Option<String>,
104    pub parent_control: Option<ExecutionControl>,
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108#[repr(u8)]
109pub enum SqlQueryPhase {
110    Queued = 0,
111    Planning = 1,
112    Executing = 2,
113    Streaming = 3,
114    Serializing = 4,
115    CommitCritical = 5,
116    Cancelling = 6,
117    Completed = 7,
118    Failed = 8,
119    Cancelled = 9,
120}
121
122impl SqlQueryPhase {
123    fn from_u8(value: u8) -> Self {
124        match value {
125            0 => Self::Queued,
126            1 => Self::Planning,
127            2 => Self::Executing,
128            3 => Self::Streaming,
129            4 => Self::Serializing,
130            5 => Self::CommitCritical,
131            6 => Self::Cancelling,
132            7 => Self::Completed,
133            8 => Self::Failed,
134            9 => Self::Cancelled,
135            _ => Self::Failed,
136        }
137    }
138
139    fn is_terminal(self) -> bool {
140        matches!(self, Self::Completed | Self::Failed | Self::Cancelled)
141    }
142}
143
144#[derive(Debug, Clone, Copy, PartialEq, Eq)]
145pub enum CancelOutcome {
146    Accepted,
147    AlreadyCancelling,
148    TooLate,
149    AlreadyFinished,
150    NotFound,
151}
152
153#[derive(Debug, Clone)]
154pub struct QueryStatus {
155    pub query_id: QueryId,
156    pub owner: Option<String>,
157    pub session_id: Option<String>,
158    pub phase: SqlQueryPhase,
159    pub started_at: Instant,
160    pub deadline: Option<Instant>,
161    pub operation: String,
162    pub sql_fingerprint: [u8; 32],
163    pub cancellation_reason: CancellationReason,
164    pub committed: bool,
165    pub completed_statements: usize,
166    pub statement_index: usize,
167    pub cancel_requested_at: Option<Instant>,
168    pub queue_duration: Duration,
169    pub planning_duration: Duration,
170    pub execution_duration: Duration,
171    pub serialization_duration: Duration,
172    pub cancel_requested_phase: Option<SqlQueryPhase>,
173    pub cancel_observed_phase: Option<SqlQueryPhase>,
174    pub commit_fence_outcome: CommitFenceOutcome,
175}
176
177#[derive(Debug, Clone, Copy, PartialEq, Eq)]
178pub enum CommitFenceOutcome {
179    NotReached,
180    CancelWon,
181    CommitWon,
182}
183
184#[derive(Debug, Clone)]
185struct QueryTrace {
186    phase_started_at: Instant,
187    queue_duration: Duration,
188    planning_duration: Duration,
189    execution_duration: Duration,
190    serialization_duration: Duration,
191    cancel_requested_phase: Option<SqlQueryPhase>,
192    cancel_observed_phase: Option<SqlQueryPhase>,
193    commit_fence_outcome: CommitFenceOutcome,
194}
195
196impl QueryTrace {
197    fn new(started_at: Instant) -> Self {
198        Self {
199            phase_started_at: started_at,
200            queue_duration: Duration::ZERO,
201            planning_duration: Duration::ZERO,
202            execution_duration: Duration::ZERO,
203            serialization_duration: Duration::ZERO,
204            cancel_requested_phase: None,
205            cancel_observed_phase: None,
206            commit_fence_outcome: CommitFenceOutcome::NotReached,
207        }
208    }
209
210    fn transition(&mut self, phase: SqlQueryPhase) {
211        let now = Instant::now();
212        let elapsed = now.saturating_duration_since(self.phase_started_at);
213        match phase {
214            SqlQueryPhase::Queued => self.queue_duration += elapsed,
215            SqlQueryPhase::Planning => self.planning_duration += elapsed,
216            SqlQueryPhase::Executing | SqlQueryPhase::Streaming | SqlQueryPhase::CommitCritical => {
217                self.execution_duration += elapsed
218            }
219            SqlQueryPhase::Serializing => self.serialization_duration += elapsed,
220            SqlQueryPhase::Cancelling
221            | SqlQueryPhase::Completed
222            | SqlQueryPhase::Failed
223            | SqlQueryPhase::Cancelled => {}
224        }
225        self.phase_started_at = now;
226    }
227}
228
229#[derive(Debug)]
230struct RegisteredQuery {
231    id: QueryId,
232    owner: Option<String>,
233    session_id: Option<String>,
234    control: ExecutionControl,
235    phase: AtomicU8,
236    started_at: Instant,
237    deadline: Option<Instant>,
238    operation: Mutex<String>,
239    sql_fingerprint: Mutex<[u8; 32]>,
240    committed: AtomicBool,
241    completed_statements: AtomicUsize,
242    statement_index: AtomicUsize,
243    cancel_requested_at: Mutex<Option<Instant>>,
244    trace: Mutex<QueryTrace>,
245}
246
247impl RegisteredQuery {
248    fn phase(&self) -> SqlQueryPhase {
249        SqlQueryPhase::from_u8(self.phase.load(Ordering::Acquire))
250    }
251
252    fn status(&self) -> QueryStatus {
253        let mut trace = self.trace.lock().unwrap().clone();
254        trace.transition(self.phase());
255        QueryStatus {
256            query_id: self.id,
257            owner: self.owner.clone(),
258            session_id: self.session_id.clone(),
259            phase: self.phase(),
260            started_at: self.started_at,
261            deadline: self.deadline,
262            operation: self.operation.lock().unwrap().clone(),
263            sql_fingerprint: *self.sql_fingerprint.lock().unwrap(),
264            cancellation_reason: self.control.reason(),
265            committed: self.committed.load(Ordering::Acquire),
266            completed_statements: self.completed_statements.load(Ordering::Acquire),
267            statement_index: self.statement_index.load(Ordering::Acquire),
268            cancel_requested_at: *self.cancel_requested_at.lock().unwrap(),
269            queue_duration: trace.queue_duration,
270            planning_duration: trace.planning_duration,
271            execution_duration: trace.execution_duration,
272            serialization_duration: trace.serialization_duration,
273            cancel_requested_phase: trace.cancel_requested_phase,
274            cancel_observed_phase: trace.cancel_observed_phase,
275            commit_fence_outcome: trace.commit_fence_outcome,
276        }
277    }
278
279    fn record_transition(&self, phase: SqlQueryPhase) {
280        self.trace.lock().unwrap().transition(phase);
281    }
282}
283
284#[derive(Debug, Clone)]
285struct FinishedQuery {
286    status: QueryStatus,
287    finished_at: Instant,
288    approximate_bytes: usize,
289}
290
291#[derive(Debug, Default)]
292struct RegistryState {
293    active: HashMap<QueryId, Arc<RegisteredQuery>>,
294    finished: VecDeque<FinishedQuery>,
295    finished_bytes: usize,
296}
297
298#[derive(Debug)]
299pub struct SqlQueryRegistry {
300    state: Mutex<RegistryState>,
301    max_active: usize,
302    max_finished: usize,
303    max_finished_bytes: usize,
304    finished_ttl: Duration,
305}
306
307impl Default for SqlQueryRegistry {
308    fn default() -> Self {
309        Self::new(
310            DEFAULT_MAX_ACTIVE,
311            DEFAULT_MAX_FINISHED,
312            DEFAULT_MAX_FINISHED_BYTES,
313            DEFAULT_FINISHED_TTL,
314        )
315    }
316}
317
318impl SqlQueryRegistry {
319    pub fn new(
320        max_active: usize,
321        max_finished: usize,
322        max_finished_bytes: usize,
323        finished_ttl: Duration,
324    ) -> Self {
325        Self {
326            state: Mutex::new(RegistryState::default()),
327            max_active: max_active.max(1),
328            max_finished,
329            max_finished_bytes,
330            finished_ttl,
331        }
332    }
333
334    pub fn register(self: &Arc<Self>, options: SqlQueryOptions) -> Result<RegisteredSqlQuery> {
335        validate_metadata("owner", options.owner.as_deref())?;
336        validate_metadata("session id", options.session_id.as_deref())?;
337        let id = match options.query_id {
338            Some(id) => id,
339            None => QueryId::random()?,
340        };
341        let deadline = options.timeout.map(|timeout| Instant::now() + timeout);
342        let control = match options.parent_control {
343            Some(parent) => parent.child_with_deadline(deadline),
344            None => ExecutionControl::new(deadline),
345        };
346        let started_at = Instant::now();
347        let query = Arc::new(RegisteredQuery {
348            id,
349            owner: options.owner,
350            session_id: options.session_id,
351            deadline: control.deadline(),
352            control,
353            phase: AtomicU8::new(SqlQueryPhase::Queued as u8),
354            started_at,
355            operation: Mutex::new("UNKNOWN".into()),
356            sql_fingerprint: Mutex::new([0; 32]),
357            committed: AtomicBool::new(false),
358            completed_statements: AtomicUsize::new(0),
359            statement_index: AtomicUsize::new(0),
360            cancel_requested_at: Mutex::new(None),
361            trace: Mutex::new(QueryTrace::new(started_at)),
362        });
363        let mut state = self.state.lock().unwrap();
364        self.prune_locked(&mut state);
365        if state.active.contains_key(&id) {
366            return Err(MongrelQueryError::QueryIdConflict { query_id: id });
367        }
368        if state.active.len() >= self.max_active {
369            return Err(MongrelQueryError::QueryRegistryFull);
370        }
371        state.active.insert(id, Arc::clone(&query));
372        Ok(RegisteredSqlQuery {
373            registry: Arc::downgrade(self),
374            query,
375        })
376    }
377
378    pub fn cancel(&self, query_id: QueryId) -> CancelOutcome {
379        let query = {
380            let mut state = self.state.lock().unwrap();
381            self.prune_locked(&mut state);
382            if let Some(query) = state.active.get(&query_id) {
383                Some(Arc::clone(query))
384            } else if state
385                .finished
386                .iter()
387                .any(|finished| finished.status.query_id == query_id)
388            {
389                return CancelOutcome::AlreadyFinished;
390            } else {
391                return CancelOutcome::NotFound;
392            }
393        };
394        query
395            .unwrap()
396            .request_cancel(CancellationReason::ClientRequest)
397    }
398
399    pub fn status(&self, query_id: QueryId) -> Option<QueryStatus> {
400        let mut state = self.state.lock().unwrap();
401        self.prune_locked(&mut state);
402        state
403            .active
404            .get(&query_id)
405            .map(|query| query.status())
406            .or_else(|| {
407                state
408                    .finished
409                    .iter()
410                    .find(|finished| finished.status.query_id == query_id)
411                    .map(|finished| finished.status.clone())
412            })
413    }
414
415    pub fn cancel_session(&self, session_id: &str, reason: CancellationReason) -> usize {
416        let queries = {
417            let state = self.state.lock().unwrap();
418            state
419                .active
420                .values()
421                .filter(|query| query.session_id.as_deref() == Some(session_id))
422                .cloned()
423                .collect::<Vec<_>>()
424        };
425        let mut accepted = 0;
426        for query in queries {
427            if query.request_cancel(reason) == CancelOutcome::Accepted {
428                accepted += 1;
429            }
430        }
431        accepted
432    }
433
434    pub fn cancel_all(&self, reason: CancellationReason) -> usize {
435        let queries = self
436            .state
437            .lock()
438            .unwrap()
439            .active
440            .values()
441            .cloned()
442            .collect::<Vec<_>>();
443        let mut accepted = 0;
444        for query in queries {
445            if query.request_cancel(reason) == CancelOutcome::Accepted {
446                accepted += 1;
447            }
448        }
449        accepted
450    }
451
452    pub fn active_count(&self) -> usize {
453        self.state.lock().unwrap().active.len()
454    }
455
456    pub fn active_statuses(&self) -> Vec<QueryStatus> {
457        self.state
458            .lock()
459            .unwrap()
460            .active
461            .values()
462            .map(|query| query.status())
463            .collect()
464    }
465
466    pub fn active_for_session(&self, session_id: &str) -> usize {
467        self.state
468            .lock()
469            .unwrap()
470            .active
471            .values()
472            .filter(|query| query.session_id.as_deref() == Some(session_id))
473            .count()
474    }
475
476    pub fn queued_count(&self) -> usize {
477        self.state
478            .lock()
479            .unwrap()
480            .active
481            .values()
482            .filter(|query| query.phase() == SqlQueryPhase::Queued)
483            .count()
484    }
485
486    pub fn entry_count(&self) -> usize {
487        let mut state = self.state.lock().unwrap();
488        self.prune_locked(&mut state);
489        state.active.len() + state.finished.len()
490    }
491
492    pub fn approximate_bytes(&self) -> usize {
493        let mut state = self.state.lock().unwrap();
494        self.prune_locked(&mut state);
495        state.finished_bytes
496            + state
497                .active
498                .len()
499                .saturating_mul(std::mem::size_of::<RegisteredQuery>())
500    }
501
502    pub fn finished_count(&self) -> usize {
503        let mut state = self.state.lock().unwrap();
504        self.prune_locked(&mut state);
505        state.finished.len()
506    }
507
508    fn finish(&self, query: &Arc<RegisteredQuery>, phase: SqlQueryPhase) {
509        debug_assert!(phase.is_terminal());
510        let previous = query.phase();
511        query.phase.store(phase as u8, Ordering::Release);
512        query.record_transition(previous);
513        let status = query.status();
514        let approximate_bytes = std::mem::size_of::<FinishedQuery>()
515            + status.owner.as_ref().map_or(0, String::len)
516            + status.session_id.as_ref().map_or(0, String::len)
517            + status.operation.len();
518        let mut state = self.state.lock().unwrap();
519        if state.active.remove(&query.id).is_none() {
520            return;
521        }
522        if self.max_finished > 0 && self.max_finished_bytes > 0 {
523            state.finished_bytes = state.finished_bytes.saturating_add(approximate_bytes);
524            state.finished.push_back(FinishedQuery {
525                status,
526                finished_at: Instant::now(),
527                approximate_bytes,
528            });
529        }
530        self.prune_locked(&mut state);
531    }
532
533    fn prune_locked(&self, state: &mut RegistryState) {
534        let now = Instant::now();
535        while state.finished.front().is_some_and(|entry| {
536            now.saturating_duration_since(entry.finished_at) >= self.finished_ttl
537                || state.finished.len() > self.max_finished
538                || state.finished_bytes > self.max_finished_bytes
539        }) {
540            if let Some(entry) = state.finished.pop_front() {
541                state.finished_bytes = state.finished_bytes.saturating_sub(entry.approximate_bytes);
542            }
543        }
544    }
545}
546
547impl RegisteredQuery {
548    fn request_cancel(&self, reason: CancellationReason) -> CancelOutcome {
549        loop {
550            let phase = self.phase();
551            match phase {
552                SqlQueryPhase::Queued
553                | SqlQueryPhase::Planning
554                | SqlQueryPhase::Executing
555                | SqlQueryPhase::Streaming
556                | SqlQueryPhase::Serializing => {
557                    if self
558                        .phase
559                        .compare_exchange(
560                            phase as u8,
561                            SqlQueryPhase::Cancelling as u8,
562                            Ordering::AcqRel,
563                            Ordering::Acquire,
564                        )
565                        .is_ok()
566                    {
567                        *self.cancel_requested_at.lock().unwrap() = Some(Instant::now());
568                        let mut trace = self.trace.lock().unwrap();
569                        trace.transition(phase);
570                        trace.cancel_requested_phase = Some(phase);
571                        trace.commit_fence_outcome = CommitFenceOutcome::CancelWon;
572                        self.control.cancel(reason);
573                        return CancelOutcome::Accepted;
574                    }
575                }
576                SqlQueryPhase::Cancelling | SqlQueryPhase::Cancelled => {
577                    return CancelOutcome::AlreadyCancelling;
578                }
579                SqlQueryPhase::CommitCritical => return CancelOutcome::TooLate,
580                SqlQueryPhase::Completed | SqlQueryPhase::Failed => {
581                    return CancelOutcome::AlreadyFinished;
582                }
583            }
584        }
585    }
586}
587
588#[derive(Debug, Clone)]
589pub struct RegisteredSqlQuery {
590    registry: Weak<SqlQueryRegistry>,
591    query: Arc<RegisteredQuery>,
592}
593
594/// Query-specific state attached to a fresh DataFusion `TaskContext` for each
595/// execution. Reusable logical and physical plans never own this value.
596pub(crate) struct SqlTaskContext {
597    query: RegisteredSqlQuery,
598    test_hook: Option<SqlTestHook>,
599}
600
601impl SqlTaskContext {
602    pub(crate) fn new(query: RegisteredSqlQuery, test_hook: Option<SqlTestHook>) -> Self {
603        Self { query, test_hook }
604    }
605
606    pub(crate) fn query(&self) -> &RegisteredSqlQuery {
607        &self.query
608    }
609
610    pub(crate) fn test_hook(&self) -> Option<&SqlTestHook> {
611        self.test_hook.as_ref()
612    }
613}
614
615impl RegisteredSqlQuery {
616    pub fn id(&self) -> QueryId {
617        self.query.id
618    }
619
620    pub fn control(&self) -> &ExecutionControl {
621        &self.query.control
622    }
623
624    pub fn phase(&self) -> SqlQueryPhase {
625        self.query.phase()
626    }
627
628    pub fn status(&self) -> QueryStatus {
629        self.query.status()
630    }
631
632    pub fn set_sql_metadata(&self, sql: &str) {
633        let normalized = sql.split_whitespace().collect::<Vec<_>>().join(" ");
634        let operation = normalized
635            .split_whitespace()
636            .next()
637            .unwrap_or("UNKNOWN")
638            .to_ascii_uppercase();
639        *self.query.operation.lock().unwrap() = operation;
640        *self.query.sql_fingerprint.lock().unwrap() = Sha256::digest(normalized.as_bytes()).into();
641    }
642
643    pub fn transition(&self, expected: SqlQueryPhase, next: SqlQueryPhase) -> Result<()> {
644        self.query
645            .phase
646            .compare_exchange(
647                expected as u8,
648                next as u8,
649                Ordering::AcqRel,
650                Ordering::Acquire,
651            )
652            .map(|_| self.query.record_transition(expected))
653            .map_err(|actual| {
654                let actual = SqlQueryPhase::from_u8(actual);
655                if actual == SqlQueryPhase::Cancelling {
656                    self.cancellation_error()
657                } else {
658                    MongrelQueryError::InvalidQueryState(format!(
659                        "query {} expected {expected:?}, found {actual:?}",
660                        self.id()
661                    ))
662                }
663            })
664    }
665
666    pub fn enter_commit_critical(&self) -> Result<()> {
667        if let Err(error) = self.checkpoint() {
668            self.query.trace.lock().unwrap().commit_fence_outcome = CommitFenceOutcome::CancelWon;
669            return Err(error);
670        }
671        self.transition(SqlQueryPhase::Executing, SqlQueryPhase::CommitCritical)?;
672        self.query.trace.lock().unwrap().commit_fence_outcome = CommitFenceOutcome::CommitWon;
673        Ok(())
674    }
675
676    pub fn exit_commit_critical(&self) -> Result<()> {
677        self.transition(SqlQueryPhase::CommitCritical, SqlQueryPhase::Executing)
678    }
679
680    pub fn begin_serialization(&self) -> Result<()> {
681        match self.phase() {
682            SqlQueryPhase::Executing => {
683                self.transition(SqlQueryPhase::Executing, SqlQueryPhase::Serializing)
684            }
685            SqlQueryPhase::Streaming => {
686                self.transition(SqlQueryPhase::Streaming, SqlQueryPhase::Serializing)
687            }
688            // A durable COMMIT keeps its fence through response generation.
689            SqlQueryPhase::CommitCritical => Ok(()),
690            phase => Err(MongrelQueryError::InvalidQueryState(format!(
691                "query {} cannot serialize from {phase:?}",
692                self.id()
693            ))),
694        }
695    }
696
697    pub fn mark_committed(&self) {
698        self.query.committed.store(true, Ordering::Release);
699    }
700
701    pub fn begin_statement(&self, index: usize) {
702        self.query.statement_index.store(index, Ordering::Release);
703    }
704
705    pub fn complete_statement(&self, index: usize) {
706        self.query
707            .completed_statements
708            .store(index.saturating_add(1), Ordering::Release);
709    }
710
711    pub fn complete_current_statement(&self) {
712        let index = self.query.statement_index.load(Ordering::Acquire);
713        self.complete_statement(index);
714    }
715
716    pub fn request_cancel(&self, reason: CancellationReason) -> CancelOutcome {
717        self.query.request_cancel(reason)
718    }
719
720    pub fn checkpoint(&self) -> Result<()> {
721        let result = self
722            .query
723            .control
724            .checkpoint()
725            .map_err(|error| match error {
726                mongreldb_core::MongrelError::DeadlineExceeded => {
727                    MongrelQueryError::DeadlineExceeded {
728                        query_id: self.id(),
729                        timeout_ms: self.query.deadline.map(|deadline| {
730                            deadline
731                                .saturating_duration_since(self.query.started_at)
732                                .as_millis()
733                                .min(u128::from(u64::MAX)) as u64
734                        }),
735                        completed_statements: self
736                            .query
737                            .completed_statements
738                            .load(Ordering::Acquire),
739                        cancelled_statement_index: self
740                            .query
741                            .statement_index
742                            .load(Ordering::Acquire),
743                    }
744                }
745                mongreldb_core::MongrelError::Cancelled => self.cancellation_error(),
746                other => MongrelQueryError::Core(other),
747            });
748        if result.is_err() && self.query.control.is_cancelled() {
749            let mut trace = self.query.trace.lock().unwrap();
750            trace.cancel_observed_phase = trace.cancel_requested_phase.or(Some(self.phase()));
751        }
752        result
753    }
754
755    pub fn complete(&self) {
756        self.finish(SqlQueryPhase::Completed);
757    }
758
759    pub fn fail(&self) {
760        let phase = if self.query.control.is_cancelled() {
761            SqlQueryPhase::Cancelled
762        } else {
763            SqlQueryPhase::Failed
764        };
765        self.finish(phase);
766    }
767
768    fn finish(&self, phase: SqlQueryPhase) {
769        if let Some(registry) = self.registry.upgrade() {
770            registry.finish(&self.query, phase);
771        } else {
772            self.query.phase.store(phase as u8, Ordering::Release);
773        }
774    }
775
776    fn cancellation_error(&self) -> MongrelQueryError {
777        match self.query.control.reason() {
778            CancellationReason::Deadline => MongrelQueryError::DeadlineExceeded {
779                query_id: self.id(),
780                timeout_ms: self.query.deadline.map(|deadline| {
781                    deadline
782                        .saturating_duration_since(self.query.started_at)
783                        .as_millis()
784                        .min(u128::from(u64::MAX)) as u64
785                }),
786                completed_statements: self.query.completed_statements.load(Ordering::Acquire),
787                cancelled_statement_index: self.query.statement_index.load(Ordering::Acquire),
788            },
789            reason => MongrelQueryError::QueryCancelled {
790                query_id: self.id(),
791                reason,
792                completed_statements: self.query.completed_statements.load(Ordering::Acquire),
793                cancelled_statement_index: self.query.statement_index.load(Ordering::Acquire),
794            },
795        }
796    }
797}
798
799pub struct RegisteredQueryGuard {
800    query: Option<RegisteredSqlQuery>,
801}
802
803impl RegisteredQueryGuard {
804    pub fn new(query: RegisteredSqlQuery) -> Self {
805        Self { query: Some(query) }
806    }
807
808    pub fn query(&self) -> &RegisteredSqlQuery {
809        self.query
810            .as_ref()
811            .expect("registered query guard consumed")
812    }
813
814    pub fn complete(mut self) {
815        if let Some(query) = self.query.take() {
816            query.complete();
817        }
818    }
819
820    pub fn fail(mut self) {
821        if let Some(query) = self.query.take() {
822            query.fail();
823        }
824    }
825
826    pub fn into_query(mut self) -> RegisteredSqlQuery {
827        self.query.take().expect("registered query guard consumed")
828    }
829}
830
831impl Drop for RegisteredQueryGuard {
832    fn drop(&mut self) {
833        if let Some(query) = self.query.take() {
834            query.fail();
835        }
836    }
837}
838
839fn validate_metadata(name: &str, value: Option<&str>) -> Result<()> {
840    if value.is_some_and(|value| value.len() > MAX_METADATA_BYTES) {
841        return Err(MongrelQueryError::Core(
842            mongreldb_core::MongrelError::InvalidArgument(format!(
843                "{name} exceeds {MAX_METADATA_BYTES} bytes"
844            )),
845        ));
846    }
847    Ok(())
848}
849
850#[cfg(test)]
851mod tests {
852    use super::*;
853
854    #[test]
855    fn query_ids_are_random_strict_and_round_trip() {
856        let first = QueryId::random().unwrap();
857        let second = QueryId::random().unwrap();
858        assert_ne!(first, second);
859        assert_eq!(first.to_string().parse::<QueryId>().unwrap(), first);
860        assert!("abc".parse::<QueryId>().is_err());
861        assert!("0000000000000000000000000000000z"
862            .parse::<QueryId>()
863            .is_err());
864    }
865
866    #[test]
867    fn duplicate_active_ids_are_rejected_and_cleanup_is_bounded() {
868        let registry = Arc::new(SqlQueryRegistry::new(1, 1, 1024, Duration::from_secs(60)));
869        let id = QueryId::random().unwrap();
870        let query = registry
871            .register(SqlQueryOptions {
872                query_id: Some(id),
873                ..SqlQueryOptions::default()
874            })
875            .unwrap();
876        assert!(matches!(
877            registry.register(SqlQueryOptions {
878                query_id: Some(id),
879                ..SqlQueryOptions::default()
880            }),
881            Err(MongrelQueryError::QueryIdConflict { .. })
882        ));
883        query.complete();
884        assert_eq!(registry.active_count(), 0);
885        assert_eq!(registry.finished_count(), 1);
886    }
887
888    #[test]
889    fn cancel_and_commit_fence_have_one_winner() {
890        for cancel_first in [true, false] {
891            let registry = Arc::new(SqlQueryRegistry::default());
892            let query = registry.register(SqlQueryOptions::default()).unwrap();
893            query
894                .transition(SqlQueryPhase::Queued, SqlQueryPhase::Executing)
895                .unwrap();
896            if cancel_first {
897                assert_eq!(
898                    query.request_cancel(CancellationReason::ClientRequest),
899                    CancelOutcome::Accepted
900                );
901                assert!(query.enter_commit_critical().is_err());
902                let status = query.status();
903                assert_eq!(
904                    status.cancel_requested_phase,
905                    Some(SqlQueryPhase::Executing)
906                );
907                assert_eq!(status.cancel_observed_phase, Some(SqlQueryPhase::Executing));
908                assert_eq!(status.commit_fence_outcome, CommitFenceOutcome::CancelWon);
909            } else {
910                query.enter_commit_critical().unwrap();
911                assert_eq!(
912                    query.request_cancel(CancellationReason::ClientRequest),
913                    CancelOutcome::TooLate
914                );
915                query.mark_committed();
916                query.complete();
917                let status = registry.status(query.id()).unwrap();
918                assert!(status.committed);
919                assert_eq!(status.commit_fence_outcome, CommitFenceOutcome::CommitWon);
920            }
921        }
922    }
923
924    #[test]
925    fn guard_cleans_up_dropped_execution_without_raw_sql() {
926        let registry = Arc::new(SqlQueryRegistry::default());
927        let query = registry
928            .register(SqlQueryOptions {
929                owner: Some("alice".into()),
930                session_id: Some("session".into()),
931                ..SqlQueryOptions::default()
932            })
933            .unwrap();
934        query.set_sql_metadata("SELECT secret FROM docs WHERE token = 'private'");
935        let id = query.id();
936        drop(RegisteredQueryGuard::new(query));
937        let status = registry.status(id).unwrap();
938        assert_eq!(status.phase, SqlQueryPhase::Failed);
939        assert_eq!(status.operation, "SELECT");
940        assert_ne!(status.sql_fingerprint, [0; 32]);
941    }
942}
mongreldb_query/query_registry.rs

mongreldb_query/
query_registry.rs
