tenshift-core 0.1.3

//! Worker-thread transform execution for tenshift.
//!
//! Worker threads are the parallel middle of the architecture, applying
//! stateless transforms to source samples before ordered stages run in the
//! collector.

#![allow(clippy::module_name_repetitions)]

use super::{panic_message, record_first_fatal_error};
use crate::error::{Error, Result};
use crate::pipeline::ErrorPolicy;
use crate::pipeline::SampleChunk;
use crate::sample::Sample;
use crate::transform::{Transform, TransformResult};
use crossbeam_channel::{Receiver, Sender};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::thread::JoinHandle;

/// Apply stateless transforms to a sample.
///
/// When `error_policy` is `Skip`, transform errors and panics do NOT abort
/// processing; instead, the failing item is dropped and processing continues
/// with the remaining items in the buffer. This prevents data loss where
/// a single bad sample would cause the entire buffer to be dropped.
pub(crate) fn apply_stateless(
    transforms: &[Box<dyn Transform>],
    sample: Sample,
    current: &mut Vec<Sample>,
    next: &mut Vec<Sample>,
    error_policy: ErrorPolicy,
) -> Result<u64> {
    let mut errors: u64 = 0;
    current.clear();
    current.push(sample);

    for transform in transforms {
        next.clear();
        for item in current.drain(..) {
            let index = item.metadata().map_or(0, |metadata| metadata.index);
            let outcome = catch_unwind(AssertUnwindSafe(|| transform.apply(item)));
            match outcome {
                Ok(TransformResult::Sample(s)) => next.push(s),
                Ok(TransformResult::Samples(many)) => next.extend(many),
                Ok(TransformResult::Skip) => {}
                Ok(TransformResult::Error(error)) => {
                    if error_policy == ErrorPolicy::Fail {
                        return Err(error);
                    }
                    // Skip policy: count error and continue with remaining items
                    errors += 1;
                    tracing::warn!(
                        "skipping failed transform item at index {}: {}",
                        index,
                        error
                    );
                }
                Err(payload) => {
                    let error = Error::TransformFailed {
                        index,
                        reason: format!(
                            "transform '{}' panicked: {}",
                            transform.name(),
                            panic_message(payload)
                        ),
                    };
                    if error_policy == ErrorPolicy::Fail {
                        return Err(error);
                    }
                    // Skip policy: count error and continue with remaining items
                    errors += 1;
                    tracing::warn!(
                        "skipping panicked transform item at index {}: {}",
                        index,
                        error
                    );
                }
            }
        }
        std::mem::swap(current, next);
        if current.is_empty() {
            break;
        }
    }

    Ok(errors)
}


