tenshift_core/

source.rs

//! Source trait  -  the extension point for data origins.
//!
//! A [`Source`] produces [`Sample`]s from some origin (filesystem, network, database, etc.).
//! Community contributors add new sources by implementing this single trait.
//!
//! Built-in sources include [`GlobSource`][crate::sources::GlobSource],
//! [`MemorySource`][crate::sources::MemorySource], and
//! [`DistributedSampler`][crate::sources::DistributedSampler], with optional
//! feature-gated sources such as CSV and `io_uring` readers.
//!
//! # Async Sources
//!
//! For async data origins, implement [`AsyncSource`] and wrap with
//! [`Pipeline::from_async_source`](crate::pipeline::Pipeline::from_async_source).
//! This runs the async source in a dedicated Tokio runtime without blocking
//! the pipeline's worker threads.

use crate::error::Result;
use crate::sample::Sample;
use std::future::Future;
use std::pin::Pin;

/// A source of training data.
///
/// Sources are lazy  -  they describe WHERE data comes from, not the data itself.
/// Data is only loaded when the pipeline pulls from the source.
///
/// # Implementing a Source
///
/// ```rust
/// use tenshift_core::source::{Source, SourceIterator};
/// use tenshift_core::sample::Sample;
/// use tenshift_core::error::Result;
///
/// struct MyDatabase {
///     connection_string: String,
/// }
///
/// impl Source for MyDatabase {
///     fn open(&self) -> Result<Box<dyn SourceIterator>> {
///         // Open connection, return iterator over rows
///         # Ok(Box::new(std::iter::empty::<Result<Sample>>()))
///     }
///
///     fn len_hint(&self) -> Option<u64> {
///         None // unknown size
///     }
///
///     fn name(&self) -> &str {
///         "my_database"
///     }
/// }
/// ```
pub trait Source: Send + Sync {
    /// Create an iterator over samples from this source.
    ///
    /// Called once per epoch. The iterator is consumed by the pipeline's
    /// worker threads.
    ///
    /// # Errors
    ///
    /// Returns an error if the source cannot be opened (e.g., file not found,
    /// network unreachable, invalid credentials).
    fn open(&self) -> Result<Box<dyn SourceIterator>>;

    /// Optional hint about the total number of samples.
    ///
    /// Used for progress reporting and pre-allocation. Return `None` if unknown.
    fn len_hint(&self) -> Option<u64> {
        None
    }

    /// Human-readable name for this source (for logging and error messages).
    fn name(&self) -> &str;
}

impl<T> Source for Box<T>
where
    T: Source + ?Sized,
{
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        (**self).open()
    }

    fn len_hint(&self) -> Option<u64> {
        (**self).len_hint()
    }

    fn name(&self) -> &str {
        (**self).name()
    }
}

/// Boxed future used by async source extension traits.
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;

/// Async-capable source of training data.
pub trait AsyncSource: Send + Sync {
    /// Open the source asynchronously and return an async iterator.
    fn open_async(&self) -> BoxFuture<'_, Result<Box<dyn AsyncSourceIterator>>>;

    /// Optional hint about the total number of samples.
    fn len_hint(&self) -> Option<u64> {
        None
    }

    /// Human-readable name for this source.
    fn name(&self) -> &str;
}

/// Iterator over samples from a source.
///
/// This is a regular iterator that yields `Result<Sample>`. Errors on individual
/// items (e.g., corrupt file) are returned as `Err` without stopping the iterator.
pub trait SourceIterator: Send {
    /// Get the next sample, or `None` if exhausted.
    fn next_sample(&mut self) -> Option<Result<Sample>>;
}

/// Blanket implementation: any `Iterator<Item = Result<Sample>> + Send` is a `SourceIterator`.
impl<I> SourceIterator for I
where
    I: Iterator<Item = Result<Sample>> + Send,
{
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        self.next()
    }
}

/// Async iterator over samples from a source.
pub trait AsyncSourceIterator: Send {
    /// Get the next sample asynchronously, or `None` if exhausted.
    fn next_sample_async(&mut self) -> BoxFuture<'_, Option<Result<Sample>>>;
}

struct SyncSourceAdapter {
    inner: Box<dyn SourceIterator>,
}

impl AsyncSourceIterator for SyncSourceAdapter {
    fn next_sample_async(&mut self) -> BoxFuture<'_, Option<Result<Sample>>> {
        Box::pin(async move { self.inner.next_sample() })
    }
}

impl<T> AsyncSource for T
where
    T: Source,
{
    fn open_async(&self) -> BoxFuture<'_, Result<Box<dyn AsyncSourceIterator>>> {
        Box::pin(async move {
            self.open()
                .map(|inner| Box::new(SyncSourceAdapter { inner }) as Box<dyn AsyncSourceIterator>)
        })
    }

    fn len_hint(&self) -> Option<u64> {
        Source::len_hint(self)
    }

    fn name(&self) -> &str {
        Source::name(self)
    }
}

/// Adapts an [`AsyncSource`] into a synchronous [`Source`] by running a Tokio runtime.
pub struct AsyncToSyncAdapter<A> {
    inner: A,
}

impl<A: AsyncSource> AsyncToSyncAdapter<A> {
    /// Create a new adapter.
    pub fn new(inner: A) -> Self {
        Self { inner }
    }
}

impl<A: AsyncSource + 'static> Source for AsyncToSyncAdapter<A> {
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        let rt = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(1)
            .enable_all()
            .build()
            .map_err(|e| crate::error::Error::InvalidConfig {
                reason: format!("failed to build tokio runtime for async source: {e}"),
            })?;

        // Block to initialize the iter, since the caller requires `Result`
        let mut inner_iter = rt.block_on(self.inner.open_async())?;

        // Construct high-speed queue bridging async I/O worker to sync transform workers
        let (tx, rx) = crossbeam_channel::bounded(128);

        // Native prefetching: Decouple the async source polling from the CPU pipeline
        // to immediately start fetching the next item in the background.
        rt.spawn(async move {
            loop {
                match inner_iter.next_sample_async().await {
                    Some(Ok(sample)) => {
                        if tx.send(Some(Ok(sample))).is_err() {
                            break;
                        }
                    }
                    Some(Err(e)) => {
                        let _ = tx.send(Some(Err(e)));
                        break;
                    }
                    None => {
                        let _ = tx.send(None);
                        break;
                    }
                }
            }
        });

        Ok(Box::new(AsyncIteratorAdapter { rt, rx }))
    }

    fn len_hint(&self) -> Option<u64> {
        self.inner.len_hint()
    }

    fn name(&self) -> &str {
        self.inner.name()
    }
}

/// Runs [`AsyncSourceIterator`] asynchronously block by block.
pub struct AsyncIteratorAdapter {
    #[allow(dead_code)]
    rt: tokio::runtime::Runtime,
    rx: crossbeam_channel::Receiver<Option<Result<Sample>>>,
}

impl SourceIterator for AsyncIteratorAdapter {
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        match self.rx.recv() {
            Ok(item) => item,
            // Channel disconnected  -  async source dropped without sending None sentinel.
            // Surface as an error rather than silently ending the stream.
            Err(_) => Some(Err(crate::error::Error::SourceFailed {
                source_name: "async_adapter".into(),
                reason: "async source channel disconnected unexpectedly".into(),
            })),
        }
    }
}