irontide-session 1.0.1

//! IOCP disk I/O backend -- wraps [`PosixDiskIo`], overrides `write_block_direct()`.
//!
//! Uses a single shared IOCP completion port per session. Writes and reads are
//! submitted synchronously within `block_in_place` context -- NOT async-integrated
//! with tokio.
//!
//! The entire module is gated with `#[cfg(all(target_os = "windows", feature = "iocp"))]`
//! at the `mod` declaration site in `lib.rs`.

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use bytes::Bytes;
use irontide_core::{Id20, Id32};
use irontide_storage::{FileMap, IocpConfig, IocpStorageState, TorrentStorage};
use parking_lot::{Mutex, RwLock};
use tracing::warn;
use windows_sys::Win32::Foundation::{
    CloseHandle, ERROR_IO_PENDING, GetLastError, HANDLE, INVALID_HANDLE_VALUE,
};
use windows_sys::Win32::Storage::FileSystem::{ReadFile, WriteFile};
use windows_sys::Win32::System::IO::{
    CreateIoCompletionPort, GetQueuedCompletionStatusEx, OVERLAPPED, OVERLAPPED_ENTRY,
};
use windows_sys::Win32::System::Threading::INFINITE;

use crate::disk::DiskConfig;
use crate::disk_backend::{DiskIoBackend, DiskIoStats, PosixDiskIo};

/// Per-torrent IOCP state: pre-opened HANDLEs + file map.
pub(crate) struct IocpTorrentState {
    handles: IocpStorageState,
    file_map: FileMap,
}

/// IOCP disk I/O backend.
///
/// Wraps [`PosixDiskIo`] for cache and hashing. Overrides `write_block_direct()`
/// and volatile `read_chunk()` to submit I/O via Windows overlapped IOCP.
///
/// Falls back to the inner [`PosixDiskIo`] when:
/// - Storage does not implement `filesystem_info()` (memory, mmap backends)
/// - IOCP handle opening fails at registration time
/// - An IOCP I/O submission fails at write/read time
pub(crate) struct IocpDiskIo {
    inner: PosixDiskIo,
    iocp: HANDLE,
    /// Serialize submit+reap to prevent completion interleaving between threads.
    iocp_guard: Mutex<()>,
    pub(crate) iocp_states: RwLock<HashMap<Id20, IocpTorrentState>>,
    config: IocpConfig,
    /// Cumulative bytes written via the IOCP path (lock-free stat tracking).
    iocp_write_bytes: AtomicU64,
    /// Cumulative bytes read via the IOCP path (lock-free stat tracking).
    iocp_read_bytes: AtomicU64,
}

/// Extract contiguous bytes from two slices at a given offset and length.
///
/// Returns a `Vec<u8>` containing the data from `s0` and `s1` as if they were
/// concatenated, starting at byte `pos` for `len` bytes.
fn extract_segment_data(s0: &[u8], s1: &[u8], pos: usize, len: usize) -> Vec<u8> {
    let end = pos.saturating_add(len);
    if end <= s0.len() {
        // Entirely within s0.
        s0[pos..end].to_vec()
    } else if pos >= s0.len() {
        // Entirely within s1.
        let s1_start = pos.saturating_sub(s0.len());
        let s1_end = end.saturating_sub(s0.len());
        s1[s1_start..s1_end].to_vec()
    } else {
        // Straddle: part from s0, part from s1.
        let mut buf = Vec::with_capacity(len);
        buf.extend_from_slice(&s0[pos..]);
        let s1_need = len.saturating_sub(s0.len().saturating_sub(pos));
        buf.extend_from_slice(&s1[..s1_need]);
        buf
    }
}

