feat: Add into_bytes_sink for Writer

core/src/raw/oio/buf/flex_buf.rs

@@ -103,10 +103,16 @@ impl FlexBuf {
         bs.advance(cnt);
 
         if bs.is_empty() {
-            self.frozen = None;
-            // This reserve cloud be cheap since we can reuse already allocated memory.
-            // (if all references to the frozen buffer are dropped)
-            self.buf.reserve(self.cap);
+            self.clean()
         }
     }
+
+    /// Cleanup the buffer, reset to the initial state.
+    #[inline]
+    pub fn clean(&mut self) {
+        self.frozen = None;
+        // This reserve cloud be cheap since we can reuse already allocated memory.
+        // (if all references to the frozen buffer are dropped)
+        self.buf.reserve(self.cap);
+    }
 }

core/src/types/read/buffer_stream.rs

@@ -118,7 +118,7 @@ pub struct BufferStream(Buffered<Iter<FutureIterator>>);
 impl BufferStream {
     /// Create a new buffer stream from given reader.
     #[inline]
-    pub fn new(r: oio::Reader, options: OpReader, range: impl RangeBounds<u64>) -> Self {
+    pub(crate) fn new(r: oio::Reader, options: OpReader, range: impl RangeBounds<u64>) -> Self {
         let iter = FutureIterator::new(r, options.chunk(), range);
         let stream = stream::iter(iter).buffered(options.concurrent());
 
@@ -150,7 +150,6 @@ mod tests {
 
     use super::*;
 
-    /// Make sure BufferStream implements `Unpin` and `Send`.
     trait AssertTrait: Unpin + MaybeSend + 'static {}
     impl AssertTrait for BufferStream {}
 

core/src/types/read/futures_async_reader.rs

@@ -51,6 +51,9 @@ pub struct FuturesAsyncReader {
     pos: u64,
 }
 
+/// Safety: FuturesAsyncReader only exposes `&mut self` to the outside world,
+unsafe impl Sync for FuturesAsyncReader {}
+
 impl FuturesAsyncReader {
     /// NOTE: don't allow users to create FuturesAsyncReader directly.
     ///
@@ -179,6 +182,9 @@ mod tests {
 
     use super::*;
 
+    trait AssertTrait: Unpin + MaybeSend + Sync + 'static {}
+    impl AssertTrait for FuturesAsyncReader {}
+
     #[tokio::test]
     async fn test_futures_async_read() {
         let r: oio::Reader = Arc::new(Buffer::from(vec![

core/src/types/read/futures_bytes_stream.rs

@@ -31,15 +31,18 @@ use crate::*;
 
 /// FuturesBytesStream is the adapter of [`Stream`] generated by [`Reader::into_bytes_stream`].
 ///
-/// Users can use this adapter in cases where they need to use [`Stream`] trait. FuturesBytesStream reuses the same concurrent and chunk
-/// settings from [`Reader`].
+/// Users can use this adapter in cases where they need to use [`Stream`] trait. FuturesBytesStream
+/// reuses the same concurrent and chunk settings from [`Reader`].
 ///
 /// FuturesStream also implements [`Unpin`], [`Send`] and [`Sync`].
 pub struct FuturesBytesStream {
     stream: BufferStream,
     buf: Buffer,
 }
 
+/// Safety: FuturesBytesStream only exposes `&mut self` to the outside world,
+unsafe impl Sync for FuturesBytesStream {}
+
 impl FuturesBytesStream {
     /// NOTE: don't allow users to create FuturesStream directly.
     #[inline]
@@ -84,6 +87,9 @@ mod tests {
 
     use super::*;
 
+    trait AssertTrait: Unpin + MaybeSend + Sync + 'static {}
+    impl AssertTrait for FuturesBytesStream {}
+
     #[tokio::test]
     async fn test_futures_bytes_stream() {
         let r: oio::Reader = Arc::new(Buffer::from(vec![

core/src/types/read/reader.rs

@@ -285,9 +285,14 @@ mod tests {
     use rand::Rng;
     use rand::RngCore;
 
+    use super::*;
+    use crate::raw::MaybeSend;
     use crate::services;
     use crate::Operator;
 
+    trait AssertTrait: Unpin + MaybeSend + Sync + 'static {}
+    impl AssertTrait for Reader {}
+
     fn gen_random_bytes() -> Vec<u8> {
         let mut rng = ThreadRng::default();
         // Generate size between 1B..16MB.

core/src/types/write/buffer_sink.rs

@@ -0,0 +1,167 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::raw::*;
+use crate::*;
+use bytes::Buf;
+use std::pin::Pin;
+use std::task::{ready, Context, Poll};
+
+/// BufferSink is the adapter of [`Sink`] generated by [`Writer`].
+///
+/// Users can use this adapter in cases where they need to use [`Sink`]
+pub struct BufferSink {
+    state: State,
+    buf: Buffer,
+}
+
+enum State {
+    Idle(Option<oio::Writer>),
+    Writing(BoxedStaticFuture<(oio::Writer, Result<usize>)>),
+    Closing(BoxedStaticFuture<(oio::Writer, Result<()>)>),
+}
+
+/// # Safety
+///
+/// FuturesReader only exposes `&mut self` to the outside world, so it's safe to be `Sync`.
+unsafe impl Sync for State {}
+
+impl BufferSink {
+    /// Create a new sink from a [`oio::Writer`].
+    #[inline]
+    pub(crate) fn new(w: oio::Writer) -> Self {
+        BufferSink {
+            state: State::Idle(Some(w)),
+            buf: Buffer::new(),
+        }
+    }
+}
+
+impl futures::Sink<Buffer> for BufferSink {
+    type Error = Error;
+
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        self.poll_flush(cx)
+    }
+
+    fn start_send(mut self: Pin<&mut Self>, item: Buffer) -> Result<()> {
+        self.buf = item;
+        Ok(())
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        let this = self.get_mut();
+        loop {
+            match &mut this.state {
+                State::Idle(w) => {
+                    if this.buf.is_empty() {
+                        return Poll::Ready(Ok(()));
+                    }
+                    let Some(mut w) = w.take() else {
+                        return Poll::Ready(Err(Error::new(
+                            ErrorKind::Unexpected,
+                            "state invalid: sink has been closed",
+                        )));
+                    };
+                    let buf = this.buf.clone();
+                    let fut = async move {
+                        let res = w.write_dyn(buf).await;
+                        (w, res)
+                    };
+                    this.state = State::Writing(Box::pin(fut));
+                }
+                State::Writing(fut) => {
+                    let (w, res) = ready!(fut.as_mut().poll(cx));
+                    this.state = State::Idle(Some(w));
+                    match res {
+                        Ok(n) => {
+                            this.buf.advance(n);
+                        }
+                        Err(err) => return Poll::Ready(Err(err)),
+                    }
+                }
+                State::Closing(_) => {
+                    return Poll::Ready(Err(Error::new(
+                        ErrorKind::Unexpected,
+                        "state invalid: sink is closing",
+                    )))
+                }
+            }
+        }
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+        let this = self.get_mut();
+        loop {
+            match &mut this.state {
+                State::Idle(w) => {
+                    let Some(mut w) = w.take() else {
+                        return Poll::Ready(Err(Error::new(
+                            ErrorKind::Unexpected,
+                            "state invalid: sink has been closed",
+                        )));
+                    };
+
+                    if this.buf.is_empty() {
+                        let fut = async move {
+                            let res = w.close_dyn().await;
+                            (w, res)
+                        };
+                        this.state = State::Closing(Box::pin(fut));
+                    } else {
+                        let buf = this.buf.clone();
+                        let fut = async move {
+                            let res = w.write_dyn(buf).await;
+                            (w, res)
+                        };
+                        this.state = State::Writing(Box::pin(fut));
+                    }
+                }
+                State::Writing(fut) => {
+                    let (w, res) = ready!(fut.as_mut().poll(cx));
+                    this.state = State::Idle(Some(w));
+                    match res {
+                        Ok(n) => {
+                            this.buf.advance(n);
+                        }
+                        Err(err) => return Poll::Ready(Err(err)),
+                    }
+                }
+                State::Closing(fut) => {
+                    let (w, res) = ready!(fut.as_mut().poll(cx));
+                    this.state = State::Idle(Some(w));
+                    match res {
+                        Ok(_) => {
+                            this.state = State::Idle(None);
+                            return Poll::Ready(Ok(()));
+                        }
+                        Err(err) => return Poll::Ready(Err(err)),
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::raw::MaybeSend;
+    use crate::BufferSink;
+
+    trait AssertTrait: Unpin + MaybeSend + Sync + 'static {}
+    impl AssertTrait for BufferSink {}
+}