liminal_sdk/remote/

read_wake.rs

//! An explicit wake for an idle socket read. It never closes the TCP socket.

use std::io::{self, Read, Write};
use std::os::fd::BorrowedFd;
use std::os::unix::net::UnixStream;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

use rustix::event::{PollFd, PollFlags, Timespec, poll};

use crate::SdkError;

/// Wakes a participant receive that is waiting on the socket with a budget.
///
/// # Which reads it interrupts
///
/// A wake ends, with [`SdkError::ReadInterrupted`], a read made through a
/// budgeted participant door while it waits for bytes:
/// `RemoteParticipantHandle::receive_within` with a nonzero budget, whatever
/// the budget (including one at or above the 60 s response deadline that a
/// caller uses to await a correlated answer), and the wait for the liveness
/// `Pong` inside `RemoteParticipantHandle::try_receive`. A frame already
/// decoded is returned without waiting and so without consuming the wake, and
/// a zero budget never waits at all. The interrupted read keeps any
/// partly-arrived frame and the socket; a `Pong` still owed to an interrupted
/// `try_receive` is consumed by a later read.
///
/// The reply-owed door `RemoteParticipantHandle::receive`, and every
/// request/response exchange the transport makes for itself (publish,
/// subscribe, conversations, the handshake), never observe a wake. "Idle"
/// here means the door the caller chose, not whether a reply is owed: a caller
/// that awaits a reply through `receive_within` can be interrupted.
///
/// # Sharing
///
/// There is one `ReadWake` per transport, not per handle. Every clone of it,
/// every `RemoteConfig` sharing the transport, and every participant built
/// over that transport reach the same wake, and it carries over a participant
/// reconnect. A wake from any of them interrupts whichever interruptible read
/// on that transport observes it first, through whichever handle is reading;
/// it is not addressed to the caller that sent it.
///
/// Wake requests coalesce: a wake sent while no interruptible read is waiting
/// stays pending and interrupts the next one, and several wakes before that
/// read are one interrupt. A wake that lands while a reader is claiming an
/// earlier one is never lost; at worst it costs one spurious interrupt of a
/// later read. The handle holds only the wake pipe, so retaining it does not
/// retain the transport socket or participant binding.
#[derive(Clone, Debug)]
pub struct ReadWake {
    inner: Arc<Pipe>,
}

#[derive(Debug)]
struct Pipe {
    read: UnixStream,
    write: UnixStream,
    pending: AtomicBool,
    #[cfg(test)]
    probe: Probe,
}

pub(super) enum Ready {
    Socket,
    TimedOut,
}

impl ReadWake {
    pub(super) fn new() -> Result<Self, SdkError> {
        let (read, write) = UnixStream::pair().map_err(|error| connection_error(&error))?;
        read.set_nonblocking(true)
            .map_err(|error| connection_error(&error))?;
        write
            .set_nonblocking(true)
            .map_err(|error| connection_error(&error))?;
        Ok(Self {
            inner: Arc::new(Pipe {
                read,
                write,
                pending: AtomicBool::new(false),
                #[cfg(test)]
                probe: Probe::default(),
            }),
        })
    }

    /// Wake the interruptible read waiting now on this transport, or the next
    /// one to wait (see the type's documentation). This never waits for the
    /// reader.
    ///
    /// # Errors
    /// Returns a connection error if the local wake pipe cannot be written.
    pub fn wake(&self) -> Result<(), SdkError> {
        if self.inner.pending.swap(true, Ordering::AcqRel) {
            return Ok(());
        }
        match (&self.inner.write).write(&[1]) {
            Ok(1) => Ok(()),
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(()),
            Ok(_) => {
                self.inner.pending.store(false, Ordering::Release);
                Err(connection_error(&io::Error::new(
                    io::ErrorKind::WriteZero,
                    "empty wake write",
                )))
            }
            Err(error) => {
                self.inner.pending.store(false, Ordering::Release);
                Err(connection_error(&error))
            }
        }
    }

    pub(super) fn wait(&self, socket: BorrowedFd<'_>, budget: Duration) -> Result<Ready, SdkError> {
        let started = Instant::now();
        loop {
            let remaining = budget.saturating_sub(started.elapsed());
            let timeout = Timespec::try_from(remaining).map_err(|error| SdkError::Connection {
                description: format!("idle wake deadline cannot be represented: {error}"),
            })?;
            #[cfg(test)]
            self.observe(WaitStep::Polling {
                pending: self.inner.pending.load(Ordering::Acquire),
            });
            let mut descriptors = [
                PollFd::new(&socket, PollFlags::IN),
                PollFd::new(&self.inner.read, PollFlags::IN),
            ];
            match poll(&mut descriptors, Some(&timeout)) {
                Ok(0) => return Ok(Ready::TimedOut),
                Err(rustix::io::Errno::INTR) => continue,
                Err(error) => return Err(connection_error(&io::Error::from(error))),
                Ok(_) => {}
            }
            if !descriptors[1].revents().is_empty() {
                // Clear the flag BEFORE taking the byte. A wake landing between
                // the two steps then finds it clear and writes a byte of its own,
                // so the worst case is one spurious interrupt of a later receive;
                // the other order absorbs that wake into this one and loses it.
                self.inner.pending.swap(false, Ordering::AcqRel);
                #[cfg(test)]
                self.observe(WaitStep::Claiming);
                let mut byte = [0];
                (&self.inner.read)
                    .read_exact(&mut byte)
                    .map_err(|error| connection_error(&error))?;
                return Err(SdkError::ReadInterrupted);
            }
            if !descriptors[0].revents().is_empty() {
                return Ok(Ready::Socket);
            }
        }
    }
}

/// Where a reader stands inside [`ReadWake::wait`], reported to a test probe.
#[cfg(test)]
pub(super) enum WaitStep {
    /// About to block in `poll`; `pending` is the wake flag as it enters.
    Polling { pending: bool },
    /// Between the two steps that claim an observed wake.
    Claiming,
}

#[cfg(test)]
type ProbeFn = Box<dyn FnMut(WaitStep, &ReadWake) + Send>;

/// A test's view into the wait, so a test proves where the reader stands
/// instead of sleeping and hoping it got there.
#[cfg(test)]
#[derive(Default)]
struct Probe(std::sync::Mutex<Option<ProbeFn>>);

#[cfg(test)]
impl std::fmt::Debug for Probe {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str("Probe")
    }
}

#[cfg(test)]
impl ReadWake {
    /// Reports every step of every later wait on this pipe to `probe`, on the
    /// reader's own thread.
    pub(super) fn probe_waits(&self, probe: impl FnMut(WaitStep, &Self) + Send + 'static) {
        if let Ok(mut installed) = self.inner.probe.0.lock() {
            *installed = Some(Box::new(probe));
        }
    }

    fn observe(&self, step: WaitStep) {
        if let Ok(mut installed) = self.inner.probe.0.lock() {
            if let Some(probe) = installed.as_mut() {
                probe(step, self);
            }
        }
    }
}

fn connection_error(error: &io::Error) -> SdkError {
    SdkError::Connection {
        description: format!("idle receive wake pipe: {error}"),
    }
}