#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
pub(crate) fn spawn_worker_threads(
    num_workers: usize,
    pin_threads: bool,
    error_policy: ErrorPolicy,
    raw_rx: Receiver<SampleChunk>,
    proc_tx: Sender<SampleChunk>,
    transforms: Arc<Vec<Box<dyn Transform>>>,
    shutdown: Arc<AtomicBool>,
    errors_skipped: Arc<AtomicU64>,
    fatal_error: Arc<std::sync::Mutex<Option<crate::error::Error>>>,
) -> Result<Vec<JoinHandle<()>>> {
    let mut handles = Vec::with_capacity(num_workers);
    let core_ids = if pin_threads {
        core_affinity::get_core_ids()
    } else {
        None
    };

    for worker_id in 0..num_workers {
        let rx = raw_rx.clone();
        let tx = proc_tx.clone();
        let transforms = Arc::clone(&transforms);
        let w_shutdown = Arc::clone(&shutdown);
        let w_errors = Arc::clone(&errors_skipped);
        let w_fatal = Arc::clone(&fatal_error);
        let worker_core = core_ids.as_ref().map(|ids| {
            // If there are fewer cores than workers, wrap around
            ids[worker_id % ids.len()]
        });

        let handle = std::thread::Builder::new()
            .name(format!("tenshift-worker-{worker_id}"))
            .spawn(move || {
                if let Some(core_id) = worker_core {
                    if !core_affinity::set_for_current(core_id) {
                        tracing::warn!("failed to set thread affinity for worker {}", worker_id);
                    }
                }

                let mut local_errors = 0_u64;
                let mut output = Vec::with_capacity(64);
                // Scratch buffers reused across every sample this worker processes
                // (apply_stateless ping-pongs between them). Hoisted out of the
                // per-sample loop so each sample no longer pays two fresh 64-slot
                // Vec allocations in the hot transform path.
                let mut buf1: Vec<Sample> = Vec::with_capacity(64);
                let mut buf2: Vec<Sample> = Vec::with_capacity(64);

                while let Ok(chunk) = rx.recv() {
                    if w_shutdown.load(Ordering::Relaxed) {
                        break;
                    }

                    output.clear();
                    output.reserve(chunk.samples.len());

                    for sample in chunk.samples {
                        // Wrap each sample processing in catch_unwind to isolate panics.
                        // buf1/buf2 (and `transforms`) are borrowed into the closure via
                        // AssertUnwindSafe rather than moved, so the reused allocations
                        // survive to the next sample. apply_stateless clears both on entry,
                        // so any partial state a panicked sample leaves behind is harmless.
                        let process_result =
                            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                                apply_stateless(
                                    &transforms,
                                    sample,
                                    &mut buf1,
                                    &mut buf2,
                                    error_policy,
                                )
                            }));

                        match process_result {
                            Ok(Ok(err_count)) => {
                                // apply_stateless leaves the finished samples in `current`,
                                // which is buf1; drain them out and keep buf1's capacity.
                                output.append(&mut buf1);
                                local_errors += err_count;
                            }
                            Ok(Err(error)) => {
                                // Non-transform error from apply_stateless (e.g., allocation failure)
                                match error_policy {
                                    ErrorPolicy::Skip => {
                                        local_errors += 1;
                                        tracing::warn!("skipping failed transform item: {error}");
                                    }
                                    ErrorPolicy::Fail => {
                                        tracing::error!("worker transform error: {error}");
                                        record_first_fatal_error(&w_fatal, error);
                                        // Send partial output before shutting down to avoid data loss
                                        let _ = tx.send(SampleChunk {
                                            sequence: chunk.sequence,
                                            samples: std::mem::take(&mut output),
                                        });
                                        w_shutdown.store(true, Ordering::Relaxed);
                                        break;
                                    }
                                }
                            }
                            Err(payload) => {
                                // Panic caught during sample processing (e.g., in Vec::with_capacity)
                                let msg = super::panic_message(payload);
                                tracing::error!("worker thread panicked during transform: {}", msg);
                                match error_policy {
                                    ErrorPolicy::Skip => {
                                        local_errors += 1;
                                        tracing::warn!("skipping panicked transform item");
                                        // Continue with next sample
                                    }
                                    ErrorPolicy::Fail => {
                                        record_first_fatal_error(
                                            &w_fatal,
                                            crate::error::Error::TransformFailed {
                                                index: 0,
                                                reason: format!("worker thread panicked: {msg}"),
                                            },
                                        );
                                        // Send partial output before shutting down to avoid data loss
                                        let _ = tx.send(SampleChunk {
                                            sequence: chunk.sequence,
                                            samples: std::mem::take(&mut output),
                                        });
                                        w_shutdown.store(true, Ordering::Relaxed);
                                        break;
                                    }
                                }
                            }
                        }
                    }

                    if w_shutdown.load(Ordering::Relaxed) && error_policy == ErrorPolicy::Fail {
                        break;
                    }

                    if tx
                        .send(SampleChunk {
                            sequence: chunk.sequence,
                            samples: std::mem::take(&mut output),
                        })
                        .is_err()
                    {
                        break;
                    }
                }

                if local_errors > 0 {
                    w_errors.fetch_add(local_errors, Ordering::Relaxed);
                }
            })?;
        handles.push(handle);
    }

    Ok(handles)
}

#[cfg(test)]
mod tests {
    use super::{apply_stateless, record_first_fatal_error};
    use crate::error::Error;
    use crate::pipeline::ErrorPolicy;
    use crate::sample::Sample;
    use crate::transform::{Transform, TransformResult};
    use std::sync::{Arc, Mutex};

