(opendal) branch provided-buf created (now 2f52a85bad)

[xu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xuanwo pushed a change to branch provided-buf
in repository https://gitbox.apache.org/repos/asf/opendal.git


      at 2f52a85bad feat(raw/oio): Add PooledBuf to allow reuse buffer

This branch includes the following new commits:

     new 2f52a85bad feat(raw/oio): Add PooledBuf to allow reuse buffer

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.

(opendal) 01/01: feat(raw/oio): Add PooledBuf to allow reuse buffer

[xu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xuanwo pushed a commit to branch provided-buf
in repository https://gitbox.apache.org/repos/asf/opendal.git

commit 2f52a85bad1681cb84dea65b716290466e31803a
Author: Xuanwo <gi...@xuanwo.io>
AuthorDate: Wed Apr 24 17:02:26 2024 +0800

    feat(raw/oio): Add PooledBuf to allow reuse buffer
    
    Signed-off-by: Xuanwo <gi...@xuanwo.io>
---
 core/src/raw/oio/buf/mod.rs        |   3 +
 core/src/raw/oio/buf/pooled_buf.rs | 114 +++++++++++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+)

diff --git a/core/src/raw/oio/buf/mod.rs b/core/src/raw/oio/buf/mod.rs
index c50a264f27..d60be171e0 100644
--- a/core/src/raw/oio/buf/mod.rs
+++ b/core/src/raw/oio/buf/mod.rs
@@ -20,3 +20,6 @@ pub use flex_buf::FlexBuf;
 
 mod queue_buf;
 pub use queue_buf::QueueBuf;
+
+mod pooled_buf;
+pub use pooled_buf::PooledBuf;
diff --git a/core/src/raw/oio/buf/pooled_buf.rs b/core/src/raw/oio/buf/pooled_buf.rs
new file mode 100644
index 0000000000..750ea82238
--- /dev/null
+++ b/core/src/raw/oio/buf/pooled_buf.rs
@@ -0,0 +1,114 @@
+// 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 std::collections::VecDeque;
+use std::sync::Mutex;
+
+use bytes::BytesMut;
+
+/// PooledBuf is a buffer pool that designed for reusing already allocated bufs.
+///
+/// It works as best-effort that tries to reuse the buffer if possible. It
+/// won't block the thread if the pool is locked, just returning a new buffer
+/// or dropping existing buffer.
+pub struct PooledBuf {
+    pool: Mutex<VecDeque<BytesMut>>,
+    size: usize,
+    initial_capacity: usize,
+}
+
+impl PooledBuf {
+    /// Create a new buffer pool with a given size.
+    pub fn new(size: usize) -> Self {
+        Self {
+            pool: Mutex::new(VecDeque::with_capacity(size)),
+            size,
+            initial_capacity: 0,
+        }
+    }
+
+    /// Set the initial capacity of the buffer.
+    ///
+    /// The default value is 0.
+    pub fn with_initial_capacity(mut self, initial_capacity: usize) -> Self {
+        self.initial_capacity = initial_capacity;
+        self
+    }
+
+    /// Get a [`BytesMut`] from the pool.
+    ///
+    /// It's guaranteed that the buffer is empty.
+    pub fn get(&self) -> BytesMut {
+        // We don't want to block the thread if the pool is locked.
+        //
+        // Just returning a new buffer in this case.
+        let Ok(mut pool) = self.pool.try_lock() else {
+            return BytesMut::with_capacity(self.initial_capacity);
+        };
+
+        if let Some(buf) = pool.pop_front() {
+            buf
+        } else {
+            BytesMut::with_capacity(self.initial_capacity)
+        }
+    }
+
+    /// Put a [`BytesMut`] back to the pool.
+    pub fn put(&self, mut buf: BytesMut) {
+        // We don't want to block the thread if the pool is locked.
+        //
+        // Just dropping the buffer in this case.
+        let Ok(mut pool) = self.pool.try_lock() else {
+            return;
+        };
+
+        if pool.len() < self.size {
+            // Clean the buffer before putting it back to the pool.
+            buf.clear();
+            pool.push_back(buf);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use bytes::BufMut;
+
+    use super::*;
+
+    #[test]
+    fn test_pooled_buf() {
+        let pool = PooledBuf::new(2);
+
+        let mut buf1 = pool.get();
+        buf1.put_slice(b"hello, world!");
+
+        let mut buf2 = pool.get();
+        buf2.reserve(1024);
+
+        pool.put(buf1);
+        pool.put(buf2);
+
+        let buf3 = pool.get();
+        assert_eq!(buf3.len(), 0);
+        assert_eq!(buf3.capacity(), 13);
+
+        let buf4 = pool.get();
+        assert_eq!(buf4.len(), 0);
+        assert_eq!(buf4.capacity(), 1024);
+    }
+}