datum/io/

tokio_io.rs

use crate::stream::{BoxStream, Flow, NotUsed, Sink, Source, StreamCompletion};
use crate::{StreamError, StreamResult};
use futures::{FutureExt, channel::oneshot};
use std::future::Future;
use std::net::SocketAddr;
use std::panic::AssertUnwindSafe;
use std::path::PathBuf;
use std::sync::{
    Arc, Mutex,
    atomic::{AtomicBool, Ordering},
    mpsc as std_mpsc,
};
use std::thread::{self, Thread};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, ToSocketAddrs};
use tokio::sync::{mpsc, watch};

const DEFAULT_CHUNK_SIZE: usize = 8192;
const FILE_READ_AHEAD_CHUNKS: usize = 8;
const FILE_INTERNAL_READ_SIZE: usize = 256 * 1024;
const TCP_READ_AHEAD_CHUNKS: usize = 1;
const PARK_INTERVAL: Duration = Duration::from_millis(1);
const READ_READY_SPINS: usize = 256;
const BACKPRESSURE_READY_SPINS: usize = 64;
const BACKPRESSURE_PARK: Duration = Duration::from_micros(10);

/// Shared wake-on-send handoff between the consumer poll loop and the
/// producer task. The consumer captures its `Thread` on the first `next()`
/// call; the producer calls `unpark()` after every successful channel send
/// so the consumer's `park_timeout` returns immediately instead of waiting
/// for the next timeout tick. The per-iteration `stream_cancelled` check
/// and the cancellation `watch` on the producer side are preserved, so
/// neither side loses its termination signal.
#[derive(Clone)]
struct ConsumerWaker {
    thread: Arc<Mutex<Option<Thread>>>,
}

impl ConsumerWaker {
    fn new() -> Self {
        Self {
            thread: Arc::new(Mutex::new(None)),
        }
    }

    fn capture_current(&self) {
        let mut slot = self.thread.lock().expect("consumer waker poisoned");
        if slot.is_none() {
            *slot = Some(thread::current());
        }
    }

    fn unpark(&self) {
        let slot = self.thread.lock().expect("consumer waker poisoned");
        if let Some(t) = slot.as_ref() {
            t.unpark();
        }
    }
}

fn io_error(error: std::io::Error) -> StreamError {
    StreamError::Failed(error.to_string())
}

fn write_zero_error() -> StreamError {
    StreamError::Failed("async writer returned zero bytes".to_owned())
}

/// Akka-style materialized IO result for Tokio-backed file and TCP byte streams.
///
/// `bytes` records the number of bytes successfully read or written before the
/// terminal status. IO failures are also propagated through the stream as
/// [`StreamError`]; the materialized result preserves the byte count for
/// partial reads/writes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IoResult {
    pub bytes: u64,
    pub status: StreamResult<()>,
}

impl IoResult {
    /// A successful result that transferred `bytes` bytes.
    #[must_use]
    pub fn succeeded(bytes: u64) -> Self {
        Self {
            bytes,
            status: Ok(()),
        }
    }

    /// A failed result preserving the `bytes` transferred before `error`.
    #[must_use]
    pub fn failed(bytes: u64, error: StreamError) -> Self {
        Self {
            bytes,
            status: Err(error),
        }
    }

    /// Bytes successfully transferred before the terminal status.
    #[must_use]
    pub fn bytes(&self) -> u64 {
        self.bytes
    }

    /// Terminal status: `Ok(())` on success, `Err(StreamError)` on failure.
    #[must_use = "the terminal status should be inspected"]
    pub fn status(&self) -> StreamResult<()> {
        self.status.clone()
    }

    /// `true` when the transfer completed without error.
    #[must_use]
    pub fn is_success(&self) -> bool {
        self.status.is_ok()
    }
}

/// A byte source whose materialized value resolves to an [`IoResult`] (bytes + terminal status).
pub type TokioByteSource = Source<Vec<u8>, StreamCompletion<IoResult>>;
/// A byte sink whose materialized value resolves to an [`IoResult`] (bytes + terminal status).
pub type TokioByteSink = Sink<Vec<u8>, StreamCompletion<IoResult>>;

#[derive(Clone)]
enum DemandTerminal {
    Complete,
    Error(StreamError),
}

enum DemandResponse<T> {
    Item(T),
    Complete,
    Error(StreamError),
}

struct DemandSourceStream<T> {
    demands: mpsc::Sender<std_mpsc::Sender<DemandResponse<T>>>,
    cancel: watch::Sender<bool>,
    terminal: Arc<Mutex<Option<DemandTerminal>>>,
    done: bool,
}

impl<T> DemandSourceStream<T> {
    fn terminal_response(&self) -> Option<Option<StreamResult<T>>> {
        self.terminal
            .lock()
            .expect("tokio source terminal poisoned")
            .clone()
            .map(|terminal| match terminal {
                DemandTerminal::Complete => None,
                DemandTerminal::Error(error) => Some(Err(error)),
            })
    }

    fn mark_done(&mut self) {
        self.done = true;
        let _ = self.cancel.send(true);
    }
}

impl<T: Send + 'static> Iterator for DemandSourceStream<T> {
    type Item = StreamResult<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }

        let stream_cancelled = crate::stream::current_stream_cancelled();
        let (reply_sender, reply_receiver) = std_mpsc::channel();
        if !send_bounded_demand(&self.demands, reply_sender, &stream_cancelled) {
            self.mark_done();
            return self
                .terminal_response()
                .unwrap_or(Some(Err(StreamError::AbruptTermination)));
        }

        loop {
            if stream_cancelled
                .as_ref()
                .is_some_and(|cancelled| cancelled.load(Ordering::SeqCst))
            {
                self.mark_done();
                return Some(Err(StreamError::Cancelled));
            }

            match reply_receiver.recv_timeout(PARK_INTERVAL) {
                Ok(DemandResponse::Item(item)) => return Some(Ok(item)),
                Ok(DemandResponse::Complete) => {
                    self.mark_done();
                    return None;
                }
                Ok(DemandResponse::Error(error)) => {
                    self.mark_done();
                    return Some(Err(error));
                }
                Err(std_mpsc::RecvTimeoutError::Timeout) => {}
                Err(std_mpsc::RecvTimeoutError::Disconnected) => {
                    self.mark_done();
                    return self
                        .terminal_response()
                        .unwrap_or(Some(Err(StreamError::AbruptTermination)));
                }
            }
        }
    }
}

impl<T> Drop for DemandSourceStream<T> {
    fn drop(&mut self) {
        let _ = self.cancel.send(true);
    }
}

struct BoundedByteSourceStream {
    receiver: mpsc::Receiver<DemandResponse<Vec<u8>>>,
    cancel: watch::Sender<bool>,
    terminal: Arc<Mutex<Option<DemandTerminal>>>,
    done: bool,
    waker: ConsumerWaker,
}

impl BoundedByteSourceStream {
    fn terminal_response(&self) -> Option<Option<StreamResult<Vec<u8>>>> {
        self.terminal
            .lock()
            .expect("tokio source terminal poisoned")
            .clone()
            .map(|terminal| match terminal {
                DemandTerminal::Complete => None,
                DemandTerminal::Error(error) => Some(Err(error)),
            })
    }

    fn mark_done(&mut self) {
        self.done = true;
        let _ = self.cancel.send(true);
    }
}

impl Iterator for BoundedByteSourceStream {
    type Item = StreamResult<Vec<u8>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }

        // Register this thread with the producer so it can unpark us
        // immediately after each successful send instead of waiting for
        // the next park_timeout tick.
        self.waker.capture_current();

        let stream_cancelled = crate::stream::current_stream_cancelled();
        let mut spins = 0usize;
        loop {
            if stream_cancelled
                .as_ref()
                .is_some_and(|cancelled| cancelled.load(Ordering::SeqCst))
            {
                self.mark_done();
                return Some(Err(StreamError::Cancelled));
            }

            match self.receiver.try_recv() {
                Ok(DemandResponse::Item(item)) => return Some(Ok(item)),
                Ok(DemandResponse::Complete) => {
                    self.mark_done();
                    return None;
                }
                Ok(DemandResponse::Error(error)) => {
                    self.mark_done();
                    return Some(Err(error));
                }
                Err(mpsc::error::TryRecvError::Empty) => read_wait(&mut spins),
                Err(mpsc::error::TryRecvError::Disconnected) => {
                    self.mark_done();
                    return self
                        .terminal_response()
                        .unwrap_or(Some(Err(StreamError::AbruptTermination)));
                }
            }
        }
    }
}

impl Drop for BoundedByteSourceStream {
    fn drop(&mut self) {
        let _ = self.cancel.send(true);
    }
}

