opendal/raw/oio/buf/
flex_buf.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use bytes::Buf;
19use bytes::BufMut;
20use bytes::Bytes;
21use bytes::BytesMut;
22
23/// FlexBuf is a buffer that support frozen bytes and reuse existing allocated memory.
24///
25/// It's useful when we want to freeze the buffer and reuse the memory for the next buffer.
26pub struct FlexBuf {
27    /// Already allocated memory size of `buf`.
28    cap: usize,
29    /// Already written bytes length inside `buf`.
30    len: usize,
31    buf: BytesMut,
32    frozen: Option<Bytes>,
33}
34
35impl FlexBuf {
36    /// Initializes a new `FlexBuf` with the given capacity.
37    pub fn new(cap: usize) -> Self {
38        FlexBuf {
39            cap,
40            len: 0,
41
42            buf: BytesMut::with_capacity(cap),
43            frozen: None,
44        }
45    }
46
47    /// Put slice into flex buf.
48    ///
49    /// Return 0 means the buffer is frozen.
50    pub fn put(&mut self, bs: &[u8]) -> usize {
51        if self.frozen.is_some() {
52            return 0;
53        }
54
55        let n = (self.cap - self.len).min(bs.len());
56        self.buf.put_slice(&bs[..n]);
57        self.len += n;
58
59        if self.len >= self.cap {
60            let frozen = self.buf.split();
61            self.len = 0;
62            self.frozen = Some(frozen.freeze());
63        }
64
65        n
66    }
67
68    /// Freeze the buffer no matter it's full or not.
69    ///
70    /// It's a no-op if the buffer has already been frozen.
71    pub fn freeze(&mut self) {
72        if self.len == 0 {
73            return;
74        }
75        let frozen = self.buf.split();
76        self.len = 0;
77        self.frozen = Some(frozen.freeze());
78    }
79
80    /// Get the frozen buffer.
81    ///
82    /// Return `None` if the buffer is not frozen.
83    ///
84    /// # Notes
85    ///
86    /// This operation did nothing to the buffer. We use `&mut self` just for make
87    /// the API consistent with other APIs.
88    pub fn get(&mut self) -> Option<Bytes> {
89        self.frozen.clone()
90    }
91
92    // Advance the frozen buffer.
93    ///
94    /// # Panics
95    ///
96    /// Panic if the buffer is not frozen.
97    pub fn advance(&mut self, cnt: usize) {
98        debug_assert!(self.len == 0, "The buffer must be empty during advance");
99
100        let Some(bs) = self.frozen.as_mut() else {
101            unreachable!("It must be a bug to advance on not frozen buffer")
102        };
103        bs.advance(cnt);
104
105        if bs.is_empty() {
106            self.clean()
107        }
108    }
109
110    /// Cleanup the buffer, reset to the initial state.
111    #[inline]
112    pub fn clean(&mut self) {
113        self.frozen = None;
114        // This reserve cloud be cheap since we can reuse already allocated memory.
115        // (if all references to the frozen buffer are dropped)
116        self.buf.reserve(self.cap);
117    }
118}
opendal/raw/oio/buf/flex_buf.rs

opendal/raw/oio/buf/
flex_buf.rs
