linkrs_query/
query_manager.rs

1//! Query Manager
2//!
3//! Responsible for tracking and managing the queries that are currently in progress.
4
5use dashmap::DashMap;
6use log::{info, warn};
7use parking_lot::RwLock;
8use std::sync::atomic::{AtomicI64, Ordering};
9use std::sync::Arc;
10use std::time::{SystemTime, UNIX_EPOCH};
11
12use linkrs_core::error::{ManagerError, ManagerResult};
13
14use super::session_events::{SessionEvent, SessionEventCallback};
15use crate::executor::streaming::query_registry::{QueryId, QueryRegistry};
16use crate::executor::streaming::transaction_scope::CancelReason;
17use linkrs_core::event_dispatch::{EventFilter, EventSubscriptions, SubscriptionId};
18
19/// Row-count progress notification for a running query.
20///
21/// Independent lightweight hook: deliberately not a `SessionEvent`
22/// variant, so high-frequency progress never touches the low-frequency
23/// session-lifecycle matches. Zero overhead when no observer is registered.
24/// The executor calls `emit_progress` only at its configured row interval.
25/// The notification is valid only for the dispatch call; observers must
26/// not retain it beyond the callback.
27#[derive(Debug, Clone)]
28pub struct QueryProgress {
29    pub session_id: i64,
30    pub query_id: i64,
31    pub rows_processed: u64,
32}
33
34/// Runtime observer for query progress.
35pub type QueryProgressCallback = Arc<dyn Fn(&QueryProgress) + Send + Sync>;
36
37/// Query status
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum QueryStatus {
40    Running,
41    Finished,
42    Failed,
43    Killed,
44}
45
46/// Query information
47#[derive(Debug, Clone)]
48pub struct QueryInfo {
49    pub query_id: i64,
50    pub session_id: i64,
51    pub user_name: String,
52    pub space_name: Option<String>,
53    pub query_text: String,
54    pub status: QueryStatus,
55    pub start_time: SystemTime,
56    pub duration_ms: Option<i64>,
57    pub execution_plan: Option<String>,
58}
59
60impl QueryInfo {
61    pub fn new(
62        query_id: i64,
63        session_id: i64,
64        user_name: String,
65        space_name: Option<String>,
66        query_text: String,
67    ) -> Self {
68        Self {
69            query_id,
70            session_id,
71            user_name,
72            space_name,
73            query_text,
74            status: QueryStatus::Running,
75            start_time: SystemTime::now(),
76            duration_ms: None,
77            execution_plan: None,
78        }
79    }
80
81    pub fn finish(&mut self) {
82        self.status = QueryStatus::Finished;
83        self.duration_ms = Some(
84            SystemTime::now()
85                .duration_since(self.start_time)
86                .unwrap_or_default()
87                .as_millis() as i64,
88        );
89    }
90
91    pub fn fail(&mut self) {
92        self.status = QueryStatus::Failed;
93        self.duration_ms = Some(
94            SystemTime::now()
95                .duration_since(self.start_time)
96                .unwrap_or_default()
97                .as_millis() as i64,
98        );
99    }
100
101    pub fn kill(&mut self) {
102        self.status = QueryStatus::Killed;
103        self.duration_ms = Some(
104            SystemTime::now()
105                .duration_since(self.start_time)
106                .unwrap_or_default()
107                .as_millis() as i64,
108        );
109    }
110}
111
112/// Query Manager
113///
114/// Single enum dual emission source: query lifecycle halves
115/// (`QueryStarted` / `QueryCompleted` / `SlowQueryDetected`) are emitted here,
116/// session lifecycle halves (`SessionCreated` / `SessionDestroyed`) by
117/// `GraphSessionManager`. Both sides can share one
118/// `Arc<EventSubscriptions<SessionEvent>>` via `new_with_shared` so observers
119/// subscribe once and receive both halves.
120pub struct QueryManager {
121    queries: DashMap<i64, QueryInfo>,
122    next_query_id: AtomicI64,
123    session_callbacks: Arc<EventSubscriptions<SessionEvent>>,
124    progress_callbacks: Arc<EventSubscriptions<QueryProgress>>,
125    // Reverse bridge for KILL QUERY: when the assembly registers the
126    // executor's QueryRegistry here, kill_query also cancels the token.
127    // Absent mapping degrades to status-marking only.
128    query_registry: RwLock<Option<Arc<QueryRegistry>>>,
129    slow_query_threshold_ms: AtomicI64,
130}
131
132impl std::fmt::Debug for QueryManager {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        f.debug_struct("QueryManager")
135            .field("queries_count", &self.queries.len())
136            .field("session_callbacks", &self.session_callbacks.len())
137            .field("progress_callbacks", &self.progress_callbacks.len())
138            .field(
139                "slow_query_threshold_ms",
140                &self.slow_query_threshold_ms.load(Ordering::SeqCst),
141            )
142            .finish()
143    }
144}
145
146impl QueryManager {
147    pub fn new() -> Self {
148        Self {
149            queries: DashMap::new(),
150            next_query_id: AtomicI64::new(1),
151            session_callbacks: Arc::new(EventSubscriptions::new()),
152            progress_callbacks: Arc::new(EventSubscriptions::new()),
153            query_registry: RwLock::new(None),
154            slow_query_threshold_ms: AtomicI64::new(1000),
155        }
156    }
157
158    /// Build a manager sharing one session-event registry with
159    /// `GraphSessionManager`.
160    pub fn new_with_shared(shared: Arc<EventSubscriptions<SessionEvent>>) -> Self {
161        Self {
162            queries: DashMap::new(),
163            next_query_id: AtomicI64::new(1),
164            session_callbacks: shared,
165            progress_callbacks: Arc::new(EventSubscriptions::new()),
166            query_registry: RwLock::new(None),
167            slow_query_threshold_ms: AtomicI64::new(1000),
168        }
169    }
170
171    /// Shared session-event registry behind this manager.
172    pub fn shared_session_callbacks(&self) -> Arc<EventSubscriptions<SessionEvent>> {
173        Arc::clone(&self.session_callbacks)
174    }
175
176    /// Point this manager at a shared session-event registry.
177    pub fn set_shared_session_callbacks(&mut self, shared: Arc<EventSubscriptions<SessionEvent>>) {
178        self.session_callbacks = shared;
179    }
180
181    /// Register a runtime observer for session/query events.
182    pub fn register_session_callback(&self, callback: SessionEventCallback) -> SubscriptionId {
183        self.session_callbacks.add(callback)
184    }
185
186    /// Register a filtered observer invoked only when `filter` returns true.
187    pub fn register_session_callback_filtered(
188        &self,
189        callback: SessionEventCallback,
190        filter: EventFilter<SessionEvent>,
191    ) -> SubscriptionId {
192        self.session_callbacks.add_filtered(callback, Some(filter))
193    }
194
195    /// Remove a previously registered observer. Returns true if present.
196    pub fn unregister_session_callback(&self, id: SubscriptionId) -> bool {
197        self.session_callbacks.remove(id)
198    }
199
200    /// Number of registered session observers.
201    pub fn session_callback_count(&self) -> usize {
202        self.session_callbacks.len()
203    }
204
205    /// Attach the executor's query registry so `kill_query` also cancels the
206    /// execution token. Without it, `kill_query` only marks the query status.
207    pub fn set_query_registry(&self, registry: Arc<QueryRegistry>) {
208        *self.query_registry.write() = Some(registry);
209    }
210
211    /// Detach the executor's query registry; `kill_query` returns to
212    /// status-marking only.
213    pub fn clear_query_registry(&self) {
214        *self.query_registry.write() = None;
215    }
216
217    /// Register a progress observer. The executor invokes it at its own row
218    /// interval; unregistered by default, hence zero hot-path overhead.
219    pub fn register_progress_callback(&self, callback: QueryProgressCallback) -> SubscriptionId {
220        self.progress_callbacks.add(callback)
221    }
222
223    /// Register a progress observer thinned to every `rows_interval` rows.
224    ///
225    /// Observer-side thinning: deliveries whose `rows_processed` is not a
226    /// multiple of `rows_interval` are dropped before reaching `callback`.
227    /// `rows_interval == 0` disables thinning (every emission is delivered).
228    /// The executor still decides when to call `emit_progress`; pairing a
229    /// coarse executor cadence with a fine observer interval costs nothing
230    /// extra, while a fine executor cadence with a coarse observer interval
231    /// only pays for the dropped filter checks.
232    pub fn register_progress_callback_with_interval(
233        &self,
234        callback: QueryProgressCallback,
235        rows_interval: u64,
236    ) -> SubscriptionId {
237        if rows_interval == 0 {
238            return self.progress_callbacks.add(callback);
239        }
240        self.progress_callbacks.add_filtered(
241            callback,
242            Some(Arc::new(move |progress: &QueryProgress| {
243                progress.rows_processed.is_multiple_of(rows_interval)
244            })),
245        )
246    }
247
248    /// Register a filtered progress observer.
249    pub fn register_progress_callback_filtered(
250        &self,
251        callback: QueryProgressCallback,
252        filter: EventFilter<QueryProgress>,
253    ) -> SubscriptionId {
254        self.progress_callbacks.add_filtered(callback, Some(filter))
255    }
256
257    /// Remove a previously registered progress observer.
258    pub fn unregister_progress_callback(&self, id: SubscriptionId) -> bool {
259        self.progress_callbacks.remove(id)
260    }
261
262    /// Number of registered progress observers.
263    pub fn progress_callback_count(&self) -> usize {
264        self.progress_callbacks.len()
265    }
266
267    /// Emit a progress notification. No-op when nobody listens.
268    ///
269    /// Explicit executor entry point: streaming operators call this when
270    /// their processed row count reaches the configured cadence. Callers
271    /// outside the executor (tests, manual drivers) may invoke it directly.
272    pub fn emit_progress(&self, session_id: i64, query_id: i64, rows_processed: u64) {
273        if self.progress_callbacks.is_empty() {
274            return;
275        }
276        self.progress_callbacks.dispatch(
277            "progress",
278            &QueryProgress {
279                session_id,
280                query_id,
281                rows_processed,
282            },
283        );
284    }
285
286    /// Configure the threshold used to emit `SlowQueryDetected`.
287    pub fn set_slow_query_threshold_ms(&self, threshold_ms: i64) {
288        self.slow_query_threshold_ms
289            .store(threshold_ms, Ordering::SeqCst);
290    }
291
292    /// Current slow-query threshold in milliseconds.
293    pub fn slow_query_threshold_ms(&self) -> i64 {
294        self.slow_query_threshold_ms.load(Ordering::SeqCst)
295    }
296
297    fn emit_session_event(&self, event: SessionEvent) {
298        if self.session_callbacks.is_empty() {
299            // Slow queries are still worth a log line even without observers.
300            if let SessionEvent::SlowQueryDetected {
301                session_id,
302                query_id,
303                duration_ms,
304                threshold_ms,
305            } = &event
306            {
307                warn!(
308                    "Slow query detected: session={}, query={}, duration={}ms threshold={}ms",
309                    session_id, query_id, duration_ms, threshold_ms
310                );
311            }
312            return;
313        }
314        self.session_callbacks.dispatch("session", &event);
315    }
316
317    fn check_slow_query(&self, session_id: i64, query_id: i64, duration_ms: i64) {
318        let threshold = self.slow_query_threshold_ms.load(Ordering::SeqCst);
319        if duration_ms >= threshold {
320            self.emit_session_event(SessionEvent::SlowQueryDetected {
321                session_id,
322                query_id,
323                duration_ms,
324                threshold_ms: threshold,
325            });
326        }
327    }
328
329    /// Generate a new query ID.
330    fn generate_query_id(&self) -> i64 {
331        self.next_query_id.fetch_add(1, Ordering::SeqCst)
332    }
333
334    /// Register a new query
335    pub fn register_query(
336        &self,
337        session_id: i64,
338        user_name: String,
339        space_name: Option<String>,
340        query_text: String,
341    ) -> i64 {
342        let query_id = self.generate_query_id();
343        let query_info = QueryInfo::new(
344            query_id,
345            session_id,
346            user_name,
347            space_name,
348            query_text.clone(),
349        );
350
351        self.queries.insert(query_id, query_info);
352
353        info!(
354            "Query registered: id={}, session_id={}, query={}",
355            query_id, session_id, query_text
356        );
357
358        self.emit_session_event(SessionEvent::QueryStarted {
359            session_id,
360            query_id,
361            query_text,
362        });
363
364        query_id
365    }
366
367    /// Register a new query under a caller-assigned id.
368    ///
369    /// The network service assigns one monotonic id per execution and uses
370    /// it for the engine request, this registry, and the session tracker,
371    /// so SHOW QUERIES rows double as KILL QUERY targets. The id lane
372    /// advances past the fixed id to keep generated ids collision-free.
373    /// An id already held by a running query is rejected; a stale finished
374    /// entry under the same id is replaced.
375    pub fn register_query_with_id(
376        &self,
377        query_id: i64,
378        session_id: i64,
379        user_name: String,
380        space_name: Option<String>,
381        query_text: String,
382    ) -> ManagerResult<()> {
383        if query_id < 0 {
384            return Err(ManagerError::InvalidInput(format!(
385                "invalid query id {query_id}"
386            )));
387        }
388        if let Some(existing) = self.queries.get(&query_id) {
389            if existing.status == QueryStatus::Running {
390                return Err(ManagerError::InvalidInput(format!(
391                    "query id {query_id} is already running"
392                )));
393            }
394        }
395        self.next_query_id
396            .fetch_max(query_id.saturating_add(1), Ordering::SeqCst);
397        let query_info = QueryInfo::new(
398            query_id,
399            session_id,
400            user_name,
401            space_name,
402            query_text.clone(),
403        );
404        self.queries.insert(query_id, query_info);
405        info!(
406            "Query registered: id={}, session_id={}, query={}",
407            query_id, session_id, query_text
408        );
409        self.emit_session_event(SessionEvent::QueryStarted {
410            session_id,
411            query_id,
412            query_text,
413        });
414        Ok(())
415    }
416
417    /// Complete the query.
418    pub fn finish_query(&self, query_id: i64) -> ManagerResult<()> {
419        if let Some(mut query) = self.queries.get_mut(&query_id) {
420            query.finish();
421            let duration_ms = query.duration_ms.unwrap_or(0);
422            let session_id = query.session_id;
423            info!(
424                "Query finished: id={}, duration={}ms",
425                query_id, duration_ms
426            );
427            drop(query);
428            self.emit_session_event(SessionEvent::QueryCompleted {
429                session_id,
430                query_id,
431                duration_ms,
432                success: true,
433            });
434            self.check_slow_query(session_id, query_id, duration_ms);
435            Ok(())
436        } else {
437            Err(ManagerError::NotFound(format!(
438                "Query {} not found",
439                query_id
440            )))
441        }
442    }
443
444    /// The marker query failed.
445    pub fn fail_query(&self, query_id: i64) -> ManagerResult<()> {
446        if let Some(mut query) = self.queries.get_mut(&query_id) {
447            query.fail();
448            let duration_ms = query.duration_ms.unwrap_or(0);
449            let session_id = query.session_id;
450            warn!("Query failed: id={}, duration={}ms", query_id, duration_ms);
451            drop(query);
452            self.emit_session_event(SessionEvent::QueryCompleted {
453                session_id,
454                query_id,
455                duration_ms,
456                success: false,
457            });
458            self.check_slow_query(session_id, query_id, duration_ms);
459            Ok(())
460        } else {
461            Err(ManagerError::NotFound(format!(
462                "Query {} not found",
463                query_id
464            )))
465        }
466    }
467
468    /// Terminate the query.
469    ///
470    /// Marks the query Killed and, when a query registry is attached via
471    /// `set_query_registry`, cancels its execution token with
472    /// `CancelReason::UserKill` so a running query actually stops.
473    /// Without an attached registry this only marks the status; the
474    /// executor-side `QueryRegistry` remains the source of truth for
475    /// cancellation. Query ids are generated as positive `i64` values, so a
476    /// negative id is rejected as invalid input instead of wrapping on the
477    /// `i64` to executor `u64` id conversion.
478    pub fn kill_query(&self, query_id: i64) -> ManagerResult<()> {
479        if query_id < 0 {
480            return Err(ManagerError::InvalidInput(format!(
481                "invalid query id {query_id}"
482            )));
483        }
484        if let Some(mut query) = self.queries.get_mut(&query_id) {
485            query.kill();
486            let duration_ms = query.duration_ms.unwrap_or(0);
487            let session_id = query.session_id;
488            warn!("Query killed: id={}", query_id);
489            drop(query);
490            if let Some(registry) = self.query_registry.read().as_ref() {
491                registry.cancel(QueryId(query_id as u64), CancelReason::UserKill);
492            }
493            self.emit_session_event(SessionEvent::QueryCompleted {
494                session_id,
495                query_id,
496                duration_ms,
497                success: false,
498            });
499            self.check_slow_query(session_id, query_id, duration_ms);
500            Ok(())
501        } else {
502            Err(ManagerError::NotFound(format!(
503                "Query {} not found",
504                query_id
505            )))
506        }
507    }
508
509    /// Obtain the query information
510    pub fn get_query(&self, query_id: i64) -> Option<QueryInfo> {
511        self.queries.get(&query_id).map(|v| v.clone())
512    }
513
514    /// Retrieve all queries
515    pub fn get_all_queries(&self) -> Vec<QueryInfo> {
516        self.queries.iter().map(|v| v.value().clone()).collect()
517    }
518
519    /// Obtain the queries that are currently running.
520    pub fn get_running_queries(&self) -> Vec<QueryInfo> {
521        self.queries
522            .iter()
523            .filter(|q| q.value().status == QueryStatus::Running)
524            .map(|v| v.value().clone())
525            .collect()
526    }
527
528    /// Clean up the completed queries (retaining only the last N of them).
529    pub fn cleanup_finished_queries(&self, keep_count: usize) {
530        let mut finished_queries: Vec<_> = self
531            .queries
532            .iter()
533            .filter(|q| q.value().status != QueryStatus::Running)
534            .map(|q| *q.key())
535            .collect();
536
537        // Sort by start time, keeping the most recent items at the top.
538        finished_queries.sort_by_key(|id| {
539            self.queries
540                .get(id)
541                .map(|q| q.start_time)
542                .unwrap_or(UNIX_EPOCH)
543        });
544
545        let to_remove = finished_queries.len().saturating_sub(keep_count);
546        for id in finished_queries.into_iter().take(to_remove) {
547            self.queries.remove(&id);
548        }
549    }
550}
551
552impl Default for QueryManager {
553    fn default() -> Self {
554        Self::new()
555    }
556}
557
558#[cfg(test)]
559mod tests {
560    use super::*;
561    use crate::executor::streaming::query_registry::QueryMetadata;
562    use std::sync::atomic::AtomicUsize;
563    use std::time::Instant;
564
565    #[test]
566    fn kill_query_cancels_registry_token_when_attached() {
567        let manager = QueryManager::new();
568        let registry = Arc::new(QueryRegistry::new());
569        manager.set_query_registry(Arc::clone(&registry));
570
571        let query_id = manager.register_query(
572            7,
573            "tester".to_string(),
574            None,
575            "MATCH (n) RETURN n".to_string(),
576        );
577        // Mirror the executor assembly: registry entry under the same id.
578        let (registry_id, token) = registry.register_with_id(
579            QueryId(query_id as u64),
580            QueryMetadata {
581                query_id: QueryId(query_id as u64),
582                session_id: Some(7),
583                user_name: Some("tester".to_string()),
584                space_name: None,
585                query_text: Some("MATCH (n) RETURN n".to_string()),
586                start_time: Instant::now(),
587            },
588        );
589        assert_eq!(registry_id, QueryId(query_id as u64));
590
591        manager.kill_query(query_id).expect("kill must succeed");
592        assert!(token.is_cancelled());
593        assert_eq!(
594            token.reason(),
595            Some(CancelReason::UserKill),
596            "kill must propagate the typed reason"
597        );
598        assert_eq!(
599            manager.get_query(query_id).map(|q| q.status),
600            Some(QueryStatus::Killed)
601        );
602    }
603
604    #[test]
605    fn kill_query_without_registry_only_marks_status() {
606        let manager = QueryManager::new();
607        let query_id = manager.register_query(
608            7,
609            "tester".to_string(),
610            None,
611            "MATCH (n) RETURN n".to_string(),
612        );
613        manager.kill_query(query_id).expect("kill must succeed");
614        assert_eq!(
615            manager.get_query(query_id).map(|q| q.status),
616            Some(QueryStatus::Killed)
617        );
618    }
619
620    #[test]
621    fn progress_hook_receives_emissions_and_unsubscribes() {
622        let manager = QueryManager::new();
623        assert_eq!(manager.progress_callback_count(), 0);
624        // No listeners: no-op, must not panic.
625        manager.emit_progress(1, 1, 100);
626
627        let rows = Arc::new(AtomicUsize::new(0));
628        let probe = Arc::clone(&rows);
629        let id = manager.register_progress_callback(Arc::new(move |progress| {
630            assert_eq!(progress.session_id, 1);
631            assert_eq!(progress.query_id, 2);
632            probe.fetch_add(progress.rows_processed as usize, Ordering::SeqCst);
633        }));
634        manager.emit_progress(1, 2, 100);
635        manager.emit_progress(1, 2, 50);
636        assert_eq!(rows.load(Ordering::SeqCst), 150);
637        assert!(manager.unregister_progress_callback(id));
638        manager.emit_progress(1, 2, 100);
639        assert_eq!(rows.load(Ordering::SeqCst), 150);
640    }
641}
linkrs_query/query_manager.rs

linkrs_query/
query_manager.rs