fn send_bounded_demand<T>(
    sender: &mpsc::Sender<T>,
    mut message: T,
    stream_cancelled: &Option<Arc<AtomicBool>>,
) -> bool {
    let mut spins = 0usize;
    loop {
        if stream_cancelled
            .as_ref()
            .is_some_and(|cancelled| cancelled.load(Ordering::SeqCst))
        {
            return false;
        }

        match sender.try_send(message) {
            Ok(()) => return true,
            Err(mpsc::error::TrySendError::Full(returned)) => {
                message = returned;
                backpressure_wait(&mut spins);
            }
            Err(mpsc::error::TrySendError::Closed(_)) => return false,
        }
    }
}

fn finish_terminal(terminal: &Arc<Mutex<Option<DemandTerminal>>>, value: DemandTerminal) {
    let mut slot = terminal.lock().expect("tokio source terminal poisoned");
    if slot.is_none() {
        *slot = Some(value);
    }
}

async fn next_demand<T>(
    demands: &mut mpsc::Receiver<std_mpsc::Sender<DemandResponse<T>>>,
    cancel: &mut watch::Receiver<bool>,
) -> Option<std_mpsc::Sender<DemandResponse<T>>> {
    if *cancel.borrow() {
        return None;
    }

    tokio::select! {
        demand = demands.recv() => demand,
        changed = cancel.changed() => {
            let _ = changed;
            None
        }
    }
}

fn async_read_source<R, Fut>(
    open: impl FnOnce() -> Fut + Send + 'static,
    chunk_size: usize,
    internal_read_size: usize,
    read_ahead_chunks: usize,
) -> (BoxStream<Vec<u8>>, StreamCompletion<IoResult>)
where
    R: AsyncRead + Unpin + Send + 'static,
    Fut: Future<Output = std::io::Result<R>> + Send + 'static,
{
    assert!(chunk_size > 0, "chunk size must be greater than zero");
    assert!(
        read_ahead_chunks > 0,
        "read-ahead bound must be greater than zero"
    );
    let internal_read_size = internal_read_size.max(chunk_size);
    let (item_sender, item_receiver) = mpsc::channel(read_ahead_chunks);
    let (cancel_sender, cancel_receiver) = watch::channel(false);
    let (mat_sender, mat_receiver) = oneshot::channel();
    let terminal = Arc::new(Mutex::new(None));
    let terminal_for_task = Arc::clone(&terminal);
    let waker = ConsumerWaker::new();
    let producer_waker = waker.clone();

    crate::stream::stream_tokio_runtime().spawn(async move {
        let result = AssertUnwindSafe(run_async_read_task(
            open(),
            chunk_size,
            internal_read_size,
            item_sender,
            cancel_receiver,
            Arc::clone(&terminal_for_task),
            producer_waker,
        ))
        .catch_unwind()
        .await
        .unwrap_or_else(|_| {
            finish_terminal(
                &terminal_for_task,
                DemandTerminal::Error(StreamError::AbruptTermination),
            );
            Err(StreamError::AbruptTermination)
        });
        let _ = mat_sender.send(result);
    });

    (
        Box::new(BoundedByteSourceStream {
            receiver: item_receiver,
            cancel: cancel_sender,
            terminal,
            done: false,
            waker,
        }) as BoxStream<Vec<u8>>,
        StreamCompletion::from_receiver(mat_receiver, None),
    )
}

async fn run_async_read_task<R, Fut>(
    open: Fut,
    chunk_size: usize,
    internal_read_size: usize,
    items: mpsc::Sender<DemandResponse<Vec<u8>>>,
    mut cancel: watch::Receiver<bool>,
    terminal: Arc<Mutex<Option<DemandTerminal>>>,
    waker: ConsumerWaker,
) -> StreamResult<IoResult>
where
    R: AsyncRead + Unpin + Send + 'static,
    Fut: Future<Output = std::io::Result<R>> + Send + 'static,
{
    let mut bytes = 0_u64;
    let mut reader = tokio::select! {
        reader = open => match reader {
            Ok(reader) => reader,
            Err(error) => {
                let error = io_error(error);
                finish_terminal(&terminal, DemandTerminal::Error(error.clone()));
                let _ = send_read_item(&items, DemandResponse::Error(error.clone()), &mut cancel, &waker).await;
                return Ok(IoResult::failed(bytes, error));
            }
        },
        changed = cancel.changed() => {
            let _ = changed;
            finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
            return Ok(IoResult::failed(bytes, StreamError::Cancelled));
        }
    };

    let mut buffer = vec![0_u8; internal_read_size];
    let mut pending_tail = Vec::with_capacity(chunk_size);
    loop {
        let read = tokio::select! {
            read = reader.read(&mut buffer) => read,
            changed = cancel.changed() => {
                let _ = changed;
                finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
                return Ok(IoResult::failed(bytes, StreamError::Cancelled));
            }
        };

        match read {
            Ok(0) => {
                if !pending_tail.is_empty()
                    && !send_read_item(
                        &items,
                        DemandResponse::Item(std::mem::take(&mut pending_tail)),
                        &mut cancel,
                        &waker,
                    )
                    .await
                {
                    finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
                    return Ok(IoResult::failed(bytes, StreamError::Cancelled));
                }
                finish_terminal(&terminal, DemandTerminal::Complete);
                let _ = send_read_item(&items, DemandResponse::Complete, &mut cancel, &waker).await;
                return Ok(IoResult::succeeded(bytes));
            }
            Ok(read) => {
                bytes += read as u64;
                if !send_read_chunks(
                    &items,
                    chunk_size,
                    &mut pending_tail,
                    &buffer[..read],
                    &mut cancel,
                    &waker,
                )
                .await
                {
                    finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
                    return Ok(IoResult::failed(bytes, StreamError::Cancelled));
                }
            }
            Err(error) => {
                let error = io_error(error);
                finish_terminal(&terminal, DemandTerminal::Error(error.clone()));
                let _ = send_read_item(
                    &items,
                    DemandResponse::Error(error.clone()),
                    &mut cancel,
                    &waker,
                )
                .await;
                return Ok(IoResult::failed(bytes, error));
            }
        }
    }
}

async fn send_read_chunks(
    sender: &mpsc::Sender<DemandResponse<Vec<u8>>>,
    chunk_size: usize,
    pending_tail: &mut Vec<u8>,
    read_buffer: &[u8],
    cancel: &mut watch::Receiver<bool>,
    waker: &ConsumerWaker,
) -> bool {
    let mut offset = 0;
    if !pending_tail.is_empty() {
        let needed = chunk_size - pending_tail.len();
        let take = needed.min(read_buffer.len());
        pending_tail.extend_from_slice(&read_buffer[..take]);
        offset += take;
        if pending_tail.len() == chunk_size
            && !send_read_item(
                sender,
                DemandResponse::Item(std::mem::take(pending_tail)),
                cancel,
                waker,
            )
            .await
        {
            return false;
        }
    }

    while offset + chunk_size <= read_buffer.len() {
        let next = offset + chunk_size;
        if !send_read_item(
            sender,
            DemandResponse::Item(read_buffer[offset..next].to_vec()),
            cancel,
            waker,
        )
        .await
        {
            return false;
        }
        offset = next;
    }

    if offset < read_buffer.len() {
        pending_tail.extend_from_slice(&read_buffer[offset..]);
    }
    true
}

async fn send_read_item<T>(
    sender: &mpsc::Sender<DemandResponse<T>>,
    item: DemandResponse<T>,
    cancel: &mut watch::Receiver<bool>,
    waker: &ConsumerWaker,
) -> bool
where
    T: Send + 'static,
{
    let result = tokio::select! {
        result = sender.send(item) => result,
        changed = cancel.changed() => {
            let _ = changed;
            return false;
        }
    };
    if result.is_ok() {
        waker.unpark();
    }
    result.is_ok()
}

enum WriteCommand {
    Chunk(Vec<u8>),
    Finish(StreamResult<()>),
}

struct TokioCancelGuard {
    cancel: watch::Sender<bool>,
    armed: bool,
}

impl TokioCancelGuard {
    fn new(cancel: watch::Sender<bool>) -> Self {
        Self {
            cancel,
            armed: true,
        }
    }

    fn disarm(&mut self) {
        self.armed = false;
    }
}

impl Drop for TokioCancelGuard {
    fn drop(&mut self) {
        if self.armed {
            let _ = self.cancel.send(true);
        }
    }
}

fn async_write_sink<W, F, Fut>(open: F) -> TokioByteSink
where
    W: AsyncWrite + Unpin + Send + 'static,
    F: Fn() -> Fut + Send + Sync + 'static,
    Fut: Future<Output = std::io::Result<W>> + Send + 'static,
{
    let open = Arc::new(open);
    Sink::from_runner(move |input, materializer| {
        run_async_write_sink::<W, F, Fut>(input, materializer, Arc::clone(&open))
    })
}