impl IocpDiskIo {
    /// Create a new IOCP backend wrapping [`PosixDiskIo`].
    ///
    /// # Errors
    ///
    /// Returns an I/O error if `CreateIoCompletionPort` fails to create the
    /// completion port.
    pub fn new(disk_config: &DiskConfig, iocp_config: IocpConfig) -> std::io::Result<Self> {
        // SAFETY: Creating a new IOCP with no file handle association.
        // INVALID_HANDLE_VALUE as the file handle + 0 for existing port = new port.
        let iocp = unsafe {
            CreateIoCompletionPort(
                INVALID_HANDLE_VALUE,
                std::ptr::null_mut(),
                0,
                iocp_config.concurrent_threads,
            )
        };
        if iocp.is_null() {
            return Err(std::io::Error::last_os_error());
        }

        Ok(Self {
            inner: PosixDiskIo::new(disk_config),
            iocp,
            iocp_guard: Mutex::new(()),
            iocp_states: RwLock::new(HashMap::new()),
            config: iocp_config,
            iocp_write_bytes: AtomicU64::new(0),
            iocp_read_bytes: AtomicU64::new(0),
        })
    }

    /// Associate a file HANDLE with the IOCP completion port.
    ///
    /// # Errors
    ///
    /// Returns an I/O error if `CreateIoCompletionPort` fails to associate
    /// the handle.
    fn associate_handle(&self, handle: HANDLE) -> std::io::Result<()> {
        // SAFETY: `handle` is a valid file HANDLE opened with FILE_FLAG_OVERLAPPED.
        // `self.iocp` is a valid completion port from the constructor.
        let result = unsafe { CreateIoCompletionPort(handle, self.iocp, 0, 0) };
        if result.is_null() {
            return Err(std::io::Error::last_os_error());
        }
        Ok(())
    }

