tenshift_core/pipeline/

iterator.rs

//! Iterator types for consuming pipeline output.
//!
//! This module is the consumer-facing endpoint of the pipeline architecture,
//! turning collector output into a pull-based iterator with stats and timeout
//! helpers.

#![allow(clippy::module_name_repetitions)]

use super::get_fatal_error;
use crate::sample::Sample;
use crossbeam_channel::Receiver;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// A running pipeline that yields batches of samples.
pub struct PipelineIterator {
    pub(crate) receiver: Receiver<Vec<Sample>>,
    pub(crate) shutdown: Arc<AtomicBool>,
    #[allow(dead_code)]
    pub(crate) workers: Vec<std::thread::JoinHandle<()>>,
    pub(crate) items_yielded: u64,
    pub(crate) errors_skipped: Arc<AtomicU64>,
    pub(crate) started_at: Instant,
    pub(crate) fatal_error: Arc<std::sync::Mutex<Option<crate::error::Error>>>,
}

/// Error returned when waiting for the next batch with a timeout.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum NextTimeoutError {
    /// The timeout elapsed before a batch was ready.
    Timeout,
    /// The pipeline finished and disconnected.
    Disconnected,
}

impl PipelineIterator {
    /// Number of items yielded so far.
    pub fn items_yielded(&self) -> u64 {
        self.items_yielded
    }

    /// Number of errors that were skipped.
    pub fn errors_skipped(&self) -> u64 {
        self.errors_skipped.load(Ordering::Relaxed)
    }

    /// Wall-clock time since the pipeline started.
    pub fn elapsed(&self) -> Duration {
        self.started_at.elapsed()
    }

    /// Current throughput in items per second.
    ///
    /// Returns `0.0` if no items have been yielded yet.
    #[allow(clippy::cast_precision_loss)] // throughput metrics don't need sub-ulp precision
    pub fn throughput(&self) -> f64 {
        let elapsed = self.started_at.elapsed().as_secs_f64();
        if elapsed <= 0.0 {
            return 0.0;
        }
        self.items_yielded as f64 / elapsed
    }

    /// Snapshot of pipeline statistics for observability.
    #[allow(clippy::cast_precision_loss)] // throughput metrics don't need sub-ulp precision
    pub fn error(&self) -> Option<crate::error::Error> {
        get_fatal_error(&self.fatal_error)
    }

    /// Snapshot of pipeline statistics for observability.
    #[allow(clippy::cast_precision_loss)] // throughput metrics don't need sub-ulp precision
    pub fn stats(&self) -> PipelineStats {
        let elapsed = self.started_at.elapsed();
        let elapsed_secs = elapsed.as_secs_f64();
        PipelineStats {
            items_yielded: self.items_yielded,
            errors_skipped: self.errors_skipped.load(Ordering::Relaxed),
            elapsed,
            throughput: if elapsed_secs > 0.0 {
                self.items_yielded as f64 / elapsed_secs
            } else {
                0.0
            },
        }
    }

    /// Signal the pipeline to stop.
    pub fn stop(&self) {
        self.shutdown.store(true, Ordering::Relaxed);
    }

    /// Block for the next batch with a timeout.
    ///
    /// Useful for external bindings (e.g. Python `PyO3`) that must periodically
    /// wake up to check for signals like Ctrl+C.
    ///
    /// # Errors
    ///
    /// Returns [`NextTimeoutError::Timeout`] if the timeout elapses.
    /// Returns [`NextTimeoutError::Disconnected`] if the pipeline finishes.
    pub fn next_timeout(
        &mut self,
        timeout: Duration,
    ) -> std::result::Result<Vec<Sample>, NextTimeoutError> {
        let deadline = std::time::Instant::now() + timeout;
        loop {
            let remaining = deadline.saturating_duration_since(std::time::Instant::now());
            if remaining.is_zero() {
                return Err(NextTimeoutError::Timeout);
            }
            match self.receiver.recv_timeout(remaining) {
                Ok(batch) if batch.is_empty() => {
                    // Empty batches are control signals (e.g., error indicators).
                    // Continue waiting, but with the original deadline.
                }
                Ok(batch) => {
                    self.items_yielded += 1;
                    return Ok(batch);
                }
                Err(crossbeam_channel::RecvTimeoutError::Timeout) => {
                    return Err(NextTimeoutError::Timeout)
                }
                Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
                    return Err(NextTimeoutError::Disconnected);
                }
            }
        }
    }
}

