wireshift-fallback 0.1.1

use std::collections::HashMap;
use std::io::Read;
use std::os::unix::fs::FileExt;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use wireshift_core::buffer::{Buffer, Submitted};
use wireshift_core::op::CompletionPayload;
use wireshift_core::{Error, Result};

use crate::ops::retry_eintr;

/// Execute a single read at a specific file offset into a submitted buffer.
pub fn execute_read(
    mut file: std::fs::File,
    offset: u64,
    mut buffer: Buffer<Submitted>,
    canceled: Arc<AtomicBool>,
    len: Option<usize>,
) -> Result<CompletionPayload> {
    if canceled.load(Ordering::Relaxed) {
        return Err(Error::canceled(
            "read canceled before data transfer started",
            "avoid canceling the request before the backend starts transferring data",
        ));
    }
    let buf = buffer.backend_mut();
    let read_len = len.unwrap_or(buf.len());
    let target_buf = if read_len < buf.len() {
        &mut buf[..read_len]
    } else {
        buf
    };
    apply_madv_sequential(target_buf)?;

    let bytes = retry_eintr!(file.read_at(target_buf, offset))
        .or_else(|error| {
            let is_espipe = error.raw_os_error() == Some(rustix::io::Errno::SPIPE.raw_os_error());
            if is_espipe && offset == 0 {
                retry_eintr!(file.read(target_buf))
            } else {
                Err(error)
            }
        })
        .map_err(|error| {
            Error::io(
                "read failed",
                error,
                "ensure the file descriptor is readable",
            )
        })?;
    let completed = buffer.into_completed(bytes)?;
    Ok(CompletionPayload::Read {
        buffer: completed,
        bytes,
    })
}

/// Execute a vectored (scatter) read across multiple buffers.
pub fn execute_read_vectored(
    mut file: std::fs::File,
    offset: u64,
    buffers: Vec<Buffer<Submitted>>,
) -> Result<CompletionPayload> {
    let mut total = 0_usize;
    let mut current_offset = offset;
    let mut completed_buffers = Vec::with_capacity(buffers.len());
    for mut buffer in buffers {
        let (bytes, is_short) = {
            let chunk = buffer.backend_mut();
            apply_madv_sequential(chunk)?;
            let bytes = retry_eintr!(file.read_at(chunk, current_offset))
                .or_else(|error| {
                    let is_espipe =
                        error.raw_os_error() == Some(rustix::io::Errno::SPIPE.raw_os_error());
                    if is_espipe && offset == 0 {
                        retry_eintr!(file.read(chunk))
                    } else {
                        Err(error)
                    }
                })
                .map_err(|error| {
                    Error::io(
                        "readv segment read failed",
                        error,
                        "ensure the file descriptor remains readable for all readv segments",
                    )
                })?;
            (bytes, bytes < chunk.len())
        };
        total += bytes;
        current_offset += bytes as u64;
        completed_buffers.push(buffer.into_completed(bytes)?);
        if is_short {
            break;
        }
    }
    Ok(CompletionPayload::ReadVectored {
        buffers: completed_buffers,
        bytes: total,
    })
}

/// Execute a read into a page-aligned GPU staging buffer.
pub fn execute_read_gpu(
    file: std::fs::File,
    offset: u64,
    mut buffer: wireshift_core::ops::AlignedBuffer,
    canceled: Arc<AtomicBool>,
) -> Result<CompletionPayload> {
    if canceled.load(Ordering::Relaxed) {
        return Err(Error::canceled(
            "gpu staged read canceled before data transfer started",
            "avoid canceling the request before the backend starts transferring data",
        ));
    }
    let slice = buffer.as_mut_slice();
    apply_madv_sequential(slice)?;
    let bytes = retry_eintr!(file.read_at(slice, offset)).map_err(|error| {
        Error::io(
            "gpu staged read failed",
            error,
            "ensure the file descriptor is readable and supports positional reads",
        )
    })?;
    Ok(CompletionPayload::GpuRead { buffer, bytes })
}

/// Execute a read using a pre-registered (fixed) file descriptor slot.
pub fn execute_read_fixed(
    slot: u32,
    offset: u64,
    mut buffer: Buffer<Submitted>,
    canceled: Arc<AtomicBool>,
    fixed_files: &mut HashMap<u32, std::fs::File>,
) -> Result<CompletionPayload> {
    if canceled.load(Ordering::Relaxed) {
        return Err(Error::canceled(
            "read canceled before data transfer started",
            "avoid canceling the request before the backend starts transferring data",
        ));
    }
    let file = fixed_files.get(&slot).ok_or_else(|| {
        Error::completion(
            format!("direct descriptor slot {slot} was read before open"),
            "ensure linked reads open the direct descriptor before issuing a fixed read",
        )
    })?;
    let slice = buffer.backend_mut();
    apply_madv_sequential(slice)?;
    let bytes = retry_eintr!(file.read_at(slice, offset)).map_err(|error| {
        Error::io(
            "read failed",
            error,
            "ensure the file descriptor is readable",
        )
    })?;
    let completed = buffer.into_completed(bytes)?;
    Ok(CompletionPayload::Read {
        buffer: completed,
        bytes,
    })
}

/// Apply `MADV_SEQUENTIAL` to `buffer`, hinting the kernel to read ahead.
///
/// A no-op on an empty buffer.
pub fn apply_madv_sequential(buffer: &mut [u8]) -> Result<()> {
    if buffer.is_empty() {
        return Ok(());
    }

    let page_size = rustix::param::page_size();

    let start = buffer.as_mut_ptr() as usize;
    let end = start.checked_add(buffer.len()).ok_or_else(|| {
        Error::validation(
            "buffer address overflowed while preparing MADV_SEQUENTIAL",
            "reduce the read buffer length so address calculations stay in range",
        )
    })?;
    let aligned_start = start / page_size * page_size;
    let aligned_end = end
        .checked_add(page_size - 1)
        .map(|value| value / page_size * page_size)
        .ok_or_else(|| {
            Error::validation(
                "buffer range overflowed while page-aligning MADV_SEQUENTIAL",
                "reduce the read buffer length so page alignment stays in range",
            )
        })?;
    let aligned_len = aligned_end.checked_sub(aligned_start).ok_or_else(|| {
        Error::validation(
            "aligned MADV_SEQUENTIAL range underflowed",
            "reduce the read buffer length so alignment calculations stay valid",
        )
    })?;

    #[cfg(target_os = "linux")]
    {
        use rustix::mm::{madvise, Advice};
        // SAFETY: madvise requires a valid pointer and length; aligned_start and
        // aligned_len were computed from the buffer pointer and length and
        // validated for overflow, so the pointer/length pair is safe to pass.
        let ptr = aligned_start as *mut std::ffi::c_void;
        unsafe {
            madvise(ptr, aligned_len, Advice::Sequential).map_err(|error| {
                Error::io(
                    "madvise(MADV_SEQUENTIAL) failed",
                    std::io::Error::from(error),
                    "ensure the runtime allows calling madvise on the target memory range",
                )
            })?;
        }
    }

    // On non-linux platforms madvise is a no-op for this backend.
    Ok(())
}


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

    #[test]
    fn verify_apply_madv_sequential() {
        let mut buf = vec![0u8; 16 * 1024];
        assert!(apply_madv_sequential(&mut buf).is_ok());
    }
}