tenshift-core 0.1.3

//! `io_uring`-accelerated batch file reader powered by `santh-io`.
//!
//! This source batches file reads through a [`wireshift::Ring`], submitting many
//! reads in a single syscall on Linux `io_uring` and falling back to a worker
//! pool on other platforms. The public API is identical to [`super::GlobSource`]
//! but with dramatically higher throughput on `NVMe` storage.
//!
//! # Feature gate
//!
//! Requires the `uring` feature:
//! ```toml
//! tenshift-core = { version = "0.1", features = ["uring"] }
//! ```
//!
//! # Example
//!
//! ```rust,ignore
//! use tenshift_core::sources::RingSource;
//! use tenshift_core::Pipeline;
//!
//! let source = RingSource::from_glob("data/train/*.bin", 64)?;
//! let pipeline = Pipeline::from_source(source)
//!     .workers(4)
//!     .batch(32)
//!     .prefetch(4);
//!
//! for batch in pipeline {
//!     // Loaded via io_uring batch reads
//! }
//! ```

use std::collections::VecDeque;
use std::path::PathBuf;

use std::fs::File;

use wireshift::ops::{OpenAt, Read, Statx};
use wireshift::{Request, Ring, RingConfig};

use crate::error::{Error, Result};
use crate::sample::{Sample, Tensor};
use crate::source::{Source, SourceIterator};

/// Batch file reader using `santh-io` for `io_uring`-accelerated I/O.
///
/// Submits `batch_size` file reads per `io_uring` submission, drastically
/// reducing syscall overhead on NVMe/SSD storage. Falls back to a worker
/// pool on systems without `io_uring` support.
pub struct RingSource {
    paths: Vec<PathBuf>,
    pattern: String,
    batch_size: usize,
    ring_config: RingConfig,
}

impl RingSource {
    /// Create a source from a glob pattern with the specified batch size.
    ///
    /// `batch_size` controls how many file reads are submitted per `io_uring`
    /// batch. Higher values = fewer syscalls = higher throughput, but more
    /// memory. A good default is 64.
    ///
    /// # Errors
    ///
    /// Returns [`Error::InvalidPattern`] if the glob pattern is invalid, or
    /// [`Error::EmptySource`] if no files match.
    pub fn from_glob(pattern: &str, batch_size: usize) -> Result<Self> {
        let paths = super::collect_glob_files(pattern)?;

        if paths.is_empty() {
            return Err(Error::EmptySource {
                pattern: pattern.to_string(),
            });
        }

        tracing::info!(
            "RingSource: found {} files matching '{}' (batch_size={})",
            paths.len(),
            pattern,
            batch_size
        );

        Ok(Self {
            paths,
            pattern: pattern.to_string(),
            batch_size: batch_size.max(1),
            ring_config: RingConfig::default(),
        })
    }

    /// Create a source from explicit paths.
    ///
    /// # Errors
    ///
    /// Returns [`Error::EmptySource`] if the path list is empty.
    pub fn from_paths(paths: Vec<PathBuf>, batch_size: usize) -> Result<Self> {
        if paths.is_empty() {
            return Err(Error::EmptySource {
                pattern: "<explicit paths>".to_string(),
            });
        }

        Ok(Self {
            pattern: format!("{} explicit paths", paths.len()),
            paths,
            batch_size: batch_size.max(1),
            ring_config: RingConfig::default(),
        })
    }

    /// Override the `santh-io` ring configuration.
    pub fn with_ring_config(mut self, config: RingConfig) -> Self {
        self.ring_config = config;
        self
    }
}

impl Source for RingSource {
    fn open(&self) -> Result<Box<dyn SourceIterator>> {
        let ring = Ring::new(self.ring_config.clone()).map_err(|error| Error::SourceFailed {
            source_name: self.pattern.clone(),
            reason: format!("failed to create santh-io ring: {error}"),
        })?;

        tracing::debug!("RingSource opened with backend {:?}", ring.backend_kind());

        Ok(Box::new(RingIterator {
            paths: self.paths.clone(),
            ring,
            index: 0,
            batch_size: self.batch_size,
            pending: VecDeque::new(),
        }))
    }

