soothe_client/appkit/
query_gate.rs

1//! Single-flight query gate per application key.
2//!
3//! [`QueryGate`] enforces single-flight query execution per app key and
4//! **cancel-before-context ordering**: when a query is cancelled, the daemon is
5//! told to stop (`command_request` with `command:"cancel"`) **before** the local
6//! cancel callback runs, on a detached 10s timeout so the caller's cancellation
7//! cannot block the wire send.
8
9use std::collections::HashMap;
10use std::future::Future;
11use std::pin::Pin;
12use std::sync::{Arc, Mutex};
13use std::time::Duration;
14
15use crate::errors::Result;
16
17/// Local cancel callback for an in-flight query (typically a timeout context cancel).
18pub type CancelFn = Arc<dyn Fn() + Send + Sync>;
19
20/// Async daemon-cancel sender (`command_request` cancel for the loop).
21pub type SendCancelFn =
22    Arc<dyn Fn() -> Pin<Box<dyn Future<Output = Result<()>> + Send>> + Send + Sync>;
23
24/// Returned when an app key already has an in-flight query.
25#[derive(Debug, Clone, Copy, thiserror::Error)]
26#[error("appkit: query already in progress for app key")]
27pub struct ErrQueryBusy;
28
29struct GateState {
30    cancels: HashMap<String, CancelFn>,
31    send_cancels: HashMap<String, SendCancelFn>,
32}
33
34/// Enforces single-flight query execution per app key with daemon-first cancel.
35pub struct QueryGate {
36    inner: Mutex<GateState>,
37}
38
39impl QueryGate {
40    /// Construct an empty gate.
41    pub fn new() -> Self {
42        Self {
43            inner: Mutex::new(GateState {
44                cancels: HashMap::new(),
45                send_cancels: HashMap::new(),
46            }),
47        }
48    }
49
50    /// Reserve `app_key` for one agent turn.
51    ///
52    /// Returns [`ErrQueryBusy`] when a query is already in flight. Pass the local
53    /// cancel callback for the query's timeout context (the same one the turn loop
54    /// derives from). `send_cancel`, when provided, sends the daemon cancel and is
55    /// invoked from [`Self::cancel`] on a detached 10s timeout.
56    pub fn acquire(
57        &self,
58        app_key: &str,
59        cancel: CancelFn,
60        send_cancel: Option<SendCancelFn>,
61    ) -> std::result::Result<(), ErrQueryBusy> {
62        let mut state = self.inner.lock().expect("query gate mutex poisoned");
63        if state.cancels.contains_key(app_key) {
64            return Err(ErrQueryBusy);
65        }
66        state.cancels.insert(app_key.to_string(), cancel);
67        if let Some(send) = send_cancel {
68            state.send_cancels.insert(app_key.to_string(), send);
69        }
70        Ok(())
71    }
72
73    /// Cooperatively stop a running query for `app_key`.
74    ///
75    /// Sends the daemon cancel first (10s timeout, detached from the caller's
76    /// cancellation), then invokes the local cancel callback. Returns `Ok(())` when
77    /// no query is in flight (intent already satisfied).
78    pub async fn cancel(&self, app_key: &str) -> Result<()> {
79        let (cancel, send_cancel) = {
80            let mut state = self.inner.lock().expect("query gate mutex poisoned");
81            let cancel = state.cancels.remove(app_key);
82            let send_cancel = state.send_cancels.remove(app_key);
83            (cancel, send_cancel)
84        };
85
86        if cancel.is_none() && send_cancel.is_none() {
87            return Ok(());
88        }
89
90        if let Some(send_cancel) = send_cancel {
91            let fut = send_cancel();
92            let _ = tokio::time::timeout(Duration::from_secs(10), fut).await;
93        }
94
95        if let Some(cancel) = cancel {
96            cancel();
97        }
98        Ok(())
99    }
100
101    /// Clear the gate for `app_key` without sending a daemon cancel.
102    ///
103    /// Call when a query completes normally (success or local failure) so the next
104    /// turn can acquire.
105    pub fn release(&self, app_key: &str) {
106        let mut state = self.inner.lock().expect("query gate mutex poisoned");
107        state.cancels.remove(app_key);
108        state.send_cancels.remove(app_key);
109    }
110
111    /// Update the daemon-cancel sender for an already-acquired session.
112    pub fn set_send_cancel(&self, app_key: &str, send_cancel: SendCancelFn) {
113        let mut state = self.inner.lock().expect("query gate mutex poisoned");
114        if !state.cancels.contains_key(app_key) {
115            return;
116        }
117        state.send_cancels.insert(app_key.to_string(), send_cancel);
118    }
119
120    /// Update the local cancel callback for an already-acquired session.
121    pub fn replace_cancel(&self, app_key: &str, cancel: CancelFn) {
122        let mut state = self.inner.lock().expect("query gate mutex poisoned");
123        if !state.cancels.contains_key(app_key) {
124            return;
125        }
126        state.cancels.insert(app_key.to_string(), cancel);
127    }
128
129    /// Report whether a query is in flight for `app_key`.
130    pub fn is_active(&self, app_key: &str) -> bool {
131        self.inner
132            .lock()
133            .expect("query gate mutex poisoned")
134            .cancels
135            .contains_key(app_key)
136    }
137}
138
139impl Default for QueryGate {
140    fn default() -> Self {
141        Self::new()
142    }
143}
soothe_client/appkit/query_gate.rs

soothe_client/appkit/
query_gate.rs
