refactor(core/raw): Migrate oio::Write to async in trait (#4358)

apache · Mar 14, 2024 · 6ebd75c · 6ebd75c
1 parent 7533627
commit 6ebd75c
Show file tree

Hide file tree

Showing 65 changed files with 945 additions and 1,640 deletions.
diff --git a/core/benches/oio/utils.rs b/core/benches/oio/utils.rs
@@ -15,9 +15,6 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::task::Context;
-use std::task::Poll;
-
 use bytes::Bytes;
 use opendal::raw::oio;
 use rand::prelude::ThreadRng;
@@ -27,16 +24,16 @@ use rand::RngCore;
 pub struct BlackHoleWriter;
 
 impl oio::Write for BlackHoleWriter {
-    fn poll_write(&mut self, _: &mut Context<'_>, bs: Bytes) -> Poll<opendal::Result<usize>> {
-        Poll::Ready(Ok(bs.len()))
+    async fn write(&mut self, bs: Bytes) -> opendal::Result<usize> {
+        Ok(bs.len())
     }
 
-    fn poll_abort(&mut self, _: &mut Context<'_>) -> Poll<opendal::Result<()>> {
-        Poll::Ready(Ok(()))
+    async fn abort(&mut self) -> opendal::Result<()> {
+        Ok(())
     }
 
-    fn poll_close(&mut self, _: &mut Context<'_>) -> Poll<opendal::Result<()>> {
-        Poll::Ready(Ok(()))
+    async fn close(&mut self) -> opendal::Result<()> {
+        Ok(())
     }
 }
 

diff --git a/core/benches/oio/write.rs b/core/benches/oio/write.rs
@@ -19,7 +19,7 @@ use bytes::Buf;
 use criterion::Criterion;
 use once_cell::sync::Lazy;
 use opendal::raw::oio::ExactBufWriter;
-use opendal::raw::oio::WriteExt;
+use opendal::raw::oio::Write;
 use rand::thread_rng;
 use size::Size;
 

diff --git a/core/src/layers/blocking.rs b/core/src/layers/blocking.rs
@@ -18,7 +18,7 @@
 use async_trait::async_trait;
 use bytes;
 use bytes::Bytes;
-use futures::future::poll_fn;
+
 use tokio::runtime::Handle;
 
 use crate::raw::*;
@@ -299,13 +299,11 @@ impl<I: oio::Read + 'static> oio::BlockingRead for BlockingWrapper<I> {
 
 impl<I: oio::Write + 'static> oio::BlockingWrite for BlockingWrapper<I> {
     fn write(&mut self, bs: Bytes) -> Result<usize> {
-        self.handle
-            .block_on(poll_fn(|cx| self.inner.poll_write(cx, bs.clone())))
+        self.handle.block_on(self.inner.write(bs))
     }
 
     fn close(&mut self) -> Result<()> {
-        self.handle
-            .block_on(poll_fn(|cx| self.inner.poll_close(cx)))
+        self.handle.block_on(self.inner.close())
     }
 }
 

diff --git a/core/src/layers/complete.rs b/core/src/layers/complete.rs
@@ -18,10 +18,8 @@
 use std::cmp;
 use std::fmt::Debug;
 use std::fmt::Formatter;
+
 use std::sync::Arc;
-use std::task::ready;
-use std::task::Context;
-use std::task::Poll;
 
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -158,7 +156,7 @@ impl<A: Accessor> CompleteAccessor<A> {
         }
         if capability.write_can_empty && capability.list {
             let (_, mut w) = self.inner.write(path, OpWrite::default()).await?;
-            oio::WriteExt::close(&mut w).await?;
+            oio::Write::close(&mut w).await?;
             return Ok(RpCreateDir::default());
         }
 
@@ -712,35 +710,34 @@ impl<W> oio::Write for CompleteWriter<W>
 where
     W: oio::Write,
 {
-    fn poll_write(&mut self, cx: &mut Context<'_>, bs: Bytes) -> Poll<Result<usize>> {
+    async fn write(&mut self, bs: Bytes) -> Result<usize> {
         let w = self.inner.as_mut().ok_or_else(|| {
             Error::new(ErrorKind::Unexpected, "writer has been closed or aborted")
         })?;
-        let n = ready!(w.poll_write(cx, bs))?;
 
-        Poll::Ready(Ok(n))
+        w.write(bs).await
     }
 
-    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+    async fn close(&mut self) -> Result<()> {
         let w = self.inner.as_mut().ok_or_else(|| {
             Error::new(ErrorKind::Unexpected, "writer has been closed or aborted")
         })?;
 
-        ready!(w.poll_close(cx))?;
+        w.close().await?;
         self.inner = None;
 
-        Poll::Ready(Ok(()))
+        Ok(())
     }
 
-    fn poll_abort(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+    async fn abort(&mut self) -> Result<()> {
         let w = self.inner.as_mut().ok_or_else(|| {
             Error::new(ErrorKind::Unexpected, "writer has been closed or aborted")
         })?;
 
-        ready!(w.poll_abort(cx))?;
+        w.abort().await?;
         self.inner = None;
 
-        Poll::Ready(Ok(()))
+        Ok(())
     }
 }
 

diff --git a/core/src/layers/concurrent_limit.rs b/core/src/layers/concurrent_limit.rs
@@ -19,8 +19,6 @@ use std::fmt::Debug;
 
 use std::io::SeekFrom;
 use std::sync::Arc;
-use std::task::Context;
-use std::task::Poll;
 
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -278,16 +276,16 @@ impl<R: oio::BlockingRead> oio::BlockingRead for ConcurrentLimitWrapper<R> {
 }
 
 impl<R: oio::Write> oio::Write for ConcurrentLimitWrapper<R> {
-    fn poll_write(&mut self, cx: &mut Context<'_>, bs: Bytes) -> Poll<Result<usize>> {
-        self.inner.poll_write(cx, bs)
+    async fn write(&mut self, bs: Bytes) -> Result<usize> {
+        self.inner.write(bs).await
     }
 
-    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        self.inner.poll_close(cx)
+    async fn close(&mut self) -> Result<()> {
+        self.inner.close().await
     }
 
-    fn poll_abort(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        self.inner.poll_abort(cx)
+    async fn abort(&mut self) -> Result<()> {
+        self.inner.abort().await
     }
 }
 

diff --git a/core/src/layers/dtrace.rs b/core/src/layers/dtrace.rs
@@ -18,9 +18,8 @@
 use std::ffi::CString;
 use std::fmt::Debug;
 use std::fmt::Formatter;
+
 use std::io;
-use std::task::Context;
-use std::task::Poll;
 
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -408,12 +407,13 @@ impl<R: oio::BlockingRead> oio::BlockingRead for DtraceLayerWrapper<R> {
 }
 
 impl<R: oio::Write> oio::Write for DtraceLayerWrapper<R> {
-    fn poll_write(&mut self, cx: &mut Context<'_>, bs: Bytes) -> Poll<Result<usize>> {
+    async fn write(&mut self, bs: Bytes) -> Result<usize> {
         let c_path = CString::new(self.path.clone()).unwrap();
         probe_lazy!(opendal, writer_write_start, c_path.as_ptr());
         self.inner
-            .poll_write(cx, bs)
-            .map_ok(|n| {
+            .write(bs)
+            .await
+            .map(|n| {
                 probe_lazy!(opendal, writer_write_ok, c_path.as_ptr(), n);
                 n
             })
@@ -423,12 +423,13 @@ impl<R: oio::Write> oio::Write for DtraceLayerWrapper<R> {
             })
     }
 
-    fn poll_abort(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+    async fn abort(&mut self) -> Result<()> {
         let c_path = CString::new(self.path.clone()).unwrap();
         probe_lazy!(opendal, writer_poll_abort_start, c_path.as_ptr());
         self.inner
-            .poll_abort(cx)
-            .map_ok(|_| {
+            .abort()
+            .await
+            .map(|_| {
                 probe_lazy!(opendal, writer_poll_abort_ok, c_path.as_ptr());
             })
             .map_err(|err| {
@@ -437,12 +438,13 @@ impl<R: oio::Write> oio::Write for DtraceLayerWrapper<R> {
             })
     }
 
-    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
+    async fn close(&mut self) -> Result<()> {
         let c_path = CString::new(self.path.clone()).unwrap();
         probe_lazy!(opendal, writer_close_start, c_path.as_ptr());
         self.inner
-            .poll_close(cx)
-            .map_ok(|_| {
+            .close()
+            .await
+            .map(|_| {
                 probe_lazy!(opendal, writer_close_ok, c_path.as_ptr());
             })
             .map_err(|err| {

diff --git a/core/src/layers/error_context.rs b/core/src/layers/error_context.rs
@@ -19,8 +19,6 @@ use std::fmt::Debug;
 use std::fmt::Formatter;
 
 use std::io::SeekFrom;
-use std::task::Context;
-use std::task::Poll;
 
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -387,27 +385,26 @@ impl<T: oio::BlockingRead> oio::BlockingRead for ErrorContextWrapper<T> {
     }
 }
 
-#[async_trait::async_trait]
 impl<T: oio::Write> oio::Write for ErrorContextWrapper<T> {
-    fn poll_write(&mut self, cx: &mut Context<'_>, bs: Bytes) -> Poll<Result<usize>> {
-        self.inner.poll_write(cx, bs.clone()).map_err(|err| {
+    async fn write(&mut self, bs: Bytes) -> Result<usize> {
+        self.inner.write(bs.clone()).await.map_err(|err| {
             err.with_operation(WriteOperation::Write)
                 .with_context("service", self.scheme)
                 .with_context("path", &self.path)
                 .with_context("write_buf", bs.len().to_string())
         })
     }
 
-    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        self.inner.poll_close(cx).map_err(|err| {
+    async fn close(&mut self) -> Result<()> {
+        self.inner.close().await.map_err(|err| {
             err.with_operation(WriteOperation::Close)
                 .with_context("service", self.scheme)
                 .with_context("path", &self.path)
         })
     }
 
-    fn poll_abort(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        self.inner.poll_abort(cx).map_err(|err| {
+    async fn abort(&mut self) -> Result<()> {
+        self.inner.abort().await.map_err(|err| {
             err.with_operation(WriteOperation::Abort)
                 .with_context("service", self.scheme)
                 .with_context("path", &self.path)

diff --git a/core/src/layers/logging.rs b/core/src/layers/logging.rs
@@ -18,9 +18,6 @@
 use std::fmt::Debug;
 
 use std::io;
-use std::task::ready;
-use std::task::Context;
-use std::task::Poll;
 
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -1147,8 +1144,8 @@ impl<W> LoggingWriter<W> {
 }
 
 impl<W: oio::Write> oio::Write for LoggingWriter<W> {
-    fn poll_write(&mut self, cx: &mut Context<'_>, bs: Bytes) -> Poll<Result<usize>> {
-        match ready!(self.inner.poll_write(cx, bs.clone())) {
+    async fn write(&mut self, bs: Bytes) -> Result<usize> {
+        match self.inner.write(bs.clone()).await {
             Ok(n) => {
                 self.written += n as u64;
                 trace!(
@@ -1161,7 +1158,7 @@ impl<W: oio::Write> oio::Write for LoggingWriter<W> {
                     bs.len(),
                     n,
                 );
-                Poll::Ready(Ok(n))
+                Ok(n)
             }
             Err(err) => {
                 if let Some(lvl) = self.ctx.error_level(&err) {
@@ -1176,13 +1173,13 @@ impl<W: oio::Write> oio::Write for LoggingWriter<W> {
                         self.ctx.error_print(&err),
                     )
                 }
-                Poll::Ready(Err(err))
+                Err(err)
             }
         }
     }
 
-    fn poll_abort(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        match ready!(self.inner.poll_abort(cx)) {
+    async fn abort(&mut self) -> Result<()> {
+        match self.inner.abort().await {
             Ok(_) => {
                 trace!(
                     target: LOGGING_TARGET,
@@ -1192,7 +1189,7 @@ impl<W: oio::Write> oio::Write for LoggingWriter<W> {
                     self.path,
                     self.written,
                 );
-                Poll::Ready(Ok(()))
+                Ok(())
             }
             Err(err) => {
                 if let Some(lvl) = self.ctx.error_level(&err) {
@@ -1207,13 +1204,13 @@ impl<W: oio::Write> oio::Write for LoggingWriter<W> {
                         self.ctx.error_print(&err),
                     )
                 }
-                Poll::Ready(Err(err))
+                Err(err)
             }
         }
     }
 
-    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
-        match ready!(self.inner.poll_close(cx)) {
+    async fn close(&mut self) -> Result<()> {
+        match self.inner.close().await {
             Ok(_) => {
                 debug!(
                     target: LOGGING_TARGET,
@@ -1223,7 +1220,7 @@ impl<W: oio::Write> oio::Write for LoggingWriter<W> {
                     self.path,
                     self.written
                 );
-                Poll::Ready(Ok(()))
+                Ok(())
             }
             Err(err) => {
                 if let Some(lvl) = self.ctx.error_level(&err) {
@@ -1238,7 +1235,7 @@ impl<W: oio::Write> oio::Write for LoggingWriter<W> {
                         self.ctx.error_print(&err),
                     )
                 }
-                Poll::Ready(Err(err))
+                Err(err)
             }
         }
     }

diff --git a/core/src/layers/madsim.rs b/core/src/layers/madsim.rs
@@ -291,7 +291,7 @@ pub struct MadsimWriter {
 }
 
 impl oio::Write for MadsimWriter {
-    fn poll_write(&mut self, cx: &mut Context<'_>, bs: Bytes) -> Poll<crate::Result<usize>> {
+    async fn write(&mut self, bs: Bytes) -> crate::Result<usize> {
         #[cfg(madsim)]
         {
             let req = Request::Write(self.path.to_string(), bs);
@@ -307,15 +307,15 @@ impl oio::Write for MadsimWriter {
         }
     }
 
-    fn poll_abort(&mut self, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
-        Poll::Ready(Err(Error::new(
+    async fn abort(&mut self) -> crate::Result<()> {
+        Err(Error::new(
             ErrorKind::Unsupported,
             "will be supported in the future",
-        )))
+        ))
     }
 
-    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<crate::Result<()>> {
-        Poll::Ready(Ok(()))
+    async fn close(&mut self) -> crate::Result<()> {
+        Ok(())
     }
 }