kinetic/transforms/sample/

task.rs

//! Sampling transform task.
//!
//! Wires the `kinetic-sampling` library into the Kinetic pipeline topology.
//! Handles channel I/O, named outputs (`sampled` and `passthrough`), metrics
//! emission, error policy, and shutdown.

use crate::transforms::Transform;
use crate::transforms::util::TransformErrorHandler;
use async_trait::async_trait;
use kinetic_buffers::{BufferReceiver, BufferSender};
use kinetic_config::model::{ErrorPolicy, SamplingMode};
use kinetic_core::EventBatch;
use kinetic_sampling::{
    adaptive::{AdaptiveConfig, AdaptiveSampler},
    hash::{HashConfig, MissingKeyBehaviour},
    random::{RandomConfig, build_rng},
    recordbatch::RecordBatchConfig,
    reservoir::{ReservoirConfig, ReservoirSampler},
    stratified::{MissingFieldBehaviour, StratifiedConfig},
    systematic::{SystematicConfig, SystematicSampler},
    tail::{Decision, Operator, TailConfig, TailRule, TailSampler},
};
use metrics::{Label, counter, gauge};
use rand::SeedableRng;
use rand::rngs::StdRng;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::{self, Interval};
use tracing::{debug, error, info, warn};

use kinetic_common::register;
use kinetic_common::telemetry::EventDuration;
use std::time::Instant;

/// Sampling transform task.
pub struct SampleTask {
    receiver: BufferReceiver,
    senders: HashMap<String, BufferSender>,
    mode: SamplingMode,
    route_all: bool,
    component_id: String,
    labels: Arc<[Label]>,
    error_handler: TransformErrorHandler,
    event_duration: EventDuration,
}

#[async_trait]
impl Transform for SampleTask {
    async fn run(self: Box<Self>) {
        self.run_task().await;
    }
}

impl SampleTask {
    pub fn new(
        component_id: String,
        receiver: BufferReceiver,
        senders: HashMap<String, BufferSender>,
        mode: SamplingMode,
        error_policy: ErrorPolicy,
        route_all: bool,
    ) -> Self {
        let labels: Arc<[Label]> = Arc::new([
            Label::new("component_id", component_id.clone()),
            Label::new("component_type", "transform"),
            Label::new("component_kind", "sample"),
        ]);
        let error_handler = TransformErrorHandler::new(component_id.clone(), error_policy);
        let event_duration = register!(EventDuration::new(component_id.clone(), "transform"));

        Self {
            receiver,
            senders,
            mode,
            route_all,
            component_id,
            labels,
            error_handler,
            event_duration,
        }
    }

