khive_storage/
request_context.rs

1//! Backend-neutral request cancellation and deadline propagation.
2
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::sync::Arc;
7use std::task::Poll;
8use std::time::{Duration, Instant as WallInstant};
9
10use crate::{StorageError, StorageResult};
11
12const SEARCH_PHASES: [SearchPhase; 25] = [
13    SearchPhase::Authorization,
14    SearchPhase::TextSearch,
15    SearchPhase::VectorSearch,
16    SearchPhase::VectorEmbedding,
17    SearchPhase::Fusion,
18    SearchPhase::EntityHydration,
19    SearchPhase::NoteHydration,
20    SearchPhase::SupersededFilter,
21    SearchPhase::MailboxScopeHydration,
22    SearchPhase::MailboxAuthorization,
23    SearchPhase::TextReaderAdmission,
24    SearchPhase::TextReaderCheckout,
25    SearchPhase::TextBlockingRead,
26    SearchPhase::VectorReaderAdmission,
27    SearchPhase::VectorReaderCheckout,
28    SearchPhase::VectorBlockingRead,
29    SearchPhase::EntityReaderAdmission,
30    SearchPhase::EntityReaderCheckout,
31    SearchPhase::EntityBlockingRead,
32    SearchPhase::NoteReaderAdmission,
33    SearchPhase::NoteReaderCheckout,
34    SearchPhase::NoteBlockingRead,
35    SearchPhase::GraphReaderAdmission,
36    SearchPhase::GraphReaderCheckout,
37    SearchPhase::GraphBlockingRead,
38];
39
40/// A stage of one backend search, named in the deadline warning when that search times out.
41#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42#[repr(usize)]
43pub enum SearchPhase {
44    Authorization,
45    TextSearch,
46    VectorSearch,
47    VectorEmbedding,
48    Fusion,
49    EntityHydration,
50    NoteHydration,
51    SupersededFilter,
52    MailboxScopeHydration,
53    MailboxAuthorization,
54    TextReaderAdmission,
55    TextReaderCheckout,
56    TextBlockingRead,
57    VectorReaderAdmission,
58    VectorReaderCheckout,
59    VectorBlockingRead,
60    EntityReaderAdmission,
61    EntityReaderCheckout,
62    EntityBlockingRead,
63    NoteReaderAdmission,
64    NoteReaderCheckout,
65    NoteBlockingRead,
66    GraphReaderAdmission,
67    GraphReaderCheckout,
68    GraphBlockingRead,
69}
70
71impl SearchPhase {
72    /// The phase name written in log fields.
73    pub fn as_str(self) -> &'static str {
74        match self {
75            Self::Authorization => "authorization",
76            Self::TextSearch => "text_search",
77            Self::VectorSearch => "vector_search",
78            Self::VectorEmbedding => "vector_embedding",
79            Self::Fusion => "fusion",
80            Self::EntityHydration => "entity_hydration",
81            Self::NoteHydration => "note_hydration",
82            Self::SupersededFilter => "superseded_filter",
83            Self::MailboxScopeHydration => "mailbox_scope_hydration",
84            Self::MailboxAuthorization => "mailbox_authorization",
85            Self::TextReaderAdmission => "text_reader_admission",
86            Self::TextReaderCheckout => "text_reader_checkout",
87            Self::TextBlockingRead => "text_blocking_read",
88            Self::VectorReaderAdmission => "vector_reader_admission",
89            Self::VectorReaderCheckout => "vector_reader_checkout",
90            Self::VectorBlockingRead => "vector_blocking_read",
91            Self::EntityReaderAdmission => "entity_reader_admission",
92            Self::EntityReaderCheckout => "entity_reader_checkout",
93            Self::EntityBlockingRead => "entity_blocking_read",
94            Self::NoteReaderAdmission => "note_reader_admission",
95            Self::NoteReaderCheckout => "note_reader_checkout",
96            Self::NoteBlockingRead => "note_blocking_read",
97            Self::GraphReaderAdmission => "graph_reader_admission",
98            Self::GraphReaderCheckout => "graph_reader_checkout",
99            Self::GraphBlockingRead => "graph_blocking_read",
100        }
101    }
102}
103
104struct SearchPhaseTraceInner {
105    origin: std::time::Instant,
106    active_since_ns: [AtomicU64; SEARCH_PHASES.len()],
107    completed_ns: [AtomicU64; SEARCH_PHASES.len()],
108    completed_order: [AtomicU64; SEARCH_PHASES.len()],
109    entered: AtomicU64,
110    next_completion: AtomicU64,
111}
112
113/// Which phases one backend search has entered and finished, shared by the search and the code
114/// that reports its deadline. Marking a phase only touches atomics, so a reader never waits on a
115/// stalled phase.
116#[derive(Clone)]
117pub struct SearchPhaseTrace(Arc<SearchPhaseTraceInner>);
118
119/// Keeps one phase active until dropped; dropping it adds the elapsed time to that phase.
120pub struct SearchPhaseGuard {
121    trace: SearchPhaseTrace,
122    phase: SearchPhase,
123    started_ns: u64,
124}
125
126/// A trace read at one instant, formatted for the deadline warning fields.
127#[derive(Clone)]
128pub struct SearchPhaseSnapshot {
129    pub in_flight_phase: String,
130    pub phase_elapsed_ms: u64,
131    pub active_phases: String,
132    pub completed_phases: String,
133}
134
135tokio::task_local! {
136    static SEARCH_PHASE_TRACE: SearchPhaseTrace;
137}
138
139impl SearchPhaseTrace {
140    /// An empty trace with no phase entered.
141    pub fn new() -> Self {
142        Self(Arc::new(SearchPhaseTraceInner {
143            origin: std::time::Instant::now(),
144            active_since_ns: std::array::from_fn(|_| AtomicU64::new(0)),
145            completed_ns: std::array::from_fn(|_| AtomicU64::new(0)),
146            completed_order: std::array::from_fn(|_| AtomicU64::new(0)),
147            entered: AtomicU64::new(0),
148            next_completion: AtomicU64::new(0),
149        }))
150    }
151
152    /// Mark `phase` active from now until the returned guard drops. A phase entered again before
153    /// its earlier guard drops reports the later entry.
154    pub fn enter(&self, phase: SearchPhase) -> SearchPhaseGuard {
155        let started_ns = self.elapsed_ns();
156        let index = phase as usize;
157        self.0.active_since_ns[index].store(started_ns.saturating_add(1), Ordering::Release);
158        self.0.entered.fetch_or(1u64 << index, Ordering::Relaxed);
159        SearchPhaseGuard {
160            trace: self.clone(),
161            phase,
162            started_ns,
163        }
164    }
165
166    /// The most recently entered active phase and its time so far, every active phase, and the
167    /// finished phases in completion order with their total time.
168    pub fn snapshot(&self) -> SearchPhaseSnapshot {
169        let now_ns = self.elapsed_ns();
170        let mut active = Vec::new();
171        let mut completed = Vec::new();
172        for phase in SEARCH_PHASES {
173            let index = phase as usize;
174            let active_since = self.0.active_since_ns[index].load(Ordering::Acquire);
175            if active_since != 0 {
176                active.push((active_since - 1, phase));
177            }
178            let order = self.0.completed_order[index].load(Ordering::Acquire);
179            if order != 0 {
180                completed.push((
181                    order,
182                    phase,
183                    self.0.completed_ns[index].load(Ordering::Relaxed),
184                ));
185            }
186        }
187        active.sort_unstable_by_key(|(started_ns, _)| *started_ns);
188        completed.sort_unstable_by_key(|(order, _, _)| *order);
189
190        let (in_flight_phase, phase_elapsed_ms) = match active.last() {
191            Some((started_ns, phase)) => (
192                phase.as_str().to_string(),
193                nanos_to_millis(now_ns.saturating_sub(*started_ns)),
194            ),
195            None if self.0.entered.load(Ordering::Relaxed) == 0 => ("not_entered".into(), 0),
196            None => ("between_phases".into(), 0),
197        };
198        let active_phases = if active.is_empty() {
199            "none".to_string()
200        } else {
201            active
202                .iter()
203                .map(|(started_ns, phase)| {
204                    format!(
205                        "{}={}ms",
206                        phase.as_str(),
207                        nanos_to_millis(now_ns.saturating_sub(*started_ns))
208                    )
209                })
210                .collect::<Vec<_>>()
211                .join(",")
212        };
213        let completed_phases = if completed.is_empty() {
214            "none".to_string()
215        } else {
216            completed
217                .iter()
218                .map(|(_, phase, elapsed_ns)| {
219                    format!("{}={}ms", phase.as_str(), nanos_to_millis(*elapsed_ns))
220                })
221                .collect::<Vec<_>>()
222                .join(",")
223        };
224        SearchPhaseSnapshot {
225            in_flight_phase,
226            phase_elapsed_ms,
227            active_phases,
228            completed_phases,
229        }
230    }
231
232    fn elapsed_ns(&self) -> u64 {
233        u64::try_from(self.0.origin.elapsed().as_nanos()).unwrap_or(u64::MAX)
234    }
235}
236
237impl Default for SearchPhaseTrace {
238    fn default() -> Self {
239        Self::new()
240    }
241}
242
243impl Drop for SearchPhaseGuard {
244    fn drop(&mut self) {
245        let index = self.phase as usize;
246        let elapsed_ns = self.trace.elapsed_ns().saturating_sub(self.started_ns);
247        self.trace.0.completed_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
248        let order = self.trace.0.next_completion.fetch_add(1, Ordering::Relaxed) + 1;
249        self.trace.0.completed_order[index].store(order, Ordering::Release);
250        self.trace.0.active_since_ns[index].store(0, Ordering::Release);
251    }
252}
253
254fn nanos_to_millis(nanos: u64) -> u64 {
255    nanos / 1_000_000
256}
257
258/// Run `future` with `trace` as the current task's search phase trace.
259pub async fn scope_search_phase_trace<F: Future>(trace: SearchPhaseTrace, future: F) -> F::Output {
260    SEARCH_PHASE_TRACE.scope(trace, future).await
261}
262
263/// Mark `phase` on the current task's trace; `None` when no trace is in scope.
264pub fn enter_search_phase(phase: SearchPhase) -> Option<SearchPhaseGuard> {
265    SEARCH_PHASE_TRACE.try_with(|trace| trace.enter(phase)).ok()
266}
267
268/// The current task's trace, for work that runs on another thread.
269pub fn capture_search_phase_trace() -> Option<SearchPhaseTrace> {
270    SEARCH_PHASE_TRACE.try_with(Clone::clone).ok()
271}
272
273/// Default nonzero ceiling for request-owned read work, including compose.
274pub const DEFAULT_REQUEST_READ_TIMEOUT_SECS: u64 = 30;
275
276/// Largest accepted request-read ceiling: one day, so a trusted local batch
277/// (a whole-archive ingest, a long serial model run) can hold one deadline.
278pub const MAX_REQUEST_READ_TIMEOUT_SECS: u64 = 86_400;
279
280/// Resolve the operator-visible request-read ceiling from
281/// `KHIVE_REQUEST_READ_TIMEOUT_SECS`.
282///
283/// Unset uses the default. A value above [`MAX_REQUEST_READ_TIMEOUT_SECS`] is
284/// clamped to it; zero or a non-integer uses the default, so the guard is never
285/// disabled. Either correction logs a warning the first time it is made in a
286/// process, naming the value that was set and the ceiling in force.
287pub fn request_read_timeout_from_env() -> Duration {
288    let raw = std::env::var("KHIVE_REQUEST_READ_TIMEOUT_SECS").ok();
289    let (seconds, correction) = resolve_request_read_timeout_secs(raw.as_deref());
290    if let Some(correction) = correction {
291        static WARNED: std::sync::Once = std::sync::Once::new();
292        WARNED.call_once(|| {
293            tracing::warn!(
294                value = raw.as_deref().unwrap_or_default(),
295                ceiling_secs = seconds,
296                max_secs = MAX_REQUEST_READ_TIMEOUT_SECS,
297                "KHIVE_REQUEST_READ_TIMEOUT_SECS {correction}"
298            );
299        });
300    }
301    Duration::from_secs(seconds)
302}
303
304/// The ceiling in seconds for a raw setting, and the correction applied when
305/// the setting was not used as written.
306fn resolve_request_read_timeout_secs(raw: Option<&str>) -> (u64, Option<&'static str>) {
307    let Some(raw) = raw else {
308        return (DEFAULT_REQUEST_READ_TIMEOUT_SECS, None);
309    };
310    match raw.trim().parse::<u64>() {
311        Ok(seconds) if (1..=MAX_REQUEST_READ_TIMEOUT_SECS).contains(&seconds) => (seconds, None),
312        Ok(seconds) if seconds > MAX_REQUEST_READ_TIMEOUT_SECS => (
313            MAX_REQUEST_READ_TIMEOUT_SECS,
314            Some("is above the maximum; clamped to the maximum"),
315        ),
316        _ => (
317            DEFAULT_REQUEST_READ_TIMEOUT_SECS,
318            Some("is not a whole number of seconds from 1 to the maximum; using the default"),
319        ),
320    }
321}
322
323/// One absolute request deadline represented on async and blocking clocks.
324///
325/// Tokio's instant is authoritative for async timeout selection, including
326/// paused-clock tests. The wall instant lets blocking backends enforce the
327/// same non-renewing deadline without depending on Tokio's runtime clock.
328#[derive(Clone, Copy, Debug)]
329pub struct RequestReadDeadline {
330    async_at: tokio::time::Instant,
331    blocking_at: WallInstant,
332}
333
334impl RequestReadDeadline {
335    /// Create a deadline `duration` from now.
336    pub fn after(duration: Duration) -> Self {
337        Self {
338            async_at: tokio::time::Instant::now() + duration,
339            blocking_at: WallInstant::now() + duration,
340        }
341    }
342
343    /// Tokio-clock instant used by request coordinators and timeout selection.
344    pub fn async_at(self) -> tokio::time::Instant {
345        self.async_at
346    }
347
348    /// Wall-clock instant used by blocking backend cancellation checks.
349    pub fn blocking_at(self) -> WallInstant {
350        self.blocking_at
351    }
352
353    fn earlier(self, other: Self) -> Self {
354        if self.async_at <= other.async_at {
355            self
356        } else {
357            other
358        }
359    }
360}
361
362/// The first request-level condition that stopped read-only work.
363#[derive(Clone, Copy, Debug, Eq, PartialEq)]
364pub enum RequestReadStopReason {
365    /// An explicit cancellation signal fired or its sender disappeared.
366    Cancelled,
367    /// The request's absolute deadline elapsed.
368    Deadline,
369}
370
371/// Opaque snapshot of the current request read context.
372///
373/// Backend implementations capture this once before crossing into blocking
374/// work so nested scopes and spawned children share one absolute deadline.
375#[derive(Clone, Default)]
376pub struct RequestReadContext {
377    cancellations: Arc<[tokio::sync::watch::Receiver<bool>]>,
378    deadline: Option<RequestReadDeadline>,
379    store_acquisition_operation: Option<&'static str>,
380}
381
382impl RequestReadContext {
383    /// Opt one blocking store-constructor call into request-aware admission.
384    /// Ordinary accessors under a read context remain outside this scope.
385    /// The previous context is restored even if the constructor panics.
386    pub fn scope_store_acquisition<T>(
387        mut self,
388        operation: &'static str,
389        work: impl FnOnce() -> T,
390    ) -> T {
391        self.store_acquisition_operation = Some(operation);
392        REQUEST_READ_CONTEXT.sync_scope(self, work)
393    }
394
395    /// The explicitly scoped constructor operation, absent for ordinary callers.
396    pub fn store_acquisition_operation(&self) -> Option<&'static str> {
397        self.store_acquisition_operation
398    }
399
400    /// Observe cancellation and the original wall-clock deadline in blocking work.
401    pub fn blocking_stop_reason(&self) -> Option<RequestReadStopReason> {
402        if self.cancellations.iter().any(receiver_cancelled) {
403            Some(RequestReadStopReason::Cancelled)
404        } else if self
405            .deadline
406            .is_some_and(|deadline| WallInstant::now() >= deadline.blocking_at)
407        {
408            Some(RequestReadStopReason::Deadline)
409        } else {
410            None
411        }
412    }
413
414    /// Return the request's absolute deadline, when one is installed.
415    pub fn deadline(&self) -> Option<RequestReadDeadline> {
416        self.deadline
417    }
418
419    /// Return the currently observed stop cause without waiting.
420    pub fn stop_reason(&self) -> Option<RequestReadStopReason> {
421        if self.cancellations.iter().any(receiver_cancelled) {
422            Some(RequestReadStopReason::Cancelled)
423        } else if self
424            .deadline
425            .is_some_and(|deadline| tokio::time::Instant::now() >= deadline.async_at)
426        {
427            Some(RequestReadStopReason::Deadline)
428        } else {
429            None
430        }
431    }
432
433    /// Wait until one merged cancellation source fires or the deadline elapses.
434    pub async fn wait_for_stop(self) -> RequestReadStopReason {
435        let cancellation_receivers = self.cancellations;
436        let request_deadline = self.deadline;
437        let cancellation = async move {
438            match cancellation_receivers.len() {
439                0 => std::future::pending::<()>().await,
440                1 => wait_for_receiver_cancellation(cancellation_receivers[0].clone()).await,
441                2 => {
442                    tokio::select! {
443                        _ = wait_for_receiver_cancellation(cancellation_receivers[0].clone()) => {},
444                        _ = wait_for_receiver_cancellation(cancellation_receivers[1].clone()) => {},
445                    }
446                }
447                _ => wait_for_receiver_set(cancellation_receivers).await,
448            }
449        };
450        let deadline = async move {
451            match request_deadline {
452                Some(deadline) => tokio::time::sleep_until(deadline.async_at).await,
453                None => std::future::pending::<()>().await,
454            }
455        };
456        tokio::select! {
457            _ = cancellation => RequestReadStopReason::Cancelled,
458            _ = deadline => RequestReadStopReason::Deadline,
459        }
460    }
461}
462
463tokio::task_local! {
464    static REQUEST_READ_CONTEXT: RequestReadContext;
465}
466
467/// Capture the current request context for a backend operation.
468pub fn capture_request_read_context() -> RequestReadContext {
469    REQUEST_READ_CONTEXT
470        .try_with(Clone::clone)
471        .unwrap_or_default()
472}
473
474/// Constrain `candidate` by the current request's deadline without installing
475/// a new scope.
476pub fn effective_request_read_deadline(candidate: RequestReadDeadline) -> RequestReadDeadline {
477    capture_request_read_context()
478        .deadline
479        .map_or(candidate, |existing| existing.earlier(candidate))
480}
481
482/// Scope `future` to an explicit read-cancellation signal.
483///
484/// Backends observe the signal only for work they classify as read-only.
485pub async fn scope_request_read_cancellation<F>(
486    cancellation: tokio::sync::watch::Receiver<bool>,
487    future: F,
488) -> F::Output
489where
490    F: Future,
491{
492    let mut context = capture_request_read_context();
493    let mut cancellations = Vec::with_capacity(context.cancellations.len() + 1);
494    cancellations.extend(context.cancellations.iter().cloned());
495    cancellations.push(cancellation);
496    context.cancellations = cancellations.into();
497    REQUEST_READ_CONTEXT.scope(context, future).await
498}
499
500/// Scope `future` to one relative read deadline, preserving an earlier outer
501/// deadline instead of renewing it.
502pub async fn scope_request_read_deadline<F>(duration: Duration, future: F) -> F::Output
503where
504    F: Future,
505{
506    scope_request_read_deadline_at(RequestReadDeadline::after(duration), future).await
507}
508
509/// Scope `future` to an already-created absolute read deadline.
510pub async fn scope_request_read_deadline_at<F>(
511    deadline: RequestReadDeadline,
512    future: F,
513) -> F::Output
514where
515    F: Future,
516{
517    let mut context = capture_request_read_context();
518    context.deadline = Some(match context.deadline {
519        Some(existing) => existing.earlier(deadline),
520        None => deadline,
521    });
522    REQUEST_READ_CONTEXT.scope(context, future).await
523}
524
525/// Capture the current request read context for a spawned child future.
526pub fn inherit_request_read_context<F>(future: F) -> impl Future<Output = F::Output> + Send
527where
528    F: Future + Send,
529{
530    let inherited = REQUEST_READ_CONTEXT.try_with(Clone::clone).ok();
531    async move {
532        match inherited {
533            Some(context) => REQUEST_READ_CONTEXT.scope(context, future).await,
534            None => future.await,
535        }
536    }
537}
538
539/// Intention-revealing alias for cancellation-only child tasks.
540pub fn inherit_request_read_cancellation<F>(future: F) -> impl Future<Output = F::Output> + Send
541where
542    F: Future + Send,
543{
544    inherit_request_read_context(future)
545}
546
547/// Return whether the current request has stopped its read-only work.
548pub fn request_read_is_cancelled() -> bool {
549    capture_request_read_context().stop_reason().is_some()
550}
551
552/// Refuse to begin another request-owned read phase after cancellation.
553pub fn ensure_request_read_active(operation: &'static str) -> StorageResult<()> {
554    if request_read_is_cancelled() {
555        Err(StorageError::Timeout {
556            operation: operation.into(),
557        })
558    } else {
559        Ok(())
560    }
561}
562
563/// Wait until the current request's read-only work is cancelled or timed out.
564pub async fn wait_for_request_read_cancellation() {
565    let _ = capture_request_read_context().wait_for_stop().await;
566}
567
568/// Await one request-owned read phase under the current absolute deadline.
569pub async fn await_request_read_phase<F>(
570    operation: &'static str,
571    future: F,
572) -> StorageResult<F::Output>
573where
574    F: Future,
575{
576    ensure_request_read_active(operation)?;
577    tokio::pin!(future);
578    tokio::select! {
579        biased;
580        _ = wait_for_request_read_cancellation() => Err(StorageError::Timeout {
581            operation: operation.into(),
582        }),
583        output = &mut future => {
584            ensure_request_read_active(operation)?;
585            Ok(output)
586        }
587    }
588}
589
590fn receiver_cancelled(receiver: &tokio::sync::watch::Receiver<bool>) -> bool {
591    *receiver.borrow() || receiver.has_changed().is_err()
592}
593
594async fn wait_for_receiver_cancellation(mut receiver: tokio::sync::watch::Receiver<bool>) {
595    loop {
596        if *receiver.borrow_and_update() {
597            return;
598        }
599        if receiver.changed().await.is_err() {
600            return;
601        }
602    }
603}
604
605async fn wait_for_receiver_set(receivers: Arc<[tokio::sync::watch::Receiver<bool>]>) {
606    let mut waits: Vec<Pin<Box<dyn Future<Output = ()> + Send>>> = receivers
607        .iter()
608        .cloned()
609        .map(|receiver| {
610            Box::pin(wait_for_receiver_cancellation(receiver))
611                as Pin<Box<dyn Future<Output = ()> + Send>>
612        })
613        .collect();
614    std::future::poll_fn(move |cx| {
615        if waits
616            .iter_mut()
617            .any(|wait| wait.as_mut().poll(cx).is_ready())
618        {
619            Poll::Ready(())
620        } else {
621            Poll::Pending
622        }
623    })
624    .await;
625}
626
627#[cfg(test)]
628#[path = "request_context/tests.rs"]
629mod tests;
khive_storage/request_context.rs

khive_storage/
request_context.rs