async fn run_async_write_task<W, Fut>(
    open: Fut,
    mut commands: mpsc::Receiver<WriteCommand>,
    mut cancel: watch::Receiver<bool>,
) -> StreamResult<IoResult>
where
    W: AsyncWrite + Unpin + Send + 'static,
    Fut: Future<Output = std::io::Result<W>> + Send + 'static,
{
    let mut bytes = 0_u64;
    let mut writer = tokio::select! {
        writer = open => match writer {
            Ok(writer) => writer,
            Err(error) => return Ok(IoResult::failed(bytes, io_error(error))),
        },
        changed = cancel.changed() => {
            let _ = changed;
            return Ok(IoResult::failed(bytes, StreamError::Cancelled));
        }
    };

    loop {
        let command = tokio::select! {
            command = commands.recv() => command,
            changed = cancel.changed() => {
                let _ = changed;
                return Ok(IoResult::failed(bytes, StreamError::Cancelled));
            }
        };

        match command {
            Some(WriteCommand::Chunk(chunk)) => {
                if let Err(error) = write_chunk(&mut writer, &chunk, &mut cancel, &mut bytes).await
                {
                    return Ok(IoResult::failed(bytes, error));
                }
            }
            Some(WriteCommand::Finish(upstream_status)) => {
                let shutdown_status = shutdown_writer(&mut writer, &mut cancel).await;
                return Ok(IoResult {
                    bytes,
                    status: upstream_status.and(shutdown_status),
                });
            }
            None => {
                let _ = shutdown_writer(&mut writer, &mut cancel).await;
                return Ok(IoResult::failed(bytes, StreamError::Cancelled));
            }
        }
    }
}

async fn write_chunk<W>(
    writer: &mut W,
    chunk: &[u8],
    cancel: &mut watch::Receiver<bool>,
    bytes: &mut u64,
) -> StreamResult<()>
where
    W: AsyncWrite + Unpin,
{
    let mut offset = 0usize;
    while offset < chunk.len() {
        let written = tokio::select! {
            written = writer.write(&chunk[offset..]) => written.map_err(io_error)?,
            changed = cancel.changed() => {
                let _ = changed;
                return Err(StreamError::Cancelled);
            }
        };

        if written == 0 {
            return Err(write_zero_error());
        }
        offset += written;
        *bytes += written as u64;
    }
    Ok(())
}

async fn shutdown_writer<W>(writer: &mut W, cancel: &mut watch::Receiver<bool>) -> StreamResult<()>
where
    W: AsyncWrite + Unpin,
{
    tokio::select! {
        result = writer.flush() => result.map_err(io_error)?,
        changed = cancel.changed() => {
            let _ = changed;
            return Err(StreamError::Cancelled);
        }
    }

    tokio::select! {
        result = writer.shutdown() => result.map_err(io_error),
        changed = cancel.changed() => {
            let _ = changed;
            Err(StreamError::Cancelled)
        }
    }
}

fn feed_async_writer(
    mut input: BoxStream<Vec<u8>>,
    command_sender: mpsc::Sender<WriteCommand>,
    done_receiver: std_mpsc::Receiver<StreamResult<IoResult>>,
    cancelled: Arc<AtomicBool>,
    cancel_sender: watch::Sender<bool>,
) -> StreamResult<IoResult> {
    let mut terminal = Ok(());
    loop {
        if cancelled.load(Ordering::SeqCst) {
            terminal = Err(StreamError::Cancelled);
            break;
        }

        match input.next() {
            Some(Ok(chunk)) => {
                if !send_write_command(&command_sender, WriteCommand::Chunk(chunk), &cancelled) {
                    break;
                }
            }
            Some(Err(error)) => {
                terminal = Err(error);
                break;
            }
            None => break,
        }
    }

    if cancelled.load(Ordering::SeqCst) {
        let _ = cancel_sender.send(true);
    } else {
        let _ = send_write_command(&command_sender, WriteCommand::Finish(terminal), &cancelled);
    }
    drop(command_sender);

    loop {
        match done_receiver.recv_timeout(PARK_INTERVAL) {
            Ok(result) => return result,
            Err(std_mpsc::RecvTimeoutError::Timeout) => {
                if cancelled.load(Ordering::SeqCst) {
                    // Cancellation can arrive after Finish is sent while the
                    // Tokio writer is stuck in write/flush/shutdown. Forward
                    // the watch signal and keep waiting for its real IoResult.
                    let _ = cancel_sender.send(true);
                }
            }
            Err(std_mpsc::RecvTimeoutError::Disconnected) => {
                return Err(StreamError::AbruptTermination);
            }
        }
    }
}

fn send_write_command(
    sender: &mpsc::Sender<WriteCommand>,
    mut command: WriteCommand,
    cancelled: &AtomicBool,
) -> bool {
    let mut spins = 0usize;
    loop {
        if cancelled.load(Ordering::SeqCst) {
            return false;
        }

        match sender.try_send(command) {
            Ok(()) => return true,
            Err(mpsc::error::TrySendError::Full(returned)) => {
                command = returned;
                backpressure_wait(&mut spins);
            }
            Err(mpsc::error::TrySendError::Closed(_)) => return false,
        }
    }
}

fn backpressure_wait(spins: &mut usize) {
    if *spins < BACKPRESSURE_READY_SPINS {
        *spins += 1;
        thread::yield_now();
    } else {
        thread::park_timeout(BACKPRESSURE_PARK);
    }
}

fn read_wait(spins: &mut usize) {
    if *spins < READ_READY_SPINS {
        *spins += 1;
        thread::yield_now();
    } else {
        thread::park_timeout(PARK_INTERVAL);
    }
}

fn tokio_file_read_source(
    path: PathBuf,
    chunk_size: usize,
) -> (BoxStream<Vec<u8>>, StreamCompletion<IoResult>) {
    async_read_source(
        move || tokio::fs::File::open(path),
        chunk_size,
        FILE_INTERNAL_READ_SIZE,
        FILE_READ_AHEAD_CHUNKS,
    )
}

async fn tokio_file_write_open(path: Arc<PathBuf>) -> std::io::Result<tokio::fs::File> {
    tokio::fs::OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .open(path.as_ref())
        .await
}

fn run_async_write_sink<W, F, Fut>(
    input: BoxStream<Vec<u8>>,
    materializer: &crate::stream::Materializer,
    open: Arc<F>,
) -> StreamResult<StreamCompletion<IoResult>>
where
    W: AsyncWrite + Unpin + Send + 'static,
    F: Fn() -> Fut + Send + Sync + 'static,
    Fut: Future<Output = std::io::Result<W>> + Send + 'static,
{
    let (command_sender, command_receiver) = mpsc::channel(1);
    let (cancel_sender, cancel_receiver) = watch::channel(false);
    let (done_sender, done_receiver) = std_mpsc::sync_channel(1);

    crate::stream::stream_tokio_runtime().spawn(async move {
        let result = AssertUnwindSafe(run_async_write_task(
            open(),
            command_receiver,
            cancel_receiver,
        ))
        .catch_unwind()
        .await
        .unwrap_or(Err(StreamError::AbruptTermination));
        let _ = done_sender.send(result);
    });

    Ok(materializer.spawn_stream(move |cancelled| {
        let mut guard = TokioCancelGuard::new(cancel_sender.clone());
        let result = feed_async_writer(
            input,
            command_sender,
            done_receiver,
            cancelled,
            cancel_sender,
        );
        guard.disarm();
        result
    }))
}

fn tokio_file_write_sink(path: PathBuf) -> TokioByteSink {
    let path = Arc::new(path);
    async_write_sink(move || tokio_file_write_open(Arc::clone(&path)))
}

pub struct TokioFileIO;

impl TokioFileIO {
    /// Creates a Tokio-backed file source that reads `path` into byte chunks.
    ///
    /// The source uses the shared Datum stream Tokio runtime and opens the file
    /// asynchronously when materialized. File reads use a 256 KiB internal
    /// buffer but preserve `chunk_size` emissions except for the final tail.
    /// Reads have a bounded read-ahead of at most eight emitted chunks; once
    /// that buffer is full, the Tokio task stops reading until downstream
    /// consumes data. The materialized [`IoResult`] records the total bytes
    /// read and terminal status.
    #[must_use]
    pub fn from_path(path: impl Into<PathBuf>, chunk_size: usize) -> TokioByteSource {
        assert!(chunk_size > 0, "chunk size must be greater than zero");
        let path = path.into();
        Source::from_materialized_factory(move |_materializer| {
            let path = path.clone();
            Ok(tokio_file_read_source(path, chunk_size))
        })
    }