    pub async fn run_task(mut self) {
        info!("Starting Sample transform task: {}", self.component_id);

        match &self.mode {
            SamplingMode::Random { rate, seed } => {
                let config = RandomConfig {
                    rate: *rate,
                    seed: *seed,
                };
                self.run_random(config).await;
            }
            SamplingMode::Systematic { interval } => {
                match SystematicSampler::new(SystematicConfig {
                    interval: *interval,
                }) {
                    Ok(sampler) => self.run_systematic(sampler).await,
                    Err(e) => error!(
                        "Sample {}: invalid systematic config: {}",
                        self.component_id, e
                    ),
                }
            }
            SamplingMode::Hash {
                key,
                rate,
                missing_key_behaviour,
            } => {
                let behaviour = if missing_key_behaviour == "keep" {
                    MissingKeyBehaviour::Keep
                } else {
                    MissingKeyBehaviour::Drop
                };
                let config = HashConfig {
                    key: key.clone(),
                    rate: *rate,
                    missing_key_behaviour: behaviour,
                };
                self.run_hash(config).await;
            }
            SamplingMode::Recordbatch { rate, seed } => {
                let config = RecordBatchConfig {
                    rate: *rate,
                    seed: *seed,
                };
                self.run_recordbatch(config).await;
            }
            SamplingMode::Stratified {
                field,
                strata,
                default_rate,
                missing_field_behaviour,
                seed,
            } => {
                let behaviour = if missing_field_behaviour == "drop" {
                    MissingFieldBehaviour::Drop
                } else {
                    MissingFieldBehaviour::Keep
                };
                let config = StratifiedConfig {
                    field: field.clone(),
                    strata: strata.clone(),
                    default_rate: *default_rate,
                    missing_field_behaviour: behaviour,
                    seed: *seed,
                };
                self.run_stratified(config, *seed).await;
            }
            SamplingMode::Adaptive {
                base_rate,
                min_rate,
                max_rate,
                error_field,
                error_value,
                latency_field,
                latency_threshold_ms,
                window_secs,
                sensitivity,
            } => {
                let config = AdaptiveConfig {
                    base_rate: *base_rate,
                    min_rate: *min_rate,
                    max_rate: *max_rate,
                    error_field: error_field.clone(),
                    error_value: error_value.clone(),
                    latency_field: latency_field.clone(),
                    latency_threshold_ms: *latency_threshold_ms,
                    window: Duration::from_secs(*window_secs),
                    sensitivity: *sensitivity,
                    seed: None,
                };
                match AdaptiveSampler::new(config) {
                    Ok(sampler) => {
                        self.run_adaptive(sampler).await;
                    }
                    Err(e) => error!(
                        "Sample {}: invalid adaptive config: {}",
                        self.component_id, e
                    ),
                }
            }
            SamplingMode::Reservoir {
                size,
                flush_interval_secs,
            } => {
                let config = ReservoirConfig {
                    size: *size,
                    flush_interval_secs: *flush_interval_secs,
                };
                match ReservoirSampler::new(config) {
                    Ok(sampler) => {
                        self.run_reservoir(sampler).await;
                    }
                    Err(e) => error!(
                        "Sample {}: invalid reservoir config: {}",
                        self.component_id, e
                    ),
                }
            }
            SamplingMode::Tail {
                group_key,
                decision_window_secs,
                max_groups,
                default_decision,
                rules,
            } => {
                let tail_rules: Vec<TailRule> = rules
                    .iter()
                    .map(|r| TailRule {
                        field: r.field.clone(),
                        operator: parse_operator(&r.operator),
                        value: r.value.clone(),
                        decision: if r.decision == "keep" {
                            Decision::Keep
                        } else {
                            Decision::Drop
                        },
                    })
                    .collect();
                let config = TailConfig {
                    group_key: group_key.clone(),
                    decision_window_secs: *decision_window_secs,
                    max_groups: *max_groups,
                    default_decision: if default_decision == "keep" {
                        Decision::Keep
                    } else {
                        Decision::Drop
                    },
                    rules: tail_rules,
                };
                match TailSampler::new(config) {
                    Ok(sampler) => self.run_tail(sampler).await,
                    Err(e) => error!("Sample {}: invalid tail config: {}", self.component_id, e),
                }
            }
        }

        info!("Sample transform task {} shutting down", self.component_id);
    }

    async fn run_random(&mut self, config: RandomConfig) {
        let mut rng = build_rng(config.seed);
        while let Some(mut batch) = self.receiver.recv().await {
            let start = Instant::now();
            let received_rows = batch.num_rows();
            self.inc_received(received_rows, batch.estimated_size());

            match kinetic_sampling::random::sample(
                &batch.payload,
                &config,
                self.route_all,
                &mut *rng,
            ) {
                Ok(result) => {
                    self.emit_sample_rate(&config.rate.to_string());
                    self.send_sampled(result.selected, received_rows, &mut batch)
                        .await;
                    if let Some(rejected) = result.rejected {
                        self.send_passthrough(rejected).await;
                    }
                }
                Err(e) => {
                    self.handle_batch_error(format!("random sample error: {e}"), &mut batch)
                        .await;
                }
            }
            self.event_duration.emit(start.elapsed());
        }
    }

    async fn run_systematic(&mut self, sampler: SystematicSampler) {
        while let Some(mut batch) = self.receiver.recv().await {
            let start = Instant::now();
            let received_rows = batch.num_rows();
            self.inc_received(received_rows, batch.estimated_size());
            let rate = 1.0 / sampler.config.interval as f64;

            match sampler.sample(&batch.payload, self.route_all) {
                Ok(result) => {
                    self.emit_sample_rate(&rate.to_string());
                    self.send_sampled(result.selected, received_rows, &mut batch)
                        .await;
                    if let Some(rejected) = result.rejected {
                        self.send_passthrough(rejected).await;
                    }
                }
                Err(e) => {
                    self.handle_batch_error(format!("systematic sample error: {e}"), &mut batch)
                        .await;
                }
            }
            self.event_duration.emit(start.elapsed());
        }
    }