    struct Identity;
    impl Transform for Identity {
        fn apply(&self, sample: Sample) -> TransformResult {
            TransformResult::Sample(sample)
        }
        fn name(&self) -> &str {
            "identity"
        }
    }

    #[test]
    fn apply_stateless_reuses_scratch_buffers_across_samples() {
        // Two identity transforms leave the result back in buf1 (even number of
        // internal ping-pong swaps) and drop nothing. The worker hoists buf1/buf2
        // out of the per-sample loop and drains buf1 via `append`, which keeps its
        // allocation. So across many samples the scratch buffer must be REUSED:
        // same backing allocation (pointer) and capacity throughout. The old code
        // allocated two fresh `Vec::with_capacity(64)` inside the per-sample
        // closure, so this is a direct regression guard on that fix.
        let transforms: Vec<Box<dyn Transform>> = vec![Box::new(Identity), Box::new(Identity)];
        let mut buf1: Vec<Sample> = Vec::with_capacity(64);
        let mut buf2: Vec<Sample> = Vec::with_capacity(64);
        let ptr_before = buf1.as_ptr();
        let cap_before = buf1.capacity();

        let mut collected = 0_u64;
        for i in 0..500_u64 {
            let sample = Sample::new().with_metadata("s", i);
            let errs =
                apply_stateless(&transforms, sample, &mut buf1, &mut buf2, ErrorPolicy::Fail)
                    .expect("identity transform chain never errors");
            assert_eq!(errs, 0);
            assert_eq!(buf1.len(), 1, "identity chain yields exactly one sample");
            assert_eq!(
                buf1[0].metadata().expect("metadata preserved").index,
                i,
                "identity must pass the exact sample through"
            );
            // Drain like the worker does; `append` empties buf1 but keeps its allocation.
            let mut out: Vec<Sample> = Vec::new();
            out.append(&mut buf1);
            collected += out.len() as u64;
        }

        assert_eq!(collected, 500);
        assert_eq!(
            buf1.capacity(),
            cap_before,
            "scratch buffer must not have reallocated per sample"
        );
        assert_eq!(
            buf1.as_ptr(),
            ptr_before,
            "scratch buffer must be the same reused allocation, not a fresh one per sample"
        );
    }

    #[test]
    fn record_first_fatal_error_keeps_only_the_first() {
        let fatal: Mutex<Option<Error>> = Mutex::new(None);
        record_first_fatal_error(&fatal, Error::TransformFailed { index: 1, reason: "first".into() });
        record_first_fatal_error(&fatal, Error::TransformFailed { index: 2, reason: "second".into() });
        let guard = fatal.lock().unwrap();
        match guard.as_ref().expect("a fatal error was recorded") {
            Error::TransformFailed { index, reason } => {
                assert_eq!(*index, 1);
                assert_eq!(reason, "first");
            }
            other => panic!("unexpected error variant: {other:?}"),
        }
    }

    #[test]
    fn record_first_fatal_error_recovers_a_poisoned_lock() {
        // A worker panicking while another holds the fatal-error lock poisons it.
        // The recording path must still record the error (recover the poison), not
        // silently drop it (Law 10) — otherwise a failed epoch looks clean.
        let fatal: Arc<Mutex<Option<Error>>> = Arc::new(Mutex::new(None));
        let poisoner = Arc::clone(&fatal);
        let handle = std::thread::spawn(move || {
            let _guard = poisoner.lock().unwrap();
            panic!("poison the fatal-error mutex");
        });
        assert!(handle.join().is_err(), "poisoning thread must have panicked");
        assert!(fatal.is_poisoned(), "mutex must now be poisoned");

        record_first_fatal_error(
            &fatal,
            Error::TransformFailed { index: 7, reason: "recorded through poison".into() },
        );

        let guard = fatal
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        match guard.as_ref().expect("error recorded despite poison") {
            Error::TransformFailed { index, reason } => {
                assert_eq!(*index, 7);
                assert_eq!(reason, "recorded through poison");
            }
            other => panic!("unexpected error variant: {other:?}"),
        }
    }
}