    /// Submit overlapped writes via IOCP for a single block.
    ///
    /// Maps (piece, begin, s0, s1) to file segments via [`FileMap`], builds
    /// `OVERLAPPED` per segment, issues `WriteFile`, and reaps completions via
    /// `GetQueuedCompletionStatusEx`.
    ///
    /// # Errors
    ///
    /// Returns an I/O error if WriteFile submission fails (not pending), if
    /// reaping fails, or if fewer completions arrive than expected.
    fn iocp_write(
        &self,
        state: &IocpTorrentState,
        piece: u32,
        begin: u32,
        s0: &[u8],
        s1: &[u8],
    ) -> crate::Result<()> {
        let total_len = s0.len().saturating_add(s1.len());
        let total_u32 = u32::try_from(total_len).map_err(|_| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "block length exceeds u32::MAX",
            )
        })?;
        let segments = state.file_map.chunk_segments(piece, begin, total_u32);
        let num_segments = segments.len();

        if num_segments == 0 {
            return Ok(());
        }

        // Prepare per-segment data buffers and OVERLAPPED structures.
        // Buffers must remain alive until completions are reaped.
        let mut seg_bufs: Vec<Vec<u8>> = Vec::with_capacity(num_segments);
        let mut overlappeds: Vec<OVERLAPPED> = Vec::with_capacity(num_segments);

        let mut pos: usize = 0;
        for seg in &segments {
            let seg_len = seg.len as usize;
            let buf = extract_segment_data(s0, s1, pos, seg_len);
            pos = pos.saturating_add(seg_len);

            let mut ov: OVERLAPPED = unsafe { std::mem::zeroed() };
            // Set the file offset in the OVERLAPPED union fields.
            unsafe {
                ov.Anonymous.Anonymous.Offset = (seg.file_offset & 0xFFFF_FFFF) as u32;
                ov.Anonymous.Anonymous.OffsetHigh = (seg.file_offset >> 32) as u32;
            }

            seg_bufs.push(buf);
            overlappeds.push(ov);
        }

        // Lock the IOCP guard to serialize submit+reap.
        let _guard = self.iocp_guard.lock();

        for i in 0..num_segments {
            let handle = state.handles.handle(segments[i].file_index);
            let buf = &seg_bufs[i];
            let ov = &mut overlappeds[i];
            let len_u32 = u32::try_from(buf.len()).map_err(|_| {
                std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "segment length exceeds u32::MAX",
                )
            })?;

            // SAFETY: `handle` is valid and associated with IOCP. `buf` and `ov`
            // are alive for the duration of the I/O. The guard ensures no
            // concurrent submissions interleave completions.
            let ok = unsafe {
                WriteFile(
                    handle,
                    buf.as_ptr(),
                    len_u32,
                    std::ptr::null_mut(),
                    ov as *mut OVERLAPPED,
                )
            };

            if ok == 0 {
                let err = unsafe { GetLastError() };
                if err != ERROR_IO_PENDING {
                    return Err(crate::Error::Io(std::io::Error::from_raw_os_error(
                        err as i32,
                    )));
                }
            }
        }

        // Reap all completions.
        let mut entries: Vec<OVERLAPPED_ENTRY> = vec![unsafe { std::mem::zeroed() }; num_segments];
        let mut removed: u32 = 0;

        // SAFETY: `entries` has capacity for `num_segments` entries. `self.iocp`
        // is a valid completion port. INFINITE timeout blocks until all
        // completions arrive.
        let ok = unsafe {
            GetQueuedCompletionStatusEx(
                self.iocp,
                entries.as_mut_ptr(),
                u32::try_from(num_segments).unwrap_or(u32::MAX),
                &mut removed,
                INFINITE,
                0, // not alertable
            )
        };

        if ok == 0 {
            return Err(crate::Error::Io(std::io::Error::last_os_error()));
        }

        let completed = removed as usize;
        if completed < num_segments {
            return Err(crate::Error::Io(std::io::Error::other(format!(
                "iocp: expected {num_segments} write completions, got {completed}"
            ))));
        }

        // Check for per-entry errors via Internal field.
        for entry in entries.iter().take(completed) {
            // Internal == 0 means success (STATUS_SUCCESS / NTSTATUS 0).
            // Non-zero Internal values are NTSTATUS error codes.
            if entry.Internal != 0 {
                // Convert NTSTATUS to a meaningful error. The lower 16 bits
                // of many NTSTATUS codes map to Win32 error codes, but the
                // safest approach is to report the raw status.
                return Err(crate::Error::Io(std::io::Error::other(format!(
                    "iocp: write completion failed with NTSTATUS 0x{:08X}",
                    entry.Internal
                ))));
            }
        }

        Ok(())
    }

    /// Submit overlapped reads via IOCP for a contiguous chunk.
    ///
    /// Maps (piece, begin, length) to file segments via [`FileMap`], builds
    /// `OVERLAPPED` per segment, issues `ReadFile`, and reaps completions.
    ///
    /// # Errors
    ///
    /// Returns an I/O error if ReadFile submission fails (not pending), if
    /// reaping fails, or if fewer completions arrive than expected.
    fn iocp_read(
        &self,
        state: &IocpTorrentState,
        piece: u32,
        begin: u32,
        length: u32,
    ) -> crate::Result<Vec<u8>> {
        let segments = state.file_map.chunk_segments(piece, begin, length);
        let num_segments = segments.len();

        let mut buf = vec![0u8; length as usize];

        if num_segments == 0 {
            return Ok(buf);
        }

        // Build per-segment buffer slices and OVERLAPPED structures.
        // We track (offset, len) into `buf` for each segment.
        let mut seg_ranges: Vec<(usize, usize)> = Vec::with_capacity(num_segments);
        let mut overlappeds: Vec<OVERLAPPED> = Vec::with_capacity(num_segments);

        let mut pos: usize = 0;
        for seg in &segments {
            let seg_len = seg.len as usize;
            let seg_end = pos.saturating_add(seg_len);

            let mut ov: OVERLAPPED = unsafe { std::mem::zeroed() };
            unsafe {
                ov.Anonymous.Anonymous.Offset = (seg.file_offset & 0xFFFF_FFFF) as u32;
                ov.Anonymous.Anonymous.OffsetHigh = (seg.file_offset >> 32) as u32;
            }

            seg_ranges.push((pos, seg_end));
            overlappeds.push(ov);
            pos = seg_end;
        }

        // Lock the IOCP guard to serialize submit+reap.
        let _guard = self.iocp_guard.lock();

        for i in 0..num_segments {
            let handle = state.handles.handle(segments[i].file_index);
            let (start, end) = seg_ranges[i];
            let seg_buf = &mut buf[start..end];
            let ov = &mut overlappeds[i];
            let len_u32 = u32::try_from(seg_buf.len()).map_err(|_| {
                std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "segment length exceeds u32::MAX",
                )
            })?;

            // SAFETY: `handle` is valid and associated with IOCP. `seg_buf` and
            // `ov` are alive for the duration of the I/O. The guard prevents
            // concurrent interleaving.
            let ok = unsafe {
                ReadFile(
                    handle,
                    seg_buf.as_mut_ptr(),
                    len_u32,
                    std::ptr::null_mut(),
                    ov as *mut OVERLAPPED,
                )
            };

            if ok == 0 {
                let err = unsafe { GetLastError() };
                if err != ERROR_IO_PENDING {
                    return Err(crate::Error::Io(std::io::Error::from_raw_os_error(
                        err as i32,
                    )));
                }
            }
        }

        // Reap all completions.
        let mut entries: Vec<OVERLAPPED_ENTRY> = vec![unsafe { std::mem::zeroed() }; num_segments];
        let mut removed: u32 = 0;

        let ok = unsafe {
            GetQueuedCompletionStatusEx(
                self.iocp,
                entries.as_mut_ptr(),
                u32::try_from(num_segments).unwrap_or(u32::MAX),
                &mut removed,
                INFINITE,
                0,
            )
        };

        if ok == 0 {
            return Err(crate::Error::Io(std::io::Error::last_os_error()));
        }

        let completed = removed as usize;
        if completed < num_segments {
            return Err(crate::Error::Io(std::io::Error::other(format!(
                "iocp: expected {num_segments} read completions, got {completed}"
            ))));
        }

        for entry in entries.iter().take(completed) {
            if entry.Internal != 0 {
                return Err(crate::Error::Io(std::io::Error::other(format!(
                    "iocp: read completion failed with NTSTATUS 0x{:08X}",
                    entry.Internal
                ))));
            }
        }

        Ok(buf)
    }
}