impl Iterator for PipelineIterator {
    type Item = Vec<Sample>;

    // Fail-closed by design: a fatal worker error panics instead of silently
    // ending the epoch (see the `Err(_)` arm below). The fallible inherent
    // `next()` and `.error()` APIs are the non-panicking alternatives.
    #[allow(clippy::panic)]
    fn next(&mut self) -> Option<Vec<Sample>> {
        loop {
            match self.receiver.recv() {
                Ok(batch) if batch.is_empty() => {}
                Ok(batch) => {
                    self.items_yielded += 1;
                    return Some(batch);
                }
                // The channel disconnected. A clean end (all workers finished with
                // no error) yields `None`. But if a worker died on a FATAL error it
                // also disconnects the channel, and returning `None` there would let
                // a `for batch in pipeline {}` / `.collect()` consumer finish exactly
                // as on a clean epoch — silently training on a TRUNCATED epoch with
                // no signal (Law 10). Fail closed: surface the captured error so it
                // is impossible to miss. Consumers that need to recover use the
                // fallible inherent `next()` (returns a `Result`) or `.error()`.
                Err(_) => {
                    if let Some(error) = self.error() {
                        panic!(
                            "tenshift pipeline terminated early on a fatal worker error after {} batch(es): {error}. \
                             The epoch is TRUNCATED and must not be treated as a clean end; use the fallible `next()` \
                             (Result) API or check `.error()` to handle this without panicking.",
                            self.items_yielded
                        );
                    }
                    return None;
                }
            }
        }
    }
}

impl Drop for PipelineIterator {
    fn drop(&mut self) {
        let stats = self.stats();
        // A pipeline that ended on a fatal worker error must not drop quietly at
        // debug level — an operator watching at the default log level would see
        // nothing wrong even though the epoch was truncated (Law 10). Surface it
        // at error level, distinct from the clean-completion debug stats.
        if let Some(error) = self.error() {
            tracing::error!(%error, "{stats} — pipeline terminated on a FATAL worker error; the epoch was truncated");
        } else {
            tracing::debug!("{stats}");
        }

        self.shutdown.store(true, Ordering::Relaxed);
        // Clear the receiver safely to unblock any pending channels.
        while self.receiver.try_recv().is_ok() {}

        // SQLite Fix: We explicitly do NOT `join()` the worker handles here.
        // `self.receiver` doesn't drop until after this scope completes.
        // Calling `join()` while channels are still alive causes indefinite
        // cyclic deadlocks if the threads were halted on a full `out_tx.send()`.
        // By relying on the structural crossbeam drop propagation, dropping `self`
        // cleanly shuts down all senders upstream automatically.
    }
}

/// Observability snapshot from a running or completed pipeline.
#[derive(Debug, Clone)]
pub struct PipelineStats {
    /// Total items (batches) yielded to the consumer.
    pub items_yielded: u64,
    /// Total items skipped due to errors.
    pub errors_skipped: u64,
    /// Wall-clock time since pipeline started.
    pub elapsed: Duration,
    /// Throughput in items per second.
    pub throughput: f64,
}

impl std::fmt::Display for PipelineStats {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "tenshift: {} items in {:.2}s ({:.0} items/s, {} errors skipped)",
            self.items_yielded,
            self.elapsed.as_secs_f64(),
            self.throughput,
            self.errors_skipped,
        )
    }
}