    /// [`TokioFileIO::from_path`] with the default 8 KiB chunk size.
    #[must_use]
    pub fn from_path_default(path: impl Into<PathBuf>) -> TokioByteSource {
        Self::from_path(path, DEFAULT_CHUNK_SIZE)
    }

    /// Creates a Tokio-backed file sink that truncates or creates `path`.
    ///
    /// The sink writes one upstream chunk at a time through a capacity-1 Tokio
    /// channel so a slow file write backpressures the stream. Partial write
    /// failures preserve the number of bytes successfully written in the
    /// materialized [`IoResult`].
    #[must_use]
    pub fn to_path(path: impl Into<PathBuf>) -> TokioByteSink {
        tokio_file_write_sink(path.into())
    }
}

/// Opt-in `tokio-uring` backed file I/O.
///
/// `UringFileIO` is available only when the `io-uring-file` feature is enabled.
/// `TokioFileIO` remains the default implementation. On Linux, materialized
/// `UringFileIO` sources and sinks enqueue file jobs onto one feature-owned
/// `tokio-uring` current-thread runtime, so Datum's/Ractor's Tokio runtime does
/// not change. If that runtime cannot be created, or if this feature is enabled on a
/// non-Linux target, the materialization uses the same `TokioFileIO` fallback
/// path and preserves the same chunking, terminal, failure, and cancellation
/// semantics.
#[cfg(feature = "io-uring-file")]
pub struct UringFileIO;

#[cfg(feature = "io-uring-file")]
impl UringFileIO {
    /// Creates an opt-in io_uring-backed file source.
    ///
    /// The emitted chunks match [`TokioFileIO::from_path`]: `chunk_size`
    /// boundaries are preserved except for the final tail, the materialized
    /// [`IoResult`] records bytes read, and cancellation fails the materialized
    /// result with [`StreamError::Cancelled`].
    #[must_use]
    pub fn from_path(path: impl Into<PathBuf>, chunk_size: usize) -> TokioByteSource {
        assert!(chunk_size > 0, "chunk size must be greater than zero");
        let path = path.into();
        Source::from_materialized_factory(move |_materializer| {
            let path = path.clone();
            Ok(uring_file_read_source_or_fallback(path, chunk_size))
        })
    }

    /// [`UringFileIO::from_path`] with the default 8 KiB chunk size.
    #[must_use]
    pub fn from_path_default(path: impl Into<PathBuf>) -> TokioByteSource {
        Self::from_path(path, DEFAULT_CHUNK_SIZE)
    }

    /// Creates an opt-in io_uring-backed file sink.
    ///
    /// The sink truncates or creates `path` and preserves
    /// [`TokioFileIO::to_path`] materialized byte-count and terminal semantics.
    #[must_use]
    pub fn to_path(path: impl Into<PathBuf>) -> TokioByteSink {
        uring_file_write_sink_or_fallback(path.into())
    }
}

#[cfg(all(feature = "io-uring-file", not(target_os = "linux")))]
fn uring_file_read_source_or_fallback(
    path: PathBuf,
    chunk_size: usize,
) -> (BoxStream<Vec<u8>>, StreamCompletion<IoResult>) {
    tokio_file_read_source(path, chunk_size)
}

