diff --git a/crates/sail-common/src/config/application.rs b/crates/sail-common/src/config/application.rs index 02b65e0a..5474b72f 100644 --- a/crates/sail-common/src/config/application.rs +++ b/crates/sail-common/src/config/application.rs @@ -943,6 +943,7 @@ pub struct FlightConfig { #[serde(deny_unknown_fields)] pub struct TelemetryConfig { pub export_traces: bool, + pub c2_observer: bool, pub export_metrics: bool, pub export_logs: bool, pub exporter: TelemetryExporterConfig, diff --git a/crates/sail-common/src/config/application.yaml b/crates/sail-common/src/config/application.yaml index ce5091d2..7a657b9e 100644 --- a/crates/sail-common/src/config/application.yaml +++ b/crates/sail-common/src/config/application.yaml @@ -975,6 +975,15 @@ description: Whether to export traces to an OpenTelemetry collector. experimental: true +- key: telemetry.c2_observer + type: boolean + default: "false" + description: >- + Observe existing planning and execution calls for C2 through the configured + trace exporter. Requires telemetry.export_traces and an OTLP endpoint. + Durations overlap and do not form an additive execution breakdown. + experimental: true + - key: telemetry.export_metrics type: boolean default: "false" diff --git a/crates/sail-common/src/telemetry/c2.rs b/crates/sail-common/src/telemetry/c2.rs new file mode 100644 index 00000000..f6ffe12d --- /dev/null +++ b/crates/sail-common/src/telemetry/c2.rs @@ -0,0 +1,634 @@ +//! Optional C2 observations of existing calls, exported through fastrace. +//! Durations are local monotonic intervals and may overlap. They are not +//! an additive execution breakdown or proof that resources were reclaimed. +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Mutex, PoisonError}; +use std::time::Instant; + +use fastrace::Span; +use fastrace::collector::SpanContext; +use serde::Serialize; + +const ID_BYTES: usize = 256; +const PHASE_COUNT: usize = 43; +static STATE: State = State::new(); + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +#[repr(usize)] +pub enum Phase { + OperationBinding, + JobAccepted, + TaskAssigned, + Resolve, + LogicalExecute, + Optimize, + PhysicalPlan, + JobGraph, + JobTopology, + Schedule, + Definition, + PhysicalEncode, + PreparationQueue, + Preparation, + DefinitionCacheHit, + BatchDispatch, + TaskDispatched, + TaskDefinitionEncode, + TaskDefinitionDecode, + BatchPhysicalDecode, + PhysicalDecode, + BatchAdmission, + TaskAdmitted, + TaskTerminal, + ReaderDeclaration, + WriterDeclaration, + ReaderOpen, + WriterOpen, + StreamConsumption, + StreamOwner, + FirstPoll, + FirstBatch, + SinkCommit, + SinkAbort, + OutputForward, + CleanupDispatch, + CleanupQueuedAck, + CloseJobExecuted, + LocalStreamsRemoved, + ReleaseResponseBuffer, + ReleaseOperation, + ExecutorOwner, + ExecutorBufferOwner, +} + +impl Phase { + fn name(self) -> &'static str { + match self { + Self::OperationBinding => "c2.operation.binding", + Self::JobAccepted => "c2.job.accepted", + Self::TaskAssigned => "c2.task.assigned", + Self::Resolve => "c2.plan.resolve", + Self::LogicalExecute => "c2.plan.logical_execute", + Self::Optimize => "c2.plan.optimize", + Self::PhysicalPlan => "c2.plan.physical", + Self::JobGraph => "c2.job.graph", + Self::JobTopology => "c2.job.topology", + Self::Schedule => "c2.scheduler.select_assign", + Self::Definition => "c2.stage.definition", + Self::PhysicalEncode => "c2.stage.physical_encode", + Self::PreparationQueue => "c2.task.preparation_queue", + Self::Preparation => "c2.task.preparation", + Self::DefinitionCacheHit => "c2.stage.definition_cache_hit", + Self::BatchDispatch => "c2.batch.dispatch", + Self::TaskDispatched => "c2.task.dispatched", + Self::TaskDefinitionEncode => "c2.batch.definition_encode", + Self::TaskDefinitionDecode => "c2.batch.definition_decode", + Self::BatchPhysicalDecode => "c2.batch.physical_decode", + Self::PhysicalDecode => "c2.task.physical_decode", + Self::BatchAdmission => "c2.batch.admission", + Self::TaskAdmitted => "c2.task.admitted", + Self::TaskTerminal => "c2.task.terminal", + Self::ReaderDeclaration => "c2.reader.declaration", + Self::WriterDeclaration => "c2.writer.declaration", + Self::ReaderOpen => "c2.reader.open", + Self::WriterOpen => "c2.writer.open", + Self::StreamConsumption => "c2.stream.consumption", + Self::StreamOwner => "c2.stream.owner", + Self::FirstPoll => "c2.stream.first_poll", + Self::FirstBatch => "c2.stream.first_batch", + Self::SinkCommit => "c2.sink.commit", + Self::SinkAbort => "c2.sink.abort", + Self::OutputForward => "c2.output.forward", + Self::CleanupDispatch => "c2.cleanup.dispatch", + Self::CleanupQueuedAck => "c2.cleanup.queued_ack", + Self::CloseJobExecuted => "c2.cleanup.close_job_executed", + Self::LocalStreamsRemoved => "c2.cleanup.local_streams_removed", + Self::ReleaseResponseBuffer => "c2.release.response_buffer", + Self::ReleaseOperation => "c2.release.operation", + Self::ExecutorOwner => "c2.executor.owner", + Self::ExecutorBufferOwner => "c2.executor.buffer_owner", + } + } +} + +#[derive(Clone, Copy, Debug)] +pub enum Placement { + Driver, + Worker(u64), +} + +/// Borrowed identity is evaluated only in enabled mode and is not retained. +/// Native trace IDs bind pre-job planning to the actual operation root. +pub enum Identity<'a> { + CausalContext, + Operation { + session: &'a str, + operation: &'a str, + }, + Session(&'a str), + CausalJob { + job: u64, + }, + CausalStage { + job: u64, + stage: u64, + }, + PlacementJob { + session: &'a str, + job: u64, + placement: Placement, + }, + Job { + session: &'a str, + job: u64, + }, + Stage { + session: &'a str, + job: u64, + stage: usize, + }, + Task { + session: &'a str, + job: u64, + stage: usize, + partition: usize, + attempt: usize, + placement: Placement, + }, +} + +impl Identity<'_> { + fn properties(self) -> (Vec<(&'static str, String)>, bool) { + let mut values = Vec::with_capacity(10); + let mut complete = true; + let mut session = None; + let mut job = None; + match self { + Self::CausalContext => values.push(("c2.identity_kind", "causal_context".into())), + Self::Operation { + session: id, + operation, + } => { + values.push(("c2.identity_kind", "operation".into())); + session = Some(id); + let (value, intact) = bounded(operation); + values.push(("c2.operation_id", value)); + complete &= intact; + } + Self::Session(id) => { + values.push(("c2.identity_kind", "session".into())); + session = Some(id); + } + Self::CausalJob { job: number } => { + values.push(("c2.identity_kind", "causal_job".into())); + job = Some(number); + } + Self::CausalStage { job: number, stage } => { + values.push(("c2.identity_kind", "causal_stage".into())); + values.push(("execution.stage", stage.to_string())); + job = Some(number); + } + Self::PlacementJob { + session: id, + job: number, + placement, + } => { + values.push(("c2.identity_kind", "placement_job".into())); + match placement { + Placement::Driver => values.push(("c2.placement", "driver".into())), + Placement::Worker(worker) => { + values.push(("c2.placement", "worker".into())); + values.push(("cluster.worker.id", worker.to_string())); + } + } + session = Some(id); + job = Some(number); + } + Self::Job { + session: id, + job: number, + } => { + values.push(("c2.identity_kind", "job".into())); + session = Some(id); + job = Some(number); + } + Self::Stage { + session: id, + job: number, + stage, + } => { + values.push(("c2.identity_kind", "stage".into())); + values.push(("execution.stage", stage.to_string())); + session = Some(id); + job = Some(number); + } + Self::Task { + session: id, + job: number, + stage, + partition, + attempt, + placement, + } => { + values.push(("c2.identity_kind", "task".into())); + values.push(("execution.stage", stage.to_string())); + values.push(("execution.partition", partition.to_string())); + values.push(("execution.attempt", attempt.to_string())); + match placement { + Placement::Driver => values.push(("c2.placement", "driver".into())), + Placement::Worker(worker) => { + values.push(("c2.placement", "worker".into())); + values.push(("cluster.worker.id", worker.to_string())); + } + } + session = Some(id); + job = Some(number); + } + } + if let Some(id) = session { + let (value, intact) = bounded(id); + values.push(("session.id", value)); + complete &= intact; + } + if let Some(number) = job { + values.push(("execution.job.id", number.to_string())); + } + (values, complete) + } +} + +fn bounded(value: &str) -> (String, bool) { + let mut end = value.len().min(ID_BYTES); + while !value.is_char_boundary(end) { + end -= 1; + } + ( + value[..end].to_owned(), + end == value.len() && !value.is_empty(), + ) +} + +#[derive(Clone, Copy, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum Outcome { + Succeeded, + Failed, + Cancelled, + Abandoned, +} + +impl Outcome { + fn name(self) -> &'static str { + match self { + Self::Succeeded => "succeeded", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + Self::Abandoned => "abandoned", + } + } +} + +struct State { + enabled: AtomicBool, + sequence: AtomicU64, + ended: AtomicU64, + failed: AtomicU64, + cancelled: AtomicU64, + abandoned: AtomicU64, + incomplete_identity: AtomicU64, + saturated: AtomicBool, + phases: [AtomicU64; PHASE_COUNT], +} + +impl State { + const fn new() -> Self { + Self { + enabled: AtomicBool::new(false), + sequence: AtomicU64::new(0), + ended: AtomicU64::new(0), + failed: AtomicU64::new(0), + cancelled: AtomicU64::new(0), + abandoned: AtomicU64::new(0), + incomplete_identity: AtomicU64::new(0), + saturated: AtomicBool::new(false), + phases: [const { AtomicU64::new(0) }; PHASE_COUNT], + } + } + + fn increment(&self, counter: &AtomicU64) -> u64 { + match counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_add(1)) { + Ok(n) => n + 1, + Err(_) => { + self.saturated.store(true, Ordering::Relaxed); + u64::MAX + } + } + } + + fn summary(&self) -> Summary { + let created = self.sequence.load(Ordering::Relaxed); + let ended = self.ended.load(Ordering::Relaxed); + Summary { + schema_version: 1, + pid: std::process::id(), + created, + ended, + failed: self.failed.load(Ordering::Relaxed), + cancelled: self.cancelled.load(Ordering::Relaxed), + abandoned: self.abandoned.load(Ordering::Relaxed), + incomplete_identity: self.incomplete_identity.load(Ordering::Relaxed), + outstanding: created.saturating_sub(ended), + counters_complete: !self.saturated.load(Ordering::Relaxed) && ended <= created, + phase_created: self + .phases + .iter() + .map(|phase| phase.load(Ordering::Relaxed)) + .collect(), + } + } +} + +#[derive(Debug, Serialize)] +pub struct Summary { + pub schema_version: u8, + pub pid: u32, + pub created: u64, + pub ended: u64, + pub failed: u64, + pub cancelled: u64, + pub abandoned: u64, + pub incomplete_identity: u64, + pub outstanding: u64, + pub counters_complete: bool, + /// Array positions are the explicitly ordered Phase discriminants. + pub phase_created: Vec, +} + +pub fn enabled() -> bool { + STATE.enabled.load(Ordering::Relaxed) +} + +/// Unmetered RPC context for guards whose callers previously had no parent. +/// These names are outside the `c2.*` observation namespace and do not alter +/// its sequence or phase counters. +#[derive(Clone, Copy)] +pub enum Rpc { + ReleaseExecute, + RegisterWorker, +} + +impl Rpc { + fn name(self) -> &'static str { + match self { + Self::ReleaseExecute => "ReleaseExecute", + Self::RegisterWorker => "RegisterWorker", + } + } +} + +pub fn rpc_span(rpc: Rpc, properties: impl FnOnce() -> Vec<(&'static str, String)>) -> Span { + rpc_span_in(&STATE, rpc, properties) +} + +fn rpc_span_in( + state: &State, + rpc: Rpc, + properties: impl FnOnce() -> Vec<(&'static str, String)>, +) -> Span { + if !state.enabled.load(Ordering::Relaxed) { + return Span::noop(); + } + Span::root(rpc.name(), SpanContext::random()).with_properties(properties) +} + +/// Called once during telemetry initialization, after its exporter is ready. +pub fn configure(enabled: bool) { + STATE.enabled.store(enabled, Ordering::Relaxed); +} + +/// A shutdown snapshot, not a delivery acknowledgement from the collector. +pub fn summary() -> Option { + enabled().then(|| STATE.summary()) +} + +struct Active<'a> { + state: &'a State, + span: Span, + details: Mutex>, + clock: Option, + sequence: u64, +} + +/// Dropping a guard without a terminal result records Abandoned, never success. +pub struct Guard<'a> { + active: Option>, +} + +impl Guard<'static> { + pub fn point<'id>(phase: Phase, identity: impl FnOnce() -> Identity<'id>) -> Self { + Self::make(&STATE, phase, identity, false, || { + let parent = SpanContext::current_local_parent(); + let has_parent = parent.is_some(); + let span = match parent { + Some(parent) => Span::root(phase.name(), parent), + None => Span::root(phase.name(), SpanContext::random()), + }; + (span, has_parent) + }) + } + + pub fn start<'id>(phase: Phase, identity: impl FnOnce() -> Identity<'id>) -> Self { + Self::make(&STATE, phase, identity, true, || { + let parent = SpanContext::current_local_parent(); + let has_parent = parent.is_some(); + let span = match parent { + Some(parent) => Span::root(phase.name(), parent), + None => Span::root(phase.name(), SpanContext::random()), + }; + (span, has_parent) + }) + } +} + +impl<'a> Guard<'a> { + fn make<'id>( + state: &'a State, + phase: Phase, + identity: impl FnOnce() -> Identity<'id>, + timed: bool, + span: impl FnOnce() -> (Span, bool), + ) -> Self { + if !state.enabled.load(Ordering::Relaxed) { + return Self { active: None }; + } + let clock = timed.then(Instant::now); + let sequence = state.increment(&state.sequence); + state.increment(&state.phases[phase as usize]); + let (span, has_parent) = span(); + let (mut properties, complete) = identity().properties(); + let complete = complete && has_parent; + // Keep identity on the owned raw span. Separate add_property commands + // can arrive after this span is exported when it moves across threads. + let span = span.with_properties(|| { + properties.extend([ + ("c2.schema_version", "1".into()), + ("c2.observation_id", sequence.to_string()), + ("c2.pid", std::process::id().to_string()), + ("c2.causal_parent_present", has_parent.to_string()), + ("c2.identity_complete", complete.to_string()), + ("c2.kind", if timed { "interval" } else { "event" }.into()), + ]); + properties + }); + if !complete { + state.increment(&state.incomplete_identity); + } + Self { + active: Some(Active { + state, + span, + details: Mutex::new(Vec::new()), + clock, + sequence, + }), + } + } + + pub fn observation_id(&self) -> Option { + self.active.as_ref().map(|active| active.sequence) + } + + /// Enabled-only causal scope for an existing async call; no new data action. + pub fn child_span(&self, name: &'static str) -> Option { + self.active + .as_ref() + .map(|active| Span::enter_with_parent(name, &active.span)) + } + + pub fn detail(&self, key: &'static str, value: impl FnOnce() -> String) { + if let Some(active) = &self.active + && SpanContext::from_span(&active.span).is_some() + { + // Preserve noop callback laziness and run user code before locking. + let value = value(); + active + .details + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push((key, value)); + } + } + + pub fn finish_result(self, result: &Result) { + self.finish(if result.is_ok() { + Outcome::Succeeded + } else { + Outcome::Failed + }); + } + + pub fn finish(mut self, outcome: Outcome) { + self.end(outcome); + } + + fn end(&mut self, outcome: Outcome) { + if let Some(active) = self.active.take() { + let elapsed = active.clock.map(|clock| { + let elapsed = u64::try_from(clock.elapsed().as_nanos()); + if elapsed.is_err() { + active.state.saturated.store(true, Ordering::Relaxed); + } + elapsed.unwrap_or(u64::MAX) + }); + let mut properties = active + .details + .into_inner() + .unwrap_or_else(PoisonError::into_inner); + let span = active.span.with_properties(|| { + if let Some(elapsed) = elapsed { + properties.push(("c2.elapsed_ns", elapsed.to_string())); + } + properties.push(("c2.outcome", outcome.name().into())); + properties + }); + active.state.increment(&active.state.ended); + match outcome { + Outcome::Succeeded => {} + Outcome::Failed => { + active.state.increment(&active.state.failed); + } + Outcome::Cancelled => { + active.state.increment(&active.state.cancelled); + } + Outcome::Abandoned => { + active.state.increment(&active.state.abandoned); + } + } + drop(span); + } + } +} + +impl Drop for Guard<'_> { + fn drop(&mut self) { + self.end(Outcome::Abandoned); + } +} + +/// Lifetime of this named owner, not proof that its allocator returned memory. +/// The enclosing value's ordinary destruction remains unchanged. +pub struct NamedOwner<'a> { + guard: Option>, +} + +impl NamedOwner<'static> { + pub fn optional<'id>( + phase: Phase, + identity: impl FnOnce() -> Identity<'id>, + role: &'static str, + ) -> Option { + if !enabled() { + return None; + } + let guard = Guard::start(phase, identity); + guard.detail("c2.owner_role", || role.into()); + Some(Self { guard: Some(guard) }) + } +} + +impl Drop for NamedOwner<'_> { + fn drop(&mut self) { + if let Some(guard) = self.guard.take() { + guard.detail("c2.terminal_reason", || { + "named_owner_drop_not_allocator_reclaim".into() + }); + guard.finish(Outcome::Succeeded); + } + } +} + +pub fn event<'id>(phase: Phase, identity: impl FnOnce() -> Identity<'id>) { + Guard::point(phase, identity).finish(Outcome::Succeeded); +} + +/// One sequenced binding under the existing operation root, without altering it. +pub fn bind_operation(parent: &Span, session: &str, operation: &str) { + Guard::make( + &STATE, + Phase::OperationBinding, + || Identity::Operation { session, operation }, + false, + || { + ( + Span::enter_with_parent(Phase::OperationBinding.name(), parent), + SpanContext::from_span(parent).is_some(), + ) + }, + ) + .finish(Outcome::Succeeded); +} + +#[cfg(test)] +#[path = "c2_tests.rs"] +mod tests; diff --git a/crates/sail-common/src/telemetry/c2_tests.rs b/crates/sail-common/src/telemetry/c2_tests.rs new file mode 100644 index 00000000..4f888762 --- /dev/null +++ b/crates/sail-common/src/telemetry/c2_tests.rs @@ -0,0 +1,521 @@ +use std::cell::Cell; + +use fastrace::collector::{Config, SpanRecord, TestReporter}; +use fastrace::future::FutureExt; + +use super::*; + +fn state(enabled: bool) -> State { + let state = State::new(); + state.enabled.store(enabled, Ordering::Relaxed); + state +} + +fn guard(state: &State) -> Guard<'_> { + Guard::make( + state, + Phase::Preparation, + || Identity::CausalContext, + true, + || (Span::noop(), true), + ) +} + +#[test] +fn disabled_rpc_context_skips_properties_and_span_creation() { + let state = state(false); + let calls = Cell::new(0); + for rpc in [Rpc::ReleaseExecute, Rpc::RegisterWorker] { + let span = rpc_span_in(&state, rpc, || { + calls.set(calls.get() + 1); + vec![("request", "must not allocate".into())] + }); + assert!(SpanContext::from_span(&span).is_none()); + } + assert_eq!(calls.get(), 0); + assert_eq!(state.summary().created, 0); +} + +#[tokio::test] +async fn release_and_initial_worker_registration_have_real_rpc_parents() -> Result<(), &'static str> +{ + for (rpc, phases) in [ + ( + Rpc::ReleaseExecute, + vec![Phase::ReleaseResponseBuffer, Phase::ReleaseOperation], + ), + (Rpc::RegisterWorker, vec![Phase::Schedule]), + ] { + let state = state(true); + let span = rpc_span_in(&state, rpc, || vec![("request", "actual request".into())]); + let rpc_context = SpanContext::from_span(&span).ok_or("enabled RPC span")?; + let original: Result<(), &str> = async { + // Yield across a poll boundary: the RPC future must restore its + // parent on each poll, including an error return. + tokio::task::yield_now().await; + let parent = SpanContext::current_local_parent().ok_or("RPC local parent")?; + assert_eq!(parent.trace_id, rpc_context.trace_id); + assert_eq!(parent.span_id, rpc_context.span_id); + for phase in phases { + let parent = SpanContext::current_local_parent().ok_or("guard parent")?; + let observation = Guard::make( + &state, + phase, + || match rpc { + Rpc::ReleaseExecute => Identity::Operation { + session: "session", + operation: "operation", + }, + Rpc::RegisterWorker => Identity::Session("session"), + }, + true, + || (Span::root(phase.name(), parent), true), + ); + observation.finish(Outcome::Succeeded); + } + Err("original RPC error") + } + .in_span(span) + .await; + assert_eq!(original, Err("original RPC error")); + let summary = state.summary(); + assert_eq!(summary.created, summary.ended); + assert_eq!(summary.incomplete_identity, 0); + assert_eq!(summary.outstanding, 0); + assert_eq!( + summary.created, + match rpc { + Rpc::ReleaseExecute => 2, + Rpc::RegisterWorker => 1, + } + ); + assert!(SpanContext::current_local_parent().is_none()); + } + Ok(()) +} + +#[test] +fn disabled_guard_skips_identity_span_and_detail_callbacks() { + let state = state(false); + let calls = Cell::new(0); + let observation = Guard::make( + &state, + Phase::Resolve, + || { + calls.set(calls.get() + 1); + Identity::CausalContext + }, + true, + || { + calls.set(calls.get() + 1); + (Span::noop(), true) + }, + ); + assert!(observation.active.is_none()); + observation.detail("extra", || { + calls.set(calls.get() + 1); + String::new() + }); + observation.finish(Outcome::Succeeded); + assert_eq!(calls.get(), 0); + assert_eq!(state.summary().created, 0); +} + +#[test] +fn observed_call_delegates_once_and_preserves_original_error() { + for enabled in [false, true] { + let state = state(enabled); + let calls = Cell::new(0); + let observation = guard(&state); + let original: Result = { + calls.set(calls.get() + 1); + Err("original error") + }; + observation.finish_result(&original); + assert_eq!(calls.get(), 1); + assert_eq!(original, Err("original error")); + assert_eq!(state.summary().failed, u64::from(enabled)); + } +} + +#[test] +fn all_terminal_outcomes_close_exactly_one_guard() { + let state = state(true); + guard(&state).finish(Outcome::Succeeded); + let error: Result<(), &str> = Err("source error"); + guard(&state).finish_result(&error); + guard(&state).finish(Outcome::Cancelled); + drop(guard(&state)); + let summary = state.summary(); + assert_eq!(summary.created, 4); + assert_eq!(summary.ended, 4); + assert_eq!(summary.failed, 1); + assert_eq!(summary.cancelled, 1); + assert_eq!(summary.abandoned, 1); + assert_eq!(summary.outstanding, 0); + assert!(summary.counters_complete); +} + +#[test] +fn pending_guard_is_outstanding_until_its_owner_drops() { + let state = state(true); + let observation = guard(&state); + assert_eq!(state.summary().outstanding, 1); + drop(observation); + assert_eq!(state.summary().outstanding, 0); + assert_eq!(state.summary().abandoned, 1); +} + +#[test] +fn point_binding_has_no_duration_clock() { + let state = state(true); + let observation = Guard::make( + &state, + Phase::JobAccepted, + || Identity::Job { + session: "s", + job: 17, + }, + false, + || (Span::noop(), true), + ); + assert!( + observation + .active + .as_ref() + .is_some_and(|active| active.clock.is_none()) + ); + observation.finish(Outcome::Succeeded); + assert_eq!( + state.summary().phase_created[Phase::JobAccepted as usize], + 1 + ); +} + +#[test] +fn operation_and_job_identities_have_no_mutable_current_request_slot() { + let (first, first_complete) = Identity::Operation { + session: "s1", + operation: "op1", + } + .properties(); + let (second, second_complete) = Identity::Operation { + session: "s2", + operation: "op2", + } + .properties(); + let (job, job_complete) = Identity::Job { + session: "s1", + job: 29, + } + .properties(); + assert!(first_complete && second_complete && job_complete); + assert!( + first + .iter() + .any(|(key, value)| *key == "c2.operation_id" && value == "op1") + ); + assert!( + second + .iter() + .any(|(key, value)| *key == "c2.operation_id" && value == "op2") + ); + assert!( + job.iter() + .any(|(key, value)| *key == "session.id" && value == "s1") + ); + assert!( + job.iter() + .any(|(key, value)| *key == "execution.job.id" && value == "29") + ); + assert!(!job.iter().any(|(key, _)| *key == "c2.operation_id")); +} + +#[test] +fn task_identity_preserves_actual_attempt_and_worker_placement() { + let (values, complete) = Identity::Task { + session: "s", + job: u64::MAX, + stage: 4, + partition: 31, + attempt: 2, + placement: Placement::Worker(9), + } + .properties(); + assert!(complete); + for (key, expected) in [ + ("execution.job.id", u64::MAX.to_string()), + ("execution.partition", "31".into()), + ("execution.attempt", "2".into()), + ("cluster.worker.id", "9".into()), + ] { + assert!( + values + .iter() + .any(|(actual_key, value)| *actual_key == key && value == &expected) + ); + } +} + +#[test] +fn oversized_or_empty_identity_is_bounded_and_explicitly_incomplete() { + let large = "λ".repeat(ID_BYTES); + let (properties, complete) = Identity::Operation { + session: "", + operation: &large, + } + .properties(); + assert!(!complete); + assert!(properties.iter().all(|(_, value)| value.len() <= ID_BYTES)); + assert!( + properties + .iter() + .all(|(_, value)| value.is_char_boundary(value.len())) + ); +} + +#[test] +fn missing_causal_parent_cannot_qualify_as_complete_identity() { + let state = state(true); + Guard::make( + &state, + Phase::JobAccepted, + || Identity::Job { + session: "s", + job: 1, + }, + false, + || (Span::noop(), false), + ) + .finish(Outcome::Succeeded); + assert_eq!(state.summary().incomplete_identity, 1); +} + +#[test] +fn saturated_sequence_is_explicit_and_never_wraps_to_a_fresh_id() { + let state = state(true); + state.sequence.store(u64::MAX, Ordering::Relaxed); + guard(&state).finish(Outcome::Succeeded); + assert_eq!(state.summary().created, u64::MAX); + assert!(!state.summary().counters_complete); +} + +#[test] +fn named_owner_drop_closes_its_lifetime_without_claiming_allocator_release() { + let state = state(true); + let owner = NamedOwner { + guard: Some(Guard::make( + &state, + Phase::ExecutorOwner, + || Identity::Operation { + session: "s", + operation: "op", + }, + true, + || (Span::noop(), true), + )), + }; + assert_eq!(state.summary().outstanding, 1); + drop(owner); + let summary = state.summary(); + assert_eq!(summary.created, 1); + assert_eq!(summary.ended, 1); + assert_eq!(summary.abandoned, 0); + assert_eq!(summary.outstanding, 0); +} + +#[test] +fn job_scoped_cleanup_does_not_invent_a_stage() { + let (values, complete) = Identity::CausalJob { job: 17 }.properties(); + assert!(complete); + assert!( + values + .iter() + .any(|(key, value)| *key == "execution.job.id" && value == "17") + ); + assert!(!values.iter().any(|(key, _)| *key == "execution.stage")); +} + +#[test] +fn every_phase_has_a_counter_slot() { + assert_eq!(Phase::ExecutorBufferOwner as usize + 1, PHASE_COUNT); +} + +fn require_send_sync() {} + +#[test] +fn observation_and_named_owner_preserve_send_and_sync() { + require_send_sync::>(); + require_send_sync::>(); +} + +#[test] +fn enabled_noop_guard_skips_detail_callback() -> Result<(), &'static str> { + let state = state(true); + let calls = Cell::new(0); + let observation = guard(&state); + observation.detail("noop", || { + calls.set(calls.get() + 1); + "must not evaluate".into() + }); + let child = observation + .child_span("noop child") + .ok_or("enabled guard retains its child-span API")?; + assert!(SpanContext::from_span(&child).is_none()); + drop(observation); + assert_eq!(calls.get(), 0); + assert_eq!(state.summary().abandoned, 1); + assert_eq!(state.summary().outstanding, 0); + Ok(()) +} + +fn property<'a>(record: &'a SpanRecord, key: &str) -> Result<&'a str, &'static str> { + let mut values = record.properties.iter().filter(|(name, _)| name == key); + let (_, value) = values.next().ok_or("exported property is missing")?; + assert!( + values.next().is_none(), + "duplicate exported property: {key}" + ); + Ok(value) +} + +#[test] +fn cross_thread_guard_export_preserves_initial_details_and_terminal_fields() +-> Result<(), &'static str> { + let (reporter, records) = TestReporter::new(); + fastrace::set_reporter(reporter, Config::default()); + let state = state(true); + let parent = Span::root("guard_export_parent", SpanContext::random()); + let context = SpanContext::from_span(&parent).ok_or("recording parent")?; + let make = |phase: Phase, timed: bool| { + Guard::make( + &state, + phase, + || Identity::Task { + session: "export-session", + job: 17, + stage: 4, + partition: 31, + attempt: 2, + placement: Placement::Worker(9), + }, + timed, + || (Span::root(phase.name(), context), true), + ) + }; + let owner = make(Phase::StreamOwner, true); + let consumption = make(Phase::StreamConsumption, true); + let abandoned = make(Phase::ReaderOpen, true); + let cancelled = make(Phase::TaskTerminal, false); + owner.detail("c2.owner_role", || { + // A detail callback may add another detail without holding our mutex. + owner.detail("c2.open_partition", || "31".into()); + "opened_shuffle_source".into() + }); + consumption.detail("c2.owner_role", || "opened_shuffle_source".into()); + let owner_context = owner + .active + .as_ref() + .and_then(|active| SpanContext::from_span(&active.span)) + .ok_or("recording owner")?; + let child = owner.child_span("guard_export_child").ok_or("child span")?; + std::thread::scope(|scope| -> Result<(), &'static str> { + scope + .spawn(move || { + drop(child); + owner.detail("c2.terminal_reason", || "named_wrapper_owner_drop".into()); + owner.finish(Outcome::Succeeded); + consumption.detail("c2.rows", || "0".into()); + consumption.detail("c2.terminal_reason", || "inner_stream_error".into()); + let original: Result<(), &str> = Err("original stream error"); + consumption.finish_result(&original); + assert_eq!(original, Err("original stream error")); + drop(abandoned); + cancelled.finish(Outcome::Cancelled); + }) + .join() + .map_err(|_| "guard export thread panicked")?; + Ok(()) + })?; + drop(parent); + fastrace::flush(); + let records = records.lock(); + let observed: Vec<_> = records + .iter() + .filter(|record| record.trace_id == context.trace_id && record.name.starts_with("c2.")) + .collect(); + assert_eq!(observed.len(), 4); + let pid = std::process::id().to_string(); + for (phase, sequence, outcome, timed) in [ + (Phase::StreamOwner, "1", "succeeded", true), + (Phase::StreamConsumption, "2", "failed", true), + (Phase::ReaderOpen, "3", "abandoned", true), + (Phase::TaskTerminal, "4", "cancelled", false), + ] { + let record = observed + .iter() + .copied() + .find(|record| record.name == phase.name()) + .ok_or("exported observation is missing")?; + assert_eq!(record.parent_id, context.span_id); + for (key, expected) in [ + ("c2.identity_kind", "task"), + ("session.id", "export-session"), + ("execution.job.id", "17"), + ("execution.stage", "4"), + ("execution.partition", "31"), + ("execution.attempt", "2"), + ("c2.placement", "worker"), + ("cluster.worker.id", "9"), + ("c2.schema_version", "1"), + ("c2.observation_id", sequence), + ("c2.pid", pid.as_str()), + ("c2.causal_parent_present", "true"), + ("c2.identity_complete", "true"), + ("c2.kind", if timed { "interval" } else { "event" }), + ("c2.outcome", outcome), + ] { + assert_eq!(property(record, key)?, expected, "{key}"); + } + if timed { + property(record, "c2.elapsed_ns")? + .parse::() + .map_err(|_| "exported duration is invalid")?; + } else { + assert!( + !record + .properties + .iter() + .any(|(key, _)| key == "c2.elapsed_ns") + ); + } + if record.name == Phase::StreamOwner.name() { + assert_eq!(property(record, "c2.owner_role")?, "opened_shuffle_source"); + assert_eq!(property(record, "c2.open_partition")?, "31"); + assert_eq!( + property(record, "c2.terminal_reason")?, + "named_wrapper_owner_drop" + ); + } else if record.name == Phase::StreamConsumption.name() { + assert_eq!(property(record, "c2.owner_role")?, "opened_shuffle_source"); + assert_eq!(property(record, "c2.rows")?, "0"); + assert_eq!( + property(record, "c2.terminal_reason")?, + "inner_stream_error" + ); + } + } + let child = records + .iter() + .find(|record| record.trace_id == context.trace_id && record.name == "guard_export_child") + .ok_or("exported child is missing")?; + assert_eq!(child.parent_id, owner_context.span_id); + let summary = state.summary(); + assert_eq!(summary.created, 4); + assert_eq!(summary.ended, 4); + assert_eq!(summary.failed, 1); + assert_eq!(summary.cancelled, 1); + assert_eq!(summary.abandoned, 1); + assert_eq!(summary.outstanding, 0); + Ok(()) +} diff --git a/crates/sail-common/src/telemetry/mod.rs b/crates/sail-common/src/telemetry/mod.rs index 3fd26c69..35d633c0 100644 --- a/crates/sail-common/src/telemetry/mod.rs +++ b/crates/sail-common/src/telemetry/mod.rs @@ -1,3 +1,4 @@ +pub mod c2; mod common; mod futures; mod layers; diff --git a/crates/sail-execution/src/driver/actor/handler.rs b/crates/sail-execution/src/driver/actor/handler.rs index eea02a4e..7e130667 100644 --- a/crates/sail-execution/src/driver/actor/handler.rs +++ b/crates/sail-execution/src/driver/actor/handler.rs @@ -8,6 +8,7 @@ use futures::TryStreamExt; use log::{debug, error, info, warn}; use sail_celeborn::lifecycle::{LifecycleManagerActor, LocalLifecycleManager}; use sail_common::actor::{ActorAction, ActorContext, ActorHandle}; +use sail_common::telemetry::c2::{self, Guard, Identity, Outcome, Phase, Placement}; use sail_common_datafusion::error::CommonErrorCause; use sail_python_udf::error::PyErrExtractor; use tokio::sync::oneshot; @@ -215,6 +216,10 @@ impl DriverActor { ) -> ActorAction { let out = self.job_scheduler.accept_job(ctx, plan, context); if let Ok((job_id, _)) = &out { + c2::event(Phase::JobAccepted, || Identity::Job { + session: &self.options.session_id, + job: (*job_id).into(), + }); self.refresh_job(ctx, *job_id); self.run_tasks(ctx); self.reconcile_worker_demands(ctx); @@ -605,7 +610,12 @@ impl DriverActor { /// Reserve complete regions before resolving any routing. Definitions are scoped to this /// scheduling snapshot; batches additionally preserve region and worker boundaries. fn run_tasks(&mut self, ctx: &mut ActorContext) { + let observation = Guard::start(Phase::Schedule, || { + Identity::Session(&self.options.session_id) + }); let assignments = self.task_assigner.assign_tasks(); + observation.detail("c2.assignment_sets", || assignments.len().to_string()); + observation.finish(Outcome::Succeeded); self.task_assigner.track_streams(&assignments); let mut batches = indexmap::IndexMap::<_, Vec>::new(); for assignment in assignments { @@ -633,7 +643,20 @@ impl DriverActor { let mut definitions = std::collections::HashMap::new(); for ((job_id, _, stage, worker_id), keys) in batches { let Some(first) = keys.first() else { continue }; + if c2::enabled() && definitions.contains_key(&(job_id, stage)) { + c2::event(Phase::DefinitionCacheHit, || Identity::Stage { + session: &self.options.session_id, + job: job_id.into(), + stage, + }); + } let definition = definitions.entry((job_id, stage)).or_insert_with(|| { + let observation = Guard::start(Phase::Definition, || Identity::Stage { + session: &self.options.session_id, + job: job_id.into(), + stage, + }); + observation.detail("c2.cache", || "miss_in_scheduling_invocation".into()); let started = Instant::now(); let result = self .job_scheduler @@ -645,6 +668,7 @@ impl DriverActor { CommonErrorCause::new::(&error), ) }); + observation.finish_result(&result); debug!( "job {job_id} stage {stage} definition construction {:?}", started.elapsed() @@ -667,6 +691,17 @@ impl DriverActor { } }; for key in &keys { + c2::event(Phase::TaskAssigned, || Identity::Task { + session: &self.options.session_id, + job: key.job_id.into(), + stage: key.stage, + partition: key.partition, + attempt: key.attempt, + placement: match worker_id { + Some(worker_id) => Placement::Worker(worker_id.into()), + None => Placement::Driver, + }, + }); self.job_scheduler .update_task(key, TaskState::Scheduled, None, None); } diff --git a/crates/sail-execution/src/driver/job_scheduler/core.rs b/crates/sail-execution/src/driver/job_scheduler/core.rs index f73a01dc..52faa8f9 100644 --- a/crates/sail-execution/src/driver/job_scheduler/core.rs +++ b/crates/sail-execution/src/driver/job_scheduler/core.rs @@ -10,6 +10,7 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec; use indexmap::{IndexMap, IndexSet}; use log::{debug, warn}; use sail_common::actor::ActorContext; +use sail_common::telemetry::c2::{Guard, Identity, Phase}; use sail_common_datafusion::driver_extension::{ contains_driver_extension, release_driver_extensions, }; @@ -67,17 +68,34 @@ impl JobScheduler { "job {job_id} execution plan\n{}", DisplayableExecutionPlan::new(plan.as_ref()).indent(true) ); + let observation = Guard::start(Phase::JobGraph, || Identity::Job { + session: &self.options.session_id, + job: job_id.into(), + }); let graph = JobGraph::try_new( plan, crate::job_graph::JobGraphOptions { shuffle_backend: self.options.shuffle_backend.clone(), }, - )?; + ); + observation.finish_result(&graph); + let graph = graph?; debug!("job {job_id} job graph \n{graph}"); - let topology = JobTopology::try_new(&graph)?; + let observation = Guard::start(Phase::JobTopology, || Identity::Job { + session: &self.options.session_id, + job: job_id.into(), + }); + let topology = JobTopology::try_new(&graph); + observation.finish_result(&topology); + let topology = topology?; worker_topology::validate(&graph, &topology)?; - let (output, stream) = build_job_output(ctx, job_id, graph.schema().clone()); + let (output, stream) = build_job_output( + ctx, + job_id, + graph.schema().clone(), + &self.options.session_id, + ); let descriptor = JobDescriptor::new(graph, topology, JobState::Running { output }, context); self.jobs.insert(job_id, descriptor); @@ -721,7 +739,17 @@ impl JobScheduler { ))); }; - let plan = encode_remote_physical_plan(self.codec.as_ref(), stage.plan.clone())?; + let observation = Guard::start(Phase::PhysicalEncode, || Identity::Stage { + session: &self.options.session_id, + job: key.job_id.into(), + stage: key.stage, + }); + let plan = encode_remote_physical_plan(self.codec.as_ref(), stage.plan.clone()); + if let Ok(plan) = &plan { + observation.detail("c2.encoded_plan_bytes", || plan.len().to_string()); + } + observation.finish_result(&plan); + let plan = plan?; let inputs = stage .inputs .iter() diff --git a/crates/sail-execution/src/driver/output.rs b/crates/sail-execution/src/driver/output.rs index fd0a57a4..0e0eb778 100644 --- a/crates/sail-execution/src/driver/output.rs +++ b/crates/sail-execution/src/driver/output.rs @@ -11,6 +11,7 @@ use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use futures::stream::SelectAll; use futures::{Stream, StreamExt}; use sail_common::actor::ActorContext; +use sail_common::telemetry::c2::{Guard, Identity, Outcome, Phase}; use sail_common_datafusion::error::CommonErrorCause; use tokio::sync::mpsc; use tonic::codegen::tokio_stream::wrappers::ReceiverStream; @@ -94,12 +95,22 @@ pub fn build_job_output( ctx: &mut ActorContext, job_id: JobId, schema: SchemaRef, + session_id: &str, ) -> (JobOutputManager, SendableRecordBatchStream) { let (sender, receiver) = mpsc::channel(JobOutputItem::CHANNEL_SIZE); let (tx, rx) = mpsc::channel(1); let handle = ctx.handle().clone(); + let forward = Guard::start(Phase::OutputForward, || Identity::Job { + session: session_id, + job: job_id.into(), + }); ctx.spawn(async move { let outcome = forward_job_output(JobOutputStream::new(receiver), &tx).await; + forward.finish(match outcome { + JobOutputOutcome::Completed => Outcome::Succeeded, + JobOutputOutcome::Failed => Outcome::Failed, + JobOutputOutcome::Canceled => Outcome::Cancelled, + }); // Output errors and consumer cancellation can precede terminal task updates. // Tell the scheduler why output ended so cleanup records the correct job status. // Keep `tx` alive until cleanup has been sent, before signaling EOF to the client. @@ -109,10 +120,14 @@ pub fn build_job_output( }); ( JobOutputManager { sender }, - Box::pin(RecordBatchStreamAdapter::new( - schema, - ReceiverStream::new(rx), - )), + crate::stream::observation::job_stream( + Box::pin(RecordBatchStreamAdapter::new( + schema, + ReceiverStream::new(rx), + )), + session_id, + job_id.into(), + ), ) } diff --git a/crates/sail-execution/src/driver/server.rs b/crates/sail-execution/src/driver/server.rs index dcbcda2d..f78863d8 100644 --- a/crates/sail-execution/src/driver/server.rs +++ b/crates/sail-execution/src/driver/server.rs @@ -1,8 +1,10 @@ use std::sync::Arc; +use fastrace::future::FutureExt; use log::debug; use opentelemetry_proto::tonic::metrics::v1::ResourceMetrics; use prost::Message; +use sail_common::telemetry::c2::{self, Rpc}; use tokio::sync::oneshot; use tonic::{Request, Response, Status}; @@ -33,33 +35,43 @@ impl DriverService for DriverServer { request: Request, ) -> Result, Status> { let request = request.into_inner(); - debug!("{request:?}"); - let RegisterWorkerRequest { - driver_id, - worker_id, - host, - port, - } = request; - let port = u16::try_from(port).map_err(|_| { - Status::invalid_argument("port must be a valid 16-bit unsigned integer") - })?; - let (tx, rx) = oneshot::channel(); - let message = DriverMessage::RegisterWorker { - worker_id: WorkerId::from(worker_id), - host, - port, - result: tx, - }; - self.registry - .get(DriverId::from(driver_id)) - .await? - .send(message) - .await - .map_err(ExecutionError::from)?; - rx.await.map_err(ExecutionError::from)??; - let response = RegisterWorkerResponse {}; - debug!("{response:?}"); - Ok(Response::new(response)) + let span = c2::rpc_span(Rpc::RegisterWorker, || { + vec![ + ("cluster.driver.id", request.driver_id.to_string()), + ("cluster.worker.id", request.worker_id.to_string()), + ] + }); + async { + debug!("{request:?}"); + let RegisterWorkerRequest { + driver_id, + worker_id, + host, + port, + } = request; + let port = u16::try_from(port).map_err(|_| { + Status::invalid_argument("port must be a valid 16-bit unsigned integer") + })?; + let (tx, rx) = oneshot::channel(); + let message = DriverMessage::RegisterWorker { + worker_id: WorkerId::from(worker_id), + host, + port, + result: tx, + }; + self.registry + .get(DriverId::from(driver_id)) + .await? + .send(message) + .await + .map_err(ExecutionError::from)?; + rx.await.map_err(ExecutionError::from)??; + let response = RegisterWorkerResponse {}; + debug!("{response:?}"); + Ok(Response::new(response)) + } + .in_span(span) + .await } async fn report_worker_heartbeat( diff --git a/crates/sail-execution/src/driver/worker_pool/core.rs b/crates/sail-execution/src/driver/worker_pool/core.rs index d988ae5b..d41ce934 100644 --- a/crates/sail-execution/src/driver/worker_pool/core.rs +++ b/crates/sail-execution/src/driver/worker_pool/core.rs @@ -9,6 +9,7 @@ use futures::TryStreamExt; use log::{debug, error, info, warn}; use sail_common::actor::ActorContext; use sail_common::telemetry::SpanAttribute; +use sail_common::telemetry::c2::{self, Guard, Identity, Phase, Placement}; use sail_common_datafusion::error::CommonErrorCause; use sail_python_udf::error::PyErrExtractor; use sail_system_store::SystemEvent; @@ -359,11 +360,28 @@ impl WorkerPool { } }; let handle = ctx.handle().clone(); + let observation_session = c2::enabled().then(|| self.options.session_id.clone()); ctx.spawn(async move { - if let Err(error) = client + let dispatched = client .run_task_batch(job_id, stage, tasks.clone(), definition, peers) - .await - { + .await; + if let Some(session) = observation_session { + for task in &tasks { + let guard = Guard::point(Phase::TaskDispatched, || Identity::Task { + session: &session, + job: job_id.into(), + stage, + partition: task.partition, + attempt: task.attempt, + placement: Placement::Worker(worker_id.into()), + }); + guard.detail("c2.scope", || { + "remote_admission_response_not_execution_completion".into() + }); + guard.finish_result(&dispatched); + } + } + if let Err(error) = dispatched { for task in tasks { let _ = handle .send(DriverMessage::UpdateTask { @@ -503,8 +521,18 @@ impl WorkerPool { return; } }; + let cleanup = Guard::start(Phase::CleanupDispatch, || Identity::PlacementJob { + session: &options.session_id, + job: job_id.into(), + placement: Placement::Worker(worker_id.into()), + }); + cleanup.detail("c2.scope", || { + "rpc_dispatch_to_mailbox_ack_not_executed_cleanup".into() + }); ctx.spawn(async move { - if let Err(e) = client.clean_up_job(job_id, stage).await { + let output = client.clean_up_job(job_id, stage).await; + cleanup.finish_result(&output); + if let Err(e) = output { Self::log_worker_control_error("clean up job", worker_id, &e); } }); diff --git a/crates/sail-execution/src/lib.rs b/crates/sail-execution/src/lib.rs index f01782d6..f64ee11f 100644 --- a/crates/sail-execution/src/lib.rs +++ b/crates/sail-execution/src/lib.rs @@ -15,4 +15,5 @@ mod worker; pub mod worker_manager; pub use id::{DriverId, IdGenerator}; +pub use stream::observation::operation_stream; pub use worker::entrypoint::run_worker; diff --git a/crates/sail-execution/src/stream/accessor.rs b/crates/sail-execution/src/stream/accessor.rs index 1a749a96..09210180 100644 --- a/crates/sail-execution/src/stream/accessor.rs +++ b/crates/sail-execution/src/stream/accessor.rs @@ -11,6 +11,7 @@ use tokio::sync::oneshot; use crate::error::ExecutionResult; use crate::id::{JobId, TaskKey, TaskStreamKey, WorkerId}; use crate::stream::merge::merged_stream; +use crate::stream::observation::{self, ObserverContext}; use crate::stream::reader::{TaskStreamReader, TaskStreamSource}; use crate::stream::writer::{ MultiChannelTaskStreamSink, TaskStreamChannelSink, TaskStreamSink, TaskStreamWriter, @@ -23,6 +24,7 @@ pub struct TaskStreamFactory { context: Arc, celeborn: bool, mappers: usize, + observation: Option, } impl Clone for TaskStreamFactory { @@ -32,22 +34,25 @@ impl Clone for TaskStreamFactory { context: self.context.clone(), celeborn: self.celeborn, mappers: self.mappers, + observation: self.observation.clone(), } } } impl TaskStreamFactory { - pub fn new( + pub(crate) fn new( handle: ActorHandle, context: Arc, celeborn: bool, mappers: usize, + observation: Option, ) -> Self { Self { handle, context, celeborn, mappers, + observation, } } @@ -57,13 +62,15 @@ impl TaskStreamFactory { input: TaskInput, schema: SchemaRef, ) -> Arc { - Arc::new(MultiChannelTaskStreamReader::new( + let scope = self.observation.as_ref().map(|context| context.task(&key)); + let inner = Arc::new(MultiChannelTaskStreamReader::new( self.handle.clone(), self.context.clone(), key, input, schema, - )) + )); + observation::reader(inner, scope) } pub fn writer( @@ -72,24 +79,27 @@ impl TaskStreamFactory { output: TaskOutput, schema: SchemaRef, ) -> Arc { - if self.celeborn && matches!(output.locator, TaskOutputLocator::Blocking) { - Arc::new(CelebornTaskStreamWriter::new( - self.handle.clone(), - self.context.clone(), - key, - output.channels(), - schema, - self.mappers, - )) - } else { - Arc::new(MultiChannelTaskStreamWriter::new( - self.handle.clone(), - self.context.clone(), - key, - output, - schema, - )) - } + let scope = self.observation.as_ref().map(|context| context.task(&key)); + let inner: Arc = + if self.celeborn && matches!(output.locator, TaskOutputLocator::Blocking) { + Arc::new(CelebornTaskStreamWriter::new( + self.handle.clone(), + self.context.clone(), + key, + output.channels(), + schema, + self.mappers, + )) + } else { + Arc::new(MultiChannelTaskStreamWriter::new( + self.handle.clone(), + self.context.clone(), + key, + output, + schema, + )) + }; + observation::writer(inner, scope) } } diff --git a/crates/sail-execution/src/stream/mod.rs b/crates/sail-execution/src/stream/mod.rs index 783c874f..013c04ff 100644 --- a/crates/sail-execution/src/stream/mod.rs +++ b/crates/sail-execution/src/stream/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod celeborn; pub mod error; pub(crate) mod local; pub mod merge; +pub mod observation; pub mod reader; pub(crate) mod service; pub(crate) mod storage; diff --git a/crates/sail-execution/src/stream/observation.rs b/crates/sail-execution/src/stream/observation.rs new file mode 100644 index 00000000..220d2d41 --- /dev/null +++ b/crates/sail-execution/src/stream/observation.rs @@ -0,0 +1,424 @@ +//! Default-off observations of named stream owners, never allocator reclamation. +use std::fmt; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::Result; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::physical_plan::RecordBatchStream; +use futures::Stream; +use sail_common::telemetry::c2::{self, Guard, Identity, Outcome, Phase, Placement}; + +use crate::id::TaskKey; +use crate::stream::reader::{TaskStreamReader, TaskStreamSource}; +use crate::stream::writer::{TaskStreamSink, TaskStreamWriteState, TaskStreamWriter}; + +#[derive(Clone)] +pub(crate) struct ObserverContext { + session: Arc, + placement: Placement, +} + +impl ObserverContext { + pub(crate) fn optional(session: &str, placement: Placement) -> Option { + c2::enabled().then(|| Self { + session: Arc::from(session), + placement, + }) + } + + pub(crate) fn task(&self, key: &TaskKey) -> Scope { + Scope::Task { + context: self.clone(), + key: key.clone(), + } + } +} + +#[derive(Clone)] +pub(crate) enum Scope { + Task { + context: ObserverContext, + key: TaskKey, + }, + Job { + session: Arc, + job: u64, + }, + Operation { + session: Arc, + operation: Arc, + }, +} + +impl Scope { + fn identity(&self) -> Identity<'_> { + match self { + Self::Task { context, key } => Identity::Task { + session: &context.session, + job: key.job_id.into(), + stage: key.stage, + partition: key.partition, + attempt: key.attempt, + placement: context.placement, + }, + Self::Job { session, job } => Identity::Job { session, job: *job }, + Self::Operation { session, operation } => Identity::Operation { session, operation }, + } + } + + fn start(&self, phase: Phase, role: &'static str) -> Guard<'static> { + let guard = Guard::start(phase, || self.identity()); + guard.detail("c2.owner_role", || role.into()); + guard + } + + fn point(&self, phase: Phase, role: &'static str) -> Guard<'static> { + let guard = Guard::point(phase, || self.identity()); + guard.detail("c2.owner_role", || role.into()); + guard + } +} + +/// Bounded records: one first poll/batch, one consumption terminal and one owner drop. +/// Errors do not replace inner errors; dropped open/commit/abort futures are abandoned. +/// Writes count offered inputs; they make no per-write cancellation assertion. +struct Flow { + scope: Scope, + role: &'static str, + owner: Option>, + consumption: Option>, + polled: bool, + batches: u64, + rows: u64, + counts_complete: bool, +} + +impl Flow { + fn new(scope: Scope, role: &'static str) -> Self { + Self { + owner: Some(scope.start(Phase::StreamOwner, role)), + consumption: Some(scope.start(Phase::StreamConsumption, role)), + scope, + role, + polled: false, + batches: 0, + rows: 0, + counts_complete: true, + } + } + + fn declaration(scope: Scope, role: &'static str) -> Self { + // A declaration is an owner, not a stream that failed to reach EOF. + Self { + owner: Some(scope.start(Phase::StreamOwner, role)), + consumption: None, + scope, + role, + polled: false, + batches: 0, + rows: 0, + counts_complete: true, + } + } + + fn first_poll(&mut self) { + if !self.polled { + self.polled = true; + self.scope + .point(Phase::FirstPoll, self.role) + .finish(Outcome::Succeeded); + } + } + + fn batch(&mut self, rows: usize) { + if self.batches == 0 { + self.scope + .point(Phase::FirstBatch, self.role) + .finish(Outcome::Succeeded); + } + let rows = u64::try_from(rows).unwrap_or(u64::MAX); + match (self.batches.checked_add(1), self.rows.checked_add(rows)) { + (Some(batches), Some(rows)) => { + self.batches = batches; + self.rows = rows; + } + _ => self.counts_complete = false, + } + } + + fn counts(&self, guard: &Guard<'_>) { + guard.detail("c2.count_semantics", || { + if self.role == "opened_shuffle_sink" { + "offered_input_before_inner_write_returns".into() + } else { + "observed_successful_output_batches".into() + } + }); + guard.detail("c2.batches", || self.batches.to_string()); + guard.detail("c2.rows", || self.rows.to_string()); + guard.detail("c2.counts_complete", || self.counts_complete.to_string()); + guard.detail("c2.polled", || self.polled.to_string()); + } + + fn terminal(&mut self, reason: &'static str, outcome: Outcome) { + if let Some(guard) = self.consumption.take() { + self.counts(&guard); + guard.detail("c2.terminal_reason", || reason.into()); + guard.finish(outcome); + } + } +} + +impl Drop for Flow { + fn drop(&mut self) { + self.terminal("owner_drop_before_terminal", Outcome::Abandoned); + if let Some(guard) = self.owner.take() { + self.counts(&guard); + guard.detail("c2.terminal_reason", || "named_wrapper_owner_drop".into()); + guard.finish(Outcome::Succeeded); + } + } +} + +struct ObservedStream { + inner: S, + flow: Flow, +} + +impl Stream for ObservedStream +where + S: Stream> + Unpin, +{ + type Item = std::result::Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + this.flow.first_poll(); + // Exactly one original poll per outer poll, including after a terminal result. + let result = Pin::new(&mut this.inner).poll_next(cx); + match &result { + Poll::Ready(Some(Ok(batch))) => this.flow.batch(batch.num_rows()), + Poll::Ready(Some(Err(_))) => this.flow.terminal("observed_error", Outcome::Failed), + Poll::Ready(None) => this.flow.terminal("eof", Outcome::Succeeded), + Poll::Pending => {} + } + result + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} + +// The public functions receive a pinned boxed trait object. It implements Stream, +// but RecordBatchStream's schema method is available through deref, not a blanket +// RecordBatchStream implementation on Pin>. +impl RecordBatchStream for ObservedStream { + fn schema(&self) -> SchemaRef { + self.inner.schema() + } +} + +pub fn operation_stream( + inner: SendableRecordBatchStream, + session: &str, + operation: &str, +) -> SendableRecordBatchStream { + if !c2::enabled() { + return inner; + } + Box::pin(ObservedStream { + inner, + flow: Flow::new( + Scope::Operation { + session: Arc::from(session), + operation: Arc::from(operation), + }, + "executor_input_stream", + ), + }) +} + +pub(crate) fn job_stream( + inner: SendableRecordBatchStream, + session: &str, + job: u64, +) -> SendableRecordBatchStream { + if !c2::enabled() { + return inner; + } + Box::pin(ObservedStream { + inner, + flow: Flow::new( + Scope::Job { + session: Arc::from(session), + job, + }, + "driver_output_receiver", + ), + }) +} + +pub(crate) fn task_stream( + inner: SendableRecordBatchStream, + context: Option, + key: &TaskKey, +) -> SendableRecordBatchStream { + let Some(context) = context else { + return inner; + }; + Box::pin(ObservedStream { + inner, + flow: Flow::new(context.task(key), "task_completion_stream"), + }) +} + +pub(crate) fn reader( + inner: Arc, + scope: Option, +) -> Arc { + let Some(scope) = scope else { + return inner; + }; + let declaration = scope.point(Phase::ReaderDeclaration, "shuffle_reader"); + declaration.finish(Outcome::Succeeded); + Arc::new(ObservedReader { + inner, + owner: Flow::declaration(scope, "reader_declaration_owner"), + }) +} + +struct ObservedReader { + inner: Arc, + owner: Flow, +} + +impl fmt::Debug for ObservedReader { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ObservedReader") + .field("inner", &self.inner) + .finish() + } +} + +#[tonic::async_trait] +impl TaskStreamReader for ObservedReader { + async fn open(&self, partition: usize) -> Result { + let guard = self.owner.scope.start(Phase::ReaderOpen, "reader_open"); + guard.detail("c2.open_partition", || partition.to_string()); + let result = self.inner.open(partition).await; + guard.finish_result(&result); + let inner = result?; + Ok(Box::pin(ObservedStream { + inner, + flow: Flow::new(self.owner.scope.clone(), "opened_shuffle_reader"), + })) + } +} + +pub(crate) fn writer( + inner: Arc, + scope: Option, +) -> Arc { + let Some(scope) = scope else { + return inner; + }; + scope + .point(Phase::WriterDeclaration, "shuffle_writer") + .finish(Outcome::Succeeded); + Arc::new(ObservedWriter { + inner, + owner: Flow::declaration(scope, "writer_declaration_owner"), + }) +} + +struct ObservedWriter { + inner: Arc, + owner: Flow, +} + +impl fmt::Debug for ObservedWriter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ObservedWriter") + .field("inner", &self.inner) + .finish() + } +} + +#[tonic::async_trait] +impl TaskStreamWriter for ObservedWriter { + async fn open(&self, partition: usize) -> Result> { + let guard = self.owner.scope.start(Phase::WriterOpen, "writer_open"); + guard.detail("c2.open_partition", || partition.to_string()); + let result = self.inner.open(partition).await; + guard.finish_result(&result); + Ok(Box::new(ObservedSink { + inner: result?, + flow: Flow::new(self.owner.scope.clone(), "opened_shuffle_sink"), + })) + } +} + +struct ObservedSink { + inner: Box, + flow: Flow, +} + +#[tonic::async_trait] +impl TaskStreamSink for ObservedSink { + async fn write(&mut self, batches: Vec>) -> Result { + self.flow.first_poll(); + for batch in batches.iter().flatten() { + self.flow.batch(batch.num_rows()); + } + let result = self.inner.write(batches).await; + match &result { + Err(_) => self.flow.terminal("write_error", Outcome::Failed), + Ok(TaskStreamWriteState::Closed) => { + self.flow.terminal("write_closed", Outcome::Succeeded) + } + Ok(TaskStreamWriteState::Active) => {} + } + result + } + + async fn commit(self: Box) -> Result<()> { + let Self { inner, mut flow } = *self; + let guard = flow.scope.start(Phase::SinkCommit, "sink_commit"); + let result = inner.commit().await; + guard.finish_result(&result); + flow.terminal( + "commit_returned", + if result.is_ok() { + Outcome::Succeeded + } else { + Outcome::Failed + }, + ); + result + } + + async fn abort(self: Box) -> Result<()> { + let Self { inner, mut flow } = *self; + let guard = flow.scope.start(Phase::SinkAbort, "sink_abort"); + let result = inner.abort().await; + guard.finish_result(&result); + flow.terminal( + "abort_returned", + if result.is_ok() { + Outcome::Succeeded + } else { + Outcome::Failed + }, + ); + result + } +} + +#[cfg(test)] +#[path = "observation_tests.rs"] +mod tests; diff --git a/crates/sail-execution/src/stream/observation_tests.rs b/crates/sail-execution/src/stream/observation_tests.rs new file mode 100644 index 00000000..4f3291a4 --- /dev/null +++ b/crates/sail-execution/src/stream/observation_tests.rs @@ -0,0 +1,221 @@ +use std::collections::VecDeque; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::arrow::array::{ArrayRef, Int64Array}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::DataFusionError; +use futures::task::noop_waker; + +use super::*; + +type Item = Option>; + +struct Inner { + schema: SchemaRef, + polls: Arc, + outputs: VecDeque>, +} + +impl Stream for Inner { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + let inner = self.get_mut(); + inner.polls.fetch_add(1, Ordering::SeqCst); + inner.outputs.pop_front().unwrap_or(Poll::Ready(None)) + } + + fn size_hint(&self) -> (usize, Option) { + (0, Some(self.outputs.len())) + } +} + +impl RecordBatchStream for Inner { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +fn scope() -> Scope { + Scope::Job { + session: Arc::from("session"), + job: 23, + } +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) +} + +#[test] +fn observed_stream_preserves_each_poll_backpressure_batch_error_and_eof() -> Result<()> { + let schema = schema(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int64Array::from(vec![i64::MIN, 1, i64::MAX])) as ArrayRef], + )?; + let polls = Arc::new(AtomicUsize::new(0)); + let inner = Inner { + schema: schema.clone(), + polls: polls.clone(), + outputs: VecDeque::from([ + Poll::Pending, + Poll::Ready(Some(Ok(batch.clone()))), + Poll::Ready(Some(Err(DataFusionError::Execution( + "original sentinel".into(), + )))), + Poll::Ready(None), + ]), + }; + let inner: SendableRecordBatchStream = Box::pin(inner); + let mut observed = ObservedStream { + inner, + flow: Flow::new(scope(), "test_reader"), + }; + assert_eq!(observed.schema(), schema); + assert_eq!(observed.size_hint(), (0, Some(4))); + let waker = noop_waker(); + let mut context = Context::from_waker(&waker); + assert!(Pin::new(&mut observed).poll_next(&mut context).is_pending()); + assert_eq!(polls.load(Ordering::SeqCst), 1); + assert_eq!(observed.size_hint(), (0, Some(3))); + assert!(observed.flow.polled); + assert_eq!(observed.flow.batches, 0); + assert!(matches!(Pin::new(&mut observed).poll_next(&mut context), + Poll::Ready(Some(Ok(actual))) if actual == batch)); + assert_eq!(observed.flow.rows, 3); + assert_eq!(observed.flow.batches, 1); + assert!(matches!(Pin::new(&mut observed).poll_next(&mut context), + Poll::Ready(Some(Err(error))) if error.to_string().contains("original sentinel"))); + assert!(observed.flow.consumption.is_none()); + assert!(matches!( + Pin::new(&mut observed).poll_next(&mut context), + Poll::Ready(None) + )); + assert_eq!(polls.load(Ordering::SeqCst), 4); + Ok(()) +} + +#[test] +fn pinned_record_batch_wrapper_keeps_sendable_schema_and_unpolled_drop() { + let original_schema = schema(); + let polls = Arc::new(AtomicUsize::new(0)); + let inner: SendableRecordBatchStream = Box::pin(Inner { + schema: original_schema.clone(), + polls: polls.clone(), + outputs: VecDeque::from([Poll::Pending]), + }); + let wrapped: SendableRecordBatchStream = Box::pin(ObservedStream { + inner, + flow: Flow::new(scope(), "test_pinned_reader"), + }); + assert!(Arc::ptr_eq(&wrapped.schema(), &original_schema)); + assert_eq!(wrapped.size_hint(), (0, Some(1))); + drop(wrapped); + assert_eq!(polls.load(Ordering::SeqCst), 0); +} + +#[test] +fn declaration_owner_has_no_consumption_or_first_poll_to_abandon() { + let flow = Flow::declaration(scope(), "reader_declaration_owner"); + assert!(flow.owner.is_some()); + assert!(flow.consumption.is_none()); + assert!(!flow.polled); + assert_eq!(flow.batches, 0); +} + +#[test] +fn unpolled_stream_drop_does_not_poll_original() { + let polls = Arc::new(AtomicUsize::new(0)); + let stream = ObservedStream { + inner: Inner { + schema: schema(), + polls: polls.clone(), + outputs: VecDeque::new(), + }, + flow: Flow::new(scope(), "test_reader"), + }; + drop(stream); + assert_eq!(polls.load(Ordering::SeqCst), 0); +} + +#[test] +fn saturating_counts_are_explicitly_incomplete_and_never_wrap() { + let mut flow = Flow::new(scope(), "test_reader"); + flow.rows = u64::MAX; + flow.batches = u64::MAX; + flow.batch(1); + assert!(!flow.counts_complete); + assert_eq!(flow.rows, u64::MAX); + assert_eq!(flow.batches, u64::MAX); +} + +struct InnerSink { + writes: Arc, + commits: Arc, + aborts: Arc, +} + +#[tonic::async_trait] +impl TaskStreamSink for InnerSink { + async fn write(&mut self, _: Vec>) -> Result { + self.writes.fetch_add(1, Ordering::SeqCst); + Err(DataFusionError::Execution("original write error".into())) + } + + async fn commit(self: Box) -> Result<()> { + self.commits.fetch_add(1, Ordering::SeqCst); + Err(DataFusionError::Execution("original commit error".into())) + } + + async fn abort(self: Box) -> Result<()> { + self.aborts.fetch_add(1, Ordering::SeqCst); + Err(DataFusionError::Execution("original abort error".into())) + } +} + +#[tokio::test] +async fn observed_sink_preserves_write_error_and_following_abort_once() { + let writes = Arc::new(AtomicUsize::new(0)); + let commits = Arc::new(AtomicUsize::new(0)); + let aborts = Arc::new(AtomicUsize::new(0)); + let mut sink = ObservedSink { + inner: Box::new(InnerSink { + writes: writes.clone(), + commits: commits.clone(), + aborts: aborts.clone(), + }), + flow: Flow::new(scope(), "opened_shuffle_sink"), + }; + let result = sink + .write(vec![Some(RecordBatch::new_empty(schema()))]) + .await; + assert!(matches!(result, Err(error) if error.to_string().contains("original write error"))); + assert!(sink.flow.consumption.is_none()); + assert_eq!(sink.flow.batches, 1); + let result = Box::new(sink).abort().await; + assert!(matches!(result, Err(error) if error.to_string().contains("original abort error"))); + assert_eq!(writes.load(Ordering::SeqCst), 1); + assert_eq!(commits.load(Ordering::SeqCst), 0); + assert_eq!(aborts.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn observed_sink_preserves_commit_error_without_adding_abort() { + let writes = Arc::new(AtomicUsize::new(0)); + let commits = Arc::new(AtomicUsize::new(0)); + let aborts = Arc::new(AtomicUsize::new(0)); + let sink = ObservedSink { + inner: Box::new(InnerSink { + writes: writes.clone(), + commits: commits.clone(), + aborts: aborts.clone(), + }), + flow: Flow::new(scope(), "opened_shuffle_sink"), + }; + let result = Box::new(sink).commit().await; + assert!(matches!(result, Err(error) if error.to_string().contains("original commit error"))); + assert_eq!(writes.load(Ordering::SeqCst), 0); + assert_eq!(commits.load(Ordering::SeqCst), 1); + assert_eq!(aborts.load(Ordering::SeqCst), 0); +} diff --git a/crates/sail-execution/src/task_runner/actor/handler.rs b/crates/sail-execution/src/task_runner/actor/handler.rs index 3eabbc53..5b779e1d 100644 --- a/crates/sail-execution/src/task_runner/actor/handler.rs +++ b/crates/sail-execution/src/task_runner/actor/handler.rs @@ -7,6 +7,7 @@ use datafusion_proto::protobuf::PhysicalPlanNode; use log::{debug, error, warn}; use prost::Message; use sail_common::actor::{ActorAction, ActorContext}; +use sail_common::telemetry::c2::{self, Guard, Identity, Outcome, Phase, Placement}; use sail_common_datafusion::error::CommonErrorCause; use tokio::sync::oneshot; @@ -23,6 +24,13 @@ use crate::task_runner::{TaskRunnerActor, TaskRunnerPlacement}; use crate::worker::{WorkerLocation, WorkerMessage}; impl TaskRunnerActor { + fn observation_placement(&self) -> Placement { + match &self.placement { + TaskRunnerPlacement::Driver { .. } => Placement::Driver, + TaskRunnerPlacement::Worker { worker_id, .. } => Placement::Worker((*worker_id).into()), + } + } + pub(super) fn handle_run_task_batch( &mut self, ctx: &mut ActorContext, @@ -60,13 +68,31 @@ impl TaskRunnerActor { } // Validate shared descriptions before admitting any task. Only descriptions are shared; // every task gets a fresh converter, executable plan, and shuffle reader/writer. - let proto = Arc::new(PhysicalPlanNode::decode(definition.plan.as_ref()).map_err( - |error| ExecutionError::InvalidArgument(format!("invalid physical plan: {error}")), - )?); + let decode = Guard::start(Phase::BatchPhysicalDecode, || Identity::Stage { + session: &self.session_id, + job: job_id.into(), + stage, + }); + decode.detail("c2.encoded_physical_bytes", || { + definition.plan.len().to_string() + }); + let decoded = PhysicalPlanNode::decode(definition.plan.as_ref()).map_err(|error| { + ExecutionError::InvalidArgument(format!("invalid physical plan: {error}")) + }); + decode.finish_result(&decoded); + let proto = Arc::new(decoded?); self.extension_jobs.admit(job_id, &context)?; self.tasks.record_batch(job_id, stage, &tasks); for task in tasks { let key = task.task_key(job_id, stage); + c2::event(Phase::TaskAdmitted, || Identity::Task { + session: &self.session_id, + job: key.job_id.into(), + stage: key.stage, + partition: key.partition, + attempt: key.attempt, + placement: self.observation_placement(), + }); let stream = TaskPreparation { session_id: self.session_id.clone(), handle: ctx.handle().clone(), @@ -98,10 +124,19 @@ impl TaskRunnerActor { } pub(super) fn handle_close_job(&mut self, job_id: JobId) -> ActorAction { + let closed = Guard::start(Phase::CloseJobExecuted, || Identity::PlacementJob { + session: &self.session_id, + job: job_id.into(), + placement: self.observation_placement(), + }); self.tasks.close_job(job_id); self.extension_jobs.close_job(&self.session_id, job_id); self.extensions.local_streams.remove_streams(job_id, None); self.signals.retain(|key, _| key.job_id != job_id); + closed.detail("c2.scope", || { + "actor_close_job_returned_named_owners_only".into() + }); + closed.finish(Outcome::Succeeded); ActorAction::Continue } @@ -129,7 +164,22 @@ impl TaskRunnerActor { .unwrap_or_else(|| "".into()) ); } - if !matches!(status, TaskStatus::Running) { + let terminal_outcome = match status { + TaskStatus::Running => None, + TaskStatus::Succeeded => Some(Outcome::Succeeded), + TaskStatus::Failed => Some(Outcome::Failed), + TaskStatus::Canceled => Some(Outcome::Cancelled), + }; + if let Some(outcome) = terminal_outcome { + let terminal = Guard::point(Phase::TaskTerminal, || Identity::Task { + session: &self.session_id, + job: key.job_id.into(), + stage: key.stage, + partition: key.partition, + attempt: key.attempt, + placement: self.observation_placement(), + }); + terminal.finish(outcome); self.signals.remove(&key); } match &mut self.placement { @@ -422,7 +472,16 @@ impl TaskRunnerActor { job_id: JobId, stage: Option, ) -> ActorAction { + let removed = Guard::start(Phase::LocalStreamsRemoved, || Identity::PlacementJob { + session: &self.session_id, + job: job_id.into(), + placement: self.observation_placement(), + }); self.extensions.local_streams.remove_streams(job_id, stage); + removed.detail("c2.scope", || { + "local_registry_remove_returned_not_allocator_reclamation".into() + }); + removed.finish(Outcome::Succeeded); ActorAction::Continue } diff --git a/crates/sail-execution/src/task_runner/preparation.rs b/crates/sail-execution/src/task_runner/preparation.rs index 2dc07c8b..917e7c3e 100644 --- a/crates/sail-execution/src/task_runner/preparation.rs +++ b/crates/sail-execution/src/task_runner/preparation.rs @@ -15,6 +15,7 @@ use datafusion_proto::protobuf::PhysicalPlanNode; use futures::TryStreamExt; use log::debug; use sail_common::actor::ActorHandle; +use sail_common::telemetry::c2::{self, Guard, Identity, Outcome, Phase, Placement}; use sail_common_datafusion::schema_evolution::SchemaEvolutionPhysicalExprAdapterFactory; use sail_common_datafusion::worker_extension::{ WorkerExtensionRegistry, WorkerJobIdentity, WorkerTaskScope, @@ -28,6 +29,7 @@ use crate::id::{TaskKey, TaskKeyDisplay}; use crate::plan::{ShuffleReadExec, ShuffleWriteExec, StageInputExec}; use crate::proto::{RemoteExecutionCodec, proto_to_physical_plan}; use crate::stream::accessor::TaskStreamFactory; +use crate::stream::observation::{self, ObserverContext}; use crate::task::definition::{TaskDefinition, TaskInput, TaskOutput}; use crate::task_runner::TaskRunnerActor; @@ -38,6 +40,24 @@ pub(super) struct TaskPreparation { pub worker_id: Option, } +struct PreparationObservation { + session: String, + placement: Placement, +} + +impl PreparationObservation { + fn identity(&self, key: &TaskKey) -> Identity<'_> { + Identity::Task { + session: &self.session, + job: key.job_id.into(), + stage: key.stage, + partition: key.partition, + attempt: key.attempt, + placement: self.placement, + } + } +} + impl TaskPreparation { /// The returned schema belongs to ShuffleWriteExec's completion stream, not the stage's data. pub fn stream( @@ -47,7 +67,14 @@ impl TaskPreparation { proto: Arc, context: Arc, ) -> SendableRecordBatchStream { - preparation_stream(key.clone(), move |canceled| { + let observation = c2::enabled().then(|| PreparationObservation { + session: self.session_id.clone(), + placement: match self.worker_id { + Some(worker_id) => Placement::Worker(worker_id), + None => Placement::Driver, + }, + }); + preparation_stream(key.clone(), observation, move |canceled| { self.execute_plan(&key, &definition, &proto, &canceled, context) }) } @@ -83,7 +110,17 @@ impl TaskPreparation { } _ => context, }; - let plan = proto_to_physical_plan(&context, &RemoteExecutionCodec, proto)?; + let decode = Guard::start(Phase::PhysicalDecode, || Identity::Task { + session: &self.session_id, + job: key.job_id.into(), + stage: key.stage, + partition: key.partition, + attempt: key.attempt, + placement: self.worker_id.map_or(Placement::Driver, Placement::Worker), + }); + let decoded = proto_to_physical_plan(&context, &RemoteExecutionCodec, proto); + decode.finish_result(&decoded); + let plan = decoded?; let plan = self.rewrite_file_scans(plan)?; let plan = self.rewrite_shuffle( key, @@ -113,7 +150,15 @@ impl TaskPreparation { "task canceled during preparation".into(), )); } - Ok(plan.execute(key.partition, context)?) + let stream = plan.execute(key.partition, context)?; + Ok(observation::task_stream( + stream, + ObserverContext::optional( + &self.session_id, + self.worker_id.map_or(Placement::Driver, Placement::Worker), + ), + key, + )) } fn rewrite_file_scans( @@ -157,8 +202,16 @@ impl TaskPreparation { context: Arc, ) -> ExecutionResult> { let mappers = plan.output_partitioning().partition_count(); - let streams = - TaskStreamFactory::new(self.handle.clone(), context.clone(), self.celeborn, mappers); + let streams = TaskStreamFactory::new( + self.handle.clone(), + context.clone(), + self.celeborn, + mappers, + ObserverContext::optional( + &self.session_id, + self.worker_id.map_or(Placement::Driver, Placement::Worker), + ), + ); let result = { let streams = streams.clone(); plan.transform(move |node| { @@ -190,6 +243,7 @@ impl TaskPreparation { /// its inputs and result, so Tokio drops an abandoned result even after the monitor exits. fn preparation_stream( key: TaskKey, + observation: Option, prepare: impl FnOnce(CancellationToken) -> ExecutionResult + Send + 'static, @@ -198,16 +252,36 @@ fn preparation_stream( let canceled = CancellationToken::new(); let _cancel_on_drop = canceled.clone().drop_guard(); let queued = Instant::now(); + let queued_observation = Guard::start(Phase::PreparationQueue, || { + observation + .as_ref() + .map_or(Identity::CausalContext, |o| o.identity(&key)) + }); let span = fastrace::Span::enter_with_local_parent("TaskPreparation"); tokio::task::spawn_blocking(move || { let _parent = span.set_local_parent(); let started = Instant::now(); + queued_observation.finish(Outcome::Succeeded); + let work_observation = Guard::start(Phase::Preparation, || { + observation + .as_ref() + .map_or(Identity::CausalContext, |o| o.identity(&key)) + }); if canceled.is_cancelled() { + work_observation.finish(Outcome::Cancelled); return Err(ExecutionError::InternalError( "task canceled before preparation".into(), )); } + let observed_cancellation = c2::enabled().then(|| canceled.clone()); let result = prepare(canceled); + work_observation.detail("c2.cancelled_at_completion", || { + observed_cancellation + .as_ref() + .is_some_and(CancellationToken::is_cancelled) + .to_string() + }); + work_observation.finish_result(&result); debug!( "{} preparation wait={:?} duration={:?}", TaskKeyDisplay(&key), @@ -250,7 +324,7 @@ mod tests { async fn dropping_an_unpolled_stream_does_not_prepare() { let started = Arc::new(AtomicBool::new(false)); let flag = started.clone(); - let stream = preparation_stream(key(), move |_| { + let stream = preparation_stream(key(), None, move |_| { flag.store(true, Ordering::SeqCst); Err(ExecutionError::InternalError("should never start".into())) }); @@ -262,7 +336,7 @@ mod tests { #[tokio::test] async fn preparation_errors_and_panics_are_stream_errors() -> Result<(), Box> { - let mut stream = preparation_stream(key(), |_| { + let mut stream = preparation_stream(key(), None, |_| { Err(ExecutionError::InvalidArgument("bad plan".into())) }); assert!( @@ -273,7 +347,7 @@ mod tests { .is_err() ); assert!(stream.next().await.is_none()); - let mut stream = preparation_stream(key(), |_| { + let mut stream = preparation_stream(key(), None, |_| { std::panic::resume_unwind(Box::new("construction panic")) }); assert!(stream.next().await.ok_or("missing panic error")?.is_err()); @@ -297,7 +371,7 @@ mod tests { let (release, blocked) = std::sync::mpsc::channel(); let (observed, canceled) = oneshot::channel(); let (dropped, result_dropped) = oneshot::channel(); - let mut stream = preparation_stream(key(), move |token| { + let mut stream = preparation_stream(key(), None, move |token| { let _ = started.send(()); blocked .recv() diff --git a/crates/sail-execution/src/worker/client.rs b/crates/sail-execution/src/worker/client.rs index d6924594..3498b8ad 100644 --- a/crates/sail-execution/src/worker/client.rs +++ b/crates/sail-execution/src/worker/client.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use prost::Message; +use sail_common::telemetry::c2::{Guard, Identity, Outcome, Phase}; use crate::error::ExecutionResult; use crate::id::{JobId, TaskAttempt, TaskKey}; @@ -52,8 +53,17 @@ impl WorkerClient { definition: Arc, peers: Vec, ) -> ExecutionResult<()> { + let encode = Guard::start(Phase::TaskDefinitionEncode, || Identity::CausalStage { + job: job_id.into(), + stage: stage as u64, + }); let definition = crate::task::r#gen::TaskDefinition::from(definition.as_ref().clone()).encode_to_vec(); + encode.detail("c2.encoded_definition_bytes", || { + definition.len().to_string() + }); + encode.detail("c2.task_attempts", || tasks.len().to_string()); + encode.finish(Outcome::Succeeded); log::debug!( "task batch tasks={} encoded definition bytes={}", tasks.len(), @@ -72,9 +82,18 @@ impl WorkerClient { definition, peers: peers.into_iter().map(|x| x.into()).collect(), }; - let response = self.inner.get().await?.run_task_batch(request).await?; - let RunTaskBatchResponse {} = response.into_inner(); - Ok(()) + let dispatch = Guard::start(Phase::BatchDispatch, || Identity::CausalStage { + job: job_id.into(), + stage: stage as u64, + }); + let result = async { + let response = self.inner.get().await?.run_task_batch(request).await?; + let RunTaskBatchResponse {} = response.into_inner(); + Ok(()) + } + .await; + dispatch.finish_result(&result); + result } pub async fn stop_task(&self, key: TaskKey) -> ExecutionResult<()> { diff --git a/crates/sail-execution/src/worker/server.rs b/crates/sail-execution/src/worker/server.rs index 81fcb160..1f077a83 100644 --- a/crates/sail-execution/src/worker/server.rs +++ b/crates/sail-execution/src/worker/server.rs @@ -4,10 +4,11 @@ use datafusion::execution::TaskContext; use log::debug; use prost::Message; use sail_common::actor::ActorHandle; +use sail_common::telemetry::c2::{Guard, Identity, Phase}; use tonic::{Request, Response, Status}; use crate::error::{ExecutionError, ExecutionResult}; -use crate::id::{TaskAttempt, TaskKey}; +use crate::id::{JobId, TaskAttempt, TaskKey}; use crate::task::definition::TaskDefinition; use crate::task_runner::{TaskRunnerActor, TaskRunnerMessage}; use crate::worker::r#gen::worker_service_server::WorkerService; @@ -56,8 +57,17 @@ impl WorkerService for WorkerServer { .into_iter() .map(|x| x.try_into()) .collect::>>()?; - let definition = crate::task::r#gen::TaskDefinition::decode(definition.as_slice()) - .map_err(|e| Status::invalid_argument(format!("invalid task definition: {e}")))?; + let decode = Guard::start(Phase::TaskDefinitionDecode, || Identity::CausalStage { + job: job_id, + stage, + }); + decode.detail("c2.encoded_definition_bytes", || { + definition.len().to_string() + }); + let decoded = crate::task::r#gen::TaskDefinition::decode(definition.as_slice()) + .map_err(|e| Status::invalid_argument(format!("invalid task definition: {e}"))); + decode.finish_result(&decoded); + let definition = decoded?; let stage = usize::try_from(stage).map_err(|_| Status::invalid_argument("stage overflow"))?; let tasks = tasks @@ -71,21 +81,31 @@ impl WorkerService for WorkerServer { }) }) .collect::, Status>>()?; - let (result, rx) = tokio::sync::oneshot::channel(); - self.task_runner - .send(TaskRunnerMessage::RunTaskBatch { - job_id: job_id.into(), - stage, - tasks, - definition: Arc::new(TaskDefinition::try_from(definition)?), - context: self.context.clone(), - peers, - result, - }) - .await - .map_err(ExecutionError::from)?; - rx.await - .map_err(|_| Status::unavailable("task runner stopped before admission"))??; + let admission = Guard::start(Phase::BatchAdmission, || Identity::CausalStage { + job: job_id, + stage: stage as u64, + }); + let admitted: Result<(), Status> = async { + let (result, rx) = tokio::sync::oneshot::channel(); + self.task_runner + .send(TaskRunnerMessage::RunTaskBatch { + job_id: job_id.into(), + stage, + tasks, + definition: Arc::new(TaskDefinition::try_from(definition)?), + context: self.context.clone(), + peers, + result, + }) + .await + .map_err(ExecutionError::from)?; + rx.await + .map_err(|_| Status::unavailable("task runner stopped before admission"))??; + Ok(()) + } + .await; + admission.finish_result(&admitted); + admitted?; let response = RunTaskBatchResponse {}; debug!("{response:?}"); Ok(Response::new(response)) @@ -126,22 +146,38 @@ impl WorkerService for WorkerServer { let request = request.into_inner(); debug!("{request:?}"); let CleanUpJobRequest { job_id, stage } = request; - let job_id = job_id.into(); + let job_id: JobId = job_id.into(); let stage = stage.map(|x| x as usize); - if stage.is_none() { + let queued = Guard::start(Phase::CleanupQueuedAck, || match stage { + Some(stage) => Identity::CausalStage { + job: job_id.into(), + stage: stage as u64, + }, + None => Identity::CausalJob { job: job_id.into() }, + }); + queued.detail("c2.scope", || { + "mailbox_send_ack_only_not_execution_or_deallocation".into() + }); + let queued_result: Result<(), Status> = async { + if stage.is_none() { + self.task_runner + .send(TaskRunnerMessage::CloseJob { job_id }) + .await + .map_err(ExecutionError::from)?; + } self.task_runner - .send(TaskRunnerMessage::CloseJob { job_id }) + .send(TaskRunnerMessage::CleanUpLocalStreams { job_id, stage }) .await .map_err(ExecutionError::from)?; + self.task_runner + .send(TaskRunnerMessage::CleanUpCelebornStreams { job_id, stage }) + .await + .map_err(ExecutionError::from)?; + Ok(()) } - self.task_runner - .send(TaskRunnerMessage::CleanUpLocalStreams { job_id, stage }) - .await - .map_err(ExecutionError::from)?; - self.task_runner - .send(TaskRunnerMessage::CleanUpCelebornStreams { job_id, stage }) - .await - .map_err(ExecutionError::from)?; + .await; + queued.finish_result(&queued_result); + queued_result?; let response = CleanUpJobResponse {}; debug!("{response:?}"); Ok(Response::new(response)) diff --git a/crates/sail-plan/src/lib.rs b/crates/sail-plan/src/lib.rs index bb87e6d3..823d9b33 100644 --- a/crates/sail-plan/src/lib.rs +++ b/crates/sail-plan/src/lib.rs @@ -7,6 +7,7 @@ use datafusion_common::Result; use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan}; use datafusion_expr::LogicalPlan; use sail_common::spec; +use sail_common::telemetry::c2::{Guard, Identity, Phase}; use sail_common_datafusion::rename::physical_plan::rename_physical_plan; use crate::config::PlanConfig; @@ -38,21 +39,33 @@ pub async fn resolve_and_execute_plan( ) -> PlanResult<(Arc, Vec)> { let mut info = vec![]; let resolver = PlanResolver::new(ctx, config); - let NamedPlan { plan, fields } = resolver.resolve_named_plan(plan).await?; + let observation = Guard::start(Phase::Resolve, || Identity::CausalContext); + let resolved = resolver.resolve_named_plan(plan).await; + observation.finish_result(&resolved); + let NamedPlan { plan, fields } = resolved?; info.push(plan.to_stringified(PlanType::InitialLogicalPlan)); - let df = execute_logical_plan(ctx, plan).await?; + let observation = Guard::start(Phase::LogicalExecute, || Identity::CausalContext); + let executed = execute_logical_plan(ctx, plan).await; + observation.finish_result(&executed); + let df = executed?; let (session_state, plan) = df.into_parts(); - let plan = session_state.optimize(&plan)?; + let observation = Guard::start(Phase::Optimize, || Identity::CausalContext); + let optimized = session_state.optimize(&plan); + observation.finish_result(&optimized); + let plan = optimized?; let plan = if is_streaming_plan(&plan)? { rewrite_streaming_plan(plan)? } else { plan }; info.push(plan.to_stringified(PlanType::FinalLogicalPlan)); - let plan = session_state + let observation = Guard::start(Phase::PhysicalPlan, || Identity::CausalContext); + let planned = session_state .query_planner() .create_physical_plan(&plan, &session_state) - .await?; + .await; + observation.finish_result(&planned); + let plan = planned?; let plan = if let Some(fields) = fields { rename_physical_plan(plan, &fields)? } else { diff --git a/crates/sail-spark-connect/src/executor.rs b/crates/sail-spark-connect/src/executor.rs index 82fb2deb..e43053f2 100644 --- a/crates/sail-spark-connect/src/executor.rs +++ b/crates/sail-spark-connect/src/executor.rs @@ -14,6 +14,7 @@ use fastrace::Span; use fastrace::future::FutureExt; use futures::Stream; use futures::stream::StreamExt; +use sail_common::telemetry::c2::{self, Guard, Identity, NamedOwner, Phase}; use tokio::sync::{mpsc, oneshot}; use tokio::task::JoinHandle; use tonic::codegen::tokio_stream::wrappers::ReceiverStream; @@ -64,6 +65,8 @@ pub type ExecutorOutputStream = Pin, + // A named owner's lifetime; it makes no allocator-reclamation assertion. + _observer_owner: Option>, } // TODO: use "spark.connect.execute.reattachable.observerRetryBufferSize" @@ -75,6 +78,7 @@ impl ExecutorBuffer { Self { capacity, inner: VecDeque::with_capacity(capacity), + _observer_owner: None, } } @@ -109,6 +113,8 @@ pub(crate) struct ExecutorMetadata { pub(crate) struct Executor { pub(crate) metadata: ExecutorMetadata, state: Mutex, + observer_session: Option>, + _observer_owner: Option>, } enum ExecutorState { @@ -217,6 +223,7 @@ enum ExecutorTaskResult { impl Executor { pub(crate) fn new( + session_id: &str, metadata: ExecutorMetadata, stream: SendableRecordBatchStream, heartbeat_interval: Duration, @@ -229,11 +236,28 @@ impl Executor { completion: Some(completion), }, }; - let buffer = if metadata.reattachable { + let mut buffer = if metadata.reattachable { ExecutorBuffer::new(EXECUTOR_BUFFER_CAPACITY) } else { ExecutorBuffer::new(0) }; + buffer._observer_owner = NamedOwner::optional( + Phase::ExecutorBufferOwner, + || Identity::Operation { + session: session_id, + operation: &metadata.operation_id, + }, + "executor_response_buffer_owner", + ); + let observer_session = c2::enabled().then(|| Arc::from(session_id)); + let observer_owner = NamedOwner::optional( + Phase::ExecutorOwner, + || Identity::Operation { + session: session_id, + operation: &metadata.operation_id, + }, + "spark_executor_owner", + ); Self { metadata, state: Mutex::new(ExecutorState::Pending { @@ -245,6 +269,8 @@ impl Executor { }, span: Span::enter_with_local_parent("Executor::new"), }), + observer_session, + _observer_owner: observer_owner, } } @@ -482,6 +508,19 @@ impl Executor { } pub(crate) fn release(&self, response_id: String) -> SparkResult<()> { + let released = Guard::start(Phase::ReleaseResponseBuffer, || Identity::Operation { + session: self.observer_session.as_deref().unwrap_or(""), + operation: &self.metadata.operation_id, + }); + released.detail("c2.scope", || { + "remove_matching_buffer_prefix_not_executor_or_allocator_reclaim".into() + }); + let result = self.release_internal(response_id); + released.finish_result(&result); + result + } + + fn release_internal(&self, response_id: String) -> SparkResult<()> { let state = self.state.lock()?; let buffer = match state.deref() { ExecutorState::Running { task, span: _ } => &task.buffer, diff --git a/crates/sail-spark-connect/src/executor/tests.rs b/crates/sail-spark-connect/src/executor/tests.rs index fe8c666d..013c6905 100644 --- a/crates/sail-spark-connect/src/executor/tests.rs +++ b/crates/sail-spark-connect/src/executor/tests.rs @@ -33,6 +33,7 @@ impl Drop for PendingStream { fn executor() -> (Executor, Arc) { let dropped = Arc::new(AtomicUsize::new(0)); let executor = Executor::new( + "test-session", ExecutorMetadata { operation_id: "interrupted-operation".into(), tags: vec![], @@ -56,6 +57,23 @@ fn assert_terminal(executor: &Executor, dropped: &AtomicUsize) { } } +#[test] +fn response_release_preserves_the_unreleased_buffer_suffix_and_unknown_ids() { + let mut buffer = ExecutorBuffer::new(3); + let first = ExecutorOutput::new(ExecutorBatch::Heartbeat); + let second = ExecutorOutput::new(ExecutorBatch::Heartbeat); + let last = ExecutorOutput::complete(); + for output in [&first, &second, &last] { + buffer.add(output.clone()); + } + buffer.remove_until("not-a-response"); + assert_eq!(buffer.iter().count(), 3); + buffer.remove_until(&second.id); + let remaining = buffer.iter().collect::>(); + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].id, last.id); +} + #[tokio::test] async fn interrupt_drops_pending_and_running_contexts_but_preserves_terminal_identity() -> SparkResult<()> { diff --git a/crates/sail-spark-connect/src/server.rs b/crates/sail-spark-connect/src/server.rs index 9455ed6e..b569baf1 100644 --- a/crates/sail-spark-connect/src/server.rs +++ b/crates/sail-spark-connect/src/server.rs @@ -1,6 +1,8 @@ use async_stream; use datafusion::prelude::SessionContext; +use fastrace::future::FutureExt; use log::debug; +use sail_common::telemetry::c2::{self, Rpc}; use sail_session::session_manager::SessionManager; use tonic::codegen::tokio_stream::StreamExt; use tonic::{Request, Response, Status, Streaming}; @@ -428,26 +430,37 @@ impl SparkConnectService for SparkConnectServer { request: Request, ) -> Result, Status> { let request = request.into_inner(); - debug!("{request:?}"); - let session_id = request.session_id.clone(); - let user_id = request.user_context.map(|u| u.user_id).unwrap_or_default(); - let ctx = self - .session_manager - .get_or_create_session_context(session_id, user_id) - .await - .map_err(SparkError::from)?; - let response_id = match request.release.required("release")? { - Release::ReleaseAll(ReleaseAll {}) => None, - Release::ReleaseUntil(ReleaseUntil { response_id }) => Some(response_id), - }; - service::handle_release_execute(&ctx, request.operation_id.clone(), response_id).await?; - let response = ReleaseExecuteResponse { - session_id: request.session_id.clone(), - server_side_session_id: request.session_id, - operation_id: Some(request.operation_id), - }; - debug!("{response:?}"); - Ok(Response::new(response)) + let span = c2::rpc_span(Rpc::ReleaseExecute, || { + vec![ + ("session.id", request.session_id.clone()), + ("operation.id", request.operation_id.clone()), + ] + }); + async { + debug!("{request:?}"); + let session_id = request.session_id.clone(); + let user_id = request.user_context.map(|u| u.user_id).unwrap_or_default(); + let ctx = self + .session_manager + .get_or_create_session_context(session_id, user_id) + .await + .map_err(SparkError::from)?; + let response_id = match request.release.required("release")? { + Release::ReleaseAll(ReleaseAll {}) => None, + Release::ReleaseUntil(ReleaseUntil { response_id }) => Some(response_id), + }; + service::handle_release_execute(&ctx, request.operation_id.clone(), response_id) + .await?; + let response = ReleaseExecuteResponse { + session_id: request.session_id.clone(), + server_side_session_id: request.session_id, + operation_id: Some(request.operation_id), + }; + debug!("{response:?}"); + Ok(Response::new(response)) + } + .in_span(span) + .await } async fn release_session( diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index d0e8d103..d06e1ca6 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -9,6 +9,7 @@ use fastrace::future::FutureExt; use futures::stream; use log::debug; use sail_common::spec; +use sail_common::telemetry::c2; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; use sail_plan::resolve_and_execute_plan; @@ -120,13 +121,22 @@ async fn handle_execute_plan( let spark = ctx.extension::()?; let service = ctx.extension::()?; let operation_id = metadata.operation_id.clone(); - let (plan, _) = resolve_and_execute_plan(ctx, spark.plan_config()?, plan).await?; + c2::bind_operation(&span, spark.session_id(), &operation_id); + let (plan, _) = if c2::enabled() { + resolve_and_execute_plan(ctx, spark.plan_config()?, plan) + .in_span(Span::enter_with_parent("c2.resolve_scope", &span)) + .await? + } else { + resolve_and_execute_plan(ctx, spark.plan_config()?, plan).await? + }; let stream = { let span = Span::enter_with_parent("JobRunner::execute", &span); service.runner().execute(ctx, plan).in_span(span).await? }; let _guard = span.set_local_parent(); + let stream = sail_execution::operation_stream(stream, spark.session_id(), &operation_id); let executor = Executor::new( + spark.session_id(), metadata, stream, spark.options().execution_heartbeat_interval, diff --git a/crates/sail-spark-connect/src/session.rs b/crates/sail-spark-connect/src/session.rs index c2e98ff3..f1ea6b10 100644 --- a/crates/sail-spark-connect/src/session.rs +++ b/crates/sail-spark-connect/src/session.rs @@ -5,6 +5,7 @@ use std::time::Duration; use datafusion::execution::SendableRecordBatchStream; use datafusion::logical_expr::StringifiedPlan; +use sail_common::telemetry::c2::{Guard, Identity, Phase}; use sail_common::utils::datetime::get_system_timezone; use sail_common_datafusion::extension::SessionExtension; use sail_common_datafusion::session::lifecycle::SessionResource; @@ -214,11 +215,25 @@ impl SparkSession { } pub(crate) fn remove_executor(&self, id: &str) -> SparkResult>> { - let mut state = self.state.lock()?; - Ok(state - .executors - .remove_entry(id) - .map(|(_, executor)| executor)) + let released = Guard::start(Phase::ReleaseOperation, || Identity::Operation { + session: self.session_id(), + operation: id, + }); + released.detail("c2.scope", || { + "detach_session_map_entry_not_final_arc_owner_or_allocator_reclaim".into() + }); + let result = self + .state + .lock() + .map(|mut state| { + state + .executors + .remove_entry(id) + .map(|(_, executor)| executor) + }) + .map_err(SparkError::from); + released.finish_result(&result); + result } pub(crate) fn all_executors(&self) -> SparkResult>> { diff --git a/crates/sail-telemetry/src/telemetry.rs b/crates/sail-telemetry/src/telemetry.rs index 1d500afe..bdbce3ee 100644 --- a/crates/sail-telemetry/src/telemetry.rs +++ b/crates/sail-telemetry/src/telemetry.rs @@ -18,6 +18,7 @@ use opentelemetry_sdk::logs::{BatchConfigBuilder, BatchLogProcessor, SdkLoggerPr use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider, Temporality}; use sail_common::actor::ActorSystem; use sail_common::config::{OtlpProtocol, SystemCatalogConfig, SystemCatalogStore, TelemetryConfig}; +use sail_common::telemetry::c2; use sail_system_store::{SystemStoreHandle, SystemStoreReader}; use crate::error::{TelemetryError, TelemetryResult}; @@ -62,6 +63,13 @@ pub fn init_telemetry( match *status { TelemetryStatus::Uninitialized => { + if config.c2_observer + && (!config.export_traces || config.exporter.otlp.endpoint.is_none()) + { + return Err(TelemetryError::invalid( + "telemetry.c2_observer requires trace export and an OTLP endpoint", + )); + } let mut state = TelemetryState::default(); match init_traces(config, &mut state, &resource) .and_then(|()| init_system_store(system_config, &mut state, &resource)) @@ -70,6 +78,7 @@ pub fn init_telemetry( .and_then(|()| init_datafusion_telemetry()) { Ok(()) => { + c2::configure(config.c2_observer); debug!("OpenTelemetry initialized"); *status = TelemetryStatus::Initialized(Box::new(state)); Ok(()) @@ -272,6 +281,11 @@ fn init_datafusion_telemetry() -> TelemetryResult<()> { pub fn shutdown_telemetry() { debug!("Shutting down OpenTelemetry..."); + if let Some(summary) = c2::summary() + && let Ok(summary) = serde_json::to_string(&summary) + { + log::info!("c2_observer_summary={summary}"); + } fastrace::flush(); let state = TELEMETRY_STATUS.lock().ok().and_then(|mut status| { let previous = std::mem::replace(&mut *status, TelemetryStatus::Finalized);