rustfs/crates/io-core/src/reader.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Bytes-backed object reader implementation.

use bytes::Bytes;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};

/// Errors that can occur during Bytes-backed read operations.
#[derive(Debug, Clone)]
pub enum ZeroCopyReadError {
    /// I/O error occurred.
    Io(String),
    /// Memory mapping error.
    Mmap(String),
    /// Invalid offset or size.
    InvalidRange,
}

impl std::fmt::Display for ZeroCopyReadError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(msg) => write!(f, "I/O error: {}", msg),
            Self::Mmap(msg) => write!(f, "Mmap error: {}", msg),
            Self::InvalidRange => write!(f, "Invalid offset or size"),
        }
    }
}

impl std::error::Error for ZeroCopyReadError {}

impl From<io::Error> for ZeroCopyReadError {
    fn from(err: io::Error) -> Self {
        Self::Io(err.to_string())
    }
}

/// Bytes-backed object reader.
///
/// `from_bytes` wraps existing `Bytes` without copying, but file constructors
/// copy file data into owned `Bytes` after mmap or normal reads.
///
/// # Example
///
/// ```ignore
/// use bytes::Bytes;
/// use rustfs_io_core::BytesBufferedReader;
///
/// // Create from bytes without copying the `Bytes` buffer
/// let data = Bytes::from("hello world");
/// let reader = BytesBufferedReader::from_bytes(data);
///
/// // Read using AsyncRead trait
/// let mut buf = vec![0u8; 1024];
/// let n = reader.read(&mut buf[..]).await?;
/// ```
pub struct BytesBufferedReader {
    /// Internal data source (could be mmap or owned bytes)
    data: Bytes,
    /// Current read position
    pos: usize,
}

/// Historical name for the bytes-backed object reader.
#[deprecated(
    since = "1.0.0-beta.8",
    note = "use BytesBufferedReader; file constructors copy into owned Bytes"
)]
pub type ZeroCopyObjectReader = BytesBufferedReader;

impl BytesBufferedReader {
    /// Create a reader from existing bytes.
    ///
    /// This is a true zero-copy operation - the Bytes are wrapped
    /// without any allocation or copying.
    ///
    /// # Arguments
    ///
    /// * `data` - Bytes to wrap
    ///
    /// # Example
    ///
    /// ```ignore
    /// let data = Bytes::from("hello world");
    /// let reader = BytesBufferedReader::from_bytes(data);
    /// ```
    pub fn from_bytes(data: Bytes) -> Self {
        Self { data, pos: 0 }
    }

    /// Create a Bytes-backed reader from a file using mmap-then-copy.
    ///
    /// This maps the requested file range and copies it into owned `Bytes`
    /// before returning. It does not expose the mmap as a zero-copy buffer.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the file to memory map
    /// * `offset` - Offset within the file to start reading
    /// * `size` - Number of bytes to read
    ///
    /// # Returns
    ///
    /// A reader backed by copied file data.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be memory mapped.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let reader = BytesBufferedReader::from_file_mmap_path("large_file.bin", 0, 1024).await?;
    /// ```
    #[cfg(unix)]
    // SAFETY: The mmap is created from a read-only file handle for the
    // caller-provided range, then copied into owned `Bytes` before the file and
    // mapping are dropped.
    #[allow(unsafe_code)]
    pub async fn from_file_mmap_path(path: &std::path::Path, offset: u64, size: usize) -> Result<Self, ZeroCopyReadError> {
        use memmap2::MmapOptions;

        let path = path.to_path_buf();
        let (offset, size) = (offset, size);

        tokio::task::spawn_blocking(move || {
            // Open the file in sync context
            let std_file = std::fs::File::open(&path).map_err(|e| ZeroCopyReadError::Io(e.to_string()))?;

            // SAFETY: `std_file` remains open while the mapping is created and
            // copied, and the mapped bytes are not exposed beyond this closure.
            let mmap = unsafe { MmapOptions::new().offset(offset).len(size).map(&std_file) }
                .map_err(|e| ZeroCopyReadError::Mmap(e.to_string()))?;

            // Convert to Bytes (this is a copy, but only done once)
            Ok(Self {
                data: Bytes::copy_from_slice(&mmap),
                pos: 0,
            })
        })
        .await
        .map_err(|e| ZeroCopyReadError::Io(e.to_string()))?
    }

