Skip to main content

query_runtime/observability/recorder/
handle.rs

1//! Record-path handle methods and bounded detail retention.
2
3use std::{
4    collections::BTreeMap,
5    fmt, mem,
6    panic::{self, AssertUnwindSafe},
7    sync::{
8        Arc, Condvar, Mutex, MutexGuard,
9        atomic::{AtomicBool, Ordering},
10    },
11};
12
13use query_planning::correlation::{QueryId, TraceId};
14
15use super::{
16    AttemptReportSealProvider, BudgetWindow, BudgetWindowPeaks, CleanupFailureSummary, DeferredRuntimeObservationSeal,
17    LiveRuntimeObservationSnapshot, MetricsCaptureMode, MetricsCompleteness, OsSourceLocationDigestKeyProvider,
18    PlanExecutionId, PlanExecutionMetrics, QueryMetrics, QueryMetricsCheckpoint, RUNTIME_OBSERVATION_JSONL_SCHEMA,
19    RuntimeClockDomain, RuntimeEventContext, RuntimeObservationClock, RuntimeObservationConfig,
20    RuntimeObservationHandle, RuntimeObservationLimits, RuntimeObservationLossReason, RuntimeObservationManifest,
21    RuntimeObservationPayload, RuntimeObservationRecordClass, RuntimeObservationRecordDisposition,
22    RuntimeObservationRecorderState, RuntimeObservationSealObserver, RuntimeObservationSessionInner,
23    RuntimeSourceLocationId, RuntimeTerminalOutcome, RuntimeTopologyRecord, SealedRuntimeObservationSnapshot,
24    SourceLocationDigestKeyProvider, SourceLocationDigestKeyState, SourceLocationDigestKeyUnavailable, SourceReadFact,
25    SourceReadRuntimeMetrics, SystemRuntimeObservationClock,
26};
27use crate::{
28    metrics::BoundaryMetrics,
29    observability::{LakehouseCacheFillEvent, LakehouseCacheFillEventSnapshot, recorder::ExecutionProgress},
30};
31
32/// Growing detail families and the loss class each is retained under.
33///
34/// Families bounded by plan topology; graph calibration rows, which a later
35/// measurement updates in place, task mappings, and the keyed task summaries;
36/// are not listed: they cannot grow with data, and retaining them in full keeps
37/// in-place updates correct.
38fn detail_lengths(metrics: &QueryMetrics) -> [usize; 5] {
39    [
40        metrics.operators.len(),
41        metrics.join_strategy_decisions.len(),
42        metrics.aggregate_strategy_decisions.len(),
43        metrics.dynamic_filter_runtime_rows.len(),
44        metrics.budgeted_execution_events.len(),
45    ]
46}
47
48const DETAIL_CLASSES: [RuntimeObservationRecordClass; 5] = [
49    RuntimeObservationRecordClass::OperatorDetail,
50    RuntimeObservationRecordClass::JoinStrategyDetail,
51    RuntimeObservationRecordClass::AggregateStrategyDetail,
52    RuntimeObservationRecordClass::DynamicFilterDetail,
53    RuntimeObservationRecordClass::BudgetDetail,
54];
55
56impl RuntimeObservationHandle {
57    /// Creates a query-bound session using a system monotonic clock.
58    ///
59    /// A disabled config still allocates the session. Trace capture is off,
60    /// and the query's metrics reduction continues. Only [`Self::disabled`]
61    /// produces a session with no recorder at all.
62    pub(super) fn new(
63        query_id: impl Into<QueryId>,
64        trace_id: Option<TraceId>,
65        config: RuntimeObservationConfig,
66    ) -> Self {
67        Self::with_clock(query_id, trace_id, config, Arc::new(SystemRuntimeObservationClock::new()))
68    }
69
70    pub(super) fn with_clock(
71        query_id: impl Into<QueryId>,
72        trace_id: Option<TraceId>,
73        config: RuntimeObservationConfig,
74        clock: Arc<dyn RuntimeObservationClock>,
75    ) -> Self {
76        Self::with_clock_and_source_location_digest_key(
77            query_id,
78            trace_id,
79            config,
80            clock,
81            &OsSourceLocationDigestKeyProvider,
82        )
83    }
84
85    pub(super) fn with_clock_and_source_location_digest_key(
86        query_id: impl Into<QueryId>,
87        trace_id: Option<TraceId>,
88        config: RuntimeObservationConfig,
89        clock: Arc<dyn RuntimeObservationClock>,
90        source_location_digest_key_provider: &dyn SourceLocationDigestKeyProvider,
91    ) -> Self {
92        let manifest = RuntimeObservationManifest {
93            schema: RUNTIME_OBSERVATION_JSONL_SCHEMA.to_string(),
94            query_id: query_id.into(),
95            trace_id,
96            unix_epoch_ns: clock.unix_epoch_ns(),
97            clock_domain: RuntimeClockDomain::HostMonotonic,
98            limits: config.limits(),
99        };
100        let (source_location_digest_key, state) = match source_location_digest_key_provider.source_location_digest_key()
101        {
102            Ok(key) => (
103                SourceLocationDigestKeyState::Available(key),
104                RuntimeObservationRecorderState { next_sequence: 1, ..Default::default() },
105            ),
106            Err(SourceLocationDigestKeyUnavailable) => {
107                let mut dropped_by_reason = BTreeMap::new();
108                dropped_by_reason.insert(RuntimeObservationLossReason::SourceLocationKeyUnavailable, 1);
109                (
110                    SourceLocationDigestKeyState::Unavailable,
111                    RuntimeObservationRecorderState {
112                        next_sequence: 1,
113                        dropped_by_reason,
114                        degraded: true,
115                        ..Default::default()
116                    },
117                )
118            },
119        };
120        Self {
121            config,
122            inner: Some(Arc::new(RuntimeObservationSessionInner {
123                manifest,
124                source_location_digest_key,
125                clock,
126                state: Mutex::new(state),
127                sealed: AtomicBool::new(false),
128                seal_published: Condvar::new(),
129            })),
130        }
131    }
132
133    /// Returns a disabled session.
134    pub const fn disabled() -> Self {
135        Self { config: RuntimeObservationConfig::disabled(), inner: None }
136    }
137
138    /// Returns whether this session constructs and records trace payloads.
139    pub(crate) const fn is_enabled(&self) -> bool {
140        self.config.is_enabled() && self.inner.is_some()
141    }
142
143    // Every pre-existing handle method is a trace operation and keeps its
144    // exact disabled behaviour by going through this accessor: no lazy payload
145    // runs, and no trace snapshot is produced. Only the metrics methods below
146    // read `inner` directly, because metrics reduction outlives trace capture.
147    fn trace_inner(&self) -> Option<&Arc<RuntimeObservationSessionInner>> {
148        self.inner.as_ref().filter(|_| self.config.is_enabled())
149    }
150
151    /// Returns the immutable recorder configuration.
152    pub const fn config(&self) -> RuntimeObservationConfig {
153        self.config
154    }
155
156    /// Counts one cleanup failure in the session's exact tally, which the
157    /// attempt report reads at seal, before the caller attempts an optional
158    /// timeline projection. Failures after the seal are ignored.
159    #[doc(hidden)]
160    pub fn record_attempt_cleanup_failure(&self, summary: CleanupFailureSummary) {
161        let Some(inner) = &self.inner else {
162            return;
163        };
164        let mut state = inner.lock_state();
165        // Sealing flips this flag while holding the same lock.
166        if inner.sealed.load(Ordering::Acquire) {
167            return;
168        }
169        state.cleanup_failures.record(summary);
170    }
171
172    /// Returns the immutable query identity when recording is enabled.
173    pub fn query_id(&self) -> Option<&QueryId> {
174        self.trace_inner().map(|inner| &inner.manifest.query_id)
175    }
176
177    /// Returns the immutable optional parent trace identity.
178    pub fn trace_id(&self) -> Option<&TraceId> {
179        self.trace_inner().and_then(|inner| inner.manifest.trace_id.as_ref())
180    }
181
182    /// Returns a keyed, query-scoped source-location identity when recording is enabled.
183    pub fn source_location_id(&self, value: &str) -> Option<RuntimeSourceLocationId> {
184        self.trace_inner().and_then(|inner| inner.source_location_digest_key.source_location_id(value))
185    }
186
187    /// Lazily constructs and records a query-level event.
188    pub fn record<F>(&self, build_payload: F) -> RuntimeObservationRecordDisposition
189    where F: FnOnce() -> RuntimeObservationPayload {
190        self.record_with_context(RuntimeEventContext::default(), build_payload)
191    }
192
193    /// Lazily constructs and records an event for one plan execution.
194    pub fn record_for_plan<F>(
195        &self,
196        plan_execution_id: PlanExecutionId,
197        build_payload: F,
198    ) -> RuntimeObservationRecordDisposition
199    where
200        F: FnOnce() -> RuntimeObservationPayload,
201    {
202        self.record_with_context(RuntimeEventContext::for_plan(plan_execution_id), build_payload)
203    }
204
205    /// Lazily constructs an event before acquiring recorder order and records it infallibly.
206    pub fn record_with_context<F>(
207        &self,
208        context: RuntimeEventContext,
209        build_payload: F,
210    ) -> RuntimeObservationRecordDisposition
211    where
212        F: FnOnce() -> RuntimeObservationPayload,
213    {
214        let Some(inner) = self.trace_inner() else {
215            return RuntimeObservationRecordDisposition::Disabled;
216        };
217        if inner.sealed.load(Ordering::Acquire) {
218            return RuntimeObservationRecordDisposition::Sealed;
219        }
220        inner.record_event_lazy(context, build_payload)
221    }
222
223    /// Records an event after an exact borrowed-data size preflight.
224    ///
225    /// The payload closure is not invoked when the record or byte cap already
226    /// rejects `logical_bytes`, so variable-size strings and row-group vectors
227    /// are never cloned merely to discover that capture is exhausted.
228    #[cfg(test)]
229    pub(crate) fn record_with_context_sized<F>(
230        &self,
231        context: RuntimeEventContext,
232        logical_bytes: u64,
233        build_payload: F,
234    ) -> RuntimeObservationRecordDisposition
235    where
236        F: FnOnce() -> RuntimeObservationPayload,
237    {
238        let Some(inner) = self.trace_inner() else {
239            return RuntimeObservationRecordDisposition::Disabled;
240        };
241        if inner.sealed.load(Ordering::Acquire) {
242            return RuntimeObservationRecordDisposition::Sealed;
243        }
244        inner.record_event_sized(context, logical_bytes, build_payload)
245    }
246
247    /// Records an event whose payload depends on its exact recorder timestamp.
248    ///
249    /// The payload builder runs while recorder order is held and receives the
250    /// same query-relative monotonic timestamp retained on the event. The
251    /// timestamp is returned even when bounded capture drops the event, and
252    /// `None` is returned only when recording is disabled or already sealed.
253    #[doc(hidden)]
254    pub fn record_with_context_timestamped<F>(&self, context: RuntimeEventContext, build_payload: F) -> Option<u64>
255    where F: FnOnce(u64) -> RuntimeObservationPayload {
256        let inner = self.trace_inner()?;
257        if inner.sealed.load(Ordering::Acquire) {
258            return None;
259        }
260        inner.record_event_timestamped(context, build_payload).1
261    }
262
263    /// Lazily constructs and records one static topology row.
264    pub fn record_topology<F>(
265        &self,
266        class: RuntimeObservationRecordClass,
267        build_record: F,
268    ) -> RuntimeObservationRecordDisposition
269    where
270        F: FnOnce() -> RuntimeTopologyRecord,
271    {
272        let Some(inner) = self.trace_inner() else {
273            return RuntimeObservationRecordDisposition::Disabled;
274        };
275        if inner.sealed.load(Ordering::Acquire) {
276            return RuntimeObservationRecordDisposition::Sealed;
277        }
278        inner.record_topology_lazy(class, build_record)
279    }
280
281    /// Captures an immutable incomplete view. Returns none after sealing or when disabled.
282    pub fn live_snapshot(&self) -> Option<LiveRuntimeObservationSnapshot> {
283        let inner = self.trace_inner()?;
284        let state = inner.lock_state();
285        if state.sealed_snapshot.is_some() {
286            return None;
287        }
288        Some(LiveRuntimeObservationSnapshot { data: Arc::new(inner.snapshot_data(&state, false, None)) })
289    }
290
291    pub(super) fn seal_owned(&self, outcome: RuntimeTerminalOutcome) -> Option<SealedRuntimeObservationSnapshot> {
292        self.seal_owned_deferred(Some(outcome), None).invoke_observers()
293    }
294
295    pub(super) fn seal_owned_deferred(
296        &self,
297        outcome: Option<RuntimeTerminalOutcome>,
298        attempt_report_provider: Option<&AttemptReportSealProvider>,
299    ) -> DeferredRuntimeObservationSeal {
300        let Some(inner) = self.inner.as_ref() else {
301            return DeferredRuntimeObservationSeal { snapshot: None, observers: Vec::new() };
302        };
303        let mut state = inner.lock_state();
304        if !self.config.is_enabled() {
305            // No trace snapshot exists to seal, but the metrics reduction must
306            // still stop changing once the query's terminal outcome is fixed.
307            inner.sealed.store(true, Ordering::Release);
308            return DeferredRuntimeObservationSeal { snapshot: None, observers: Vec::new() };
309        }
310        loop {
311            if let Some(snapshot) = &state.sealed_snapshot {
312                return DeferredRuntimeObservationSeal { snapshot: Some(snapshot.clone()), observers: Vec::new() };
313            }
314            if !inner.sealed.load(Ordering::Acquire) {
315                break;
316            }
317            state = inner.wait_for_seal(state);
318        }
319        inner.sealed.store(true, Ordering::Release);
320        // Every exit after `sealed` publishes a snapshot, so a seal waiting in
321        // `wait_for_seal` never waits on one that unwound.
322        match panic::catch_unwind(AssertUnwindSafe(|| {
323            self.publish_seal(inner, state, outcome, attempt_report_provider)
324        })) {
325            Ok(deferred) => deferred,
326            Err(panic) => {
327                let mut state = inner.lock_state();
328                if state.sealed_snapshot.is_none() {
329                    let mut data = inner.take_sealed_snapshot_data(&mut state, Some(RuntimeTerminalOutcome::Panicked));
330                    data.footer.degraded = true;
331                    state.sealed_snapshot = Some(SealedRuntimeObservationSnapshot { data: Arc::new(data) });
332                }
333                let deferred = DeferredRuntimeObservationSeal {
334                    snapshot: state.sealed_snapshot.clone(),
335                    observers: mem::take(&mut state.seal_observers),
336                };
337                inner.seal_published.notify_all();
338                // Observers may re-enter the session, which locks the recorder state.
339                drop(state);
340                deferred.invoke_observers();
341                panic::resume_unwind(panic);
342            },
343        }
344    }
345
346    /// Builds and publishes the sealed snapshot of a seal that has set
347    /// `sealed`.
348    fn publish_seal<'a>(
349        &self,
350        inner: &'a RuntimeObservationSessionInner,
351        mut state: MutexGuard<'a, RuntimeObservationRecorderState>,
352        outcome: Option<RuntimeTerminalOutcome>,
353        attempt_report_provider: Option<&AttemptReportSealProvider>,
354    ) -> DeferredRuntimeObservationSeal {
355        settle_cache_fill_events(&mut state, inner.manifest.limits);
356        let mut data = inner.take_sealed_snapshot_data(&mut state, outcome);
357        data.plans = self.sealed_plan_metrics(&state, &inner.manifest);
358        let mut provider_panic = None;
359        if let Some(provider) = attempt_report_provider {
360            let cleanup_failures = state.cleanup_failures.clone();
361            // The provider runs the attempt's report sinks, which may read or
362            // record through this session; recording is already closed.
363            drop(state);
364            // The provider builds the exact immutable report that completion
365            // receivers publish, from the complete footer.
366            match panic::catch_unwind(AssertUnwindSafe(|| provider(&data.footer, &cleanup_failures))) {
367                Ok(report) => {
368                    debug_assert_eq!(data.footer.terminal.as_ref(), Some(&report.outcome));
369                    data.footer.attempt_report = Some(report);
370                },
371                Err(panic) => {
372                    data.footer.terminal = Some(RuntimeTerminalOutcome::Panicked);
373                    data.footer.degraded = true;
374                    provider_panic = Some(panic);
375                },
376            }
377            state = inner.lock_state();
378        }
379        let snapshot = SealedRuntimeObservationSnapshot { data: Arc::new(data) };
380        state.sealed_snapshot = Some(snapshot.clone());
381        let observers = mem::take(&mut state.seal_observers);
382        inner.seal_published.notify_all();
383        let deferred = DeferredRuntimeObservationSeal { snapshot: Some(snapshot), observers };
384        if let Some(panic) = provider_panic {
385            // Observers run outside the state lock. The snapshot is already
386            // published, so a recovery seal returns it without waiting.
387            drop(state);
388            deferred.invoke_observers();
389            panic::resume_unwind(panic);
390        }
391        deferred
392    }
393
394    pub(super) fn mark_degraded(&self) {
395        if let Some(inner) = self.trace_inner() {
396            inner.lock_state().degraded = true;
397        }
398    }
399
400    pub(super) fn register_seal_observer(
401        &self,
402        observer: RuntimeObservationSealObserver,
403    ) -> Option<(RuntimeObservationSealObserver, SealedRuntimeObservationSnapshot)> {
404        let inner = self.trace_inner()?;
405        let mut state = inner.lock_state();
406        if let Some(snapshot) = &state.sealed_snapshot {
407            Some((observer, snapshot.clone()))
408        } else {
409            state.seal_observers.push(observer);
410            None
411        }
412    }
413
414    /// Returns the sealed immutable snapshot when a terminal outcome was already fixed.
415    pub fn sealed_snapshot(&self) -> Option<SealedRuntimeObservationSnapshot> {
416        let inner = self.trace_inner()?;
417        inner.lock_state().sealed_snapshot.clone()
418    }
419
420    /// Folds one source-read fact into a plan execution's exact totals.
421    ///
422    /// This runs whether or not trace capture is enabled, and before any
423    /// retention decision, so the totals never depend on what was retained.
424    /// Facts arriving after the owner sealed are ignored.
425    pub fn record_source_read(&self, plan_execution_id: &PlanExecutionId, fact: SourceReadFact) {
426        let Some(inner) = &self.inner else {
427            return;
428        };
429        let mut state = inner.lock_state();
430        // Sealing flips this flag while holding the same lock, so reading it
431        // here cannot race a concurrent seal.
432        if inner.sealed.load(Ordering::Acquire) {
433            return;
434        }
435        let execution = state.metrics.entry(plan_execution_id.clone()).or_default();
436        if execution.is_frozen() {
437            return;
438        }
439        if !execution.summary.source_read.apply(&fact) {
440            execution.totals_saturated = true;
441        }
442        state.metrics_revision = state.metrics_revision.saturating_add(1);
443    }
444
445    /// Offers one source-read detail row for bounded retention.
446    ///
447    /// Rows never feed totals; record the matching [`SourceReadFact`]s for
448    /// those. With trace capture disabled the row is intentionally not kept;
449    /// otherwise it is retained under the recorder's record and byte limits,
450    /// which are shared by the whole query, and a declined row is counted as a
451    /// `SourceReadDetail` loss.
452    pub fn record_source_read_row(&self, plan_execution_id: &PlanExecutionId, row: SourceReadRuntimeMetrics) {
453        let Some(inner) = self.trace_inner() else {
454            return;
455        };
456        let mut state = inner.lock_state();
457        if inner.sealed.load(Ordering::Acquire)
458            || state.metrics.get(plan_execution_id).is_some_and(PlanExecutionMetrics::is_frozen)
459        {
460            return;
461        }
462        let logical_bytes = debug_logical_bytes(&row);
463        let sequence = state.take_sequence();
464        let admitted = state
465            .admit(inner.manifest.limits, RuntimeObservationRecordClass::SourceReadDetail, sequence, logical_bytes)
466            .is_ok();
467        let execution = state.metrics.entry(plan_execution_id.clone()).or_default();
468        if admitted {
469            execution.source_read_rows.push(row);
470        } else {
471            execution.source_read_rows_dropped = execution.source_read_rows_dropped.saturating_add(1);
472        }
473        state.metrics_revision = state.metrics_revision.saturating_add(1);
474    }
475
476    /// Offers one background lakehouse cache-fill completion for bounded
477    /// retention.
478    ///
479    /// With trace capture disabled the event is not kept. Otherwise the next
480    /// plan execution to freeze its metrics offers it under the query's record
481    /// and byte limits, and a declined event is a `LakehouseCacheFillDetail`
482    /// loss in that plan's completeness. The seal settles events no plan froze
483    /// with: an unfrozen plan execution carries them, or each is an
484    /// `Unassigned` loss.
485    #[doc(hidden)]
486    pub fn record_lakehouse_cache_fill(&self, event: &LakehouseCacheFillEvent) {
487        let Some(inner) = self.trace_inner() else {
488            return;
489        };
490        let mut state = inner.lock_state();
491        if inner.sealed.load(Ordering::Acquire) {
492            return;
493        }
494        let waiting = u64::try_from(state.cache_fill_events.len()).unwrap_or(u64::MAX);
495        if waiting < inner.manifest.limits.max_retained_records {
496            state.cache_fill_events.push(LakehouseCacheFillEventSnapshot::from(event));
497        } else {
498            state.cache_fill_events_overflow = state.cache_fill_events_overflow.saturating_add(1);
499        }
500    }
501
502    /// Returns whether this handle reaches a query-bound owner, which is
503    /// where metrics are reduced even when trace capture is off.
504    pub(crate) const fn has_owner(&self) -> bool {
505        self.inner.is_some()
506    }
507
508    /// Returns the data-boundary counters of a plan execution.
509    pub(crate) fn boundary_metrics(&self, plan_execution_id: &PlanExecutionId) -> BoundaryMetrics {
510        let Some(inner) = &self.inner else {
511            return BoundaryMetrics::default();
512        };
513        let state = inner.lock_state();
514        state.metrics.get(plan_execution_id).map(|execution| execution.reduction.boundary.clone()).unwrap_or_default()
515    }
516
517    /// Returns the progress counters of a plan execution.
518    pub(crate) fn execution_progress(&self, plan_execution_id: &PlanExecutionId) -> super::ExecutionProgress {
519        let Some(inner) = &self.inner else {
520            return ExecutionProgress::default();
521        };
522        let state = inner.lock_state();
523        state.metrics.get(plan_execution_id).map_or_else(Default::default, |execution| super::ExecutionProgress {
524            source_chunks_planned: execution.summary.source_read.source_chunks_planned,
525            source_chunks_committed: execution.summary.source_read.source_chunks_committed,
526            partial_aggregate_count: execution.reduction.partial_aggregate_count,
527        })
528    }
529
530    /// Returns one internally consistent view of a plan execution's metrics.
531    ///
532    /// Available for any query-bound session, including one whose trace
533    /// capture is disabled. A plan execution that recorded nothing yet reports
534    /// empty totals. Returns `None` only for [`Self::disabled`].
535    pub fn metrics_checkpoint(&self, plan_execution_id: &PlanExecutionId) -> Option<QueryMetricsCheckpoint> {
536        let inner = self.inner.as_ref()?;
537        let state = inner.lock_state();
538        Some(self.plan_checkpoint(&state, plan_execution_id, inner.sealed.load(Ordering::Acquire)))
539    }
540
541    /// Freezes a plan execution's metrics and returns them. Facts recorded
542    /// for that plan afterwards are ignored; freezing again returns the same
543    /// checkpoint. Returns `None` only for [`Self::disabled`].
544    pub(crate) fn freeze_plan_metrics(&self, plan_execution_id: &PlanExecutionId) -> Option<QueryMetricsCheckpoint> {
545        let inner = self.inner.as_ref()?;
546        let mut state = inner.lock_state();
547        let mut execution = state.metrics.remove(plan_execution_id).unwrap_or_default();
548        // After the seal a plan's metrics are the sealed snapshot's value; the
549        // seal already settled every waiting cache-fill event.
550        if !execution.is_frozen() && !inner.sealed.load(Ordering::Acquire) {
551            assign_cache_fill_events(&mut state, &mut execution, inner.manifest.limits);
552        }
553        state.metrics.insert(plan_execution_id.clone(), execution);
554        let checkpoint = self.plan_checkpoint(&state, plan_execution_id, true);
555        state.metrics.entry(plan_execution_id.clone()).or_default().frozen = Some(checkpoint.clone());
556        Some(checkpoint)
557    }
558
559    /// Every plan execution's metrics as a sealed snapshot carries them: the
560    /// frozen value, or the reduction at seal.
561    fn sealed_plan_metrics(
562        &self,
563        state: &RuntimeObservationRecorderState,
564        manifest: &RuntimeObservationManifest,
565    ) -> BTreeMap<PlanExecutionId, QueryMetrics> {
566        state
567            .metrics
568            .keys()
569            .map(|plan_execution_id| {
570                let mut metrics = QueryMetrics {
571                    query_id: QueryId::new(manifest.query_id.as_str()),
572                    trace_id: manifest.trace_id.as_ref().map(|trace_id| TraceId::new(trace_id.as_str())),
573                    ..QueryMetrics::default()
574                };
575                metrics.apply_owner_checkpoint(&self.plan_checkpoint(state, plan_execution_id, true));
576                (plan_execution_id.clone(), metrics)
577            })
578            .collect()
579    }
580
581    fn plan_checkpoint(
582        &self,
583        state: &RuntimeObservationRecorderState,
584        plan_execution_id: &PlanExecutionId,
585        sealed: bool,
586    ) -> QueryMetricsCheckpoint {
587        let empty = PlanExecutionMetrics::default();
588        let execution = state.metrics.get(plan_execution_id).unwrap_or(&empty);
589        if let Some(frozen) = &execution.frozen {
590            return frozen.clone();
591        }
592        let detail_rows_dropped = execution.source_read_rows_dropped.saturating_add(execution.detail_rows_dropped);
593        QueryMetricsCheckpoint {
594            revision: state.metrics_revision,
595            sealed,
596            summary: execution.summary.clone(),
597            source_read_rows: execution.source_read_rows.clone(),
598            metrics: execution.reduction.clone(),
599            completeness: MetricsCompleteness {
600                capture_mode: if self.config.is_enabled() {
601                    MetricsCaptureMode::Bounded
602                } else {
603                    MetricsCaptureMode::Disabled
604                },
605                detail_rows_dropped,
606                truncated: detail_rows_dropped > 0,
607                totals_saturated: execution.totals_saturated || state.counter_overflow,
608            },
609        }
610    }
611
612    /// Applies one typed metrics fact to a plan execution's reduction.
613    ///
614    /// Only `MetricsRecorder` calls this, with a closure that invokes one
615    /// existing `QueryMetrics` reduction; execution code never receives the
616    /// reduction itself. Rows the fact appended to a growing detail family are
617    /// then offered for bounded retention, after open budget windows have seen
618    /// them, so no total depends on what was kept.
619    pub(crate) fn apply_metrics(&self, plan_execution_id: &PlanExecutionId, reduce: impl FnOnce(&mut QueryMetrics)) {
620        let Some(inner) = &self.inner else {
621            return;
622        };
623        let mut state = inner.lock_state();
624        if inner.sealed.load(Ordering::Acquire)
625            || state.metrics.get(plan_execution_id).is_some_and(PlanExecutionMetrics::is_frozen)
626        {
627            return;
628        }
629        let mut execution = state.metrics.remove(plan_execution_id).unwrap_or_default();
630        let before = detail_lengths(&execution.reduction);
631        reduce(&mut execution.reduction);
632        let new_budget_events = execution.reduction.budgeted_execution_events.get(before[4]..).unwrap_or_default();
633        for event in new_budget_events {
634            for peaks in execution.budget_windows.values_mut() {
635                peaks.observed_peak_bytes = peaks.observed_peak_bytes.max(event.allocator_high_water_bytes);
636                peaks.retained_bytes = peaks.retained_bytes.max(event.retained_bytes);
637            }
638        }
639        retain_new_detail_rows(&mut state, &mut execution, before, inner.manifest.limits, self.config.is_enabled());
640        state.metrics.insert(plan_execution_id.clone(), execution);
641        state.metrics_revision = state.metrics_revision.saturating_add(1);
642    }
643
644    /// Opens a budget window on a plan execution.
645    pub(crate) fn open_budget_window(&self, plan_execution_id: &PlanExecutionId) -> BudgetWindow {
646        let Some(inner) = &self.inner else {
647            return BudgetWindow(0);
648        };
649        let mut state = inner.lock_state();
650        let execution = state.metrics.entry(plan_execution_id.clone()).or_default();
651        execution.next_budget_window = execution.next_budget_window.saturating_add(1);
652        let id = execution.next_budget_window;
653        execution.budget_windows.insert(id, BudgetWindowPeaks::default());
654        BudgetWindow(id)
655    }
656
657    /// Closes a budget window and returns the peaks it saw.
658    pub(crate) fn take_budget_window(
659        &self,
660        plan_execution_id: &PlanExecutionId,
661        window: BudgetWindow,
662    ) -> BudgetWindowPeaks {
663        let Some(inner) = &self.inner else {
664            return BudgetWindowPeaks::default();
665        };
666        let mut state = inner.lock_state();
667        state
668            .metrics
669            .get_mut(plan_execution_id)
670            .and_then(|execution| execution.budget_windows.remove(&window.0))
671            .unwrap_or_default()
672    }
673
674    /// Closes a budget window, folding the peaks it saw into one task.
675    pub(crate) fn close_budget_window(
676        &self,
677        plan_execution_id: &PlanExecutionId,
678        window: BudgetWindow,
679        compiled_task_id: &str,
680    ) {
681        let Some(inner) = &self.inner else {
682            return;
683        };
684        let mut state = inner.lock_state();
685        let sealed = inner.sealed.load(Ordering::Acquire);
686        let Some(execution) = state.metrics.get_mut(plan_execution_id) else {
687            return;
688        };
689        let Some(peaks) = execution.budget_windows.remove(&window.0) else {
690            return;
691        };
692        if sealed || execution.is_frozen() || (peaks.observed_peak_bytes.is_none() && peaks.retained_bytes.is_none()) {
693            return;
694        }
695        execution.reduction.record_task_budget_observation(
696            compiled_task_id,
697            peaks.observed_peak_bytes,
698            peaks.retained_bytes,
699        );
700        state.metrics_revision = state.metrics_revision.saturating_add(1);
701    }
702}
703
704/// Offers the waiting cache-fill events for retention as `execution`'s rows. A
705/// declined or overflowed event counts against `execution`'s completeness.
706fn assign_cache_fill_events(
707    state: &mut RuntimeObservationRecorderState,
708    execution: &mut PlanExecutionMetrics,
709    limits: RuntimeObservationLimits,
710) {
711    for event in mem::take(&mut state.cache_fill_events) {
712        let logical_bytes = debug_logical_bytes(&event);
713        let sequence = state.take_sequence();
714        if state
715            .admit_detail(limits, RuntimeObservationRecordClass::LakehouseCacheFillDetail, sequence, logical_bytes)
716            .is_ok()
717        {
718            execution.reduction.lakehouse_cache_fill_events.push(event);
719        } else {
720            execution.detail_rows_dropped = execution.detail_rows_dropped.saturating_add(1);
721        }
722    }
723    let overflow = mem::take(&mut state.cache_fill_events_overflow);
724    drop_cache_fill_events(state, overflow, RuntimeObservationLossReason::RecordLimit);
725    execution.detail_rows_dropped = execution.detail_rows_dropped.saturating_add(overflow);
726}
727
728/// Settles the cache-fill events no plan froze with when the session seals:
729/// the first plan execution still unfrozen carries them, and with none left
730/// each is an `Unassigned` loss.
731fn settle_cache_fill_events(state: &mut RuntimeObservationRecorderState, limits: RuntimeObservationLimits) {
732    if state.cache_fill_events.is_empty() && state.cache_fill_events_overflow == 0 {
733        return;
734    }
735    let carrier = state.metrics.iter().find(|(_, execution)| !execution.is_frozen()).map(|(id, _)| id.clone());
736    match carrier {
737        Some(carrier) => {
738            let mut execution = state.metrics.remove(&carrier).unwrap_or_default();
739            assign_cache_fill_events(state, &mut execution, limits);
740            state.metrics.insert(carrier, execution);
741        },
742        None => {
743            let unassigned = u64::try_from(mem::take(&mut state.cache_fill_events).len()).unwrap_or(u64::MAX);
744            drop_cache_fill_events(state, unassigned, RuntimeObservationLossReason::Unassigned);
745            let overflow = mem::take(&mut state.cache_fill_events_overflow);
746            drop_cache_fill_events(state, overflow, RuntimeObservationLossReason::RecordLimit);
747        },
748    }
749}
750
751fn drop_cache_fill_events(
752    state: &mut RuntimeObservationRecorderState,
753    count: u64,
754    reason: RuntimeObservationLossReason,
755) {
756    for _ in 0..count {
757        let sequence = state.take_sequence();
758        state.record_drop(RuntimeObservationRecordClass::LakehouseCacheFillDetail, sequence, reason);
759    }
760}
761
762fn retain_new_detail_rows(
763    state: &mut RuntimeObservationRecorderState,
764    execution: &mut PlanExecutionMetrics,
765    before: [usize; 5],
766    limits: RuntimeObservationLimits,
767    trace_enabled: bool,
768) {
769    let after = detail_lengths(&execution.reduction);
770    for family in 0..before.len() {
771        let (start, end) = (before[family], after[family]);
772        if end <= start {
773            continue;
774        }
775        let mut keep = Vec::with_capacity(end - start);
776        for row in start..end {
777            // With trace capture off, detail rows are intentionally not kept
778            // and are not losses: nothing was attempted.
779            let admitted = trace_enabled && {
780                let logical_bytes = detail_row_logical_bytes(&execution.reduction, family, row);
781                let sequence = state.take_sequence();
782                state.admit_detail(limits, DETAIL_CLASSES[family], sequence, logical_bytes).is_ok()
783            };
784            if trace_enabled && !admitted {
785                execution.detail_rows_dropped = execution.detail_rows_dropped.saturating_add(1);
786            }
787            keep.push(admitted);
788        }
789        retain_detail_family(&mut execution.reduction, family, start, &keep);
790    }
791}
792
793fn detail_row_logical_bytes(metrics: &QueryMetrics, family: usize, row: usize) -> u64 {
794    match family {
795        0 => debug_logical_bytes(&metrics.operators[row]),
796        1 => debug_logical_bytes(&metrics.join_strategy_decisions[row]),
797        2 => debug_logical_bytes(&metrics.aggregate_strategy_decisions[row]),
798        3 => debug_logical_bytes(&metrics.dynamic_filter_runtime_rows[row]),
799        _ => debug_logical_bytes(&metrics.budgeted_execution_events[row]),
800    }
801}
802
803fn retain_detail_family(metrics: &mut QueryMetrics, family: usize, start: usize, keep: &[bool]) {
804    fn retain<T>(rows: &mut Vec<T>, start: usize, keep: &[bool]) {
805        let appended = rows.split_off(start);
806        rows.extend(appended.into_iter().zip(keep).filter_map(|(row, kept)| kept.then_some(row)));
807    }
808    match family {
809        0 => retain(&mut metrics.operators, start, keep),
810        1 => retain(&mut metrics.join_strategy_decisions, start, keep),
811        2 => retain(&mut metrics.aggregate_strategy_decisions, start, keep),
812        3 => retain(&mut metrics.dynamic_filter_runtime_rows, start, keep),
813        _ => retain(&mut metrics.budgeted_execution_events, start, keep),
814    }
815}
816
817/// Deterministic logical size of a detail row: the length of its `Debug`
818/// rendering, measured without allocating it. This is the same measure the
819/// DataFusion metrics projection used when it selected rows on its own.
820fn debug_logical_bytes(row: &impl fmt::Debug) -> u64 {
821    struct Counter(u64);
822    impl fmt::Write for Counter {
823        fn write_str(&mut self, value: &str) -> fmt::Result {
824            self.0 = self.0.saturating_add(value.len() as u64);
825            Ok(())
826        }
827    }
828    let mut counter = Counter(0);
829    let _ = fmt::write(&mut counter, format_args!("{row:?}"));
830    counter.0
831}