algeon-query-runtime 0.1.2

//! Query observation session, seal observers, and recorder event writes.

use std::{
    collections::BTreeMap,
    fmt,
    sync::{
        Arc, Condvar, Mutex, MutexGuard, PoisonError,
        atomic::{AtomicBool, Ordering},
    },
};

#[cfg(test)]
use super::RuntimeObservationDeliveryGate;
use super::{
    CleanupFailureTally, LiveRuntimeObservationSnapshot, PlanExecutionId, QueryAttemptReport, QueryId,
    RetainedTopologyRecord, RuntimeEventContext, RuntimeObservationClassLoss, RuntimeObservationClock,
    RuntimeObservationConfig, RuntimeObservationDeliveryGuard, RuntimeObservationFooter, RuntimeObservationHandle,
    RuntimeObservationManifest, RuntimeObservationOwner, RuntimeObservationPayload, RuntimeObservationReasonLoss,
    RuntimeObservationRecordClass, RuntimeObservationRecordDisposition, RuntimeObservationRecorderState,
    RuntimeObservationSnapshotData, RuntimeTerminalOutcome, RuntimeTimelineEvent, RuntimeTopologyRecord,
    SealedRuntimeObservationSnapshot, SourceLocationDigestKeyState, TraceId,
};
#[cfg(any(test, feature = "test-hooks"))]
use super::{SystemRuntimeObservationClock, TestSourceLocationDigestKeyProvider};

/// Owner-only capability for one query observation lifetime.
///
/// This type is intentionally not [`Clone`]. Cloneable execution contexts and
/// workers receive only [`RuntimeObservationHandle`], which cannot seal.
#[derive(Debug, Default)]
pub struct RuntimeObservationSession {
    pub(super) handle: RuntimeObservationHandle,
}

/// Construction-time provider of a fixed attempt summary embedded at seal, for
/// tests that seal a recorder-only owner. An attempt passes the crate-private
/// seal provider below to each seal of its session, so the embedded value and
/// completion delivery share the same immutable allocation.
#[cfg(any(test, feature = "test-hooks"))]
pub(crate) type AttemptReportProvider = Arc<dyn Fn() -> QueryAttemptReport + Send + Sync>;

/// Attempt-owned report construction after the recorder has computed its
/// complete footer, given the session's exact cleanup-failure tally.
pub(crate) type AttemptReportSealProvider =
    Arc<dyn Fn(&RuntimeObservationFooter, &CleanupFailureTally) -> Arc<QueryAttemptReport> + Send + Sync>;

pub(super) struct RuntimeObservationSessionInner {
    pub(super) manifest: RuntimeObservationManifest,
    pub(super) source_location_digest_key: SourceLocationDigestKeyState,
    pub(super) clock: Arc<dyn RuntimeObservationClock>,
    pub(super) state: Mutex<RuntimeObservationRecorderState>,
    pub(super) sealed: AtomicBool,
    /// Signals that a sealed snapshot was published. A seal that builds its
    /// attempt report outside the state lock leaves the session sealed with
    /// no snapshot until then; a racing seal waits here for it.
    pub(super) seal_published: Condvar,
}

pub(super) struct RuntimeObservationSealObserver {
    pub(super) notify: Box<dyn Fn(SealedRuntimeObservationSnapshot) + Send + Sync>,
}

impl fmt::Debug for RuntimeObservationSealObserver {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.debug_struct("RuntimeObservationSealObserver").finish_non_exhaustive()
    }
}

impl RuntimeObservationSealObserver {
    pub(super) fn invoke(&self, snapshot: SealedRuntimeObservationSnapshot) {
        let _delivery = RuntimeObservationDeliveryGuard::enter(None);
        if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (self.notify)(snapshot))).is_err() {
            eprintln!("[query-engine] warning: runtime observation seal observer panicked");
        }
    }
}

/// Deferred observer delivery for an already-sealed query cleanup.
/// The seal has published the report; delivery runs after capacity release
/// and never gates reuse.
pub(crate) struct QueryCleanupSealDelivery {
    pub(super) owner: Option<RuntimeObservationOwner>,
    pub(super) deferred: Option<DeferredRuntimeObservationSeal>,
    pub(super) already_sealed: Option<SealedRuntimeObservationSnapshot>,
    #[cfg(test)]
    pub(super) delivery_gate: Option<Arc<RuntimeObservationDeliveryGate>>,
}

