glommio 0.6.0

// Unless explicitly stated otherwise all files in this repository are licensed
// under the MIT/Apache-2.0 License, at your convenience
//
// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2021 Datadog, Inc.
//
use crate::{
    iou::sqe::{SockAddr, SockAddrStorage},
    sys::{
        DmaBuffer,
        IoBuffer,
        OsResult,
        PollableStatus,
        ReactorQueue,
        SourceId,
        TimeSpec64,
        Wakers,
    },
    GlommioError,
    IoRequirements,
    Latency,
    ReactorErrorKind,
    RingIoStats,
    TaskQueueHandle,
};
use futures_lite::{future, io};
use std::{
    cell::{Ref, RefCell, RefMut},
    convert::TryFrom,
    ffi::CString,
    fmt,
    mem::MaybeUninit,
    os::unix::io::RawFd,
    path::PathBuf,
    rc::Rc,
    task::{Poll, Waker},
    time::Duration,
};

#[derive(Debug)]
pub(crate) enum SourceType {
    Write(PollableStatus, IoBuffer),
    Read(PollableStatus, Option<IoBuffer>),
    SockSend(DmaBuffer),
    SockRecv(Option<DmaBuffer>),
    SockRecvMsg(
        Option<DmaBuffer>,
        libc::iovec,
        libc::msghdr,
        MaybeUninit<nix::sys::socket::sockaddr_storage>,
    ),
    SockSendMsg(
        DmaBuffer,
        libc::iovec,
        libc::msghdr,
        nix::sys::socket::SockAddr,
    ),
    Open(CString),
    FdataSync,
    Fallocate,
    Truncate,
    Close,
    LinkRings,
    Statx(CString, Box<RefCell<libc::statx>>),
    Timeout(TimeSpec64),
    Connect(SockAddr),
    Accept(SockAddrStorage),
    Rename(PathBuf, PathBuf),
    CreateDir(PathBuf),
    Remove(PathBuf),
    Invalid,
    #[cfg(feature = "bench")]
    Noop,
}

impl TryFrom<SourceType> for libc::statx {
    type Error = GlommioError<()>;

    fn try_from(value: SourceType) -> Result<Self, Self::Error> {
        match value {
            SourceType::Statx(_, buf) => Ok(buf.into_inner()),
            _ => Err(GlommioError::ReactorError(
                ReactorErrorKind::IncorrectSourceType,
            )),
        }
    }
}

pub struct EnqueuedSource {
    pub(crate) id: SourceId,
    pub(crate) queue: ReactorQueue,
}

pub(crate) type StatsCollectionFn = fn(&io::Result<usize>, &mut RingIoStats, waiters: u64) -> ();
#[derive(Copy, Clone)]
pub(crate) struct StatsCollection {
    /// fulfilled runs when the source exits the reactor
    pub(crate) fulfilled: Option<StatsCollectionFn>,
    /// reused runs when a fulfilled source is reused by another consumer
    pub(crate) reused: Option<StatsCollectionFn>,
}

/// A registered source of I/O events.
pub(crate) struct InnerSource {
    /// Raw file descriptor on Unix platforms.
    pub(crate) raw: RawFd,

    /// Tasks interested in events on this source.
    pub(crate) wakers: Wakers,

    pub(crate) source_type: SourceType,

    pub(crate) io_requirements: IoRequirements,

    pub(crate) timeout: Option<TimeSpec64>,

    pub(crate) enqueued: Option<EnqueuedSource>,

    pub(crate) stats_collection: Option<StatsCollection>,

    pub(crate) task_queue: Option<TaskQueueHandle>,
}

impl InnerSource {
    pub(crate) fn update_source_type(&mut self, source_type: SourceType) -> SourceType {
        std::mem::replace(&mut self.source_type, source_type)
    }
}

impl fmt::Debug for InnerSource {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("InnerSource")
            .field("raw", &self.raw)
            .field("wakers", &self.wakers)
            .field("source_type", &self.source_type)
            .field("io_requirements", &self.io_requirements)
            .finish()
    }
}

#[derive(Debug)]
pub struct Source {
    pub(crate) inner: Rc<RefCell<InnerSource>>,
}

