linkrs_query/executor/streaming/runtime/
resources.rs

1//! Runtime resource lifecycle.
2//!
3//! [`ResourceOwner`] accumulates cleanup callbacks (cursors, temp files,
4//! spill reservations) released together at query teardown.
5//! [`QueryFinishGuard`] is an RAII guard that marks a query finished in the
6//! [`QueryManager`] on drop, covering explicit-finish omission and panics.
7
8use std::sync::Arc;
9
10use crate::query_manager::QueryManager;
11
12/// Manages cleanup of runtime resources (cursors, temp files, etc.)
13pub struct ResourceOwner {
14    cleanup: Vec<Box<dyn FnOnce() + Send>>,
15}
16
17impl std::fmt::Debug for ResourceOwner {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        f.debug_struct("ResourceOwner")
20            .field("cleanup_count", &self.cleanup.len())
21            .finish()
22    }
23}
24
25impl Default for ResourceOwner {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31impl ResourceOwner {
32    pub fn new() -> Self {
33        Self {
34            cleanup: Vec::new(),
35        }
36    }
37
38    pub fn add(&mut self, cleanup: Box<dyn FnOnce() + Send>) {
39        self.cleanup.push(cleanup);
40    }
41
42    pub fn release_all(&mut self) {
43        for f in self.cleanup.drain(..) {
44            f();
45        }
46    }
47}
48
49/// RAII guard that marks a query as finished in the QueryManager on drop.
50///
51/// Created by [`crate::executor::streaming::runtime::ExecutionRuntime::finish_guard`].
52/// Ensures the query lifecycle is tracked even when the caller forgets to
53/// call finish explicitly, or when execution panics mid-flight.
54#[derive(Debug)]
55pub struct QueryFinishGuard {
56    query_manager: Arc<QueryManager>,
57    query_id: i64,
58    finished: bool,
59}
60
61impl QueryFinishGuard {
62    pub fn new(query_manager: Arc<QueryManager>, query_id: i64) -> Self {
63        Self {
64            query_manager,
65            query_id,
66            finished: false,
67        }
68    }
69
70    /// Mark the query as finished immediately without waiting for Drop.
71    pub fn finish(&mut self) {
72        if !self.finished {
73            self.finished = true;
74            let _ = self.query_manager.finish_query(self.query_id);
75        }
76    }
77}
78
79impl Drop for QueryFinishGuard {
80    fn drop(&mut self) {
81        if !self.finished {
82            let _ = self.query_manager.finish_query(self.query_id);
83        }
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn test_resource_owner() {
93        let mut owner = ResourceOwner::new();
94        let flag = Arc::new(std::sync::atomic::AtomicBool::new(false));
95        let f = flag.clone();
96        owner.add(Box::new(move || {
97            f.store(true, std::sync::atomic::Ordering::Relaxed);
98        }));
99        owner.release_all();
100        assert!(flag.load(std::sync::atomic::Ordering::Relaxed));
101    }
102}
linkrs_query/executor/streaming/runtime/resources.rs

linkrs_query/executor/streaming/runtime/
resources.rs