impl QueryCleanupSealDelivery {
    pub(super) fn already_sealed(snapshot: Option<SealedRuntimeObservationSnapshot>) -> Self {
        Self {
            owner: None,
            deferred: None,
            already_sealed: snapshot,
            #[cfg(test)]
            delivery_gate: None,
        }
    }

    /// Holds observer delivery at `gate` until the test opens it.
    #[cfg(test)]
    pub(crate) fn with_delivery_gate(mut self, gate: Option<Arc<RuntimeObservationDeliveryGate>>) -> Self {
        self.delivery_gate = gate;
        self
    }

    pub(crate) fn deliver(self) -> Option<SealedRuntimeObservationSnapshot> {
        let Some(deferred) = self.deferred else { return self.already_sealed };
        if let Some(owner) = self.owner {
            return owner.complete_observer_delivery(deferred);
        }
        // The test gate sits before every observer; the seal has already
        // committed by this point.
        #[cfg(test)]
        if let Some(gate) = self.delivery_gate {
            gate.hold();
        }
        deferred.invoke_observers()
    }
}

pub(super) struct DeferredRuntimeObservationSeal {
    pub(super) snapshot: Option<SealedRuntimeObservationSnapshot>,
    pub(super) observers: Vec<RuntimeObservationSealObserver>,
}

impl DeferredRuntimeObservationSeal {
    pub(super) fn invoke_observers(self) -> Option<SealedRuntimeObservationSnapshot> {
        if let Some(snapshot) = &self.snapshot {
            for observer in self.observers {
                observer.invoke(snapshot.clone());
            }
        }
        self.snapshot
    }
}

impl fmt::Debug for RuntimeObservationSessionInner {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("RuntimeObservationSessionInner")
            .field("manifest", &self.manifest)
            .field("sealed", &self.sealed.load(Ordering::Acquire))
            .finish_non_exhaustive()
    }
}

impl RuntimeObservationSession {
    /// Creates an owner and recording handle using a system monotonic clock.
    pub fn new(query_id: impl Into<QueryId>, trace_id: Option<TraceId>, config: RuntimeObservationConfig) -> Self {
        Self { handle: RuntimeObservationHandle::new(query_id, trace_id, config) }
    }

    /// Creates an owner and recording handle using an injected clock.
    pub fn with_clock(
        query_id: impl Into<QueryId>,
        trace_id: Option<TraceId>,
        config: RuntimeObservationConfig,
        clock: Arc<dyn RuntimeObservationClock>,
    ) -> Self {
        Self { handle: RuntimeObservationHandle::with_clock(query_id, trace_id, config, clock) }
    }

    #[cfg(test)]
    pub(crate) fn with_source_location_test_key(
        query_id: impl Into<QueryId>,
        trace_id: Option<TraceId>,
        config: RuntimeObservationConfig,
        key: [u8; blake3::KEY_LEN],
    ) -> Self {
        Self {
            handle: RuntimeObservationHandle::with_clock_and_source_location_digest_key(
                query_id,
                trace_id,
                config,
                Arc::new(SystemRuntimeObservationClock::new()),
                &TestSourceLocationDigestKeyProvider { key: Some(key) },
            ),
        }
    }

    #[cfg(any(test, feature = "test-hooks"))]
    #[doc(hidden)]
    pub fn with_unavailable_source_location_test_key(
        query_id: impl Into<QueryId>,
        trace_id: Option<TraceId>,
        config: RuntimeObservationConfig,
    ) -> Self {
        Self {
            handle: RuntimeObservationHandle::with_clock_and_source_location_digest_key(
                query_id,
                trace_id,
                config,
                Arc::new(SystemRuntimeObservationClock::new()),
                &TestSourceLocationDigestKeyProvider { key: None },
            ),
        }
    }

    /// Creates a disabled owner and handle.
    pub const fn disabled() -> Self {
        Self { handle: RuntimeObservationHandle::disabled() }
    }

    /// Returns a cloneable recording capability without seal authority.
    pub fn handle(&self) -> RuntimeObservationHandle {
        self.handle.clone()
    }

    /// Returns the immutable recorder configuration.
    pub const fn config(&self) -> RuntimeObservationConfig {
        self.handle.config()
    }

    /// Returns the immutable query identity when recording is enabled.
    pub fn query_id(&self) -> Option<&QueryId> {
        self.handle.query_id()
    }

    /// Returns the immutable optional parent trace identity.
    pub fn trace_id(&self) -> Option<&TraceId> {
        self.handle.trace_id()
    }