    /// Create a Bytes-backed reader from a file using normal reads.
    ///
    /// This path reads the requested range into an owned buffer and wraps it in
    /// `Bytes`. It does not perform mmap or zero-copy file I/O.
    ///
    /// # Arguments
    ///
    /// * `file` - File to read from
    /// * `offset` - Offset within the file to start reading
    /// * `size` - Number of bytes to map
    ///
    /// # Returns
    ///
    /// A reader backed by copied file data.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let file = tokio::fs::File::open("large_file.bin").await?;
    /// let reader = BytesBufferedReader::from_file_read(&file, 0, 1024).await?;
    /// ```
    #[cfg(unix)]
    pub async fn from_file_read(file: &tokio::fs::File, offset: u64, size: usize) -> Result<Self, ZeroCopyReadError> {
        use tokio::io::{AsyncReadExt, AsyncSeekExt, SeekFrom};

        let mut cloned = file.try_clone().await?;
        cloned.seek(SeekFrom::Start(offset)).await?;

        let mut buffer = vec![0u8; size];
        cloned.read_exact(&mut buffer).await?;

        Ok(Self {
            data: Bytes::from(buffer),
            pos: 0,
        })
    }

    /// Create a Bytes-backed reader from a file (non-Unix fallback).
    ///
    /// On platforms that don't support mmap, this falls back to regular file I/O.
    #[cfg(not(unix))]
    pub async fn from_file_read(file: &tokio::fs::File, offset: u64, size: usize) -> Result<Self, ZeroCopyReadError> {
        use tokio::io::{AsyncReadExt, AsyncSeekExt, SeekFrom};

        let mut cloned = file.try_clone().await?;
        cloned.seek(SeekFrom::Start(offset)).await?;

        let mut buffer = vec![0u8; size];
        cloned.read_exact(&mut buffer).await?;

        Ok(Self {
            data: Bytes::from(buffer),
            pos: 0,
        })
    }

    /// Historical name for `from_file_read`.
    #[deprecated(
        since = "1.0.0-beta.8",
        note = "use from_file_read; this method performs normal reads into owned Bytes"
    )]
    pub async fn from_file_mmap(file: &tokio::fs::File, offset: u64, size: usize) -> Result<Self, ZeroCopyReadError> {
        Self::from_file_read(file, offset, size).await
    }

    /// Get the remaining data as Bytes (zero-copy).
    ///
    /// This returns a slice of the remaining data without copying.
    /// The returned Bytes shares the underlying memory with this reader.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let remaining = reader.remaining_bytes();
    /// println!("Remaining: {} bytes", remaining.len());
    /// ```
    pub fn remaining_bytes(&self) -> Bytes {
        self.data.slice(self.pos..)
    }

    /// Get the total length of the data.
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Check if the reader has reached the end.
    pub fn is_empty(&self) -> bool {
        self.pos >= self.data.len()
    }

    /// Get the current read position.
    pub fn position(&self) -> usize {
        self.pos
    }
}

impl AsyncRead for BytesBufferedReader {
    fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let remaining = self.data.len() - self.pos;
        if remaining == 0 {
            return Poll::Ready(Ok(()));
        }

        let to_read = std::cmp::min(remaining, buf.remaining());
        let slice = &self.data[self.pos..self.pos + to_read];
        buf.put_slice(slice);
        self.pos += to_read;

        Poll::Ready(Ok(()))
    }
}