impl Source {
    /// Registers an I/O source in the reactor.
    pub(crate) fn new(
        ioreq: IoRequirements,
        raw: RawFd,
        source_type: SourceType,
        stats_collection: Option<StatsCollection>,
        task_queue: Option<TaskQueueHandle>,
    ) -> Source {
        Source {
            inner: Rc::new(RefCell::new(InnerSource {
                raw,
                wakers: Wakers::new(),
                source_type,
                io_requirements: ioreq,
                enqueued: None,
                timeout: None,
                stats_collection,
                task_queue,
            })),
        }
    }

    pub(crate) fn set_timeout(&self, d: Duration) -> Option<Duration> {
        let mut inner = self.inner.borrow_mut();
        let t = &mut inner.timeout;
        let old = *t;
        *t = Some(TimeSpec64::from(d));
        old.map(Duration::from)
    }

    pub(super) fn timeout_ref(&self) -> Ref<'_, Option<TimeSpec64>> {
        Ref::map(self.inner.borrow(), |x| &x.timeout)
    }

    pub(crate) fn latency_req(&self) -> Latency {
        self.inner.borrow().io_requirements.latency_req
    }

    pub(super) fn source_type(&self) -> Ref<'_, SourceType> {
        Ref::map(self.inner.borrow(), |x| &x.source_type)
    }

    pub(crate) fn source_type_mut(&self) -> RefMut<'_, SourceType> {
        RefMut::map(self.inner.borrow_mut(), |x| &mut x.source_type)
    }

    pub(crate) fn extract_source_type(self) -> SourceType {
        self.inner
            .borrow_mut()
            .update_source_type(SourceType::Invalid)
    }

    pub(crate) fn extract_buffer(self) -> IoBuffer {
        let stype = self.extract_source_type();
        match stype {
            SourceType::Read(_, Some(buffer)) => buffer,
            SourceType::Write(_, buffer) => buffer,
            x => panic!("Could not extract buffer. Source: {:?}", x),
        }
    }

    pub(crate) fn buffer(&self) -> Ref<'_, IoBuffer> {
        Ref::map(self.source_type(), |stype| match &*stype {
            SourceType::Read(_, Some(buffer)) => buffer,
            SourceType::Write(_, buffer) => buffer,
            x => panic!("Could not extract buffer. Source: {:?}", x),
        })
    }

    pub(crate) fn result(&self) -> Option<io::Result<usize>> {
        self.inner
            .borrow()
            .wakers
            .result
            .as_ref()
            .map(|x| OsResult::from(x).into())
    }

    // adds a single waiter to the list, replacing any waiter that may already
    // exist. Should be used for single streams that map a future 1:1 to their I/O
    // source
    pub(crate) fn add_waiter_single(&self, waker: Waker) {
        let mut inner = self.inner.borrow_mut();
        inner.wakers.waiters.pop();
        inner.wakers.waiters.push(waker);
        debug_assert_eq!(inner.wakers.waiters.len(), 1)
    }

    // adds a waiter to the list. Useful for streams that have many futures waiting
    // on a single I/O source
    pub(crate) fn add_waiter_many(&self, waker: Waker) {
        self.inner.borrow_mut().wakers.waiters.push(waker)
    }

    // used for eventfd, and other internal sources that reuse the same source
    // across many invocations
    pub(super) fn take_result(&self) -> Option<io::Result<usize>> {
        self.inner
            .borrow_mut()
            .wakers
            .result
            .take()
            .map(|x| OsResult::from(x).into())
    }

    pub(super) fn raw(&self) -> RawFd {
        self.inner.borrow().raw
    }

    pub(crate) fn stats_collection(&self) -> Option<StatsCollection> {
        self.inner.borrow().stats_collection
    }

    pub(crate) async fn collect_rw(&self) -> io::Result<usize> {
        future::poll_fn(|cx| {
            if let Some(result) = self.result() {
                return Poll::Ready(result);
            }

            self.add_waiter_many(cx.waker().clone());
            Poll::Pending
        })
        .await
    }
}

impl Drop for Source {
    fn drop(&mut self) {
        let enqueued = self.inner.borrow_mut().enqueued.take();
        if let Some(EnqueuedSource { id, queue }) = enqueued {
            queue.borrow_mut().cancel_request(id);
        }
    }
}