#[cfg(all(feature = "io-uring-file", not(target_os = "linux")))]
fn uring_file_write_sink_or_fallback(path: PathBuf) -> TokioByteSink {
    tokio_file_write_sink(path)
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
type UringJob =
    Box<dyn FnOnce() -> std::pin::Pin<Box<dyn Future<Output = ()> + 'static>> + Send + 'static>;

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
#[derive(Clone)]
struct UringRuntimeHandle {
    sender: mpsc::UnboundedSender<UringJob>,
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
enum UringRuntimeInit {
    Ready,
    Failed(std::io::Error),
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
fn uring_runtime_handle() -> Result<&'static UringRuntimeHandle, String> {
    static HANDLE: std::sync::OnceLock<Result<UringRuntimeHandle, String>> =
        std::sync::OnceLock::new();
    HANDLE
        .get_or_init(start_uring_runtime)
        .as_ref()
        .map_err(Clone::clone)
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
fn start_uring_runtime() -> Result<UringRuntimeHandle, String> {
    let (init_sender, init_receiver) = std_mpsc::sync_channel(1);
    let (sender, receiver) = mpsc::unbounded_channel::<UringJob>();
    thread::Builder::new()
        .name("datum-uring-file".to_owned())
        .spawn(
            move || match tokio_uring::Runtime::new(&tokio_uring::builder()) {
                Ok(runtime) => {
                    let _ = init_sender.send(UringRuntimeInit::Ready);
                    runtime.block_on(run_uring_jobs(receiver));
                }
                Err(error) => {
                    let _ = init_sender.send(UringRuntimeInit::Failed(error));
                }
            },
        )
        .map_err(|error| error.to_string())?;

    match init_receiver.recv() {
        Ok(UringRuntimeInit::Ready) => Ok(UringRuntimeHandle { sender }),
        Ok(UringRuntimeInit::Failed(error)) => Err(error.to_string()),
        Err(_) => Err("tokio-uring runtime thread exited".to_owned()),
    }
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
async fn run_uring_jobs(mut receiver: mpsc::UnboundedReceiver<UringJob>) {
    while let Some(job) = receiver.recv().await {
        tokio_uring::spawn(job());
    }
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
fn spawn_uring_job(job: UringJob) -> Result<(), String> {
    uring_runtime_handle()?
        .sender
        .send(job)
        .map_err(|_| "tokio-uring runtime thread exited".to_owned())
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
fn uring_file_read_source_or_fallback(
    path: PathBuf,
    chunk_size: usize,
) -> (BoxStream<Vec<u8>>, StreamCompletion<IoResult>) {
    let internal_read_size = FILE_INTERNAL_READ_SIZE.max(chunk_size);
    let (item_sender, item_receiver) = mpsc::channel(FILE_READ_AHEAD_CHUNKS);
    let (cancel_sender, cancel_receiver) = watch::channel(false);
    let (mat_sender, mat_receiver) = oneshot::channel();
    let terminal = Arc::new(Mutex::new(None));
    let terminal_for_task = Arc::clone(&terminal);
    let waker = ConsumerWaker::new();
    let producer_waker = waker.clone();
    let uring_path = path.clone();

    let result = spawn_uring_job(Box::new(move || {
        Box::pin(async move {
            let task_result = AssertUnwindSafe(run_uring_read_task(
                uring_path,
                chunk_size,
                internal_read_size,
                item_sender,
                cancel_receiver,
                Arc::clone(&terminal_for_task),
                producer_waker,
            ))
            .catch_unwind()
            .await
            .unwrap_or_else(|_| {
                finish_terminal(
                    &terminal_for_task,
                    DemandTerminal::Error(StreamError::AbruptTermination),
                );
                Err(StreamError::AbruptTermination)
            });
            let _ = mat_sender.send(task_result);
        })
    }));

    if result.is_err() {
        return tokio_file_read_source(path, chunk_size);
    }

    (
        Box::new(BoundedByteSourceStream {
            receiver: item_receiver,
            cancel: cancel_sender,
            terminal,
            done: false,
            waker,
        }) as BoxStream<Vec<u8>>,
        StreamCompletion::from_receiver(mat_receiver, None),
    )
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
async fn run_uring_read_task(
    path: PathBuf,
    chunk_size: usize,
    internal_read_size: usize,
    items: mpsc::Sender<DemandResponse<Vec<u8>>>,
    mut cancel: watch::Receiver<bool>,
    terminal: Arc<Mutex<Option<DemandTerminal>>>,
    waker: ConsumerWaker,
) -> StreamResult<IoResult> {
    let mut bytes = 0_u64;
    let file = tokio::select! {
        file = tokio_uring::fs::File::open(path) => match file {
            Ok(file) => file,
            Err(error) => {
                let error = io_error(error);
                finish_terminal(&terminal, DemandTerminal::Error(error.clone()));
                let _ = send_read_item(&items, DemandResponse::Error(error.clone()), &mut cancel, &waker).await;
                return Ok(IoResult::failed(bytes, error));
            }
        },
        changed = cancel.changed() => {
            let _ = changed;
            finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
            return Ok(IoResult::failed(bytes, StreamError::Cancelled));
        }
    };

    let mut offset = 0_u64;
    let mut buffer = Vec::with_capacity(internal_read_size);
    let mut pending_tail = Vec::with_capacity(chunk_size);
    loop {
        let read_buffer = std::mem::take(&mut buffer);
        let (read, returned_buffer) = tokio::select! {
            result = file.read_at(read_buffer, offset) => result,
            changed = cancel.changed() => {
                let _ = changed;
                finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
                return Ok(IoResult::failed(bytes, StreamError::Cancelled));
            }
        };
        buffer = returned_buffer;

        match read {
            Ok(0) => {
                if !pending_tail.is_empty()
                    && !send_read_item(
                        &items,
                        DemandResponse::Item(std::mem::take(&mut pending_tail)),
                        &mut cancel,
                        &waker,
                    )
                    .await
                {
                    finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
                    return Ok(IoResult::failed(bytes, StreamError::Cancelled));
                }
                finish_terminal(&terminal, DemandTerminal::Complete);
                let _ = send_read_item(&items, DemandResponse::Complete, &mut cancel, &waker).await;
                return Ok(IoResult::succeeded(bytes));
            }
            Ok(read) => {
                bytes += read as u64;
                offset += read as u64;
                if !send_read_chunks(
                    &items,
                    chunk_size,
                    &mut pending_tail,
                    &buffer[..read],
                    &mut cancel,
                    &waker,
                )
                .await
                {
                    finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
                    return Ok(IoResult::failed(bytes, StreamError::Cancelled));
                }
                buffer.clear();
            }
            Err(error) => {
                let error = io_error(error);
                finish_terminal(&terminal, DemandTerminal::Error(error.clone()));
                let _ = send_read_item(
                    &items,
                    DemandResponse::Error(error.clone()),
                    &mut cancel,
                    &waker,
                )
                .await;
                return Ok(IoResult::failed(bytes, error));
            }
        }
    }
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
fn uring_file_write_sink_or_fallback(path: PathBuf) -> TokioByteSink {
    let path = Arc::new(path);
    Sink::from_runner(move |input, materializer| {
        let path = Arc::clone(&path);
        let (command_sender, command_receiver) = mpsc::channel(1);
        let (cancel_sender, cancel_receiver) = watch::channel(false);
        let (done_sender, done_receiver) = std_mpsc::sync_channel(1);
        let uring_path = Arc::clone(&path);

        let result = spawn_uring_job(Box::new(move || {
            Box::pin(async move {
                let task_result = AssertUnwindSafe(run_uring_write_task(
                    uring_path,
                    command_receiver,
                    cancel_receiver,
                ))
                .catch_unwind()
                .await
                .unwrap_or(Err(StreamError::AbruptTermination));
                let _ = done_sender.send(task_result);
            })
        }));

        if result.is_err() {
            let fallback_path = Arc::clone(&path);
            return run_async_write_sink::<tokio::fs::File, _, _>(
                input,
                materializer,
                Arc::new(move || tokio_file_write_open(Arc::clone(&fallback_path))),
            );
        }

        Ok(materializer.spawn_stream(move |cancelled| {
            let mut guard = TokioCancelGuard::new(cancel_sender.clone());
            let result = feed_async_writer(
                input,
                command_sender,
                done_receiver,
                cancelled,
                cancel_sender,
            );
            guard.disarm();
            result
        }))
    })
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
async fn run_uring_write_task(
    path: Arc<PathBuf>,
    mut commands: mpsc::Receiver<WriteCommand>,
    mut cancel: watch::Receiver<bool>,
) -> StreamResult<IoResult> {
    let mut bytes = 0_u64;
    let file = tokio::select! {
        file = tokio_uring::fs::File::create(path.as_ref()) => match file {
            Ok(file) => file,
            Err(error) => return Ok(IoResult::failed(bytes, io_error(error))),
        },
        changed = cancel.changed() => {
            let _ = changed;
            return Ok(IoResult::failed(bytes, StreamError::Cancelled));
        }
    };

    let mut offset = 0_u64;
    loop {
        let command = tokio::select! {
            command = commands.recv() => command,
            changed = cancel.changed() => {
                let _ = changed;
                return Ok(IoResult::failed(bytes, StreamError::Cancelled));
            }
        };

        match command {
            Some(WriteCommand::Chunk(chunk)) => {
                if let Err(error) =
                    write_uring_chunk(&file, chunk, &mut offset, &mut cancel, &mut bytes).await
                {
                    return Ok(IoResult::failed(bytes, error));
                }
            }
            Some(WriteCommand::Finish(upstream_status)) => {
                let close_status = tokio::select! {
                    result = file.close() => result.map_err(io_error),
                    changed = cancel.changed() => {
                        let _ = changed;
                        Err(StreamError::Cancelled)
                    }
                };
                return Ok(IoResult {
                    bytes,
                    status: upstream_status.and(close_status),
                });
            }
            None => {
                let _ = file.close().await;
                return Ok(IoResult::failed(bytes, StreamError::Cancelled));
            }
        }
    }
}

#[cfg(all(feature = "io-uring-file", target_os = "linux"))]
async fn write_uring_chunk(
    file: &tokio_uring::fs::File,
    chunk: Vec<u8>,
    offset: &mut u64,
    cancel: &mut watch::Receiver<bool>,
    bytes: &mut u64,
) -> StreamResult<()> {
    use tokio_uring::buf::BoundedBuf;

    let len = chunk.len();
    if len == 0 {
        return Ok(());
    }

    let mut written = 0usize;
    let mut buffer = chunk;
    while written < len {
        let slice = buffer.slice(written..len);
        let (result, returned) = tokio::select! {
            result = file.write_at(slice, *offset).submit() => result,
            changed = cancel.changed() => {
                let _ = changed;
                return Err(StreamError::Cancelled);
            }
        };
        buffer = returned.into_inner();

        match result {
            Ok(0) => return Err(write_zero_error()),
            Ok(n) => {
                written += n;
                *offset += n as u64;
                *bytes += n as u64;
            }
            Err(error) => return Err(io_error(error)),
        }
    }

    Ok(())
}

/// Address metadata for an established TCP connection, carried as the materialized value of
/// [`TokioTcp::outgoing_connection`] and accessible from a [`TcpIncomingConnection`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TcpConnection {
    /// Local (this side) socket address.
    pub local_addr: SocketAddr,
    /// Remote (peer) socket address.
    pub remote_addr: SocketAddr,
}

impl TcpConnection {
    #[must_use]
    pub fn local_addr(&self) -> SocketAddr {
        self.local_addr
    }

    #[must_use]
    pub fn remote_addr(&self) -> SocketAddr {
        self.remote_addr
    }
}

/// The bound address of a TCP listener, available as the materialized value of [`TokioTcp::bind`]
/// as soon as the listener is bound — before any connection is accepted.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TcpBinding {
    /// The address the listener is bound to (e.g. the OS-assigned port for `:0`).
    pub local_addr: SocketAddr,
}

impl TcpBinding {
    #[must_use]
    pub fn local_addr(&self) -> SocketAddr {
        self.local_addr
    }
}

/// A TCP connection accepted by [`TokioTcp::bind`].
///
/// `TokioTcp` carries plain TCP only; TLS, UDP, and QUIC live in the `datum-net` satellite crate.
/// The connection can be consumed as independent byte source/sink halves
/// ([`into_parts`](TcpIncomingConnection::into_parts)) or as a coupled
/// `Flow<Vec<u8>, Vec<u8>>` ([`into_flow`](TcpIncomingConnection::into_flow)).
pub struct TcpIncomingConnection {
    connection: TcpConnection,
    source: TokioByteSource,
    sink: TokioByteSink,
}

impl TcpIncomingConnection {
    #[must_use]
    pub fn local_addr(&self) -> SocketAddr {
        self.connection.local_addr
    }

    #[must_use]
    pub fn remote_addr(&self) -> SocketAddr {
        self.connection.remote_addr
    }

    #[must_use]
    pub fn connection(&self) -> TcpConnection {
        self.connection
    }

    /// Splits the connection into independent read (`TokioByteSource`) and write
    /// (`TokioByteSink`) halves. The halves are decoupled: completing the write half
    /// (TCP FIN) does not tear down the read half.
    #[must_use]
    pub fn into_parts(self) -> (TokioByteSource, TokioByteSink) {
        (self.source, self.sink)
    }

    /// Consumes the connection as a single coupled echo `Flow`: upstream bytes are written to the
    /// socket and bytes received from the socket are emitted downstream.
    #[must_use]
    pub fn into_flow(self) -> Flow<Vec<u8>, Vec<u8>, NotUsed> {
        Flow::from_sink_and_source_coupled(self.sink, self.source)
            .map_materialized_value(|_| NotUsed)
    }
}

pub struct TokioTcp;

impl TokioTcp {
    /// Opens a plain TCP outgoing connection as a coupled byte flow.
    ///
    /// TLS and QUIC carriers live in the `datum-net` satellite crate. Reads are bounded by a
    /// capacity-1 channel: once one chunk is waiting for downstream demand, the Tokio task stops
    /// reading until the stream consumes it. The flow does not hide an unbounded receive buffer.
    /// Panics if `chunk_size == 0`.
    #[must_use]
    pub fn outgoing_connection<A>(
        addr: A,
        chunk_size: usize,
    ) -> Flow<Vec<u8>, Vec<u8>, StreamCompletion<TcpConnection>>
    where
        A: ToSocketAddrs + Clone + Send + Sync + 'static,
    {
        assert!(chunk_size > 0, "chunk size must be greater than zero");
        Flow::future_flow(move || {
            let addr = addr.clone();
            async move {
                let stream = TcpStream::connect(addr).await.map_err(io_error)?;
                Ok(tcp_flow_from_stream(stream, chunk_size))
            }
        })
    }

    /// [`TokioTcp::outgoing_connection`] with the default 8 KiB chunk size.
    #[must_use]
    pub fn outgoing_connection_default<A>(
        addr: A,
    ) -> Flow<Vec<u8>, Vec<u8>, StreamCompletion<TcpConnection>>
    where
        A: ToSocketAddrs + Clone + Send + Sync + 'static,
    {
        Self::outgoing_connection(addr, DEFAULT_CHUNK_SIZE)
    }

    /// Binds a plain TCP listener and emits accepted incoming connections.
    ///
    /// Each accepted connection carries independent byte source/sink halves and can be converted
    /// into a coupled flow. The accept loop is demand-gated: one downstream pull permits one Tokio
    /// `accept`. The materialized [`TcpBinding`] is available as soon as the listener is bound.
    /// Panics if `chunk_size == 0`.
    #[must_use]
    pub fn bind<A>(
        addr: A,
        chunk_size: usize,
    ) -> Source<TcpIncomingConnection, StreamCompletion<TcpBinding>>
    where
        A: ToSocketAddrs + Clone + Send + Sync + 'static,
    {
        assert!(chunk_size > 0, "chunk size must be greater than zero");
        Source::from_materialized_factory(move |_materializer| {
            let (demand_sender, demand_receiver) = mpsc::channel(1);
            let (cancel_sender, cancel_receiver) = watch::channel(false);
            let (binding_sender, binding_receiver) = oneshot::channel();
            let terminal = Arc::new(Mutex::new(None));
            let terminal_for_task = Arc::clone(&terminal);
            let addr = addr.clone();

            crate::stream::stream_tokio_runtime().spawn(async move {
                let result = AssertUnwindSafe(run_tcp_bind_task(
                    addr,
                    chunk_size,
                    demand_receiver,
                    cancel_receiver,
                    binding_sender,
                    Arc::clone(&terminal_for_task),
                ))
                .catch_unwind()
                .await;
                if result.is_err() {
                    finish_terminal(
                        &terminal_for_task,
                        DemandTerminal::Error(StreamError::AbruptTermination),
                    );
                }
            });

            Ok((
                Box::new(DemandSourceStream {
                    demands: demand_sender,
                    cancel: cancel_sender,
                    terminal,
                    done: false,
                }) as BoxStream<TcpIncomingConnection>,
                StreamCompletion::from_receiver(binding_receiver, None),
            ))
        })
    }

    /// [`TokioTcp::bind`] with the default 8 KiB chunk size.
    #[must_use]
    pub fn bind_default<A>(addr: A) -> Source<TcpIncomingConnection, StreamCompletion<TcpBinding>>
    where
        A: ToSocketAddrs + Clone + Send + Sync + 'static,
    {
        Self::bind(addr, DEFAULT_CHUNK_SIZE)
    }
}

fn tcp_flow_from_stream(
    stream: TcpStream,
    chunk_size: usize,
) -> Flow<Vec<u8>, Vec<u8>, TcpConnection> {
    let connection = TcpConnection {
        local_addr: stream
            .local_addr()
            .expect("connected TCP stream has local address"),
        remote_addr: stream
            .peer_addr()
            .expect("connected TCP stream has peer address"),
    };
    let (read_half, write_half) = stream.into_split();
    let source = single_use_async_read_source(read_half, chunk_size);
    let sink = single_use_async_write_sink(write_half);
    // Use `from_sink_and_source` (not coupled) so that the write half
    // completing (TCP half-close / FIN) does **not** cancel the read half.
    // TCP half-close semantics require the read direction to stay open until
    // the remote peer sends its own FIN — if the write side released the
    // source keepalive first, the read source would be torn down before
    // buffered echoed bytes arrive, causing a race under load.
    Flow::from_sink_and_source(sink, source).map_materialized_value(move |_| connection)
}

fn single_use_async_read_source<R>(reader: R, chunk_size: usize) -> TokioByteSource
where
    R: AsyncRead + Unpin + Send + 'static,
{
    let reader = Arc::new(Mutex::new(Some(reader)));
    Source::from_materialized_factory(move |_materializer| {
        let reader = Arc::clone(&reader);
        Ok(async_read_source(
            move || async move {
                reader
                    .lock()
                    .expect("single-use async reader poisoned")
                    .take()
                    .ok_or_else(|| std::io::Error::other("async reader already materialized"))
            },
            chunk_size,
            chunk_size,
            TCP_READ_AHEAD_CHUNKS,
        ))
    })
}

fn single_use_async_write_sink<W>(writer: W) -> TokioByteSink
where
    W: AsyncWrite + Unpin + Send + 'static,
{
    let writer = Arc::new(Mutex::new(Some(writer)));
    async_write_sink(move || {
        let writer = Arc::clone(&writer);
        async move {
            writer
                .lock()
                .expect("single-use async writer poisoned")
                .take()
                .ok_or_else(|| std::io::Error::other("async writer already materialized"))
        }
    })
}

async fn run_tcp_bind_task<A>(
    addr: A,
    chunk_size: usize,
    mut demands: mpsc::Receiver<std_mpsc::Sender<DemandResponse<TcpIncomingConnection>>>,
    mut cancel: watch::Receiver<bool>,
    binding_sender: oneshot::Sender<StreamResult<TcpBinding>>,
    terminal: Arc<Mutex<Option<DemandTerminal>>>,
) where
    A: ToSocketAddrs + Send + 'static,
{
    let listener = match TcpListener::bind(addr).await {
        Ok(listener) => listener,
        Err(error) => {
            let error = io_error(error);
            finish_terminal(&terminal, DemandTerminal::Error(error.clone()));
            let _ = binding_sender.send(Err(error));
            return;
        }
    };
    let local_addr = match listener.local_addr() {
        Ok(local_addr) => local_addr,
        Err(error) => {
            let error = io_error(error);
            finish_terminal(&terminal, DemandTerminal::Error(error.clone()));
            let _ = binding_sender.send(Err(error));
            return;
        }
    };
    let _ = binding_sender.send(Ok(TcpBinding { local_addr }));

    loop {
        let Some(reply) = next_demand(&mut demands, &mut cancel).await else {
            finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
            return;
        };

        let (stream, remote_addr) = loop {
            let accepted = tokio::select! {
                accepted = listener.accept() => accepted,
                changed = cancel.changed() => {
                    let _ = changed;
                    finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
                    return;
                }
            };

            match accepted {
                Ok(accepted) => break accepted,
                Err(error) if is_transient_accept_error(&error) => continue,
                Err(error) => {
                    // EMFILE/ENFILE and other persistent listener errors are
                    // fatal by choice: retrying immediately can hot-loop under
                    // fd pressure and still cannot satisfy the pending demand.
                    let error = io_error(error);
                    finish_terminal(&terminal, DemandTerminal::Error(error.clone()));
                    let _ = reply.send(DemandResponse::Error(error));
                    return;
                }
            }
        };

        let incoming = tcp_incoming_connection(stream, remote_addr, local_addr, chunk_size);
        if reply.send(DemandResponse::Item(incoming)).is_err() {
            finish_terminal(&terminal, DemandTerminal::Error(StreamError::Cancelled));
            return;
        }
    }
}

fn is_transient_accept_error(error: &std::io::Error) -> bool {
    matches!(
        error.kind(),
        std::io::ErrorKind::Interrupted
            | std::io::ErrorKind::ConnectionAborted
            | std::io::ErrorKind::ConnectionReset
    ) || error.raw_os_error().is_some_and(is_transient_accept_errno)
}

#[cfg(target_os = "linux")]
fn is_transient_accept_errno(code: i32) -> bool {
    matches!(code, 4 | 103 | 104)
}

#[cfg(not(target_os = "linux"))]
fn is_transient_accept_errno(_code: i32) -> bool {
    false
}

fn tcp_incoming_connection(
    stream: TcpStream,
    remote_addr: SocketAddr,
    local_addr: SocketAddr,
    chunk_size: usize,
) -> TcpIncomingConnection {
    let connection = TcpConnection {
        local_addr,
        remote_addr,
    };
    let (read_half, write_half) = stream.into_split();
    let source = single_use_async_read_source(read_half, chunk_size);
    let sink = single_use_async_write_sink(write_half);
    TcpIncomingConnection {
        connection,
        source,
        sink,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::testkit::TestSink;
    use crate::{Framing, Keep, Sink, Source};
    use std::pin::Pin;
    use std::sync::atomic::{AtomicBool as StdAtomicBool, Ordering as StdOrdering};
    use std::task::{Context, Poll};
    use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

    fn unique_temp_path(name: &str) -> PathBuf {
        let nanos = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("clock after epoch")
            .as_nanos();
        std::env::temp_dir().join(format!(
            "datum-wp12b-{name}-{}-{nanos}.bin",
            std::process::id()
        ))
    }

    fn wait_until(timeout: Duration, condition: impl Fn() -> bool) -> bool {
        let deadline = Instant::now() + timeout;
        while Instant::now() < deadline {
            if condition() {
                return true;
            }
            thread::sleep(Duration::from_millis(5));
        }
        condition()
    }

    struct PendingWriter {
        polled: Arc<StdAtomicBool>,
        dropped: Arc<StdAtomicBool>,
    }

    impl AsyncWrite for PendingWriter {
        fn poll_write(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            _buf: &[u8],
        ) -> Poll<std::io::Result<usize>> {
            self.polled.store(true, StdOrdering::SeqCst);
            Poll::Pending
        }

        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
            self.polled.store(true, StdOrdering::SeqCst);
            Poll::Pending
        }

        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
            self.polled.store(true, StdOrdering::SeqCst);
            Poll::Pending
        }
    }

    impl Drop for PendingWriter {
        fn drop(&mut self) {
            self.dropped.store(true, StdOrdering::SeqCst);
        }
    }

    #[test]
    fn tokio_file_io_round_trips_bytes_and_reports_counts() {
        let path = unique_temp_path("roundtrip");
        let write_completion = Source::from_iter([b"ab".to_vec(), b"cd".to_vec()])
            .run_with(TokioFileIO::to_path(path.clone()))
            .expect("tokio file sink materializes");
        let write_result = write_completion.wait().expect("tokio file write completes");
        assert_eq!(write_result.bytes(), 4);
        assert_eq!(write_result.status(), Ok(()));

        let (read_completion, collected) = TokioFileIO::from_path(path.clone(), 2)
            .to_mat(Sink::collect(), Keep::both)
            .run()
            .expect("tokio file source materializes");
        assert_eq!(
            collected.wait().expect("collect completes"),
            vec![b"ab".to_vec(), b"cd".to_vec()]
        );
        let read_result = read_completion.wait().expect("read completion available");
        assert_eq!(read_result.bytes(), 4);
        assert_eq!(read_result.status(), Ok(()));

        std::fs::remove_file(path).expect("remove roundtrip file");
    }

    #[test]
    fn tokio_file_source_surfaces_open_failure() {
        let missing = unique_temp_path("missing");
        let (read_completion, collected) = TokioFileIO::from_path(missing, 4)
            .to_mat(Sink::collect(), Keep::both)
            .run()
            .expect("tokio file source materializes despite open failure");
        let stream_error = collected.wait().expect_err("collect fails");
        assert!(matches!(stream_error, StreamError::Failed(_)));
        let read_result = read_completion.wait().expect("io result available");
        assert_eq!(read_result.bytes(), 0);
        assert!(matches!(read_result.status(), Err(StreamError::Failed(_))));
    }

    #[test]
    fn tokio_file_source_composes_with_framing_and_sink() {
        let path = unique_temp_path("framing");
        std::fs::write(&path, b"alpha\nbeta\ngamma\n").expect("write framed seed file");

        let frames = TokioFileIO::from_path(path.clone(), 5)
            .via(Framing::delimiter(b"\n".to_vec(), 64, true))
            .run_with(Sink::collect())
            .expect("framed file stream materializes")
            .wait()
            .expect("framed file stream completes");

        assert_eq!(
            frames,
            vec![b"alpha".to_vec(), b"beta".to_vec(), b"gamma".to_vec()]
        );
        std::fs::remove_file(path).expect("remove framed file");
    }

    #[test]
    fn tokio_file_source_preserves_requested_chunk_boundaries() {
        let path = unique_temp_path("chunk-boundaries");
        let chunk_size = 8192;
        let tail_size = 13;
        let data_len = FILE_INTERNAL_READ_SIZE + tail_size;
        let data: Vec<u8> = (0..data_len).map(|index| (index % 251) as u8).collect();
        std::fs::write(&path, &data).expect("write chunk boundary seed file");

        let (read_completion, chunks) = TokioFileIO::from_path(path.clone(), chunk_size)
            .to_mat(Sink::collect(), Keep::both)
            .run()
            .expect("tokio file source materializes");
        let chunks = chunks.wait().expect("chunk boundary stream completes");

        assert!(chunks.len() > 1);
        for chunk in &chunks[..chunks.len() - 1] {
            assert_eq!(chunk.len(), chunk_size);
        }
        assert_eq!(chunks.last().expect("tail chunk exists").len(), tail_size);
        let reassembled: Vec<u8> = chunks.into_iter().flatten().collect();
        assert_eq!(reassembled, data);
        assert_eq!(
            read_completion
                .wait()
                .expect("read completion available")
                .bytes(),
            data_len as u64
        );

        std::fs::remove_file(path).expect("remove chunk boundary file");
    }

    #[cfg(all(feature = "io-uring-file", target_os = "linux"))]
    #[test]
    fn uring_file_io_round_trips_bytes_and_reports_counts() {
        let path = unique_temp_path("uring-roundtrip");
        let write_completion = Source::from_iter([b"ab".to_vec(), b"cd".to_vec()])
            .run_with(UringFileIO::to_path(path.clone()))
            .expect("uring file sink materializes");
        let write_result = write_completion.wait().expect("uring file write completes");
        assert_eq!(write_result.bytes(), 4);
        assert_eq!(write_result.status(), Ok(()));

        let (read_completion, collected) = UringFileIO::from_path(path.clone(), 2)
            .to_mat(Sink::collect(), Keep::both)
            .run()
            .expect("uring file source materializes");
        assert_eq!(
            collected.wait().expect("collect completes"),
            vec![b"ab".to_vec(), b"cd".to_vec()]
        );
        let read_result = read_completion.wait().expect("read completion available");
        assert_eq!(read_result.bytes(), 4);
        assert_eq!(read_result.status(), Ok(()));

        std::fs::remove_file(path).expect("remove uring roundtrip file");
    }

    #[cfg(all(feature = "io-uring-file", target_os = "linux"))]
    #[test]
    fn uring_file_source_surfaces_open_failure() {
        let missing = unique_temp_path("uring-missing");
        let (read_completion, collected) = UringFileIO::from_path(missing, 4)
            .to_mat(Sink::collect(), Keep::both)
            .run()
            .expect("uring file source materializes despite open failure");
        let stream_error = collected.wait().expect_err("collect fails");
        assert!(matches!(stream_error, StreamError::Failed(_)));
        let read_result = read_completion.wait().expect("io result available");
        assert_eq!(read_result.bytes(), 0);
        assert!(matches!(read_result.status(), Err(StreamError::Failed(_))));
    }

    #[cfg(all(feature = "io-uring-file", target_os = "linux"))]
    #[test]
    fn uring_file_sink_surfaces_open_failure() {
        let path = unique_temp_path("uring-sink-dir");
        std::fs::create_dir(&path).expect("create sink failure directory");

        let completion = Source::single(b"blocked".to_vec())
            .run_with(UringFileIO::to_path(path.clone()))
            .expect("uring file sink materializes despite open failure");
        let result = completion.wait().expect("uring sink result available");
        assert_eq!(result.bytes(), 0);
        assert!(matches!(result.status(), Err(StreamError::Failed(_))));

        std::fs::remove_dir(path).expect("remove sink failure directory");
    }

    #[cfg(all(feature = "io-uring-file", target_os = "linux"))]
    #[test]
    fn uring_file_source_preserves_requested_chunk_boundaries() {
        let path = unique_temp_path("uring-chunk-boundaries");
        let chunk_size = 8192;
        let tail_size = 13;
        let data_len = FILE_INTERNAL_READ_SIZE + tail_size;
        let data: Vec<u8> = (0..data_len).map(|index| (index % 251) as u8).collect();
        std::fs::write(&path, &data).expect("write uring chunk boundary seed file");

        let (read_completion, chunks) = UringFileIO::from_path(path.clone(), chunk_size)
            .to_mat(Sink::collect(), Keep::both)
            .run()
            .expect("uring file source materializes");
        let chunks = chunks.wait().expect("chunk boundary stream completes");

        assert!(chunks.len() > 1);
        for chunk in &chunks[..chunks.len() - 1] {
            assert_eq!(chunk.len(), chunk_size);
        }
        assert_eq!(chunks.last().expect("tail chunk exists").len(), tail_size);
        let reassembled: Vec<u8> = chunks.into_iter().flatten().collect();
        assert_eq!(reassembled, data);
        assert_eq!(
            read_completion
                .wait()
                .expect("read completion available")
                .bytes(),
            data_len as u64
        );

        std::fs::remove_file(path).expect("remove uring chunk boundary file");
    }

    #[cfg(all(feature = "io-uring-file", target_os = "linux"))]
    #[test]
    fn uring_file_source_cancellation_reports_cancelled() {
        let path = unique_temp_path("uring-cancel");
        let data_len = FILE_INTERNAL_READ_SIZE * 4;
        std::fs::write(&path, vec![b'z'; data_len]).expect("write uring cancel seed file");

        let (completion, ignored) = UringFileIO::from_path(path.clone(), 8192)
            .take(1)
            .to_mat(Sink::ignore(), Keep::both)
            .run()
            .expect("uring cancellation source materializes");
        ignored.wait().expect("downstream ignore completes");
        let result = completion.wait().expect("uring read completion available");

        assert!(result.bytes() > 0);
        assert!(matches!(result.status(), Err(StreamError::Cancelled)));

        std::fs::remove_file(path).expect("remove uring cancel file");
    }

    #[test]
    fn tokio_sink_cancellation_unblocks_pending_writer_completion_wait() {
        let polled = Arc::new(StdAtomicBool::new(false));
        let dropped = Arc::new(StdAtomicBool::new(false));
        let completion = Source::single(b"blocked".to_vec())
            .run_with(async_write_sink({
                let polled = Arc::clone(&polled);
                let dropped = Arc::clone(&dropped);
                move || {
                    let polled = Arc::clone(&polled);
                    let dropped = Arc::clone(&dropped);
                    async move { Ok(PendingWriter { polled, dropped }) }
                }
            }))
            .expect("pending writer sink materializes");

        assert!(wait_until(Duration::from_secs(1), || {
            polled.load(StdOrdering::SeqCst)
        }));
        drop(completion);
        assert!(wait_until(Duration::from_secs(1), || {
            dropped.load(StdOrdering::SeqCst)
        }));
    }

    #[test]
    fn tokio_tcp_accept_error_classifier_retries_only_connection_races() {
        assert!(is_transient_accept_error(&std::io::Error::new(
            std::io::ErrorKind::Interrupted,
            "interrupted"
        )));
        assert!(is_transient_accept_error(&std::io::Error::new(
            std::io::ErrorKind::ConnectionAborted,
            "aborted before accept"
        )));
        assert!(is_transient_accept_error(&std::io::Error::new(
            std::io::ErrorKind::ConnectionReset,
            "reset before accept"
        )));
        assert!(!is_transient_accept_error(&std::io::Error::other(
            "fd pressure"
        )));
    }

    #[test]
    fn tokio_source_cancellation_observed_promptly_under_wake_on_send() {
        // Source that yields a few chunks then parks with a 50 ms
        // backstop.  Between polls the runtime_checked_stream wrapper
        // feeds the Datum cancelled flag into the thread‑local slot
        // that current_stream_cancelled() reads.  When the flag is
        // set, the source returns Err(Cancelled) and records the
        // detection for verification.
        //
        // The source is deliberately shaped to block (park) so it
        // never self‑completes; only Sink‑side completion drop can
        // terminate it.  This makes the test deterministic regardless
        // of disk speed — contrast a finite file that may be fully
        // consumed before the cancellation sleep fires.
        let cancelled_detected = Arc::new(AtomicBool::new(false));
        let detected = Arc::clone(&cancelled_detected);

        let source = Source::from_materialized_factory(move |_materializer| {
            let d = Arc::clone(&detected);
            let mut i = 0_u64;
            let stream: BoxStream<Vec<u8>> = Box::new(std::iter::from_fn(move || {
                // Datum cancellation: runtime_checked_stream sets this
                // thread‑local before calling into the source.
                let is_cancelled = || {
                    crate::stream::current_stream_cancelled()
                        .as_ref()
                        .is_some_and(|c| c.load(Ordering::SeqCst))
                };
                if is_cancelled() {
                    d.store(true, Ordering::SeqCst);
                    return Some(Err(StreamError::Cancelled));
                }
                if i < 4 {
                    i += 1;
                    return Some(Ok(vec![0_u8; 8192]));
                }
                // Park — stays blocked until cancellation unpark
                // wakes us, or a 50 ms backstop lets us re‑check.
                std::thread::park_timeout(Duration::from_millis(50));
                if is_cancelled() {
                    d.store(true, Ordering::SeqCst);
                    return Some(Err(StreamError::Cancelled));
                }
                // Spurious / timeout wakeup — yield a tiny element
                // so runtime_checked_stream re‑checks the flag.
                Some(Ok(vec![0_u8]))
            }));
            Ok((stream, NotUsed))
        });

        let completion = source
            .run_with(Sink::ignore())
            .expect("cancellation source materializes");

        let cancel_thread = thread::spawn(move || {
            thread::sleep(Duration::from_millis(20));
            drop(completion);
        });

        // Wait for the source to observe the cancelled flag.
        let deadline = Instant::now() + Duration::from_secs(3);
        loop {
            if cancelled_detected.load(Ordering::SeqCst) {
                break;
            }
            assert!(
                Instant::now() < deadline,
                "cancellation not observed within 3 s; the park‑wake path may be broken"
            );
            thread::park_timeout(Duration::from_millis(10));
        }
        cancel_thread.join().expect("cancellation thread joins");
    }

    #[test]
    fn tokio_tcp_bind_and_outgoing_connection_echo_round_trip() {
        let (binding_completion, incoming_completion) = TokioTcp::bind("127.0.0.1:0", 1024)
            .to_mat(Sink::head(), Keep::both)
            .run()
            .expect("tcp bind source materializes");
        let binding = binding_completion.wait().expect("tcp binding succeeds");

        let client_completion = Source::single(b"ping".to_vec())
            .via(TokioTcp::outgoing_connection(binding.local_addr(), 1024))
            .run_with(Sink::head())
            .expect("client stream materializes");

        let incoming = incoming_completion
            .wait()
            .expect("incoming connection accepted");
        let (incoming_source, incoming_sink) = incoming.into_parts();
        let server_read = incoming_source
            .run_with(Sink::head())
            .expect("server read materializes")
            .wait()
            .expect("server reads request");
        assert_eq!(server_read, b"ping".to_vec());

        let server_write = Source::single(server_read)
            .run_with(incoming_sink)
            .expect("server write materializes");
        let write_result = server_write.wait().expect("server write completes");
        assert_eq!(write_result.bytes(), 4);
        assert_eq!(write_result.status(), Ok(()));

        assert_eq!(
            client_completion.wait().expect("client receives echo"),
            b"ping".to_vec()
        );
    }

    #[test]
    fn tokio_tcp_bind_accepts_multiple_connections_on_successive_demands() {
        let (binding_completion, mut incoming_probe) = TokioTcp::bind("127.0.0.1:0", 1024)
            .to_mat(TestSink::probe(), Keep::both)
            .run()
            .expect("tcp bind source materializes");
        let binding = binding_completion.wait().expect("tcp binding succeeds");
        // The default 3s probe timeout is a bound, not a sleep — but under full-suite
        // parallel CI load the accept -> probe hop has been observed to exceed it
        // (flaked twice on CI). A generous bound removes the flake without slowing
        // the happy path.
        incoming_probe.set_timeout(Duration::from_secs(30));

        incoming_probe.request(1);
        let first_client =
            std::net::TcpStream::connect(binding.local_addr()).expect("first client connects");
        let first = incoming_probe.expect_next();
        assert_eq!(first.local_addr(), binding.local_addr());
        assert_eq!(
            first.remote_addr(),
            first_client.local_addr().expect("first client local addr")
        );

        incoming_probe.request(1);
        let second_client =
            std::net::TcpStream::connect(binding.local_addr()).expect("second client connects");
        let second = incoming_probe.expect_next();
        assert_eq!(second.local_addr(), binding.local_addr());
        assert_eq!(
            second.remote_addr(),
            second_client
                .local_addr()
                .expect("second client local addr")
        );
        assert_ne!(first.remote_addr(), second.remote_addr());

        incoming_probe.cancel();
    }
}