    fn len_hint(&self) -> Option<u64> {
        Some(self.paths.len() as u64)
    }

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

enum ReadState {
    Submitted(Request<wireshift::ops::ReadResult>),
    Error(Error),
}

struct PendingOpen {
    request: Request<File>,
    path: PathBuf,
    index: u64,
    read_len: usize,
}

struct PendingStat {
    request: Request<wireshift::FileStat>,
    path: PathBuf,
    index: u64,
}

/// Pending read: the request handle and the path/index associated with it.
struct PendingRead {
    state: ReadState,
    path: PathBuf,
    index: u64,
}

/// Iterator that submits batch reads through santh-io's Ring.
struct RingIterator {
    paths: Vec<PathBuf>,
    ring: Ring,
    index: u64,
    batch_size: usize,
    pending: VecDeque<PendingRead>,
}

impl RingIterator {
    fn submit_batch(&mut self) -> Result<()> {
        let start = self.index as usize;
        let end = (start + self.batch_size).min(self.paths.len());
        let mut pending_stats = Vec::with_capacity(end.saturating_sub(start));
        let mut pending_opens = Vec::with_capacity(end.saturating_sub(start));

        for idx in start..end {
            let path = &self.paths[idx];
            match self.ring.submit(Statx::new(path)) {
                Ok(request) => pending_stats.push(PendingStat {
                    request,
                    path: path.clone(),
                    index: idx as u64,
                }),
                Err(error) => self.pending.push_back(PendingRead {
                    state: ReadState::Error(Error::SourceFailed {
                        source_name: path.to_string_lossy().to_string(),
                        reason: format!("failed to submit statx: {error}"),
                    }),
                    path: path.clone(),
                    index: idx as u64,
                }),
            }
        }

        for pending_stat in pending_stats {
            let read_len = match pending_stat
                .request
                .wait(Some(std::time::Duration::from_secs(30)))
            {
                Ok(stat) => match usize::try_from(stat.size) {
                    Ok(size) => size,
                    Err(_) => {
                        self.pending.push_back(PendingRead {
                            state: ReadState::Error(Error::SourceFailed {
                                source_name: pending_stat.path.to_string_lossy().to_string(),
                                reason: "file is too large to fit in an in-memory read buffer. Fix: shard the file or stream it in smaller ranges.".to_string(),
                            }),
                            path: pending_stat.path,
                            index: pending_stat.index,
                        });
                        continue;
                    }
                },
                Err(error) => {
                    self.pending.push_back(PendingRead {
                        state: ReadState::Error(Error::ReadFailed {
                            path: pending_stat.path.clone(),
                            reason: error.to_string(),
                        }),
                        path: pending_stat.path,
                        index: pending_stat.index,
                    });
                    continue;
                }
            };

            match self.ring.submit(OpenAt::read_only(&pending_stat.path)) {
                Ok(request) => pending_opens.push(PendingOpen {
                    request,
                    path: pending_stat.path,
                    index: pending_stat.index,
                    read_len,
                }),
                Err(error) => self.pending.push_back(PendingRead {
                    state: ReadState::Error(Error::SourceFailed {
                        source_name: pending_stat.path.to_string_lossy().to_string(),
                        reason: format!("failed to submit openat: {error}"),
                    }),
                    path: pending_stat.path,
                    index: pending_stat.index,
                }),
            }
        }

        for pending_open in pending_opens {
            // Read buffers are sized from async statx metadata so the completed buffer can
            // stay attached to the sample without carrying a 16 MiB over-allocation.
            let state = match pending_open
                .request
                .wait(Some(std::time::Duration::from_secs(30)))
            {
                Ok(file) => match Read::new(file, 0, pending_open.read_len) {
                    Ok(op) => match self.ring.submit(op) {
                        Ok(request) => ReadState::Submitted(request),
                        Err(error) => ReadState::Error(Error::SourceFailed {
                            source_name: pending_open.path.to_string_lossy().to_string(),
                            reason: format!("failed to submit read after openat: {error}"),
                        }),
                    },
                    Err(error) => ReadState::Error(Error::SourceFailed {
                        source_name: pending_open.path.to_string_lossy().to_string(),
                        reason: format!("failed to create read operation after openat: {error}"),
                    }),
                },
                Err(error) => ReadState::Error(Error::ReadFailed {
                    path: pending_open.path.clone(),
                    reason: error.to_string(),
                }),
            };

            self.pending.push_back(PendingRead {
                state,
                path: pending_open.path,
                index: pending_open.index,
            });
        }

        self.index = end as u64;
        Ok(())
    }

    /// Drain the next completion from pending reads.
    fn next_completion(&mut self) -> Option<Result<Sample>> {
        if self.pending.is_empty() {
            return None;
        }

        let pending = self.pending.pop_front()?;
        let result = match pending.state {
            ReadState::Submitted(request) => request
                .wait(Some(std::time::Duration::from_secs(30)))
                .map_err(|error| Error::ReadFailed {
                    path: pending.path.clone(),
                    reason: error.to_string(),
                }),
            ReadState::Error(error) => Err(error),
        };

        match result {
            Ok(read_result) => {
                let filename = match pending.path.file_name() {
                    Some(n) => n.to_string_lossy().to_string(),
                    None => String::new(),
                };

                Some(Ok(Sample::new()
                    .with("data", Tensor::bytes_from_completed(read_result.buffer))
                    .with("filename", Tensor::bytes(filename.into_bytes()))
                    .with_metadata(
                        pending.path.to_string_lossy(),
                        pending.index,
                    )))
            }
            Err(error) => Some(Err(error)),
        }
    }
}

impl SourceIterator for RingIterator {
    fn next_sample(&mut self) -> Option<Result<Sample>> {
        // If no pending completions, submit next batch
        if self.pending.is_empty() {
            if (self.index as usize) >= self.paths.len() {
                return None;
            }
            if let Err(error) = self.submit_batch() {
                return Some(Err(error));
            }
        }

        self.next_completion()
    }
}

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

    #[cfg(feature = "uring")]
    #[test]
    fn ring_source_direct_iteration_reads_all_paths() {
        let dir = tempfile::tempdir().expect("tempdir");
        for idx in 0..10 {
            std::fs::write(dir.path().join(format!("f{idx}.bin")), vec![42_u8; 4096])
                .expect("write test file");
        }

        let pattern = format!("{}/*.bin", dir.path().display());
        let source = RingSource::from_glob(&pattern, 1).expect("source");
        let mut iter = source.open().expect("open");
        let mut total = 0usize;

        while let Some(sample) = iter.next_sample() {
            let sample = sample.expect("sample");
            assert_eq!(
                sample.get("data").expect("data field").as_bytes().len(),
                4096
            );
            total += 1;
        }

        assert_eq!(total, 10);
    }
}