impl Drop for IocpDiskIo {
    fn drop(&mut self) {
        // SAFETY: `self.iocp` was created by `CreateIoCompletionPort` in the
        // constructor and has not been closed. Individual file HANDLEs are
        // closed by `IocpStorageState::drop`.
        unsafe {
            CloseHandle(self.iocp);
        }
    }
}

impl DiskIoBackend for IocpDiskIo {
    fn name(&self) -> &str {
        "iocp"
    }

    fn register(&self, info_hash: Id20, storage: Arc<dyn TorrentStorage>) {
        // Register with inner backend first (for reads/cache).
        self.inner.register(info_hash, Arc::clone(&storage));

        // Try to open IOCP handles via filesystem_info().
        if let Some((base_dir, file_paths, file_map)) = storage.filesystem_info() {
            match IocpStorageState::open_files(base_dir, file_paths, self.config.enable_direct_io) {
                Ok(handles) => {
                    // Associate each handle with our IOCP port.
                    for i in 0..handles.len() {
                        if let Err(e) = self.associate_handle(handles.handle(i)) {
                            warn!(
                                %info_hash,
                                file_index = i,
                                error = %e,
                                "iocp: failed to associate handle, falling back to posix"
                            );
                            return;
                        }
                    }

                    self.iocp_states.write().insert(
                        info_hash,
                        IocpTorrentState {
                            handles,
                            file_map: file_map.clone(),
                        },
                    );
                }
                Err(e) => {
                    warn!(
                        %info_hash,
                        error = %e,
                        "iocp: failed to open handles, falling back to posix"
                    );
                }
            }
        }
    }

    fn unregister(&self, info_hash: Id20) {
        // Remove IOCP state first (Drop closes handles).
        self.iocp_states.write().remove(&info_hash);
        self.inner.unregister(info_hash);
    }

    fn write_chunk(
        &self,
        info_hash: Id20,
        piece: u32,
        begin: u32,
        data: Bytes,
        flush: bool,
    ) -> crate::Result<()> {
        self.inner.write_chunk(info_hash, piece, begin, data, flush)
    }