impl std::fmt::Debug for BytesBufferedReader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BytesBufferedReader")
            .field("data_len", &self.data.len())
            .field("pos", &self.pos)
            .field("remaining", &(self.data.len() - self.pos))
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;
    use tokio::io::AsyncReadExt;

    fn temp_file_path(test_name: &str) -> PathBuf {
        let nonce = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .expect("system time should be after unix epoch")
            .as_nanos();
        std::env::temp_dir().join(format!("rustfs-io-core-{test_name}-{}-{nonce}", std::process::id()))
    }

    #[tokio::test]
    async fn test_from_bytes() {
        let data = Bytes::from("hello world");
        let mut reader = BytesBufferedReader::from_bytes(data.clone());

        let mut buf = [0u8; 11];
        let n = reader.read(&mut buf[..]).await.unwrap();

        assert_eq!(n, 11);
        assert_eq!(&buf[..n], b"hello world");
    }

    #[tokio::test]
    async fn test_preferred_reader_alias() {
        let data = Bytes::from("hello world");
        let mut reader = BytesBufferedReader::from_bytes(data);

        let mut buf = [0u8; 5];
        let n = reader.read(&mut buf[..]).await.expect("read bytes from alias");

        assert_eq!(n, 5);
        assert_eq!(&buf[..n], b"hello");
    }

    #[tokio::test]
    async fn test_from_file_read_reads_requested_range() {
        let path = temp_file_path("from-file-read");
        tokio::fs::write(&path, b"hello world")
            .await
            .expect("write temp file for reader test");

        let file = tokio::fs::File::open(&path).await.expect("open temp file for reader test");
        let mut reader = BytesBufferedReader::from_file_read(&file, 6, 5)
            .await
            .expect("read requested range into Bytes");

        let mut output = Vec::new();
        reader.read_to_end(&mut output).await.expect("drain reader output");

        assert_eq!(output, b"world");

        let _ = tokio::fs::remove_file(path).await;
    }

    #[tokio::test]
    #[allow(deprecated)]
    async fn test_from_file_mmap_legacy_alias_reads_requested_range() {
        let path = temp_file_path("from-file-mmap");
        tokio::fs::write(&path, b"hello world")
            .await
            .expect("write temp file for legacy reader test");

        let file = tokio::fs::File::open(&path)
            .await
            .expect("open temp file for legacy reader test");
        let mut reader = BytesBufferedReader::from_file_mmap(&file, 0, 5)
            .await
            .expect("read requested range through legacy alias");

        let mut output = Vec::new();
        reader.read_to_end(&mut output).await.expect("drain legacy reader output");

        assert_eq!(output, b"hello");

        let _ = tokio::fs::remove_file(path).await;
    }

    #[tokio::test]
    async fn test_remaining_bytes() {
        let data = Bytes::from("hello world");
        let reader = BytesBufferedReader::from_bytes(data);

        let remaining = reader.remaining_bytes();
        assert_eq!(remaining.len(), 11);
        assert_eq!(&remaining[..], b"hello world");
    }

    #[tokio::test]
    async fn test_position() {
        let data = Bytes::from("hello world");
        let mut reader = BytesBufferedReader::from_bytes(data);

        assert_eq!(reader.position(), 0);

        let mut buf = [0u8; 5];
        reader.read_exact(&mut buf[..]).await.unwrap();

        assert_eq!(reader.position(), 5);
    }

    #[tokio::test]
    async fn test_is_empty() {
        let data = Bytes::from("");
        let reader = BytesBufferedReader::from_bytes(data);
        assert!(reader.is_empty());

        let data = Bytes::from("hello");
        let reader = BytesBufferedReader::from_bytes(data);
        assert!(!reader.is_empty());
    }

    #[tokio::test]
    #[allow(deprecated)]
    async fn test_legacy_reader_alias() {
        let data = Bytes::from("hello world");
        let mut reader = ZeroCopyObjectReader::from_bytes(data);

        let mut buf = [0u8; 5];
        let n = reader.read(&mut buf[..]).await.expect("read bytes through legacy alias");

        assert_eq!(n, 5);
        assert_eq!(&buf[..n], b"hello");
    }
}