feat(runtime): submit_multi (#743)

* feat(runtime): submit_multi

* fix: unused import

* feat: apply suggestion

* fix: doctest

Author: George-Miao

compio-io/src/framed/codec/bytes.rs

@@ -17,7 +17,7 @@
 //! codec.encode(data.clone(), &mut buffer).unwrap();
 //!
 //! // Decoding
-//! let decoded = codec.decode(&buffer.as_slice()).unwrap();
+//! let decoded = codec.decode(&buffer.slice(..)).unwrap();
 //! assert_eq!(decoded, data);
 //! ```
 use std::io::{self, Write};

compio-runtime/src/runtime/future.rs

@@ -42,17 +42,31 @@ impl ContextExt for Context<'_> {
         Some(extra)
     }
 }
+pin_project_lite::pin_project! {
+    /// Returned [`Future`] for [`Runtime::submit`].
+    ///
+    /// When this is dropped and the operation hasn't finished yet, it will try to
+    /// cancel the operation.
+    ///
+    /// By default, this implements `Future<Output = BufResult<usize, T>>`. If
+    /// [`Extra`] is needed, call [`.with_extra()`] to get a `Submit<T, Extra>`
+    /// which implements `Future<Output = (BufResult<usize, T>, Extra)>`.
+    ///
+    /// [`.with_extra()`]: Submit::with_extra
+    pub struct Submit<T: OpCode, E = ()> {
+        runtime: Runtime,
+        state: Option<State<T, E>>,
+    }
+
+    impl<T: OpCode, E> PinnedDrop for Submit<T, E> {
+        fn drop(this: Pin<&mut Self>) {
+            let this = this.project();
+            if let Some(State::Submitted { key, .. }) = this.state.take() {
+                this.runtime.cancel(key);
+            }
+        }
+    }
 
-/// Return type for `Runtime::submit`
-///
-/// By default, this implements `Future<Output = BufResult<usize, T>>`. If
-/// [`Extra`] is needed, call [`.with_extra()`] to get a `Submit<T, Extra>`
-/// which implements `Future<Output = (BufResult<usize, T>, Extra)>`.
-///
-/// [`.with_extra()`]: Submit::with_extra
-pub struct Submit<T: OpCode, E = ()> {
-    runtime: Runtime,
-    state: Option<State<T, E>>,
 }
 
 enum State<T: OpCode, E> {
@@ -107,12 +121,13 @@ impl<T: OpCode + 'static> Future for Submit<T, ()> {
     type Output = BufResult<usize, T>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = unsafe { self.get_unchecked_mut() };
+        let this = self.project();
+
         loop {
             match this.state.take().expect("Cannot poll after ready") {
                 State::Submitted { key, .. } => match this.runtime.poll_task(cx.waker(), key) {
                     PushEntry::Pending(key) => {
-                        this.state = Some(State::submitted(key));
+                        *this.state = Some(State::submitted(key));
                         return Poll::Pending;
                     }
                     PushEntry::Ready(res) => return Poll::Ready(res),
@@ -127,7 +142,7 @@ impl<T: OpCode + 'static> Future for Submit<T, ()> {
                                 cancel.register(&key);
                             };
 
-                            this.state = Some(State::submitted(key))
+                            *this.state = Some(State::submitted(key))
                         }
                         PushEntry::Ready(res) => {
                             return Poll::Ready(res);
@@ -143,13 +158,14 @@ impl<T: OpCode + 'static> Future for Submit<T, Extra> {
     type Output = (BufResult<usize, T>, Extra);
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = unsafe { self.get_unchecked_mut() };
+        let this = self.project();
+
         loop {
             match this.state.take().expect("Cannot poll after ready") {
                 State::Submitted { key, .. } => {
                     match this.runtime.poll_task_with_extra(cx.waker(), key) {
                         PushEntry::Pending(key) => {
-                            this.state = Some(State::submitted(key));
+                            *this.state = Some(State::submitted(key));
                             return Poll::Pending;
                         }
                         PushEntry::Ready(res) => return Poll::Ready(res),
@@ -163,7 +179,7 @@ impl<T: OpCode + 'static> Future for Submit<T, Extra> {
                                 cancel.register(&key);
                             }
 
-                            this.state = Some(State::submitted(key))
+                            *this.state = Some(State::submitted(key))
                         }
                         PushEntry::Ready(res) => {
                             return Poll::Ready((res, this.runtime.default_extra()));
@@ -183,11 +199,3 @@ where
         self.state.is_none()
     }
 }
-
-impl<T: OpCode, E> Drop for Submit<T, E> {
-    fn drop(&mut self) {
-        if let Some(State::Submitted { key, .. }) = self.state.take() {
-            self.runtime.cancel(key);
-        }
-    }
-}

compio-runtime/src/runtime/mod.rs

@@ -25,6 +25,9 @@ use futures_util::FutureExt;
 mod future;
 pub use future::*;
 
+mod stream;
+pub use stream::*;
+
 #[cfg(feature = "time")]
 pub(crate) mod time;
 
@@ -284,6 +287,13 @@ impl Runtime {
         Submit::new(self.clone(), op)
     }
 
+    /// Submit a multishot operation to the runtime.
+    ///
+    /// You only need this when authoring your own [`OpCode`].
+    pub fn submit_multi<T: OpCode + 'static>(&self, op: T) -> SubmitMulti<T> {
+        SubmitMulti::new(self.clone(), op)
+    }
+
     pub(crate) fn cancel<T: OpCode>(&self, key: Key<T>) {
         self.driver.borrow_mut().cancel(key);
     }
@@ -327,6 +337,20 @@ impl Runtime {
         })
     }
 
+    pub(crate) fn poll_multishot<T: OpCode>(
+        &self,
+        waker: &Waker,
+        key: &Key<T>,
+    ) -> Option<BufResult<usize, Extra>> {
+        instrument!(compio_log::Level::DEBUG, "poll_multishot", ?key);
+        let mut driver = self.driver.borrow_mut();
+        if let Some(res) = driver.pop_multishot(key) {
+            return Some(res);
+        }
+        driver.update_waker(key, waker);
+        None
+    }
+
     #[cfg(feature = "time")]
     pub(crate) fn poll_timer(&self, cx: &mut Context, key: &TimerKey) -> Poll<()> {
         instrument!(compio_log::Level::DEBUG, "poll_timer", ?cx, ?key);
@@ -595,12 +619,24 @@ pub fn spawn_blocking<T: Send + 'static>(
 /// ## Panics
 ///
 /// This method doesn't create runtime and will panic if it's not within a
-/// runtime. It tries to obtain the current runtime by
+/// runtime. It tries to obtain the current runtime with
 /// [`Runtime::with_current`].
 pub fn submit<T: OpCode + 'static>(op: T) -> Submit<T> {
     Runtime::with_current(|r| r.submit(op))
 }
 
+/// Submit a multishot operation to the current runtime, and return a stream for
+/// it.
+///
+/// ## Panics
+///
+/// This method doesn't create runtime and will panic if it's not within a
+/// runtime. It tries to obtain the current runtime with
+/// [`Runtime::with_current`].
+pub fn submit_multi<T: OpCode + 'static>(op: T) -> SubmitMulti<T> {
+    Runtime::with_current(|r| r.submit_multi(op))
+}
+
 /// Register file descriptors for fixed-file operations with the current
 /// runtime's io_uring instance.
 ///

compio-runtime/src/runtime/stream.rs

@@ -0,0 +1,134 @@
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use compio_buf::BufResult;
+use compio_driver::{Extra, Key, OpCode, PushEntry};
+use futures_util::{Stream, stream::FusedStream};
+
+use crate::{ContextExt, Runtime};
+
+pin_project_lite::pin_project! {
+    /// Returned [`Stream`] for [`Runtime::submit_multi`].
+    ///
+    /// When this is dropped and the operation hasn't finished yet, it will try to
+    /// cancel the operation.
+    pub struct SubmitMulti<T: OpCode> {
+        runtime: Runtime,
+        state: Option<State<T>>,
+    }
+
+    impl<T: OpCode> PinnedDrop for SubmitMulti<T> {
+        fn drop(this: Pin<&mut Self>) {
+            let this = this.project();
+            if let Some(State::Submitted { key }) = this.state.take() {
+                this.runtime.cancel(key);
+            }
+        }
+    }
+}
+
+enum State<T: OpCode> {
+    Idle { op: T },
+    Submitted { key: Key<T> },
+    Finished { op: T },
+}
+
+impl<T: OpCode> State<T> {
+    fn submitted(key: Key<T>) -> Self {
+        State::Submitted { key }
+    }
+}
+
+impl<T: OpCode> SubmitMulti<T> {
+    pub(crate) fn new(runtime: Runtime, op: T) -> Self {
+        SubmitMulti {
+            runtime,
+            state: Some(State::Idle { op }),
+        }
+    }
+
+    /// Try to take the inner op from the stream.
+    ///
+    /// Returns `Ok(T)` if the stream:
+    ///
+    /// - has not been polled yet, or
+    /// - is finished and the op is returned by the driver
+    ///
+    /// Returns `Err(Self)` if it's still running.
+    pub fn try_take(mut self) -> Result<T, Self> {
+        match self.state.take() {
+            Some(State::Finished { op }) | Some(State::Idle { op }) => Ok(op),
+            state => {
+                debug_assert!(state.is_some());
+                self.state = state;
+                Err(self)
+            }
+        }
+    }
+}
+
+impl<T: OpCode + 'static> Stream for SubmitMulti<T> {
+    type Item = BufResult<usize, Extra>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+
+        loop {
+            match this.state.take().expect("State error, this is a bug") {
+                State::Idle { op } => {
+                    let extra = cx.as_extra(|| this.runtime.default_extra());
+                    match this.runtime.submit_raw(op, extra) {
+                        PushEntry::Pending(key) => {
+                            if let Some(cancel) = cx.as_cancel() {
+                                cancel.register(&key);
+                            }
+
+                            *this.state = Some(State::submitted(key))
+                        }
+                        PushEntry::Ready(BufResult(res, op)) => {
+                            *this.state = Some(State::Finished { op });
+                            let extra = this.runtime.default_extra();
+
+                            return Poll::Ready(Some(BufResult(res, extra)));
+                        }
+                    }
+                }
+
+                State::Submitted { key, .. } => {
+                    if let Some(res) = this.runtime.poll_multishot(cx.waker(), &key) {
+                        *this.state = Some(State::submitted(key));
+
+                        return Poll::Ready(Some(res));
+                    };
+
+                    match this.runtime.poll_task_with_extra(cx.waker(), key) {
+                        PushEntry::Pending(key) => {
+                            *this.state = Some(State::submitted(key));
+
+                            return Poll::Pending;
+                        }
+                        PushEntry::Ready((BufResult(res, op), extra)) => {
+                            *this.state = Some(State::Finished { op });
+
+                            return Poll::Ready(Some(BufResult(res, extra)));
+                        }
+                    }
+                }
+
+                State::Finished { op } => {
+                    *this.state = Some(State::Finished { op });
+
+                    return Poll::Ready(None);
+                }
+            }
+        }
+    }
+}
+
+impl<T: OpCode + 'static> FusedStream for SubmitMulti<T> {
+    fn is_terminated(&self) -> bool {
+        matches!(self.state, None | Some(State::Finished { .. }))
+    }
+}