    fn read_chunk(
        &self,
        info_hash: Id20,
        piece: u32,
        begin: u32,
        length: u32,
        volatile: bool,
    ) -> crate::Result<Bytes> {
        // Non-volatile reads go through inner (cache-aware path).
        if !volatile {
            return self
                .inner
                .read_chunk(info_hash, piece, begin, length, volatile);
        }
        // Volatile reads (won't be re-read) use IOCP directly.
        let iocp_states = self.iocp_states.read();
        if let Some(state) = iocp_states.get(&info_hash) {
            let data = self.iocp_read(state, piece, begin, length)?;
            drop(iocp_states);
            self.iocp_read_bytes
                .fetch_add(u64::from(length), Ordering::Relaxed);
            return Ok(Bytes::from(data));
        }
        drop(iocp_states);
        self.inner
            .read_chunk(info_hash, piece, begin, length, volatile)
    }

    fn read_piece(&self, info_hash: Id20, piece: u32) -> crate::Result<Vec<u8>> {
        let iocp_states = self.iocp_states.read();
        if let Some(state) = iocp_states.get(&info_hash) {
            // Check-before-flush: only flush if piece is in buffer pool cache.
            // The common path (write_block_direct) bypasses the pool entirely,
            // so flushing is usually unnecessary.
            if self.inner.cached_pieces(info_hash).contains(&piece) {
                self.inner.flush_piece(info_hash, piece)?;
            }
            let piece_size = state.file_map.piece_size(piece);
            let data = self.iocp_read(state, piece, 0, piece_size)?;
            drop(iocp_states);
            self.iocp_read_bytes
                .fetch_add(u64::from(piece_size), Ordering::Relaxed);
            return Ok(data);
        }
        drop(iocp_states);
        self.inner.read_piece(info_hash, piece)
    }

    fn hash_piece(&self, info_hash: Id20, piece: u32, expected: &Id20) -> crate::Result<bool> {
        self.inner.hash_piece(info_hash, piece, expected)
    }

    fn hash_piece_v2(&self, info_hash: Id20, piece: u32, expected: &Id32) -> crate::Result<bool> {
        self.inner.hash_piece_v2(info_hash, piece, expected)
    }

    fn hash_block(
        &self,
        info_hash: Id20,
        piece: u32,
        begin: u32,
        length: u32,
    ) -> crate::Result<Id32> {
        self.inner.hash_block(info_hash, piece, begin, length)
    }

    fn clear_piece(&self, info_hash: Id20, piece: u32) {
        self.inner.clear_piece(info_hash, piece)
    }

    fn flush_piece(&self, info_hash: Id20, piece: u32) -> crate::Result<()> {
        self.inner.flush_piece(info_hash, piece)
    }

    fn flush_all(&self) -> crate::Result<()> {
        self.inner.flush_all()
    }

    fn cached_pieces(&self, info_hash: Id20) -> Vec<u32> {
        self.inner.cached_pieces(info_hash)
    }

    fn stats(&self) -> DiskIoStats {
        let mut s = self.inner.stats();
        s.write_bytes = s
            .write_bytes
            .saturating_add(self.iocp_write_bytes.load(Ordering::Relaxed));
        s.read_bytes = s
            .read_bytes
            .saturating_add(self.iocp_read_bytes.load(Ordering::Relaxed));
        s
    }

    fn write_block_direct(
        &self,
        info_hash: Id20,
        piece: u32,
        begin: u32,
        s0: &[u8],
        s1: &[u8],
    ) -> crate::Result<()> {
        // If we have IOCP state for this torrent, use IOCP.
        let iocp_states = self.iocp_states.read();
        if let Some(state) = iocp_states.get(&info_hash) {
            let result = self.iocp_write(state, piece, begin, s0, s1);
            drop(iocp_states);

            if result.is_ok() {
                let total = (s0.len().saturating_add(s1.len())) as u64;
                self.iocp_write_bytes.fetch_add(total, Ordering::Relaxed);
                return Ok(());
            }

            // Fall through to inner on failure.
            warn!(
                %info_hash,
                piece,
                begin,
                error = %result.as_ref().unwrap_err(),
                "iocp write failed, falling back to posix"
            );
        } else {
            drop(iocp_states);
        }

        self.inner
            .write_block_direct(info_hash, piece, begin, s0, s1)
    }
}