feat(core): Add Write::sink API

core/src/docs/internals/accessor.rs

@@ -155,7 +155,7 @@
 //! vendors that provide s3-like RESTful APIs, and our s3 service is
 //! implemented to support all of them, not just AWS S3.
 //!
-//! Obviously, we can use `duck` as scheme, let's add a new variant in [`Scheme`], and implement all reqired functions like `Scheme::from_str` and `Scheme::into_static`:
+//! Obviously, we can use `duck` as scheme, let's add a new variant in [`Scheme`], and implement all required functions like `Scheme::from_str` and `Scheme::into_static`:
 //!
 //! ```ignore
 //! pub enum Scheme {

core/src/layers/complete.rs

@@ -568,6 +568,27 @@ where
         Ok(())
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        if let Some(total_size) = self.size {
+            if self.written + size > total_size {
+                return Err(Error::new(
+                    ErrorKind::ContentTruncated,
+                    &format!(
+                        "writer got too much data, expect: {size}, actual: {}",
+                        self.written + size
+                    ),
+                ));
+            }
+        }
+
+        let w = self.inner.as_mut().ok_or_else(|| {
+            Error::new(ErrorKind::Unexpected, "writer has been closed or aborted")
+        })?;
+        w.sink(size, s).await?;
+        self.written += size;
+        Ok(())
+    }
+
     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")

core/src/layers/concurrent_limit.rs

@@ -308,6 +308,10 @@ impl<R: oio::Write> oio::Write for ConcurrentLimitWrapper<R> {
         self.inner.abort().await
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner.sink(size, s).await
+    }
+
     async fn close(&mut self) -> Result<()> {
         self.inner.close().await
     }

core/src/layers/error_context.rs