    /// Lazily constructs and records a query-level event.
    pub fn record<F>(&self, build_payload: F) -> RuntimeObservationRecordDisposition
    where F: FnOnce() -> RuntimeObservationPayload {
        self.handle.record(build_payload)
    }

    /// Lazily constructs and records an event for one plan execution.
    pub fn record_for_plan<F>(
        &self,
        plan_execution_id: PlanExecutionId,
        build_payload: F,
    ) -> RuntimeObservationRecordDisposition
    where
        F: FnOnce() -> RuntimeObservationPayload,
    {
        self.handle.record_for_plan(plan_execution_id, build_payload)
    }

    /// Lazily constructs an event before acquiring recorder order and records it infallibly.
    pub fn record_with_context<F>(
        &self,
        context: RuntimeEventContext,
        build_payload: F,
    ) -> RuntimeObservationRecordDisposition
    where
        F: FnOnce() -> RuntimeObservationPayload,
    {
        self.handle.record_with_context(context, build_payload)
    }

    /// Lazily constructs and records one static topology row.
    pub fn record_topology<F>(
        &self,
        class: RuntimeObservationRecordClass,
        build_record: F,
    ) -> RuntimeObservationRecordDisposition
    where
        F: FnOnce() -> RuntimeTopologyRecord,
    {
        self.handle.record_topology(class, build_record)
    }

    /// Captures an immutable incomplete view. Returns none after sealing or when disabled.
    pub fn live_snapshot(&self) -> Option<LiveRuntimeObservationSnapshot> {
        self.handle.live_snapshot()
    }

    /// Fixes the first terminal outcome and returns the same immutable snapshot
    /// to every later or concurrent seal.
    pub fn seal(&self, outcome: RuntimeTerminalOutcome) -> Option<SealedRuntimeObservationSnapshot> {
        self.handle.seal_owned(outcome)
    }

    pub(super) fn seal_deferred(
        &self,
        outcome: Option<RuntimeTerminalOutcome>,
        attempt_report_provider: Option<&AttemptReportSealProvider>,
    ) -> DeferredRuntimeObservationSeal {
        self.handle.seal_owned_deferred(outcome, attempt_report_provider)
    }

    /// Seals this session with `outcome` and the report `report` builds,
    /// deferring observer delivery to the returned value. A seal that follows
    /// an earlier seal returns that seal's snapshot and delivers nothing.
    ///
    /// A panic while sealing marks the session degraded, seals it as
    /// `Panicked`, delivers its observers, and resumes unwinding.
    pub(crate) fn seal_with_report(
        &self,
        outcome: RuntimeTerminalOutcome,
        report: &AttemptReportSealProvider,
    ) -> QueryCleanupSealDelivery {
        let deferred =
            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.seal_deferred(Some(outcome), Some(report))));
        match deferred {
            Ok(deferred) => QueryCleanupSealDelivery {
                owner: None,
                deferred: Some(deferred),
                already_sealed: None,
                #[cfg(test)]
                delivery_gate: None,
            },
            Err(panic) => {
                self.handle.mark_degraded();
                let _ = self.seal_deferred(Some(RuntimeTerminalOutcome::Panicked), None).invoke_observers();
                std::panic::resume_unwind(panic);
            },
        }
    }

    /// Registers an infallible notification for this session's seal. It runs
    /// once, outside recorder locks: with the seal's deferred delivery, or
    /// immediately when the session has already sealed.
    pub(crate) fn register_seal_observer<F>(&self, observer: F)
    where F: Fn(SealedRuntimeObservationSnapshot) + Send + Sync + 'static {
        let observer = RuntimeObservationSealObserver { notify: Box::new(observer) };
        if let Some((observer, snapshot)) = self.handle.register_seal_observer(observer) {
            observer.invoke(snapshot);
        }
    }

    /// Returns the sealed immutable snapshot when a terminal outcome was already fixed.
    pub fn sealed_snapshot(&self) -> Option<SealedRuntimeObservationSnapshot> {
        self.handle.sealed_snapshot()
    }
}