    async fn run_hash(&mut self, config: HashConfig) {
        while let Some(mut batch) = self.receiver.recv().await {
            let start = Instant::now();
            let received_rows = batch.num_rows();
            self.inc_received(received_rows, batch.estimated_size());

            match kinetic_sampling::hash::sample(&batch.payload, &config, self.route_all) {
                Ok(result) => {
                    let rate = config.rate as f64 / 100.0;
                    self.emit_sample_rate(&rate.to_string());
                    self.send_sampled(result.selected, received_rows, &mut batch)
                        .await;
                    if let Some(rejected) = result.rejected {
                        self.send_passthrough(rejected).await;
                    }
                }
                Err(e) => {
                    self.handle_batch_error(format!("hash sample error: {e}"), &mut batch)
                        .await;
                }
            }
            self.event_duration.emit(start.elapsed());
        }
    }

    async fn run_recordbatch(&mut self, config: RecordBatchConfig) {
        let mut rng = build_rng(config.seed);
        while let Some(mut batch) = self.receiver.recv().await {
            let start = Instant::now();
            let received_rows = batch.num_rows();
            self.inc_received(received_rows, batch.estimated_size());

            let payload = batch.payload.clone();
            match kinetic_sampling::recordbatch::sample(payload, &config, self.route_all, &mut *rng)
            {
                Ok(result) => {
                    self.emit_sample_rate(&config.rate.to_string());
                    self.send_sampled(result.selected, received_rows, &mut batch)
                        .await;
                    if let Some(rejected) = result.rejected {
                        self.send_passthrough(rejected).await;
                    }
                }
                Err(e) => {
                    self.handle_batch_error(format!("recordbatch sample error: {e}"), &mut batch)
                        .await;
                }
            }
            self.event_duration.emit(start.elapsed());
        }
    }

    async fn run_stratified(&mut self, config: StratifiedConfig, seed: Option<u64>) {
        let mut rng = build_rng(seed);
        while let Some(mut batch) = self.receiver.recv().await {
            let start = Instant::now();
            let received_rows = batch.num_rows();
            self.inc_received(received_rows, batch.estimated_size());

            match kinetic_sampling::stratified::sample(
                &batch.payload,
                &config,
                self.route_all,
                &mut *rng,
            ) {
                Ok(result) => {
                    self.emit_sample_rate(&config.default_rate.to_string());
                    self.send_sampled(result.selected, received_rows, &mut batch)
                        .await;
                    if let Some(rejected) = result.rejected {
                        self.send_passthrough(rejected).await;
                    }
                }
                Err(e) => {
                    self.handle_batch_error(format!("stratified sample error: {e}"), &mut batch)
                        .await;
                }
            }
            self.event_duration.emit(start.elapsed());
        }
    }

    async fn run_adaptive(&mut self, sampler: AdaptiveSampler) {
        let mut rng = StdRng::from_entropy();
        while let Some(mut batch) = self.receiver.recv().await {
            let start = Instant::now();
            let received_rows = batch.num_rows();
            self.inc_received(received_rows, batch.estimated_size());

            match sampler.sample(&batch.payload, self.route_all, &mut rng) {
                Ok(result) => {
                    let rate = sampler.current_rate();
                    self.emit_sample_rate(&rate.to_string());
                    self.send_sampled(result.selected, received_rows, &mut batch)
                        .await;
                    if let Some(rejected) = result.rejected {
                        self.send_passthrough(rejected).await;
                    }
                }
                Err(e) => {
                    self.handle_batch_error(format!("adaptive sample error: {e}"), &mut batch)
                        .await;
                }
            }
            self.event_duration.emit(start.elapsed());
        }
    }

