fuchsia_async/handle/zircon/

rwhandle.rs

// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use crate::OnSignalsRef;
use crate::runtime::{EHandle, PacketReceiver, ReceiverRegistration};
use fuchsia_sync::Mutex;
use std::marker::PhantomData;
use std::task::{Context, Poll, Waker, ready};
use zx::AsHandleRef;

const OBJECT_PEER_CLOSED: zx::Signals = zx::Signals::OBJECT_PEER_CLOSED;
const OBJECT_READABLE: zx::Signals = zx::Signals::OBJECT_READABLE;
const OBJECT_WRITABLE: zx::Signals = zx::Signals::OBJECT_WRITABLE;

/// State of an object when it is ready for reading.
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub enum ReadableState {
    /// Received `OBJECT_READABLE`, or optimistically assuming the object is readable.
    Readable,
    /// Received `OBJECT_PEER_CLOSED`.  The object might also be readable.
    MaybeReadableAndClosed,
}

/// State of an object when it is ready for writing.
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub enum WritableState {
    /// Received `OBJECT_WRITABLE`, or optimistically assuming the object is writable.
    Writable,
    /// Received `OBJECT_PEER_CLOSED`.
    Closed,
}

/// A `Handle` that receives notifications when it is readable.
///
/// # Examples
///
/// ```
/// loop {
///     ready!(self.poll_readable(cx))?;
///     match /* make read syscall */ {
///         Err(zx::Status::SHOULD_WAIT) => ready!(self.need_readable(cx)?),
///         status => return Poll::Ready(status),
///     }
/// }
/// ```
pub trait ReadableHandle {
    /// If the object is ready for reading, returns `Ready` with the readable
    /// state. If the implementor returns Pending, it should first ensure that
    /// `need_readable` is called.
    ///
    /// This should be called in a poll function. If the syscall returns
    /// `SHOULD_WAIT`, you must call `need_readable` to schedule wakeup when the
    /// object is readable.
    ///
    /// The returned `ReadableState` does not necessarily reflect an observed
    /// `OBJECT_READABLE` signal. We optimistically assume the object remains
    /// readable until `need_readable` is called.
    fn poll_readable(&self, cx: &mut Context<'_>) -> Poll<Result<ReadableState, zx::Status>>;

    /// Arranges for the current task to be woken when the object receives an
    /// `OBJECT_READABLE` or `OBJECT_PEER_CLOSED` signal.  This can return
    /// Poll::Ready if the object has already been signaled in which case the
    /// waker *will* not be woken and it is the caller's responsibility to not
    /// lose the signal.
    fn need_readable(&self, cx: &mut Context<'_>) -> Poll<Result<(), zx::Status>>;
}

/// A `Handle` that receives notifications when it is writable.
///
/// # Examples
///
/// ```
/// loop {
///     ready!(self.poll_writable(cx))?;
///     match /* make write syscall */ {
///         Err(zx::Status::SHOULD_WAIT) => ready!(self.need_writable(cx)?),
///         status => Poll::Ready(status),
///     }
/// }
/// ```
pub trait WritableHandle {
    /// If the object is ready for writing, returns `Ready` with the writable
    /// state. If the implementor returns Pending, it should first ensure that
    /// `need_writable` is called.
    ///
    /// This should be called in a poll function. If the syscall returns
    /// `SHOULD_WAIT`, you must call `need_writable` to schedule wakeup when the
    /// object is writable.
    ///
    /// The returned `WritableState` does not necessarily reflect an observed
    /// `OBJECT_WRITABLE` signal. We optimistically assume the object remains
    /// writable until `need_writable` is called.
    fn poll_writable(&self, cx: &mut Context<'_>) -> Poll<Result<WritableState, zx::Status>>;

    /// Arranges for the current task to be woken when the object receives an
    /// `OBJECT_WRITABLE` or `OBJECT_PEER_CLOSED` signal. This can return
    /// Poll::Ready if the object has already been signaled in which case the
    /// waker *will* not be woken and it is the caller's responsibility to not
    /// lose the signal.
    fn need_writable(&self, cx: &mut Context<'_>) -> Poll<Result<(), zx::Status>>;
}

struct RWPacketReceiver<S: RWHandleSpec> {
    inner: Mutex<Inner>,
    _marker: PhantomData<S>,
}

struct Inner {
    signals: zx::Signals,
    read_task: Option<Waker>,
    write_task: Option<Waker>,
}

impl<S: RWHandleSpec> PacketReceiver for RWPacketReceiver<S> {
    fn receive_packet(&self, packet: zx::Packet) {
        let new = if let zx::PacketContents::SignalOne(p) = packet.contents() {
            // Only consider the signals that were part of the trigger. This
            // ensures that only the packets generated by the correct signal
            // observer (Read or Write) can raise the corresponding inner signal
            // bit.
            //
            // Without this, we can lose track of how many port observers are
            // installed.
            p.observed() & p.trigger()
        } else {
            return;
        };

        // We wake the tasks when the lock isn't held in case the wakers need the same lock.
        let mut read_task = None;
        let mut write_task = None;
        {
            let mut inner = self.inner.lock();
            let old = inner.signals;
            inner.signals |= new;

            let became_readable =
                new.intersection(S::READABLE_SIGNALS) != old.intersection(S::READABLE_SIGNALS);
            let became_writable =
                new.intersection(S::WRITABLE_SIGNALS) != old.intersection(S::WRITABLE_SIGNALS);
            let became_closed =
                new.contains(OBJECT_PEER_CLOSED) && !old.contains(OBJECT_PEER_CLOSED);
            if became_readable || became_closed {
                read_task = inner.read_task.take();
            }
            if became_writable || became_closed {
                write_task = inner.write_task.take();
            }
        }
        // *NOTE*: This is the only safe place to wake wakers.  In any other location, there is a
        // risk that locks are held which might be required when the waker is woken.  It is safe to
        // wake here because this is called from the executor when no locks are held.
        if let Some(read_task) = read_task {
            read_task.wake();
        }
        if let Some(write_task) = write_task {
            write_task.wake();
        }
    }
}

/// A `Handle` that receives notifications when it is readable/writable.
pub struct RWHandle<T, S: RWHandleSpec = DefaultRWHandleSpec> {
    handle: T,
    receiver: ReceiverRegistration<RWPacketReceiver<S>>,
}

impl<T> RWHandle<T, DefaultRWHandleSpec>
where
    T: AsHandleRef,
{
    /// Creates a new `RWHandle` object which will receive notifications when
    /// the underlying handle becomes readable, writable, or closes.
    ///
    /// # Panics
    ///
    /// If called outside the context of an active async executor.
    pub fn new(handle: T) -> Self {
        Self::new_with_spec(handle)
    }
}

impl<T, S> RWHandle<T, S>
where
    T: AsHandleRef,
    S: RWHandleSpec,
{
    /// Creates a new `RWHandle` with a non-default spec and an object which
    /// will receive notifications when the underlying handle becomes readable,
    /// writable, or closes.
    ///
    /// # Panics
    ///
    /// If called outside the context of an active async executor.
    ///
    /// If `S::READABLE_SIGNALS` or `S::WRITABLE_SIGNALS` contain
    /// [`zx::Signals::OBJECT_PEER_CLOSED`] or have a non-empty intersection.
    pub fn new_with_spec(handle: T) -> Self {
        let ehandle = EHandle::local();
        let internal_signals = OBJECT_PEER_CLOSED;
        assert!(
            !S::READABLE_SIGNALS.contains(internal_signals),
            "readable signals may not contain ({internal_signals:?})"
        );
        assert!(
            !S::WRITABLE_SIGNALS.contains(internal_signals),
            "writable signals may not contain ({internal_signals:?})"
        );
        assert!(!S::WRITABLE_SIGNALS.intersects(S::READABLE_SIGNALS), "signals may not intersect");
        let initial_signals = S::WRITABLE_SIGNALS | S::READABLE_SIGNALS;
        let receiver = ehandle.register_receiver(RWPacketReceiver {
            inner: Mutex::new(Inner {
                // Optimistically assume that the handle is readable and writable.
                // Reads and writes will be attempted before queueing a packet.
                // This makes handles slightly faster to read/write the first time
                // they're accessed after being created, provided they start off as
                // readable or writable. In return, there will be an extra wasted
                // syscall per read/write if the handle is not readable or writable.
                signals: initial_signals,
                read_task: None,
                write_task: None,
            }),
            _marker: PhantomData,
        });

        RWHandle { handle, receiver }
    }

    /// Returns a reference to the underlying handle.
    pub fn get_ref(&self) -> &T {
        &self.handle
    }

    /// Returns a mutable reference to the underlying handle.
    pub fn get_mut(&mut self) -> &mut T {
        &mut self.handle
    }

    /// Consumes `self` and returns the underlying handle.
    pub fn into_inner(self) -> T {
        self.handle
    }

    /// Returns true if the object received the `OBJECT_PEER_CLOSED` signal.
    pub fn is_closed(&self) -> bool {
        let signals = self.receiver().inner.lock().signals;
        if signals.contains(OBJECT_PEER_CLOSED) {
            return true;
        }

        // The signals bitset might not be updated if we haven't gotten around to processing the
        // packet telling us that yet. To provide an up-to-date response, we query the current
        // state of the signal.
        //
        // Note: we _could_ update the bitset with what we find here, if we're careful to also
        // update READABLE + WRITEABLE at the same time, and also wakeup the tasks as necessary.
        // But having `is_closed` wakeup tasks if it discovered a signal change seems too weird, so
        // we just leave the bitset as-is and let the regular notification mechanism get around to
        // it when it gets around to it.
        match self
            .handle
            .as_handle_ref()
            .wait_one(OBJECT_PEER_CLOSED, zx::MonotonicInstant::INFINITE_PAST)
            .to_result()
        {
            Ok(_) => true,
            Err(zx::Status::TIMED_OUT) => false,
            Err(status) => {
                // None of the other documented error statuses should be possible, either the type
                // system doesn't allow it or the wait from `RWHandle::new()` would have already
                // failed.
                unreachable!("status: {status}")
            }
        }
    }

    /// Returns a future that completes when `is_closed()` is true.
    pub fn on_closed(&self) -> OnSignalsRef<'_> {
        OnSignalsRef::new(self.handle.as_handle_ref(), OBJECT_PEER_CLOSED)
    }

    fn receiver(&self) -> &RWPacketReceiver<S> {
        self.receiver.receiver()
    }

    fn need_signal(&self, cx: &mut Context<'_>, signal: Signal) -> Poll<Result<(), zx::Status>> {
        let mut inner = self.receiver.inner.lock();
        let old = inner.signals;
        if old.contains(OBJECT_PEER_CLOSED) {
            // We don't want to return an error here because even though the peer has closed, the
            // object could still have queued messages that can be read.
            Poll::Ready(Ok(()))
        } else {
            let waker = cx.waker().clone();
            let signal = match signal {
                Signal::Read => {
                    inner.read_task = Some(waker);
                    S::READABLE_SIGNALS
                }
                Signal::Write => {
                    inner.write_task = Some(waker);
                    S::WRITABLE_SIGNALS
                }
            };
            if old.intersects(signal) {
                inner.signals &= !signal;
                std::mem::drop(inner);
                self.handle.as_handle_ref().wait_async(
                    self.receiver.port(),
                    self.receiver.key(),
                    signal | OBJECT_PEER_CLOSED,
                    zx::WaitAsyncOpts::empty(),
                )?;
            }
            Poll::Pending
        }
    }

    fn poll_signal(
        &self,
        cx: &mut Context<'_>,
        signal: Signal,
    ) -> Poll<Result<zx::Signals, zx::Status>> {
        let mask = match signal {
            Signal::Read => S::READABLE_SIGNALS,
            Signal::Write => S::WRITABLE_SIGNALS,
        } | OBJECT_PEER_CLOSED;

        loop {
            let signals = self.receiver().inner.lock().signals;
            let asserted = signals.intersection(mask);
            if !asserted.is_empty() {
                return Poll::Ready(Ok(asserted));
            }
            ready!(self.need_signal(cx, signal)?);
        }
    }
}

#[derive(Copy, Clone)]
enum Signal {
    Read,
    Write,
}

impl<T, S> ReadableHandle for RWHandle<T, S>
where
    T: AsHandleRef,
    S: RWHandleSpec,
{
    fn poll_readable(&self, cx: &mut Context<'_>) -> Poll<Result<ReadableState, zx::Status>> {
        let signals = ready!(self.poll_signal(cx, Signal::Read)?);
        if signals.contains(OBJECT_PEER_CLOSED) {
            return Poll::Ready(Ok(ReadableState::MaybeReadableAndClosed));
        }
        Poll::Ready(Ok(ReadableState::Readable))
    }

    fn need_readable(&self, cx: &mut Context<'_>) -> Poll<Result<(), zx::Status>> {
        self.need_signal(cx, Signal::Read)
    }
}

impl<T, S> WritableHandle for RWHandle<T, S>
where
    T: AsHandleRef,
    S: RWHandleSpec,
{
    fn poll_writable(&self, cx: &mut Context<'_>) -> Poll<Result<WritableState, zx::Status>> {
        let signals = ready!(self.poll_signal(cx, Signal::Write)?);
        if signals.contains(OBJECT_PEER_CLOSED) {
            return Poll::Ready(Ok(WritableState::Closed));
        }
        Poll::Ready(Ok(WritableState::Writable))
    }

    fn need_writable(&self, cx: &mut Context<'_>) -> Poll<Result<(), zx::Status>> {
        self.need_signal(cx, Signal::Write)
    }
}

/// A trait specifying the behavior of [`RWHandle`].
///
/// The default behavior, listening for [`zx::Signals::OBJECT_READABLE`] and
/// [`zx::Signals::OBJECT_WRITABLE`] is provided by [`DefaultRWHandleSpec`].
pub trait RWHandleSpec: Send + Sync + 'static {
    /// Signals asserted when the handle is readable.
    ///
    /// [`RWHandle`] installs wait for these signals and if any of them are
    /// asserted, the handle is considered readable.
    const READABLE_SIGNALS: zx::Signals;
    /// Signals asserted when the handle is writable.
    ///
    /// [`RWHandle`] installs wait for these signals and if any of them are
    /// asserted, the handle is considered writable.
    const WRITABLE_SIGNALS: zx::Signals;
}

/// The default behavior for [`RWHandle`].
///
/// Considers the handle readable when [`zx::Signals::OBJECT_READABLE`] is set,
/// and writable when [`zx::Signals::OBJECT_WRITABLE`] is set.
pub struct DefaultRWHandleSpec;

impl RWHandleSpec for DefaultRWHandleSpec {
    const READABLE_SIGNALS: zx::Signals = OBJECT_READABLE;
    const WRITABLE_SIGNALS: zx::Signals = OBJECT_WRITABLE;
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::TestExecutor;

    #[test]
    fn is_closed_immediately_after_close() {
        let mut exec = TestExecutor::new();
        let (tx, rx) = zx::Channel::create();
        let rx_rw_handle = RWHandle::new(rx);
        let mut noop_ctx = Context::from_waker(Waker::noop());
        // Clear optimistic readable state
        assert!(rx_rw_handle.need_readable(&mut noop_ctx).is_pending());
        // Starting state: the channel is not closed (because we haven't closed it yet)
        assert!(!rx_rw_handle.is_closed());
        // we will never set readable, so this should be Pending until we close
        assert_eq!(rx_rw_handle.poll_readable(&mut noop_ctx), Poll::Pending);

        drop(tx);

        // Implementation note: the cached state will not be updated yet
        assert_eq!(rx_rw_handle.poll_readable(&mut noop_ctx), Poll::Pending);
        // But is_closed should return true immediately
        assert!(rx_rw_handle.is_closed());
        // Still not updated, and won't be until we let the executor process port packets
        assert_eq!(rx_rw_handle.poll_readable(&mut noop_ctx), Poll::Pending);
        // So we do
        let _ = exec.run_until_stalled(&mut futures::future::pending::<()>());
        // And now it is updated, so we observe Closed
        assert_eq!(
            rx_rw_handle.poll_readable(&mut noop_ctx),
            Poll::Ready(Ok(ReadableState::MaybeReadableAndClosed))
        );
        // And is_closed should still be true, of course.
        assert!(rx_rw_handle.is_closed());
    }

    // Regression test for https://fxbug.dev/417333384.
    #[test]
    fn simultaneous_read_and_write() {
        let mut exec = TestExecutor::new();
        let (peer, local) = zx::Socket::create_stream();
        let mut buff = [0u8; 1024];
        while local.write(&buff[..]).is_ok() {}

        let rw_handle = RWHandle::new(local);
        let read_fut = futures::future::poll_fn(|cx| {
            let readable = ready!(rw_handle.poll_readable(cx));
            assert_eq!(readable, Ok(ReadableState::Readable));
            let mut buf = [0u8; 2];
            loop {
                match rw_handle.get_ref().read(&mut buf[..]) {
                    Ok(r) => assert_eq!(r, 1),
                    Err(e) => {
                        assert_eq!(e, zx::Status::SHOULD_WAIT);
                        break;
                    }
                }
            }
            assert_eq!(rw_handle.need_readable(cx), Poll::Pending);
            Poll::<()>::Pending
        });

        let write_fut = futures::future::poll_fn(|cx| {
            let writable = ready!(rw_handle.poll_writable(cx));
            assert_eq!(writable, Ok(WritableState::Writable));
            let buf = [0u8; 1];
            loop {
                match rw_handle.get_ref().write(&buf[..]) {
                    Ok(r) => assert_eq!(r, 1),
                    Err(e) => {
                        assert_eq!(e, zx::Status::SHOULD_WAIT);
                        break;
                    }
                }
            }
            assert_eq!(rw_handle.need_writable(cx), Poll::Pending);
            Poll::<()>::Pending
        });

        let mut fut = std::pin::pin!(futures::future::join(write_fut, read_fut));
        for _ in 0..5 {
            assert_eq!(exec.run_until_stalled(&mut fut), Poll::Pending);
            assert_eq!(peer.read(&mut buff[0..1]), Ok(1));
            assert_eq!(peer.write(&buff[0..1]), Ok(1));
        }

        let mut read_waits = 0;
        let mut write_waits = 0;
        while let Ok(p) = exec.port().wait(zx::MonotonicInstant::INFINITE_PAST) {
            if p.key() != rw_handle.receiver.key() {
                continue;
            }
            let p = match p.contents() {
                zx::PacketContents::SignalOne(p) => p,
                e => panic!("unexpected packet {e:?}"),
            };
            if p.trigger().contains(zx::Signals::OBJECT_READABLE) {
                read_waits += 1;
            }
            if p.trigger().contains(zx::Signals::OBJECT_WRITABLE) {
                write_waits += 1;
            }
        }
        // We should not have installed more than 1 waiter for each side of the
        // operation.
        assert_eq!((read_waits, write_waits), (1, 1));
    }
}