impl RuntimeObservationSessionInner {
    pub(super) fn lock_state(&self) -> MutexGuard<'_, RuntimeObservationRecorderState> {
        match self.state.lock() {
            Ok(state) => state,
            Err(poisoned) => {
                let mut state = poisoned.into_inner();
                state.degraded = true;
                state.poison_recoveries = state.poison_recoveries.checked_add(1).unwrap_or_else(|| {
                    state.counter_overflow = true;
                    u64::MAX
                });
                self.state.clear_poison();
                state
            },
        }
    }

    pub(super) fn record_event_lazy<F>(
        &self,
        context: RuntimeEventContext,
        build_payload: F,
    ) -> RuntimeObservationRecordDisposition
    where
        F: FnOnce() -> RuntimeObservationPayload,
    {
        self.record_event_timestamped(context, |_monotonic_ns| build_payload()).0
    }

    #[cfg(test)]
    pub(super) fn record_event_sized<F>(
        &self,
        context: RuntimeEventContext,
        logical_bytes: u64,
        build_payload: F,
    ) -> RuntimeObservationRecordDisposition
    where
        F: FnOnce() -> RuntimeObservationPayload,
    {
        let mut state = self.lock_state();
        if state.sealed_snapshot.is_some() {
            return RuntimeObservationRecordDisposition::Sealed;
        }
        let sequence = state.take_sequence();
        if let Some(reason) = state.rejection_reason(self.manifest.limits, logical_bytes) {
            state.record_drop(RuntimeObservationRecordClass::TimelineEvent, sequence, reason);
            return RuntimeObservationRecordDisposition::Dropped { sequence, reason };
        }
        let monotonic_ns = self.clock.monotonic_ns().max(state.last_monotonic_ns);
        state.last_monotonic_ns = monotonic_ns;
        let event = RuntimeTimelineEvent {
            sequence,
            monotonic_ns,
            plan_execution_id: context.plan_execution_id,
            span_id: context.span_id,
            parent_span_id: context.parent_span_id,
            payload: build_payload(),
        };
        let actual_bytes = event.logical_retained_bytes();
        if actual_bytes != logical_bytes {
            state.degraded = true;
        }
        match state.admit(self.manifest.limits, RuntimeObservationRecordClass::TimelineEvent, sequence, actual_bytes) {
            Ok(()) => {
                state.events.push(event);
                RuntimeObservationRecordDisposition::Retained { sequence, logical_bytes: actual_bytes }
            },
            Err(reason) => RuntimeObservationRecordDisposition::Dropped { sequence, reason },
        }
    }

    pub(super) fn record_event_timestamped<F>(
        &self,
        context: RuntimeEventContext,
        build_payload: F,
    ) -> (RuntimeObservationRecordDisposition, Option<u64>)
    where
        F: FnOnce(u64) -> RuntimeObservationPayload,
    {
        let mut state = self.lock_state();
        if state.sealed_snapshot.is_some() {
            return (RuntimeObservationRecordDisposition::Sealed, None);
        }
        let sequence = state.take_sequence();
        let monotonic_ns = self.clock.monotonic_ns().max(state.last_monotonic_ns);
        state.last_monotonic_ns = monotonic_ns;
        const MIN_TIMELINE_EVENT_BYTES: u64 = 42;
        if let Some(reason) = state.rejection_reason(self.manifest.limits, MIN_TIMELINE_EVENT_BYTES) {
            state.record_drop(RuntimeObservationRecordClass::TimelineEvent, sequence, reason);
            return (RuntimeObservationRecordDisposition::Dropped { sequence, reason }, Some(monotonic_ns));
        }
        let event = RuntimeTimelineEvent {
            sequence,
            monotonic_ns,
            plan_execution_id: context.plan_execution_id,
            span_id: context.span_id,
            parent_span_id: context.parent_span_id,
            payload: build_payload(monotonic_ns),
        };
        let logical_bytes = event.logical_retained_bytes();
        let disposition = match state.admit(
            self.manifest.limits,
            RuntimeObservationRecordClass::TimelineEvent,
            sequence,
            logical_bytes,
        ) {
            Ok(()) => {
                state.events.push(event);
                RuntimeObservationRecordDisposition::Retained { sequence, logical_bytes }
            },
            Err(reason) => RuntimeObservationRecordDisposition::Dropped { sequence, reason },
        };
        (disposition, Some(monotonic_ns))
    }

    pub(super) fn record_topology_lazy<F>(
        &self,
        class: RuntimeObservationRecordClass,
        build_record: F,
    ) -> RuntimeObservationRecordDisposition
    where
        F: FnOnce() -> RuntimeTopologyRecord,
    {
        let mut state = self.lock_state();
        if state.sealed_snapshot.is_some() {
            return RuntimeObservationRecordDisposition::Sealed;
        }
        let sequence = state.take_sequence();
        const MIN_TOPOLOGY_RECORD_BYTES: u64 = 24;
        if let Some(reason) = state.rejection_reason(self.manifest.limits, MIN_TOPOLOGY_RECORD_BYTES) {
            state.record_drop(class, sequence, reason);
            return RuntimeObservationRecordDisposition::Dropped { sequence, reason };
        }
        let record = build_record();
        if record.class() != class {
            state.degraded = true;
        }
        let logical_bytes = record.logical_retained_bytes();
        match state.admit(self.manifest.limits, record.class(), sequence, logical_bytes) {
            Ok(()) => {
                state.topology.push(RetainedTopologyRecord { sequence, record });
                RuntimeObservationRecordDisposition::Retained { sequence, logical_bytes }
            },
            Err(reason) => RuntimeObservationRecordDisposition::Dropped { sequence, reason },
        }
    }

    pub(super) fn snapshot_data(
        &self,
        state: &RuntimeObservationRecorderState,
        complete: bool,
        terminal: Option<RuntimeTerminalOutcome>,
    ) -> RuntimeObservationSnapshotData {
        let mut topology = state.topology.clone();
        topology.sort_by(|left, right| left.record.cmp(&right.record).then(left.sequence.cmp(&right.sequence)));
        let mut events = state.events.clone();
        events.sort_by_key(|event| event.sequence);
        let dropped_by_class = state
            .dropped_by_class
            .iter()
            .map(|(class, count)| RuntimeObservationClassLoss { class: *class, count: *count })
            .collect();
        let dropped_by_reason = state
            .dropped_by_reason
            .iter()
            .map(|(reason, count)| RuntimeObservationReasonLoss { reason: *reason, count: *count })
            .collect();
        RuntimeObservationSnapshotData {
            manifest: self.manifest.clone(),
            topology,
            events,
            footer: RuntimeObservationFooter {
                complete,
                terminal,
                retained_records: state.retained_records,
                retained_bytes: state.retained_bytes,
                first_dropped_sequence: state.first_dropped_sequence,
                last_attempted_sequence: state.last_attempted_sequence,
                dropped_by_class,
                dropped_by_reason,
                truncated: state.first_dropped_sequence.is_some(),
                degraded: state.degraded,
                poison_recoveries: state.poison_recoveries,
                attempt_report: None,
                counter_overflow: state.counter_overflow,
            },
            plans: BTreeMap::new(),
        }
    }

    /// Waits until a seal in flight publishes its snapshot.
    pub(super) fn wait_for_seal<'a>(
        &self,
        state: MutexGuard<'a, RuntimeObservationRecorderState>,
    ) -> MutexGuard<'a, RuntimeObservationRecorderState> {
        self.seal_published.wait(state).unwrap_or_else(PoisonError::into_inner)
    }

    /// Takes the retained records and computes the complete footer. The
    /// footer's attempt report is filled in by the caller.
    pub(super) fn take_sealed_snapshot_data(
        &self,
        state: &mut RuntimeObservationRecorderState,
        outcome: Option<RuntimeTerminalOutcome>,
    ) -> RuntimeObservationSnapshotData {
        let mut topology = std::mem::take(&mut state.topology);
        topology.sort_by(|left, right| left.record.cmp(&right.record).then(left.sequence.cmp(&right.sequence)));
        let mut events = std::mem::take(&mut state.events);
        events.sort_by_key(|event| event.sequence);
        let dropped_by_class = state
            .dropped_by_class
            .iter()
            .map(|(class, count)| RuntimeObservationClassLoss { class: *class, count: *count })
            .collect();
        let dropped_by_reason = state
            .dropped_by_reason
            .iter()
            .map(|(reason, count)| RuntimeObservationReasonLoss { reason: *reason, count: *count })
            .collect();
        let footer = RuntimeObservationFooter {
            complete: true,
            terminal: outcome,
            retained_records: state.retained_records,
            retained_bytes: state.retained_bytes,
            first_dropped_sequence: state.first_dropped_sequence,
            last_attempted_sequence: state.last_attempted_sequence,
            dropped_by_class,
            dropped_by_reason,
            truncated: state.first_dropped_sequence.is_some(),
            degraded: state.degraded,
            poison_recoveries: state.poison_recoveries,
            attempt_report: None,
            counter_overflow: state.counter_overflow,
        };
        RuntimeObservationSnapshotData {
            manifest: self.manifest.clone(),
            topology,
            events,
            footer,
            plans: BTreeMap::new(),
        }
    }
}