magi-code 0.80.2

//! Deadline-aware child stdin writes. No worker or queue is needed: pipes are nonblocking.
use crate::cancellation::AgentCancellation;
use std::{
    io::{self, Write},
    process::ChildStdin,
    sync::{Mutex, TryLockError},
    time::{Duration, Instant},
};

pub(crate) struct ChildPipeWriter(Mutex<ChildStdin>);

impl ChildPipeWriter {
    pub(crate) fn new(stdin: ChildStdin) -> io::Result<Self> {
        set_nonblocking(&stdin)?;
        Ok(Self(Mutex::new(stdin)))
    }

    pub(crate) fn write(
        &self,
        bytes: &[u8],
        deadline: Instant,
        cancellation: Option<&AgentCancellation>,
    ) -> anyhow::Result<()> {
        let mut writer = loop {
            check_deadline(deadline, cancellation)?;
            match self.0.try_lock() {
                Ok(writer) => break writer,
                Err(TryLockError::Poisoned(_)) => anyhow::bail!("child stdin lock poisoned"),
                Err(TryLockError::WouldBlock) => pause(deadline),
            }
        };
        let mut remaining = bytes;
        while !remaining.is_empty() {
            check_deadline(deadline, cancellation)?;
            match writer.write(remaining) {
                Ok(0) => {
                    #[cfg(not(windows))]
                    return Err(io::Error::from(io::ErrorKind::WriteZero).into());
                    #[cfg(windows)]
                    pause(deadline);
                }
                Ok(count) => remaining = &remaining[count..],
                Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                Err(error) if would_block(&error) => pause(deadline),
                Err(error) => return Err(error.into()),
            }
        }
        // ChildStdin is unbuffered; there is nothing to flush on drop.
        check_deadline(deadline, cancellation)
    }

    #[cfg(test)]
    pub(crate) fn lock(&self) -> std::sync::LockResult<std::sync::MutexGuard<'_, ChildStdin>> {
        self.0.lock()
    }
}

fn check_deadline(
    deadline: Instant,
    cancellation: Option<&AgentCancellation>,
) -> anyhow::Result<()> {
    if let Some(cancellation) = cancellation {
        cancellation.check()?;
    }
    if Instant::now() >= deadline {
        anyhow::bail!("child stdin write timed out");
    }
    Ok(())
}

fn pause(deadline: Instant) {
    std::thread::sleep(
        deadline
            .saturating_duration_since(Instant::now())
            .min(Duration::from_millis(10)),
    );
}

fn would_block(error: &io::Error) -> bool {
    #[cfg(windows)]
    if error.raw_os_error() == Some(232) {
        return true;
    } // ERROR_NO_DATA for PIPE_NOWAIT
    error.kind() == io::ErrorKind::WouldBlock
}

#[cfg(unix)]
fn set_nonblocking(stdin: &ChildStdin) -> io::Result<()> {
    use std::os::fd::AsRawFd;
    let fd = stdin.as_raw_fd();
    // SAFETY: stdin owns a live descriptor; fcntl only changes its status flags.
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
    if flags < 0 || unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) } < 0 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

#[cfg(windows)]
fn set_nonblocking(stdin: &ChildStdin) -> io::Result<()> {
    use std::os::windows::io::AsRawHandle;
    #[link(name = "kernel32")]
    unsafe extern "system" {
        fn SetNamedPipeHandleState(
            handle: *mut std::ffi::c_void,
            mode: *const u32,
            max_count: *const u32,
            timeout: *const u32,
        ) -> i32;
    }
    let mode = 1u32; // PIPE_NOWAIT; anonymous pipes are implemented with named pipes.
    // SAFETY: the handle is live and mode points to an initialized DWORD for this call.
    if unsafe {
        SetNamedPipeHandleState(
            stdin.as_raw_handle(),
            &mode,
            std::ptr::null(),
            std::ptr::null(),
        )
    } == 0
    {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

#[cfg(not(any(unix, windows)))]
fn set_nonblocking(_: &ChildStdin) -> io::Result<()> {
    Err(io::Error::new(
        io::ErrorKind::Unsupported,
        "nonblocking child stdin is unsupported on this platform",
    ))
}