feat(storage): Make sure all futures are spawned on global runtime (#10918)

* feat(storage): Make sure all futures are created on global runtime

Signed-off-by: Xuanwo <[email protected]>

* Allow dirty buffer

Signed-off-by: Xuanwo <[email protected]>

* format code

Signed-off-by: Xuanwo <[email protected]>

* Make sure runtime layer applied first

Signed-off-by: Xuanwo <[email protected]>

* Fix bug

Signed-off-by: Xuanwo <[email protected]>

* Use copy instead

Signed-off-by: Xuanwo <[email protected]>

---------

Signed-off-by: Xuanwo <[email protected]>

Author: Xuanwo

src/common/storage/src/operator.rs

@@ -85,6 +85,12 @@ pub fn build_operator<B: Builder>(builder: B) -> Result<Operator> {
     let ob = Operator::new(builder)?;
 
     let op = ob
+        // NOTE
+        //
+        // Magic happens here. We will add a layer upon original
+        // storage operator so that all underlying storage operations
+        // will send to storage runtime.
+        .layer(RuntimeLayer::new(GlobalIORuntime::instance().inner()))
         // Add retry
         .layer(RetryLayer::new().with_jitter())
         // Add metrics
@@ -93,12 +99,6 @@ pub fn build_operator<B: Builder>(builder: B) -> Result<Operator> {
         .layer(LoggingLayer::default())
         // Add tracing
         .layer(TracingLayer)
-        // NOTE
-        //
-        // Magic happens here. We will add a layer upon original
-        // storage operator so that all underlying storage operations
-        // will send to storage runtime.
-        .layer(RuntimeLayer::new(GlobalIORuntime::instance().inner()))
         .finish();
 
     Ok(op)

src/common/storage/src/runtime_layer.rs

@@ -13,17 +13,28 @@
 // limitations under the License.
 
 use std::env;
+use std::io::SeekFrom;
+use std::mem;
+use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::LazyLock;
+use std::task::Context;
+use std::task::Poll;
 use std::time::Duration;
 
 use async_trait::async_trait;
+use bytes::Bytes;
 use common_base::base::tokio::pin;
 use common_base::base::tokio::runtime::Handle;
 use common_base::base::tokio::select;
+use common_base::base::tokio::task::JoinHandle;
 use common_base::base::tokio::time;
 use common_base::runtime::TrackedFuture;
+use futures::ready;
+use futures::Future;
 use opendal::ops::*;
+use opendal::raw::oio;
+use opendal::raw::oio::ReadExt;
 use opendal::raw::Accessor;
 use opendal::raw::Layer;
 use opendal::raw::LayeredAccessor;
@@ -84,7 +95,7 @@ pub struct RuntimeAccessor<A> {
 #[async_trait]
 impl<A: Accessor> LayeredAccessor for RuntimeAccessor<A> {
     type Inner = A;
-    type Reader = A::Reader;
+    type Reader = RuntimeIO<A::Reader>;
     type BlockingReader = A::BlockingReader;
     type Writer = A::Writer;
     type BlockingWriter = A::BlockingWriter;
@@ -127,7 +138,14 @@ impl<A: Accessor> LayeredAccessor for RuntimeAccessor<A> {
         };
 
         let future = TrackedFuture::create(future);
-        self.runtime.spawn(future).await.expect("join must success")
+        self.runtime
+            .spawn(future)
+            .await
+            .expect("join must success")
+            .map(|(rp, r)| {
+                let r = RuntimeIO::new(r, self.runtime.clone());
+                (rp, r)
+            })
     }
 
     #[async_backtrace::framed]
@@ -191,3 +209,161 @@ impl<A: Accessor> LayeredAccessor for RuntimeAccessor<A> {
         self.inner.blocking_scan(path, args)
     }
 }
+
+pub struct RuntimeIO<R: 'static> {
+    runtime: Handle,
+    state: State<R>,
+    buf: Vec<u8>,
+}
+
+impl<R> RuntimeIO<R> {
+    fn new(inner: R, runtime: Handle) -> Self {
+        Self {
+            runtime,
+            state: State::Idle(Some(inner)),
+            buf: vec![],
+        }
+    }
+}
+
+pub enum State<R: 'static> {
+    Idle(Option<R>),
+    Read(JoinHandle<(R, Result<Vec<u8>>)>),
+    Seek(JoinHandle<(R, Result<u64>)>),
+    Next(JoinHandle<(R, Option<Result<Bytes>>)>),
+}
+
+/// Safety: State will only be accessed under &mut.
+unsafe impl<R> Sync for State<R> {}
+
+impl<R: oio::Read> oio::Read for RuntimeIO<R> {
+    /// TODO: the performance of `read` could be affected, we will improve it later.
+    fn poll_read(&mut self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> {
+        match &mut self.state {
+            State::Idle(r) => {
+                let mut r = r.take().expect("Idle must have a valid reader");
+                let mut buffer = mem::take(&mut self.buf);
+
+                buffer.reserve(buf.len());
+                // Safety: buffer is reserved with buf.len() bytes.
+                #[allow(clippy::uninit_vec)]
+                unsafe {
+                    buffer.set_len(buf.len())
+                }
+
+                let future = async move {
+                    let mut buffer = buffer;
+                    let res = r.read(&mut buffer).await;
+                    match res {
+                        Ok(size) => {
+                            // Safety: we trust our reader, the returning size is correct.
+                            unsafe { buffer.set_len(size) }
+                            (r, Ok(buffer))
+                        }
+                        Err(err) => (r, Err(err)),
+                    }
+                };
+                let future = TrackedFuture::create(future);
+                self.state = State::Read(self.runtime.spawn(future));
+
+                self.poll_read(cx, buf)
+            }
+            State::Read(future) => {
+                let (r, res) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+                match res {
+                    Ok(mut buffer) => {
+                        let size = buffer.len();
+                        buf[..size].copy_from_slice(&buffer);
+                        // Safety: set length to 0 as we don't care the remaining content.
+                        unsafe { buffer.set_len(0) }
+                        // Always reuse the same buffer
+                        self.buf = buffer;
+                        Poll::Ready(Ok(size))
+                    }
+                    Err(err) => Poll::Ready(Err(err)),
+                }
+            }
+            State::Seek(future) => {
+                let (r, _) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                self.poll_read(cx, buf)
+            }
+            State::Next(future) => {
+                let (r, _) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                self.poll_read(cx, buf)
+            }
+        }
+    }
+
+    fn poll_seek(&mut self, cx: &mut Context<'_>, pos: SeekFrom) -> Poll<Result<u64>> {
+        match &mut self.state {
+            State::Idle(r) => {
+                let mut r = r.take().expect("Idle must have a valid reader");
+                let future = async move {
+                    let res = r.seek(pos).await;
+                    (r, res)
+                };
+                let future = TrackedFuture::create(future);
+                self.state = State::Seek(self.runtime.spawn(future));
+
+                self.poll_seek(cx, pos)
+            }
+            State::Read(future) => {
+                let (r, _) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                self.poll_seek(cx, pos)
+            }
+            State::Seek(future) => {
+                let (r, res) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                Poll::Ready(res)
+            }
+            State::Next(future) => {
+                let (r, _) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                self.poll_seek(cx, pos)
+            }
+        }
+    }
+
+    fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<bytes::Bytes>>> {
+        match &mut self.state {
+            State::Idle(r) => {
+                let mut r = r.take().expect("Idle must have a valid reader");
+                let future = async move {
+                    let res = r.next().await;
+                    (r, res)
+                };
+                let future = TrackedFuture::create(future);
+                self.state = State::Next(self.runtime.spawn(future));
+
+                self.poll_next(cx)
+            }
+            State::Read(future) => {
+                let (r, _) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                self.poll_next(cx)
+            }
+            State::Seek(future) => {
+                let (r, _) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                self.poll_next(cx)
+            }
+            State::Next(future) => {
+                let (r, res) = ready!(Pin::new(future).poll(cx)).expect("join must success");
+                self.state = State::Idle(Some(r));
+
+                Poll::Ready(res)
+            }
+        }
+    }
+}