@@ -442,6 +442,14 @@ impl<T: oio::Write> oio::Write for ErrorContextWrapper<T> {
         })
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner.sink(size, s).await.map_err(|err| {
+            err.with_operation(WriteOperation::Sink)
+                .with_context("service", self.scheme)
+                .with_context("path", &self.path)
+        })
+    }
+
     async fn close(&mut self) -> Result<()> {
         self.inner.close().await.map_err(|err| {
             err.with_operation(WriteOperation::Close)

core/src/layers/logging.rs

@@ -1343,6 +1343,38 @@ impl<W: oio::Write> oio::Write for LoggingWriter<W> {
         }
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        match self.inner.sink(size, s).await {
+            Ok(_) => {
+                self.written += size;
+                trace!(
+                    target: LOGGING_TARGET,
+                    "service={} operation={} path={} written={} -> data sink {}B",
+                    self.scheme,
+                    WriteOperation::Sink,
+                    self.path,
+                    self.written,
+                    size
+                );
+                Ok(())
+            }
+            Err(err) => {
+                if let Some(lvl) = self.failure_level {
+                    log!(
+                        target: LOGGING_TARGET,
+                        lvl,
+                        "service={} operation={} path={} written={} -> data sink failed: {err:?}",
+                        self.scheme,
+                        WriteOperation::Sink,
+                        self.path,
+                        self.written,
+                    )
+                }
+                Err(err)
+            }
+        }
+    }
+
     async fn abort(&mut self) -> Result<()> {
         match self.inner.abort().await {
             Ok(_) => {

core/src/layers/madsim.rs

@@ -330,6 +330,13 @@ impl oio::Write for MadsimWriter {
         }
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> crate::Result<()> {
+        Err(Error::new(
+            ErrorKind::Unsupported,
+            "will be supported in the future",
+        ))
+    }
+
     async fn abort(&mut self) -> crate::Result<()> {
         Err(Error::new(
             ErrorKind::Unsupported,

core/src/layers/metrics.rs

@@ -864,6 +864,17 @@ impl<R: oio::Write> oio::Write for MetricWrapper<R> {
             })
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner
+            .sink(size, s)
+            .await
+            .map(|_| self.bytes += size)
+            .map_err(|err| {
+                self.handle.increment_errors_total(self.op, err.kind());
+                err
+            })
+    }
+
     async fn abort(&mut self) -> Result<()> {
         self.inner.abort().await.map_err(|err| {
             self.handle.increment_errors_total(self.op, err.kind());

core/src/layers/minitrace.rs

@@ -323,6 +323,16 @@ impl<R: oio::Write> oio::Write for MinitraceWrapper<R> {
             .await
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner
+            .sink(size, s)
+            .in_span(Span::enter_with_parent(
+                WriteOperation::Sink.into_static(),
+                &self.span,
+            ))
+            .await
+    }
+
     async fn abort(&mut self) -> Result<()> {
         self.inner
             .abort()

core/src/layers/oteltrace.rs

@@ -322,6 +322,10 @@ impl<R: oio::Write> oio::Write for OtelTraceWrapper<R> {
         self.inner.write(bs).await
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner.sink(size, s).await
+    }
+
     async fn abort(&mut self) -> Result<()> {
         self.inner.abort().await
     }

core/src/layers/prometheus.rs

@@ -684,6 +684,22 @@ impl<R: oio::Write> oio::Write for PrometheusMetricWrapper<R> {
             })
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner
+            .sink(size, s)
+            .await
+            .map(|_| {
+                self.stats
+                    .bytes_total
+                    .with_label_values(&[&self.scheme, Operation::Write.into_static()])
+                    .observe(size as f64)
+            })
+            .map_err(|err| {
+                self.stats.increment_errors_total(self.op, err.kind());
+                err
+            })
+    }
+
     async fn abort(&mut self) -> Result<()> {
         self.inner.abort().await.map_err(|err| {
             self.stats.increment_errors_total(self.op, err.kind());

core/src/layers/retry.rs

@@ -648,6 +648,11 @@ impl<R: oio::Write> oio::Write for RetryWrapper<R> {
         }
     }
 
+    /// Sink will move the input stream, so we can't retry it.
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner.sink(size, s).await
+    }
+
     async fn abort(&mut self) -> Result<()> {
         let mut backoff = self.builder.build();
 

core/src/layers/timeout.rs

@@ -349,6 +349,19 @@ impl<R: oio::Write> oio::Write for TimeoutWrapper<R> {
             })?
     }
 
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        let timeout = self.io_timeout(size);
+
+        tokio::time::timeout(timeout, self.inner.sink(size, s))
+            .await
+            .map_err(|_| {
+                Error::new(ErrorKind::Unexpected, "operation timeout")
+                    .with_operation(WriteOperation::Sink)
+                    .with_context("timeout", timeout.as_secs_f64().to_string())
+                    .set_temporary()
+            })?
+    }
+
     async fn abort(&mut self) -> Result<()> {
         tokio::time::timeout(self.timeout, self.inner.abort())
             .await

core/src/layers/tracing.rs

@@ -333,6 +333,14 @@ impl<R: oio::Write> oio::Write for TracingWrapper<R> {
         self.inner.write(bs).await
     }
 
+    #[tracing::instrument(
+        parent = &self.span,
+        level = "trace",
+        skip_all)]
+    async fn sink(&mut self, size: u64, s: oio::Streamer) -> Result<()> {
+        self.inner.sink(size, s).await
+    }
+
     #[tracing::instrument(
         parent = &self.span,
         level = "trace",

core/src/raw/adapters/kv/backend.rs

@@ -325,6 +325,13 @@ impl<S: Adapter> oio::Write for KvWriter<S> {
         Ok(())
     }
 
+    async fn sink(&mut self, _size: u64, _s: oio::Streamer) -> Result<()> {
+        Err(Error::new(
+            ErrorKind::Unsupported,
+            "Write::sink is not supported",
+        ))
+    }
+
     async fn abort(&mut self) -> Result<()> {
         Err(Error::new(
             ErrorKind::Unsupported,

core/src/raw/adapters/typed_kv/backend.rs

@@ -338,6 +338,13 @@ impl<S: Adapter> oio::Write for KvWriter<S> {
         Ok(())
     }
 
+    async fn sink(&mut self, _size: u64, _s: oio::Streamer) -> Result<()> {
+        Err(Error::new(
+            ErrorKind::Unsupported,
+            "Write::sink is not supported",
+        ))
+    }
+
     async fn abort(&mut self) -> Result<()> {
         self.buf.clear();
 

core/src/raw/http_util/body.rs

@@ -25,7 +25,6 @@ use std::task::Poll;
 use bytes::Buf;
 use bytes::BufMut;
 use bytes::Bytes;
-use futures::Stream;
 use futures::StreamExt;
 
 use crate::raw::*;
@@ -41,10 +40,10 @@ pub enum AsyncBody {
     Empty,
     /// Body with bytes.
     Bytes(Bytes),
+    /// Body with stream.
+    Stream(oio::Streamer),
 }
 
-type BytesStream = Box<dyn Stream<Item = Result<Bytes>> + Send + Sync + Unpin>;
-
 /// IncomingAsyncBody carries the content returned by remote servers.
 ///
 /// # Notes
@@ -58,15 +57,15 @@ pub struct IncomingAsyncBody {
     ///
     /// After [TAIT](https://rust-lang.github.io/rfcs/2515-type_alias_impl_trait.html)
     /// has been stable, we can change `IncomingAsyncBody` into `IncomingAsyncBody<S>`.
-    inner: BytesStream,
+    inner: oio::Streamer,
     size: Option<u64>,
     consumed: u64,
     chunk: Option<Bytes>,
 }
 
 impl IncomingAsyncBody {
     /// Construct a new incoming async body
-    pub fn new(s: BytesStream, size: Option<u64>) -> Self {
+    pub fn new(s: oio::Streamer, size: Option<u64>) -> Self {
         Self {
             inner: s,
             size,

core/src/raw/http_util/client.rs

@@ -27,6 +27,7 @@ use http::Response;
 use super::body::IncomingAsyncBody;
 use super::parse_content_length;
 use super::AsyncBody;
+use crate::raw::oio::into_stream;
 use crate::Error;
 use crate::ErrorKind;
 use crate::Result;
@@ -93,6 +94,7 @@ impl HttpClient {
         req_builder = match body {
             AsyncBody::Empty => req_builder.body(reqwest::Body::from("")),
             AsyncBody::Bytes(bs) => req_builder.body(reqwest::Body::from(bs)),
+            AsyncBody::Stream(s) => req_builder.body(reqwest::Body::wrap_stream(s)),
         };
 
         let mut resp = req_builder.send().await.map_err(|err| {
@@ -143,7 +145,10 @@ impl HttpClient {
                 .set_source(err)
         });
 
-        let body = IncomingAsyncBody::new(Box::new(stream), content_length);
+        let body = IncomingAsyncBody::new(
+            Box::new(into_stream::from_futures_stream(stream)),
+            content_length,
+        );
 
         let resp = hr.body(body).expect("response must build succeed");
 

core/src/raw/http_util/error.rs

@@ -98,7 +98,7 @@ pub fn new_request_build_error(err: http::Error) -> Error {
 pub fn new_request_credential_error(err: anyhow::Error) -> Error {
     Error::new(
         ErrorKind::Unexpected,
-        "loading credentail to sign http request",
+        "loading credential to sign http request",
     )
     .set_temporary()
     .with_operation("reqsign::LoadCredential")

core/src/raw/http_util/multipart.rs

@@ -38,6 +38,7 @@ use http::Version;
 use super::new_request_build_error;
 use super::AsyncBody;
 use super::IncomingAsyncBody;
+use crate::raw::oio::into_stream;
 use crate::*;
 
 /// Multipart is a builder for multipart/form-data.
@@ -250,7 +251,7 @@ pub struct MixedPart {
 }
 
 impl MixedPart {
-    /// Create a new mixed part with gien uri.
+    /// Create a new mixed part with given uri.
     pub fn new(uri: &str) -> Self {
         let mut part_headers = HeaderMap::new();
         part_headers.insert(CONTENT_TYPE, "application/http".parse().unwrap());
@@ -283,6 +284,7 @@ impl MixedPart {
         let content = match body {
             AsyncBody::Empty => Bytes::new(),
             AsyncBody::Bytes(bs) => bs,
+            AsyncBody::Stream(_) => panic!("multipart request can't contain stream body"),
         };
 
         Self {
@@ -317,8 +319,10 @@ impl MixedPart {
 
         let bs: Bytes = self.content;
         let length = bs.len();
-        let body =
-            IncomingAsyncBody::new(Box::new(stream::iter(vec![Ok(bs)])), Some(length as u64));
+        let body = IncomingAsyncBody::new(
+            Box::new(into_stream::from_futures_stream(stream::iter(vec![Ok(bs)]))),
+            Some(length as u64),
+        );
 
         builder
             .body(body)