query_runtime/observability/
lifecycle.rs

1//! Exactly-once query and task lifecycle guards.
2
3use super::{
4    MetricsRecorder, PlanExecutionBoundaryKind, PlanExecutionId, QueryMetricsCheckpoint, RuntimeEventContext,
5    RuntimeObservationHandle, RuntimeObservationOwner, RuntimeObservationPayload, RuntimeObservationSession,
6    RuntimeSpanId, RuntimeTaskId, RuntimeTaskLifecycleObservation, RuntimeTerminalOutcome,
7    SealedRuntimeObservationSnapshot, SourceReadFact, TaskAttemptId, TaskLifecycleState,
8};
9use crate::{attempt::AttemptHandle, metrics::SourceReadRuntimeMetrics};
10
11/// Owner-only query lifecycle. Normal completion consumes the owner; Drop
12/// terminalizes abandoned and unwinding executions without exposing seal
13/// authority to workers.
14#[derive(Debug)]
15pub struct RuntimeQueryLifecycle {
16    owner: Option<RuntimeObservationOwner>,
17    plan_execution_id: PlanExecutionId,
18    plan_span_id: RuntimeSpanId,
19}
20
21/// One plan invocation inside a query whose owner outlives native execution.
22///
23/// Finishing this lifecycle records and aggregates the plan outcome, but the
24/// admitted query owner remains responsible for cleanup boundaries and sealing.
25#[derive(Debug)]
26pub struct RuntimePlanExecutionLifecycle {
27    observation: RuntimeObservationHandle,
28    /// The standalone query owner a finished plan's outcome goes to. An
29    /// attempt's plans record only their own boundaries: the attempt commits
30    /// its own outcome.
31    owner: Option<RuntimeObservationOwner>,
32    plan_execution_id: PlanExecutionId,
33    plan_span_id: RuntimeSpanId,
34    finished: bool,
35}
36
37/// Cloneable identity and recording capability for descendants of one plan
38/// execution. It deliberately has no seal authority.
39#[derive(Debug, Clone)]
40pub struct RuntimeExecutionObservationScope {
41    observation: RuntimeObservationHandle,
42    plan_execution_id: PlanExecutionId,
43    plan_span_id: RuntimeSpanId,
44    metrics: MetricsRecorder,
45}
46
47impl RuntimeExecutionObservationScope {
48    /// Returns worker recording access.
49    pub const fn observation(&self) -> &RuntimeObservationHandle {
50        &self.observation
51    }
52
53    /// Returns the plan invocation identity.
54    pub const fn plan_execution_id(&self) -> &PlanExecutionId {
55        &self.plan_execution_id
56    }
57
58    /// Returns the parent span for task attempts.
59    pub const fn plan_span_id(&self) -> RuntimeSpanId {
60        self.plan_span_id
61    }
62
63    /// A scope on `observation` whose metrics go to `metrics`, for driving
64    /// plan stages outside a plan-execution lifecycle.
65    pub fn for_metrics(observation: RuntimeObservationHandle, metrics: MetricsRecorder) -> Self {
66        Self {
67            observation,
68            plan_execution_id: metrics.plan_execution_id().clone(),
69            plan_span_id: RuntimeSpanId::new(1),
70            metrics,
71        }
72    }
73
74    /// Where this plan execution's metrics are reduced.
75    pub const fn metrics(&self) -> &MetricsRecorder {
76        &self.metrics
77    }
78
79    /// Reports this plan execution's metrics to `metrics`, and the trace session
80    /// does not receive them. Use this for an execution whose metrics owner is
81    /// bound separately.
82    #[must_use]
83    pub fn with_metrics(mut self, metrics: MetricsRecorder) -> Self {
84        self.metrics = metrics;
85        self
86    }
87
88    /// Folds one source-read fact into this plan execution's exact totals.
89    pub fn record_source_read(&self, fact: SourceReadFact) {
90        self.metrics.record_source_read(fact);
91    }
92
93    /// Offers one source-read detail row for this plan execution's bounded retention.
94    pub fn record_source_read_row(&self, row: SourceReadRuntimeMetrics) {
95        self.metrics.record_source_read_runtime(row);
96    }
97
98    /// Returns one internally consistent view of this plan execution's metrics.
99    pub fn metrics_checkpoint(&self) -> Option<QueryMetricsCheckpoint> {
100        self.metrics.checkpoint()
101    }
102
103    fn new(
104        observation: RuntimeObservationHandle,
105        plan_execution_id: PlanExecutionId,
106        plan_span_id: RuntimeSpanId,
107    ) -> Self {
108        let metrics = MetricsRecorder::new(observation.clone(), plan_execution_id.clone());
109        Self { observation, plan_execution_id, plan_span_id, metrics }
110    }
111}
112
113impl RuntimeQueryLifecycle {
114    /// Opens one query and plan-execution lifecycle.
115    pub fn new(session: RuntimeObservationSession, plan_execution_id: PlanExecutionId) -> Self {
116        Self::from_owner(RuntimeObservationOwner::from_session(session), plan_execution_id)
117    }
118
119    /// Opens one lifecycle using a shared owner created before admission.
120    pub fn from_owner(owner: RuntimeObservationOwner, plan_execution_id: PlanExecutionId) -> Self {
121        let lifecycle = Self { owner: Some(owner), plan_execution_id, plan_span_id: RuntimeSpanId::new(1) };
122        lifecycle.record_plan(PlanExecutionBoundaryKind::Started);
123        lifecycle
124    }
125
126    /// Returns the invocation identity attached to all descendant events.
127    pub fn plan_execution_id(&self) -> &PlanExecutionId {
128        &self.plan_execution_id
129    }
130
131    /// Returns the root span for descendant task attempts.
132    pub const fn plan_span_id(&self) -> RuntimeSpanId {
133        self.plan_span_id
134    }
135
136    /// Returns cloneable recording access without terminal authority.
137    pub fn handle(&self) -> RuntimeObservationHandle {
138        self.owner.as_ref().map_or_else(RuntimeObservationHandle::disabled, RuntimeObservationOwner::handle)
139    }
140
141    /// Returns the non-owning scope passed to schedulers and reservations.
142    pub fn execution_scope(&self) -> RuntimeExecutionObservationScope {
143        RuntimeExecutionObservationScope::new(self.handle(), self.plan_execution_id.clone(), self.plan_span_id)
144    }
145
146    /// Records cleanup and exactly one terminal, seals the owner, and consumes
147    /// the lifecycle so Drop cannot emit a second terminal.
148    pub fn finish(mut self, outcome: RuntimeTerminalOutcome) -> Option<SealedRuntimeObservationSnapshot> {
149        self.terminate(outcome)
150    }
151
152    fn terminate(&mut self, outcome: RuntimeTerminalOutcome) -> Option<SealedRuntimeObservationSnapshot> {
153        let plan_boundary = match &outcome {
154            RuntimeTerminalOutcome::Succeeded => PlanExecutionBoundaryKind::Finished,
155            RuntimeTerminalOutcome::Cancelled => PlanExecutionBoundaryKind::Cancelled,
156            RuntimeTerminalOutcome::Failed { .. }
157            | RuntimeTerminalOutcome::Dropped
158            | RuntimeTerminalOutcome::Panicked => PlanExecutionBoundaryKind::Failed,
159        };
160        self.record_plan(plan_boundary);
161        self.owner.take().and_then(|owner| owner.finish_query(Some(outcome)))
162    }
163
164    fn record_plan(&self, boundary: PlanExecutionBoundaryKind) {
165        if let Some(owner) = &self.owner {
166            owner.handle().record_with_context(
167                RuntimeEventContext::for_plan(self.plan_execution_id.clone()).with_span(self.plan_span_id, None),
168                || RuntimeObservationPayload::PlanExecutionBoundary { boundary },
169            );
170        }
171    }
172}
173
174impl RuntimePlanExecutionLifecycle {
175    /// Opens a plan invocation using query-level owner authority retained by a
176    /// caller such as the query-service admission ticket.
177    pub fn from_owner(owner: RuntimeObservationOwner, plan_execution_id: PlanExecutionId) -> Self {
178        Self::open(owner.handle(), Some(owner), plan_execution_id)
179    }
180
181    /// Opens a plan invocation of an admitted attempt.
182    pub(crate) fn from_attempt(attempt: &AttemptHandle, plan_execution_id: PlanExecutionId) -> Self {
183        Self::open(attempt.observation(), None, plan_execution_id)
184    }
185
186    fn open(
187        observation: RuntimeObservationHandle,
188        owner: Option<RuntimeObservationOwner>,
189        plan_execution_id: PlanExecutionId,
190    ) -> Self {
191        let lifecycle =
192            Self { observation, owner, plan_execution_id, plan_span_id: RuntimeSpanId::new(1), finished: false };
193        lifecycle.record_plan(PlanExecutionBoundaryKind::Started);
194        lifecycle
195    }
196
197    /// Opens a plan invocation on a coordinator-owned query observation.
198    ///
199    /// The handle can record plan/task evidence but carries no terminal-seal
200    /// authority. The external query coordinator owns the sole query terminal.
201    pub fn from_handle(observation: RuntimeObservationHandle, plan_execution_id: PlanExecutionId) -> Self {
202        Self::open(observation, None, plan_execution_id)
203    }
204
205    /// Returns cloneable recording access without terminal authority.
206    pub fn handle(&self) -> RuntimeObservationHandle {
207        self.observation.clone()
208    }
209
210    /// Returns the non-owning scope passed to schedulers and reservations.
211    pub fn execution_scope(&self) -> RuntimeExecutionObservationScope {
212        RuntimeExecutionObservationScope::new(self.handle(), self.plan_execution_id.clone(), self.plan_span_id)
213    }
214
215    /// Records the plan terminal and contributes its outcome to the query
216    /// owner without sealing the shared observation session.
217    pub fn finish(mut self, outcome: RuntimeTerminalOutcome) {
218        self.terminate(outcome);
219    }
220
221    fn terminate(&mut self, outcome: RuntimeTerminalOutcome) {
222        let boundary = match &outcome {
223            RuntimeTerminalOutcome::Succeeded => PlanExecutionBoundaryKind::Finished,
224            RuntimeTerminalOutcome::Cancelled => PlanExecutionBoundaryKind::Cancelled,
225            RuntimeTerminalOutcome::Failed { .. }
226            | RuntimeTerminalOutcome::Dropped
227            | RuntimeTerminalOutcome::Panicked => PlanExecutionBoundaryKind::Failed,
228        };
229        self.record_plan(boundary);
230        self.finished = true;
231        if let Some(owner) = self.owner.take() {
232            owner.record_terminal_outcome(outcome);
233        }
234    }
235
236    fn record_plan(&self, boundary: PlanExecutionBoundaryKind) {
237        self.observation.record_with_context(
238            RuntimeEventContext::for_plan(self.plan_execution_id.clone()).with_span(self.plan_span_id, None),
239            || RuntimeObservationPayload::PlanExecutionBoundary { boundary },
240        );
241    }
242}
243
244impl Drop for RuntimePlanExecutionLifecycle {
245    fn drop(&mut self) {
246        if self.finished {
247            return;
248        }
249        let outcome =
250            if std::thread::panicking() { RuntimeTerminalOutcome::Panicked } else { RuntimeTerminalOutcome::Dropped };
251        self.terminate(outcome);
252    }
253}
254
255impl Drop for RuntimeQueryLifecycle {
256    fn drop(&mut self) {
257        if self.owner.is_none() {
258            return;
259        }
260        let outcome =
261            if std::thread::panicking() { RuntimeTerminalOutcome::Panicked } else { RuntimeTerminalOutcome::Dropped };
262        let _ = self.terminate(outcome);
263    }
264}
265
266/// Exactly-once lifecycle for one compiled task attempt.
267#[derive(Debug)]
268pub struct RuntimeTaskLifecycle {
269    observation: RuntimeObservationHandle,
270    context: RuntimeEventContext,
271    attempt_id: TaskAttemptId,
272    task_kind: String,
273    terminal_recorded: bool,
274}
275
276impl RuntimeTaskLifecycle {
277    /// Emits ready and start for one explicit task attempt.
278    pub fn new(
279        observation: RuntimeObservationHandle,
280        plan_execution_id: PlanExecutionId,
281        attempt_id: TaskAttemptId,
282        task_kind: impl Into<String>,
283        span_id: RuntimeSpanId,
284        parent_span_id: RuntimeSpanId,
285    ) -> Self {
286        let lifecycle = Self {
287            observation,
288            context: RuntimeEventContext::for_plan(plan_execution_id).with_span(span_id, Some(parent_span_id)),
289            attempt_id,
290            task_kind: task_kind.into(),
291            terminal_recorded: false,
292        };
293        lifecycle.record(TaskLifecycleState::Ready);
294        lifecycle.record(TaskLifecycleState::Started);
295        lifecycle
296    }
297
298    /// Emits the sole successful terminal for this task attempt.
299    pub fn finish(mut self) {
300        self.record_terminal(TaskLifecycleState::Finished);
301    }
302
303    /// Emits the sole structured failure or cancellation terminal.
304    pub fn fail(mut self, error_code: impl Into<String>, cancelled: bool) {
305        let lifecycle = if cancelled {
306            TaskLifecycleState::Cancelled
307        } else {
308            TaskLifecycleState::Failed { error_code: error_code.into() }
309        };
310        self.record_terminal(lifecycle);
311    }
312
313    /// Emits the sole panic terminal after a scheduler catches an unwind long
314    /// enough to release live resources before resuming it.
315    pub fn panicked(mut self) {
316        self.record_terminal(TaskLifecycleState::Panicked);
317    }
318
319    fn record_terminal(&mut self, lifecycle: TaskLifecycleState) {
320        if self.terminal_recorded {
321            return;
322        }
323        self.record(lifecycle);
324        self.terminal_recorded = true;
325    }
326
327    fn record(&self, lifecycle: TaskLifecycleState) {
328        self.observation.record_with_context(self.context.clone(), || {
329            RuntimeObservationPayload::TaskLifecycle(Box::new(RuntimeTaskLifecycleObservation {
330                attempt_id: self.attempt_id.clone(),
331                task_kind: self.task_kind.clone(),
332                lifecycle,
333            }))
334        });
335    }
336}
337
338impl Drop for RuntimeTaskLifecycle {
339    fn drop(&mut self) {
340        if self.terminal_recorded {
341            return;
342        }
343        let lifecycle =
344            if std::thread::panicking() { TaskLifecycleState::Panicked } else { TaskLifecycleState::Dropped };
345        self.record_terminal(lifecycle);
346    }
347}
348
349/// Builds one task-attempt identity from an explicit task and attempt number.
350pub fn task_attempt_id(task_id: impl Into<RuntimeTaskId>, attempt: u32) -> TaskAttemptId {
351    TaskAttemptId::new(task_id.into(), super::RuntimeAttemptId::new(attempt))
352}
353
354#[cfg(test)]
355mod tests {
356    use query_planning::correlation::QueryId;
357
358    use super::*;
359    use crate::observability::RuntimeObservationConfig;
360
361    #[test]
362    fn coordinator_handle_plan_lifecycle_records_without_sealing_query_root() {
363        let owner = RuntimeObservationOwner::new(
364            QueryId::new("coordinator-owned"),
365            None,
366            RuntimeObservationConfig::bounded(16, 4096),
367        );
368        let lifecycle =
369            RuntimePlanExecutionLifecycle::from_handle(owner.handle(), PlanExecutionId::new("coordinator-plan"));
370
371        lifecycle.finish(RuntimeTerminalOutcome::Succeeded);
372
373        assert!(owner.sealed_snapshot().is_none(), "plan completion must not seal a coordinator-owned query root");
374        let snapshot = owner.finish_query(Some(RuntimeTerminalOutcome::Succeeded)).unwrap();
375        assert!(snapshot.events().iter().any(|event| matches!(
376            event.payload,
377            RuntimeObservationPayload::PlanExecutionBoundary { boundary: PlanExecutionBoundaryKind::Finished }
378        )));
379    }
380}
query_runtime/observability/lifecycle.rs

query_runtime/observability/
lifecycle.rs