    async fn run_reservoir(&mut self, mut sampler: ReservoirSampler) {
        let mut rng = StdRng::from_entropy();
        let flush_secs = sampler.config.flush_interval_secs;
        let mut flush_ticker: Interval = time::interval(Duration::from_secs(flush_secs.max(1)));

        loop {
            tokio::select! {
                maybe_batch = self.receiver.recv() => {
                    match maybe_batch {
                        Some(batch) => {
                            let start = Instant::now();
                            let received_rows = batch.num_rows();
                            self.inc_received(received_rows, batch.estimated_size());
                            sampler.observe(&batch.payload, &mut rng);
                            let reservoir_len = sampler.len();
                            gauge!("sample_reservoir_size", self.labels.iter())
                                .set(reservoir_len as f64);
                            self.event_duration.emit(start.elapsed());
                        }
                        None => break,
                    }
                }
                _ = flush_ticker.tick() => {
                    let start = Instant::now();
                    self.flush_reservoir(&mut sampler).await;
                    self.event_duration.emit(start.elapsed());
                }
            }
        }
        self.flush_reservoir(&mut sampler).await;
    }

    async fn flush_reservoir(&self, sampler: &mut ReservoirSampler) {
        match sampler.flush() {
            Ok(Some(batch)) => {
                let rows = batch.num_rows();
                if let Some(sender) = self.senders.get("sampled") {
                    match EventBatch::new(batch, self.make_placeholder_metadata()) {
                        Ok(eb) => {
                            let bytes = eb.estimated_size();
                            match sender.send(eb).await {
                                Ok(_) => self.inc_sent(rows, bytes, "sampled"),
                                Err(e) => error!(
                                    "Sample {} reservoir flush send error: {:?}",
                                    self.component_id, e
                                ),
                            }
                        }
                        Err(e) => error!(
                            "Sample {} failed to create EventBatch: {}",
                            self.component_id, e
                        ),
                    }
                }
            }
            Ok(None) => {}
            Err(e) => error!("Sample {} reservoir flush error: {}", self.component_id, e),
        }
    }

    async fn run_tail(&mut self, mut sampler: TailSampler) {
        let drain_interval = Duration::from_secs(1);
        let mut drain_ticker = time::interval(drain_interval);

        loop {
            tokio::select! {
                maybe_batch = self.receiver.recv() => {
                    match maybe_batch {
                        Some(batch) => {
                            let start = Instant::now();
                            let received_rows = batch.num_rows();
                            self.inc_received(received_rows, batch.estimated_size());
                            let evicted = sampler.observe(&batch.payload);
                            for batch in evicted {
                                let rows = batch.num_rows();
                                self.emit_tail_gauges(&sampler);
                                if let Some(sender) = self.senders.get("sampled") {
                                    counter!("sample_tail_evictions_total", self.labels.iter()).increment(1);
                                    self.send_arrow_batch_to(batch, rows, sender).await;
                                }
                            }
                            self.event_duration.emit(start.elapsed());
                        }
                        None => break,
                    }
                }
                _ = drain_ticker.tick() => {
                    let start = Instant::now();
                    let (kept, _dropped) = sampler.drain_expired();
                    for batch in kept {
                        let rows = batch.num_rows();
                        self.emit_tail_gauges(&sampler);
                        if let Some(sender) = self.senders.get("sampled") {
                            self.send_arrow_batch_to(batch, rows, sender).await;
                        }
                    }
                    self.event_duration.emit(start.elapsed());
                }
            }
        }
        let (kept, _) = sampler.flush_all();
        for batch in kept {
            let rows = batch.num_rows();
            if let Some(sender) = self.senders.get("sampled") {
                self.send_arrow_batch_to(batch, rows, sender).await;
            }
        }
    }

    fn emit_tail_gauges(&self, sampler: &TailSampler) {
        gauge!("sample_tail_buffer_groups", self.labels.iter()).set(sampler.group_count() as f64);
        gauge!("sample_tail_buffer_events", self.labels.iter())
            .set(sampler.buffered_event_count() as f64);
    }

