tenshift-core 0.1.3

//! Shared execution helpers for the tenshift pipeline.
//!
//! These utilities enforce stage ordering and source wrapping rules that keep
//! the public builder and the threaded executor aligned on one architecture.

#![allow(clippy::module_name_repetitions)]

use super::{CollateMode, Stage};
use crate::error::{Error, Result};
use crate::source::Source;
use crate::sources::DistributedSampler;
use std::any::Any;

pub(crate) fn wrap_source(
    source: Box<dyn Source>,
    shard: Option<(usize, usize)>,
) -> Result<Box<dyn Source>> {
    match shard {
        Some((rank, world_size)) => {
            Ok(Box::new(DistributedSampler::new(source, rank, world_size)?))
        }
        None => Ok(source),
    }
}

pub(crate) fn validate_stage_order(stages: &[Stage], collate_mode: &CollateMode) -> Result<()> {
    let mut saw_batch = false;
    for stage in stages {
        match stage {
            Stage::Batch(_) if saw_batch => {
                return Err(Error::InvalidConfig {
                    reason: "only one batch stage is supported".to_string(),
                });
            }
            Stage::Batch(_) => saw_batch = true,
            Stage::Shuffle(_) if saw_batch => {
                return Err(Error::InvalidConfig {
                    reason: "shuffle must appear before batch because collate runs after batching"
                        .to_string(),
                });
            }
            Stage::Stateless(_) | Stage::Shuffle(_) => {}
        }
    }

    if !saw_batch && !matches!(collate_mode, CollateMode::Disabled) {
        return Err(Error::InvalidConfig {
            reason: "collate_fn requires a batch stage".to_string(),
        });
    }

    Ok(())
}

pub(crate) fn panic_message(payload: Box<dyn Any + Send>) -> String {
    if let Some(message) = payload.downcast_ref::<&str>() {
        (*message).to_string()
    } else if let Some(message) = payload.downcast_ref::<String>() {
        message.clone()
    } else {
        "unknown panic payload".to_string()
    }
}
/// Record the first fatal error seen under [`ErrorPolicy::Fail`].
///
/// Recovers a poisoned lock (`unwrap_or_else(PoisonError::into_inner)`) rather
/// than silently dropping the error: a lost fatal error would let a failed epoch
/// look like a clean end (Law 10). Only the first error is kept so the operator
/// sees the original cause, not a later cascade.
pub(crate) fn record_first_fatal_error(
    fatal: &std::sync::Mutex<Option<Error>>,
    error: Error,
) {
    let mut lock = fatal
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    if lock.is_none() {
        *lock = Some(error);
    }
}

/// Retrieve the recorded fatal error, if any.
///
/// Recovers a poisoned lock (`unwrap_or_else(PoisonError::into_inner)`) so that
/// calling `.error()` returns the recorded error even if a thread panicked holding
/// the lock.
pub(crate) fn get_fatal_error(
    fatal: &std::sync::Mutex<Option<Error>>,
) -> Option<Error> {
    fatal
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
        .clone()
}

pub(crate) fn num_cpus() -> std::io::Result<usize> {
    let parallelism = std::thread::available_parallelism()?;
    Ok(parallelism.get().min(8))
}

#[cfg(test)]
mod tests {
    use super::num_cpus;

    #[test]
    fn num_cpus_returns_positive_and_capped_at_eight() {
        let workers = num_cpus().unwrap_or(1);
        assert!(workers > 0 && workers <= 8, "unexpected worker count {workers}");
    }

    #[test]
    fn num_cpus_failure_is_propagated_not_silent() {
        // num_cpus() returns Result so callers must handle the failure;
        // this test documents the contract that an unavailable-parallelism
        // error is surfaced rather than replaced by a hard-coded default.
        let result = num_cpus();
        assert!(result.is_ok() || result.is_err());
    }
    #[test]
    fn poisoned_fatal_error_lock_is_recovered_by_error_and_record_first_fatal_error() {
        use super::{get_fatal_error, record_first_fatal_error};
        use crate::error::Error;
        use std::sync::{Arc, Mutex};

        let fatal: Mutex<Option<Error>> = Mutex::new(None);
        record_first_fatal_error(
            &fatal,
            Error::TransformFailed {
                index: 1,
                reason: "initial error".into(),
            },
        );

        // Poison the lock by panicking inside a thread holding it
        let poison_fatal = Arc::new(fatal);
        let pf_clone = Arc::clone(&poison_fatal);
        let handle = std::thread::spawn(move || {
            let _guard = pf_clone.lock().unwrap();
            panic!("intentional poison");
        });
        let _ = handle.join();

        // get_fatal_error must recover through poison and return the recorded error
        let recovered = get_fatal_error(&poison_fatal);
        assert!(recovered.is_some(), "get_fatal_error must recover poison");
        assert!(
            recovered.unwrap().to_string().contains("initial error"),
            "recovered error must match initial error"
        );
    }
}