spiral-sdk 0.1.11

//! Retained execution observations for SDK streams.

use std::collections::BTreeMap;
use std::sync::Arc;

use arrow_schema::Schema;
use parking_lot::Mutex;

use crate::PlanDescription;
use crate::scandal::ScandalError;
use crate::scandal::plan::PLAN_METADATA_KEY;
use crate::scandal::profile::PROFILE_METADATA_KEY;

/// Latest plan/profile snapshots observed by one stream, retained after failure or cancellation.
/// At most 64 reports and 16 MiB of their encoded descriptions are retained per source;
/// excess observations are counted as omitted.
#[derive(Clone, Default)]
pub struct ExecutionDiagnostics(Arc<Mutex<Reports>>);

#[derive(Default)]
struct Reports {
    plans: BTreeMap<String, PlanDescription>,
    omitted: u64,
    bytes: usize,
    sizes: BTreeMap<String, usize>,
    source: Option<ExecutionDiagnostics>,
}

impl ExecutionDiagnostics {
    /// Snapshots keyed by their execution IDs, without reading storage or polling execution.
    pub fn plans(&self) -> Vec<PlanDescription> {
        let (mut plans, source) = {
            let reports = self.0.lock();
            (reports.plans.clone(), reports.source.clone())
        };
        if let Some(source) = source {
            for plan in source.plans() {
                if let Some(id) = &plan.execution_id
                    && (plans.len() < 64 || plans.contains_key(id))
                {
                    plans.insert(id.clone(), plan);
                }
            }
        }
        plans.into_values().collect()
    }
    /// Number of observations omitted at a retention bound.
    pub fn omitted(&self) -> u64 {
        self.0.lock().omitted
    }
    pub(crate) fn bind(&self, source: Self) {
        if !Arc::ptr_eq(&self.0, &source.0) {
            self.0.lock().source = Some(source);
        }
    }
    pub(crate) fn observe(&self, schema: &Schema) -> Result<bool, ScandalError> {
        if !schema.metadata().contains_key(PROFILE_METADATA_KEY) {
            return Ok(false);
        }
        let Some(plan) = PlanDescription::from_schema(schema)? else {
            return Ok(false);
        };
        if plan.profile.is_none() {
            return Ok(false);
        }
        let Some(id) = plan.execution_id.clone() else {
            return Ok(false);
        };
        let size = schema
            .metadata()
            .get(PLAN_METADATA_KEY)
            .map_or(0, String::len);
        let mut reports = self.0.lock();
        let previous = reports.sizes.get(&id).copied().unwrap_or(0);
        let total = reports.bytes.saturating_sub(previous).saturating_add(size);
        if (reports.plans.len() < 64 || reports.plans.contains_key(&id))
            && total <= 16 * 1024 * 1024
        {
            reports.bytes = total;
            reports.sizes.insert(id.clone(), size);
            reports.plans.insert(id, plan);
        } else {
            reports.omitted = reports.omitted.saturating_add(1);
        }
        Ok(true)
    }
    pub(crate) fn cancelled(&self) {
        let source = self.0.lock().source.clone();
        if let Some(source) = source {
            source.cancelled();
        }
        for plan in self.0.lock().plans.values_mut() {
            if let Some(profile) = &mut plan.profile
                && profile.status == "running"
            {
                profile.consumer_cancelled = true;
            }
        }
    }
}