    async fn send_sampled(
        &self,
        payload: arrow_array::RecordBatch,
        received_rows: usize,
        original: &mut EventBatch,
    ) {
        let sent_rows = payload.num_rows();
        let dropped_rows = received_rows.saturating_sub(sent_rows);

        if dropped_rows > 0 {
            counter!("component_discarded_events_total", self.labels.iter())
                .increment(dropped_rows as u64);
            let drop_labels = [
                Label::new("component_id", self.component_id.clone()),
                Label::new("component_type", "transform"),
                Label::new("component_kind", "sample"),
                Label::new("decision", "dropped"),
            ];
            counter!("sample_decisions_total", drop_labels.iter()).increment(dropped_rows as u64);
        }

        if sent_rows == 0 {
            if let Some(token) = original.ack_token.take() {
                token.ack();
            }
            return;
        }

        let keep_labels = [
            Label::new("component_id", self.component_id.clone()),
            Label::new("component_type", "transform"),
            Label::new("component_kind", "sample"),
            Label::new("decision", "kept"),
        ];
        counter!("sample_decisions_total", keep_labels.iter()).increment(sent_rows as u64);

        let new_batch = match EventBatch::new(payload, original.metadata.clone()) {
            Ok(mut b) => {
                b.ack_token = original.ack_token.take();
                b
            }
            Err(e) => {
                error!(
                    "Sample {} failed to create EventBatch: {}",
                    self.component_id, e
                );
                return;
            }
        };

        if let Some(sender) = self.senders.get("sampled") {
            let bytes = new_batch.estimated_size();
            match sender.send(new_batch).await {
                Ok(_) => self.inc_sent(sent_rows, bytes, "sampled"),
                Err(e) => error!(
                    "Sample {} failed to send to sampled output: {:?}",
                    self.component_id, e
                ),
            }
        } else {
            debug!(
                "Sample {}: no 'sampled' sender wired — events dropped",
                self.component_id
            );
        }
    }

    async fn send_passthrough(&self, payload: arrow_array::RecordBatch) {
        let rows = payload.num_rows();
        if rows == 0 {
            return;
        }
        if let Some(sender) = self.senders.get("passthrough")
            && let Ok(batch) = EventBatch::new(payload, self.make_placeholder_metadata())
        {
            let bytes = batch.estimated_size();
            match sender.send(batch).await {
                Ok(_) => self.inc_sent(rows, bytes, "passthrough"),
                Err(e) => warn!(
                    "Sample {} failed to send to passthrough output: {:?}",
                    self.component_id, e
                ),
            }
        }
    }

    async fn send_arrow_batch_to(
        &self,
        payload: arrow_array::RecordBatch,
        rows: usize,
        sender: &BufferSender,
    ) {
        if let Ok(eb) = EventBatch::new(payload, self.make_placeholder_metadata()) {
            let bytes = eb.estimated_size();
            match sender.send(eb).await {
                Ok(_) => self.inc_sent(rows, bytes, "sampled"),
                Err(e) => error!("Sample {} send error: {:?}", self.component_id, e),
            }
        }
    }

    async fn handle_batch_error(&self, message: String, batch: &mut EventBatch) {
        counter!("component_errors_total", self.labels.iter()).increment(1);
        self.error_handler
            .handle_error(&self.senders, message, Some(batch))
            .await;
    }

    fn inc_received(&self, rows: usize, bytes: usize) {
        counter!("component_received_events_total", self.labels.iter()).increment(rows as u64);
        counter!("component_received_event_bytes_total", self.labels.iter())
            .increment(bytes as u64);
    }

    fn inc_sent(&self, rows: usize, bytes: usize, output: &str) {
        let labels = [
            Label::new("component_id", self.component_id.clone()),
            Label::new("component_type", "transform"),
            Label::new("component_kind", "sample"),
            Label::new("output", output.to_string()),
        ];
        counter!("component_sent_events_total", labels.iter()).increment(rows as u64);
        counter!("component_sent_event_bytes_total", labels.iter()).increment(bytes as u64);
    }

    fn emit_sample_rate(&self, rate: &str) {
        if let Ok(r) = rate.parse::<f64>() {
            gauge!("sample_rate", self.labels.iter()).set(r);
        }
    }

    fn make_placeholder_metadata(&self) -> kinetic_core::ArcEventMetadata {
        use kinetic_core::{ComponentId, EventMetadata};
        Arc::new(EventMetadata::new(
            "unknown",
            ComponentId(self.component_id.clone()),
        ))
    }
}

fn parse_operator(s: &str) -> Operator {
    match s {
        "==" => Operator::Eq,
        "!=" => Operator::Ne,
        ">" => Operator::Gt,
        ">=" => Operator::Gte,
        "<" => Operator::Lt,
        "<=" => Operator::Lte,
        "contains" => Operator::Contains,
        other => {
            warn!("Unknown tail rule operator '{}', defaulting to ==", other);
            Operator::Eq
        }
    }
}