query_runtime/observability/recorder/
session.rs

1//! Query observation session, seal observers, and recorder event writes.
2
3use std::{
4    collections::BTreeMap,
5    fmt,
6    sync::{
7        Arc, Condvar, Mutex, MutexGuard, PoisonError,
8        atomic::{AtomicBool, Ordering},
9    },
10};
11
12#[cfg(test)]
13use super::RuntimeObservationDeliveryGate;
14use super::{
15    CleanupFailureTally, LiveRuntimeObservationSnapshot, PlanExecutionId, QueryAttemptReport, QueryId,
16    RetainedTopologyRecord, RuntimeEventContext, RuntimeObservationClassLoss, RuntimeObservationClock,
17    RuntimeObservationConfig, RuntimeObservationDeliveryGuard, RuntimeObservationFooter, RuntimeObservationHandle,
18    RuntimeObservationManifest, RuntimeObservationOwner, RuntimeObservationPayload, RuntimeObservationReasonLoss,
19    RuntimeObservationRecordClass, RuntimeObservationRecordDisposition, RuntimeObservationRecorderState,
20    RuntimeObservationSnapshotData, RuntimeTerminalOutcome, RuntimeTimelineEvent, RuntimeTopologyRecord,
21    SealedRuntimeObservationSnapshot, SourceLocationDigestKeyState, TraceId,
22};
23#[cfg(any(test, feature = "test-hooks"))]
24use super::{SystemRuntimeObservationClock, TestSourceLocationDigestKeyProvider};
25
26/// Owner-only capability for one query observation lifetime.
27///
28/// This type is intentionally not [`Clone`]. Cloneable execution contexts and
29/// workers receive only [`RuntimeObservationHandle`], which cannot seal.
30#[derive(Debug, Default)]
31pub struct RuntimeObservationSession {
32    pub(super) handle: RuntimeObservationHandle,
33}
34
35/// Construction-time provider of a fixed attempt summary embedded at seal, for
36/// tests that seal a recorder-only owner. An attempt passes the crate-private
37/// seal provider below to each seal of its session, so the embedded value and
38/// completion delivery share the same immutable allocation.
39#[cfg(any(test, feature = "test-hooks"))]
40pub(crate) type AttemptReportProvider = Arc<dyn Fn() -> QueryAttemptReport + Send + Sync>;
41
42/// Attempt-owned report construction after the recorder has computed its
43/// complete footer, given the session's exact cleanup-failure tally.
44pub(crate) type AttemptReportSealProvider =
45    Arc<dyn Fn(&RuntimeObservationFooter, &CleanupFailureTally) -> Arc<QueryAttemptReport> + Send + Sync>;
46
47pub(super) struct RuntimeObservationSessionInner {
48    pub(super) manifest: RuntimeObservationManifest,
49    pub(super) source_location_digest_key: SourceLocationDigestKeyState,
50    pub(super) clock: Arc<dyn RuntimeObservationClock>,
51    pub(super) state: Mutex<RuntimeObservationRecorderState>,
52    pub(super) sealed: AtomicBool,
53    /// Signals that a sealed snapshot was published. A seal that builds its
54    /// attempt report outside the state lock leaves the session sealed with
55    /// no snapshot until then; a racing seal waits here for it.
56    pub(super) seal_published: Condvar,
57}
58
59pub(super) struct RuntimeObservationSealObserver {
60    pub(super) notify: Box<dyn Fn(SealedRuntimeObservationSnapshot) + Send + Sync>,
61}
62
63impl fmt::Debug for RuntimeObservationSealObserver {
64    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
65        formatter.debug_struct("RuntimeObservationSealObserver").finish_non_exhaustive()
66    }
67}
68
69impl RuntimeObservationSealObserver {
70    pub(super) fn invoke(&self, snapshot: SealedRuntimeObservationSnapshot) {
71        let _delivery = RuntimeObservationDeliveryGuard::enter(None);
72        if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (self.notify)(snapshot))).is_err() {
73            eprintln!("[query-engine] warning: runtime observation seal observer panicked");
74        }
75    }
76}
77
78/// Deferred observer delivery for an already-sealed query cleanup.
79/// The seal has published the report; delivery runs after capacity release
80/// and never gates reuse.
81pub(crate) struct QueryCleanupSealDelivery {
82    pub(super) owner: Option<RuntimeObservationOwner>,
83    pub(super) deferred: Option<DeferredRuntimeObservationSeal>,
84    pub(super) already_sealed: Option<SealedRuntimeObservationSnapshot>,
85    #[cfg(test)]
86    pub(super) delivery_gate: Option<Arc<RuntimeObservationDeliveryGate>>,
87}
88
89impl QueryCleanupSealDelivery {
90    pub(super) fn already_sealed(snapshot: Option<SealedRuntimeObservationSnapshot>) -> Self {
91        Self {
92            owner: None,
93            deferred: None,
94            already_sealed: snapshot,
95            #[cfg(test)]
96            delivery_gate: None,
97        }
98    }
99
100    /// Holds observer delivery at `gate` until the test opens it.
101    #[cfg(test)]
102    pub(crate) fn with_delivery_gate(mut self, gate: Option<Arc<RuntimeObservationDeliveryGate>>) -> Self {
103        self.delivery_gate = gate;
104        self
105    }
106
107    pub(crate) fn deliver(self) -> Option<SealedRuntimeObservationSnapshot> {
108        let Some(deferred) = self.deferred else { return self.already_sealed };
109        if let Some(owner) = self.owner {
110            return owner.complete_observer_delivery(deferred);
111        }
112        // The test gate sits before every observer; the seal has already
113        // committed by this point.
114        #[cfg(test)]
115        if let Some(gate) = self.delivery_gate {
116            gate.hold();
117        }
118        deferred.invoke_observers()
119    }
120}
121
122pub(super) struct DeferredRuntimeObservationSeal {
123    pub(super) snapshot: Option<SealedRuntimeObservationSnapshot>,
124    pub(super) observers: Vec<RuntimeObservationSealObserver>,
125}
126
127impl DeferredRuntimeObservationSeal {
128    pub(super) fn invoke_observers(self) -> Option<SealedRuntimeObservationSnapshot> {
129        if let Some(snapshot) = &self.snapshot {
130            for observer in self.observers {
131                observer.invoke(snapshot.clone());
132            }
133        }
134        self.snapshot
135    }
136}
137
138impl fmt::Debug for RuntimeObservationSessionInner {
139    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
140        formatter
141            .debug_struct("RuntimeObservationSessionInner")
142            .field("manifest", &self.manifest)
143            .field("sealed", &self.sealed.load(Ordering::Acquire))
144            .finish_non_exhaustive()
145    }
146}
147
148impl RuntimeObservationSession {
149    /// Creates an owner and recording handle using a system monotonic clock.
150    pub fn new(query_id: impl Into<QueryId>, trace_id: Option<TraceId>, config: RuntimeObservationConfig) -> Self {
151        Self { handle: RuntimeObservationHandle::new(query_id, trace_id, config) }
152    }
153
154    /// Creates an owner and recording handle using an injected clock.
155    pub fn with_clock(
156        query_id: impl Into<QueryId>,
157        trace_id: Option<TraceId>,
158        config: RuntimeObservationConfig,
159        clock: Arc<dyn RuntimeObservationClock>,
160    ) -> Self {
161        Self { handle: RuntimeObservationHandle::with_clock(query_id, trace_id, config, clock) }
162    }
163
164    #[cfg(test)]
165    pub(crate) fn with_source_location_test_key(
166        query_id: impl Into<QueryId>,
167        trace_id: Option<TraceId>,
168        config: RuntimeObservationConfig,
169        key: [u8; blake3::KEY_LEN],
170    ) -> Self {
171        Self {
172            handle: RuntimeObservationHandle::with_clock_and_source_location_digest_key(
173                query_id,
174                trace_id,
175                config,
176                Arc::new(SystemRuntimeObservationClock::new()),
177                &TestSourceLocationDigestKeyProvider { key: Some(key) },
178            ),
179        }
180    }
181
182    #[cfg(any(test, feature = "test-hooks"))]
183    #[doc(hidden)]
184    pub fn with_unavailable_source_location_test_key(
185        query_id: impl Into<QueryId>,
186        trace_id: Option<TraceId>,
187        config: RuntimeObservationConfig,
188    ) -> Self {
189        Self {
190            handle: RuntimeObservationHandle::with_clock_and_source_location_digest_key(
191                query_id,
192                trace_id,
193                config,
194                Arc::new(SystemRuntimeObservationClock::new()),
195                &TestSourceLocationDigestKeyProvider { key: None },
196            ),
197        }
198    }
199
200    /// Creates a disabled owner and handle.
201    pub const fn disabled() -> Self {
202        Self { handle: RuntimeObservationHandle::disabled() }
203    }
204
205    /// Returns a cloneable recording capability without seal authority.
206    pub fn handle(&self) -> RuntimeObservationHandle {
207        self.handle.clone()
208    }
209
210    /// Returns the immutable recorder configuration.
211    pub const fn config(&self) -> RuntimeObservationConfig {
212        self.handle.config()
213    }
214
215    /// Returns the immutable query identity when recording is enabled.
216    pub fn query_id(&self) -> Option<&QueryId> {
217        self.handle.query_id()
218    }
219
220    /// Returns the immutable optional parent trace identity.
221    pub fn trace_id(&self) -> Option<&TraceId> {
222        self.handle.trace_id()
223    }
224
225    /// Lazily constructs and records a query-level event.
226    pub fn record<F>(&self, build_payload: F) -> RuntimeObservationRecordDisposition
227    where F: FnOnce() -> RuntimeObservationPayload {
228        self.handle.record(build_payload)
229    }
230
231    /// Lazily constructs and records an event for one plan execution.
232    pub fn record_for_plan<F>(
233        &self,
234        plan_execution_id: PlanExecutionId,
235        build_payload: F,
236    ) -> RuntimeObservationRecordDisposition
237    where
238        F: FnOnce() -> RuntimeObservationPayload,
239    {
240        self.handle.record_for_plan(plan_execution_id, build_payload)
241    }
242
243    /// Lazily constructs an event before acquiring recorder order and records it infallibly.
244    pub fn record_with_context<F>(
245        &self,
246        context: RuntimeEventContext,
247        build_payload: F,
248    ) -> RuntimeObservationRecordDisposition
249    where
250        F: FnOnce() -> RuntimeObservationPayload,
251    {
252        self.handle.record_with_context(context, build_payload)
253    }
254
255    /// Lazily constructs and records one static topology row.
256    pub fn record_topology<F>(
257        &self,
258        class: RuntimeObservationRecordClass,
259        build_record: F,
260    ) -> RuntimeObservationRecordDisposition
261    where
262        F: FnOnce() -> RuntimeTopologyRecord,
263    {
264        self.handle.record_topology(class, build_record)
265    }
266
267    /// Captures an immutable incomplete view. Returns none after sealing or when disabled.
268    pub fn live_snapshot(&self) -> Option<LiveRuntimeObservationSnapshot> {
269        self.handle.live_snapshot()
270    }
271
272    /// Fixes the first terminal outcome and returns the same immutable snapshot
273    /// to every later or concurrent seal.
274    pub fn seal(&self, outcome: RuntimeTerminalOutcome) -> Option<SealedRuntimeObservationSnapshot> {
275        self.handle.seal_owned(outcome)
276    }
277
278    pub(super) fn seal_deferred(
279        &self,
280        outcome: Option<RuntimeTerminalOutcome>,
281        attempt_report_provider: Option<&AttemptReportSealProvider>,
282    ) -> DeferredRuntimeObservationSeal {
283        self.handle.seal_owned_deferred(outcome, attempt_report_provider)
284    }
285
286    /// Seals this session with `outcome` and the report `report` builds,
287    /// deferring observer delivery to the returned value. A seal that follows
288    /// an earlier seal returns that seal's snapshot and delivers nothing.
289    ///
290    /// A panic while sealing marks the session degraded, seals it as
291    /// `Panicked`, delivers its observers, and resumes unwinding.
292    pub(crate) fn seal_with_report(
293        &self,
294        outcome: RuntimeTerminalOutcome,
295        report: &AttemptReportSealProvider,
296    ) -> QueryCleanupSealDelivery {
297        let deferred =
298            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.seal_deferred(Some(outcome), Some(report))));
299        match deferred {
300            Ok(deferred) => QueryCleanupSealDelivery {
301                owner: None,
302                deferred: Some(deferred),
303                already_sealed: None,
304                #[cfg(test)]
305                delivery_gate: None,
306            },
307            Err(panic) => {
308                self.handle.mark_degraded();
309                let _ = self.seal_deferred(Some(RuntimeTerminalOutcome::Panicked), None).invoke_observers();
310                std::panic::resume_unwind(panic);
311            },
312        }
313    }
314
315    /// Registers an infallible notification for this session's seal. It runs
316    /// once, outside recorder locks: with the seal's deferred delivery, or
317    /// immediately when the session has already sealed.
318    pub(crate) fn register_seal_observer<F>(&self, observer: F)
319    where F: Fn(SealedRuntimeObservationSnapshot) + Send + Sync + 'static {
320        let observer = RuntimeObservationSealObserver { notify: Box::new(observer) };
321        if let Some((observer, snapshot)) = self.handle.register_seal_observer(observer) {
322            observer.invoke(snapshot);
323        }
324    }
325
326    /// Returns the sealed immutable snapshot when a terminal outcome was already fixed.
327    pub fn sealed_snapshot(&self) -> Option<SealedRuntimeObservationSnapshot> {
328        self.handle.sealed_snapshot()
329    }
330}
331
332impl RuntimeObservationSessionInner {
333    pub(super) fn lock_state(&self) -> MutexGuard<'_, RuntimeObservationRecorderState> {
334        match self.state.lock() {
335            Ok(state) => state,
336            Err(poisoned) => {
337                let mut state = poisoned.into_inner();
338                state.degraded = true;
339                state.poison_recoveries = state.poison_recoveries.checked_add(1).unwrap_or_else(|| {
340                    state.counter_overflow = true;
341                    u64::MAX
342                });
343                self.state.clear_poison();
344                state
345            },
346        }
347    }
348
349    pub(super) fn record_event_lazy<F>(
350        &self,
351        context: RuntimeEventContext,
352        build_payload: F,
353    ) -> RuntimeObservationRecordDisposition
354    where
355        F: FnOnce() -> RuntimeObservationPayload,
356    {
357        self.record_event_timestamped(context, |_monotonic_ns| build_payload()).0
358    }
359
360    #[cfg(test)]
361    pub(super) fn record_event_sized<F>(
362        &self,
363        context: RuntimeEventContext,
364        logical_bytes: u64,
365        build_payload: F,
366    ) -> RuntimeObservationRecordDisposition
367    where
368        F: FnOnce() -> RuntimeObservationPayload,
369    {
370        let mut state = self.lock_state();
371        if state.sealed_snapshot.is_some() {
372            return RuntimeObservationRecordDisposition::Sealed;
373        }
374        let sequence = state.take_sequence();
375        if let Some(reason) = state.rejection_reason(self.manifest.limits, logical_bytes) {
376            state.record_drop(RuntimeObservationRecordClass::TimelineEvent, sequence, reason);
377            return RuntimeObservationRecordDisposition::Dropped { sequence, reason };
378        }
379        let monotonic_ns = self.clock.monotonic_ns().max(state.last_monotonic_ns);
380        state.last_monotonic_ns = monotonic_ns;
381        let event = RuntimeTimelineEvent {
382            sequence,
383            monotonic_ns,
384            plan_execution_id: context.plan_execution_id,
385            span_id: context.span_id,
386            parent_span_id: context.parent_span_id,
387            payload: build_payload(),
388        };
389        let actual_bytes = event.logical_retained_bytes();
390        if actual_bytes != logical_bytes {
391            state.degraded = true;
392        }
393        match state.admit(self.manifest.limits, RuntimeObservationRecordClass::TimelineEvent, sequence, actual_bytes) {
394            Ok(()) => {
395                state.events.push(event);
396                RuntimeObservationRecordDisposition::Retained { sequence, logical_bytes: actual_bytes }
397            },
398            Err(reason) => RuntimeObservationRecordDisposition::Dropped { sequence, reason },
399        }
400    }
401
402    pub(super) fn record_event_timestamped<F>(
403        &self,
404        context: RuntimeEventContext,
405        build_payload: F,
406    ) -> (RuntimeObservationRecordDisposition, Option<u64>)
407    where
408        F: FnOnce(u64) -> RuntimeObservationPayload,
409    {
410        let mut state = self.lock_state();
411        if state.sealed_snapshot.is_some() {
412            return (RuntimeObservationRecordDisposition::Sealed, None);
413        }
414        let sequence = state.take_sequence();
415        let monotonic_ns = self.clock.monotonic_ns().max(state.last_monotonic_ns);
416        state.last_monotonic_ns = monotonic_ns;
417        const MIN_TIMELINE_EVENT_BYTES: u64 = 42;
418        if let Some(reason) = state.rejection_reason(self.manifest.limits, MIN_TIMELINE_EVENT_BYTES) {
419            state.record_drop(RuntimeObservationRecordClass::TimelineEvent, sequence, reason);
420            return (RuntimeObservationRecordDisposition::Dropped { sequence, reason }, Some(monotonic_ns));
421        }
422        let event = RuntimeTimelineEvent {
423            sequence,
424            monotonic_ns,
425            plan_execution_id: context.plan_execution_id,
426            span_id: context.span_id,
427            parent_span_id: context.parent_span_id,
428            payload: build_payload(monotonic_ns),
429        };
430        let logical_bytes = event.logical_retained_bytes();
431        let disposition = match state.admit(
432            self.manifest.limits,
433            RuntimeObservationRecordClass::TimelineEvent,
434            sequence,
435            logical_bytes,
436        ) {
437            Ok(()) => {
438                state.events.push(event);
439                RuntimeObservationRecordDisposition::Retained { sequence, logical_bytes }
440            },
441            Err(reason) => RuntimeObservationRecordDisposition::Dropped { sequence, reason },
442        };
443        (disposition, Some(monotonic_ns))
444    }
445
446    pub(super) fn record_topology_lazy<F>(
447        &self,
448        class: RuntimeObservationRecordClass,
449        build_record: F,
450    ) -> RuntimeObservationRecordDisposition
451    where
452        F: FnOnce() -> RuntimeTopologyRecord,
453    {
454        let mut state = self.lock_state();
455        if state.sealed_snapshot.is_some() {
456            return RuntimeObservationRecordDisposition::Sealed;
457        }
458        let sequence = state.take_sequence();
459        const MIN_TOPOLOGY_RECORD_BYTES: u64 = 24;
460        if let Some(reason) = state.rejection_reason(self.manifest.limits, MIN_TOPOLOGY_RECORD_BYTES) {
461            state.record_drop(class, sequence, reason);
462            return RuntimeObservationRecordDisposition::Dropped { sequence, reason };
463        }
464        let record = build_record();
465        if record.class() != class {
466            state.degraded = true;
467        }
468        let logical_bytes = record.logical_retained_bytes();
469        match state.admit(self.manifest.limits, record.class(), sequence, logical_bytes) {
470            Ok(()) => {
471                state.topology.push(RetainedTopologyRecord { sequence, record });
472                RuntimeObservationRecordDisposition::Retained { sequence, logical_bytes }
473            },
474            Err(reason) => RuntimeObservationRecordDisposition::Dropped { sequence, reason },
475        }
476    }
477
478    pub(super) fn snapshot_data(
479        &self,
480        state: &RuntimeObservationRecorderState,
481        complete: bool,
482        terminal: Option<RuntimeTerminalOutcome>,
483    ) -> RuntimeObservationSnapshotData {
484        let mut topology = state.topology.clone();
485        topology.sort_by(|left, right| left.record.cmp(&right.record).then(left.sequence.cmp(&right.sequence)));
486        let mut events = state.events.clone();
487        events.sort_by_key(|event| event.sequence);
488        let dropped_by_class = state
489            .dropped_by_class
490            .iter()
491            .map(|(class, count)| RuntimeObservationClassLoss { class: *class, count: *count })
492            .collect();
493        let dropped_by_reason = state
494            .dropped_by_reason
495            .iter()
496            .map(|(reason, count)| RuntimeObservationReasonLoss { reason: *reason, count: *count })
497            .collect();
498        RuntimeObservationSnapshotData {
499            manifest: self.manifest.clone(),
500            topology,
501            events,
502            footer: RuntimeObservationFooter {
503                complete,
504                terminal,
505                retained_records: state.retained_records,
506                retained_bytes: state.retained_bytes,
507                first_dropped_sequence: state.first_dropped_sequence,
508                last_attempted_sequence: state.last_attempted_sequence,
509                dropped_by_class,
510                dropped_by_reason,
511                truncated: state.first_dropped_sequence.is_some(),
512                degraded: state.degraded,
513                poison_recoveries: state.poison_recoveries,
514                attempt_report: None,
515                counter_overflow: state.counter_overflow,
516            },
517            plans: BTreeMap::new(),
518        }
519    }
520
521    /// Waits until a seal in flight publishes its snapshot.
522    pub(super) fn wait_for_seal<'a>(
523        &self,
524        state: MutexGuard<'a, RuntimeObservationRecorderState>,
525    ) -> MutexGuard<'a, RuntimeObservationRecorderState> {
526        self.seal_published.wait(state).unwrap_or_else(PoisonError::into_inner)
527    }
528
529    /// Takes the retained records and computes the complete footer. The
530    /// footer's attempt report is filled in by the caller.
531    pub(super) fn take_sealed_snapshot_data(
532        &self,
533        state: &mut RuntimeObservationRecorderState,
534        outcome: Option<RuntimeTerminalOutcome>,
535    ) -> RuntimeObservationSnapshotData {
536        let mut topology = std::mem::take(&mut state.topology);
537        topology.sort_by(|left, right| left.record.cmp(&right.record).then(left.sequence.cmp(&right.sequence)));
538        let mut events = std::mem::take(&mut state.events);
539        events.sort_by_key(|event| event.sequence);
540        let dropped_by_class = state
541            .dropped_by_class
542            .iter()
543            .map(|(class, count)| RuntimeObservationClassLoss { class: *class, count: *count })
544            .collect();
545        let dropped_by_reason = state
546            .dropped_by_reason
547            .iter()
548            .map(|(reason, count)| RuntimeObservationReasonLoss { reason: *reason, count: *count })
549            .collect();
550        let footer = RuntimeObservationFooter {
551            complete: true,
552            terminal: outcome,
553            retained_records: state.retained_records,
554            retained_bytes: state.retained_bytes,
555            first_dropped_sequence: state.first_dropped_sequence,
556            last_attempted_sequence: state.last_attempted_sequence,
557            dropped_by_class,
558            dropped_by_reason,
559            truncated: state.first_dropped_sequence.is_some(),
560            degraded: state.degraded,
561            poison_recoveries: state.poison_recoveries,
562            attempt_report: None,
563            counter_overflow: state.counter_overflow,
564        };
565        RuntimeObservationSnapshotData {
566            manifest: self.manifest.clone(),
567            topology,
568            events,
569            footer,
570            plans: BTreeMap::new(),
571        }
572    }
573}
query_runtime/observability/recorder/session.rs

query_runtime/observability/recorder/
session.rs
