io_buffer/
buffer.rs

1use super::utils::{safe_copy, set_zero};
2use libc;
3use nix::errno::Errno;
4use std::slice;
5use std::{
6    fmt,
7    ops::{Deref, DerefMut},
8};
9
10use fail::fail_point;
11
12/// Buffer is a static type,  size and cap (max to i32). Memory footprint is only 16B.
13///
14/// Can obtain from alloc (uninitialized, mutable and owned),
15///
16/// or wrap a raw pointer from c code (not owned,  mutable or immutable),
17///
18/// or convert `From<Vec<u8>>` (mutable and owned), and `To<Vec<u8>>`
19///
20/// When Clone, will copy the contain into a new Buffer.
21#[repr(C, align(1))]
22pub struct Buffer {
23    buf_ptr: *mut libc::c_void,
24    /// the highest bit of `size` represents `owned`
25    pub(crate) size: u32,
26    /// the highest bit of `cap` represents `mutable`
27    pub(crate) cap: u32,
28}
29
30impl fmt::Debug for Buffer {
31    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
32        write!(f, "buffer {:p} size {}", self.get_raw(), self.len())
33    }
34}
35
36unsafe impl Send for Buffer {}
37
38unsafe impl Sync for Buffer {}
39
40pub const MIN_ALIGN: usize = 512;
41pub const MAX_BUFFER_SIZE: usize = 1 << 31;
42
43fn is_aligned(offset: usize, size: usize) -> bool {
44    return (offset & (MIN_ALIGN - 1) == 0) && (size & (MIN_ALIGN - 1) == 0);
45}
46
47impl Buffer {
48    /// Allocate mutable and owned aligned buffer for aio by posix_memalign(),
49    /// with size set to capacity.
50    ///
51    /// **NOTE**: Be aware that buffer allocated is not initialized.
52    ///
53    /// `size`: must be larger than zero
54    #[inline]
55    pub fn aligned(size: i32) -> Result<Buffer, Errno> {
56        let mut _buf = Self::_alloc(MIN_ALIGN, size)?;
57        fail_point!("alloc_buf", |_| {
58            rand_buffer(&mut _buf);
59            return Ok(_buf);
60        });
61        return Ok(_buf);
62    }
63
64    /// Allocate mutable and owned non-aligned Buffer by malloc(),
65    /// with size set to capacity.
66    ///
67    /// **NOTE**: Be aware that buffer allocated is not initialized.
68    ///
69    /// `size`: must be larger than zero
70    #[inline]
71    pub fn alloc(size: i32) -> Result<Buffer, Errno> {
72        let mut _buf = Self::_alloc(0, size)?;
73        fail_point!("alloc_buf", |_| {
74            rand_buffer(&mut _buf);
75            return Ok(_buf);
76        });
77        return Ok(_buf);
78    }
79
80    /// Allocate a buffer.
81    ///
82    /// `size`: must be larger than zero
83    #[inline]
84    fn _alloc(align: usize, size: i32) -> Result<Self, Errno> {
85        assert!(size > 0);
86        let mut ptr: *mut libc::c_void = std::ptr::null_mut();
87        if align > 0 {
88            debug_assert!((align & (MIN_ALIGN - 1)) == 0);
89            debug_assert!((size as usize & (align - 1)) == 0);
90            unsafe {
91                let res =
92                    libc::posix_memalign(&mut ptr, align as libc::size_t, size as libc::size_t);
93                if res != 0 {
94                    return Err(Errno::ENOMEM);
95                }
96            }
97        } else {
98            ptr = unsafe { libc::malloc(size as libc::size_t) };
99            if ptr == std::ptr::null_mut() {
100                return Err(Errno::ENOMEM);
101            }
102        }
103        // owned == true
104        let _size = size as u32 | MAX_BUFFER_SIZE as u32;
105        // mutable == true
106        let _cap = _size;
107        Ok(Self { buf_ptr: ptr, size: _size, cap: _cap })
108    }
109
110    /// Wrap a mutable buffer passed from c code, without owner ship.
111    ///
112    /// **NOTE**: will not free on drop. You have to ensure the buffer valid throughout the lifecycle.
113    ///
114    /// `size`: must be larger than or equal to zero.
115    #[inline]
116    pub fn from_c_ref_mut(ptr: *mut libc::c_void, size: i32) -> Self {
117        assert!(size >= 0);
118        assert!(ptr != std::ptr::null_mut());
119        // owned == false
120        // mutable == true
121        let _cap = size as u32 | MAX_BUFFER_SIZE as u32;
122        Self { buf_ptr: ptr, size: size as u32, cap: _cap }
123    }
124
125    /// Wrap a const buffer passed from c code, without owner ship.
126    ///
127    /// **NOTE**: will not free on drop. You have to ensure the buffer valid throughout the lifecycle
128    ///
129    /// `size`: must be larger than or equal to zero.
130    #[inline]
131    pub fn from_c_ref_const(ptr: *const libc::c_void, size: i32) -> Self {
132        assert!(size >= 0);
133        assert!(ptr != std::ptr::null());
134        // owned == false
135        // mutable == false
136        Self { buf_ptr: unsafe { std::mem::transmute(ptr) }, size: size as u32, cap: size as u32 }
137    }
138
139    /// Tell whether the Buffer has true 'static lifetime.
140    #[inline(always)]
141    pub fn is_owned(&self) -> bool {
142        self.size & (MAX_BUFFER_SIZE as u32) != 0
143    }
144
145    /// Tell whether the Buffer can as_mut().
146    #[inline(always)]
147    pub fn is_mutable(&self) -> bool {
148        self.cap & (MAX_BUFFER_SIZE as u32) != 0
149    }
150
151    /// Return the buffer's size.
152    #[inline(always)]
153    pub fn len(&self) -> usize {
154        let size = self.size & (MAX_BUFFER_SIZE as u32 - 1);
155        size as usize
156    }
157
158    /// Return the memory capacity managed by buffer's ptr
159    #[inline(always)]
160    pub fn capacity(&self) -> usize {
161        let cap = self.cap & (MAX_BUFFER_SIZE as u32 - 1);
162        cap as usize
163    }
164
165    /// Change the buffer's size, the same as `Vec::set_len()`. Panics when len > capacity
166    #[inline(always)]
167    pub fn set_len(&mut self, len: usize) {
168        assert!(len < MAX_BUFFER_SIZE, "size {} >= {} is not supported", len, MAX_BUFFER_SIZE);
169        assert!(len <= self.cap as usize, "size {} must be <= {}", len, self.cap);
170        let owned: u32 = self.size & MAX_BUFFER_SIZE as u32;
171        self.size = owned | len as u32;
172    }
173
174    #[inline(always)]
175    pub fn as_ref(&self) -> &[u8] {
176        unsafe { slice::from_raw_parts(self.buf_ptr as *const u8, self.len()) }
177    }
178
179    /// On debug mode, will panic if the Buffer is not owned [Buffer::from_c_ref_const()]
180    ///
181    /// On release will skip the check for speed.
182    #[inline(always)]
183    pub fn as_mut(&mut self) -> &mut [u8] {
184        #[cfg(debug_assertions)]
185        {
186            if !self.is_mutable() {
187                panic!("Cannot change a mutable buffer")
188            }
189        }
190        unsafe { slice::from_raw_parts_mut(self.buf_ptr as *mut u8, self.len()) }
191    }
192
193    /// Check this buffer usable by aio. True when get from `Buffer::aligned()`.
194    #[inline(always)]
195    pub fn is_aligned(&self) -> bool {
196        is_aligned(self.buf_ptr as usize, self.capacity())
197    }
198
199    /// Get buffer raw pointer
200    #[inline]
201    pub fn get_raw(&self) -> *const u8 {
202        self.buf_ptr as *const u8
203    }
204
205    /// Get buffer raw mut pointer
206    #[inline]
207    pub fn get_raw_mut(&mut self) -> *mut u8 {
208        self.buf_ptr as *mut u8
209    }
210
211    /// Copy from another u8 slice into self[offset..].
212    ///
213    /// **NOTE**: will not do memset.
214    ///
215    /// Argument:
216    ///
217    ///  * offset: Address of this buffer to start filling.
218    #[inline]
219    pub fn copy_from(&mut self, offset: usize, other: &[u8]) {
220        let size = self.len();
221        let dst = self.as_mut();
222        if offset > 0 {
223            assert!(offset < size);
224            safe_copy(&mut dst[offset..], other);
225        } else {
226            safe_copy(dst, other);
227        }
228    }
229
230    /// Copy from another u8 slice into self[offset..], and memset the rest part.
231    ///
232    /// Argument:
233    ///
234    ///  * offset: Address of this buffer to start filling.
235    #[inline]
236    pub fn copy_and_clean(&mut self, offset: usize, other: &[u8]) {
237        let end: usize;
238        let size = self.len();
239        let dst = self.as_mut();
240        assert!(offset < size);
241        if offset > 0 {
242            set_zero(&mut dst[0..offset]);
243            end = offset + safe_copy(&mut dst[offset..], other);
244        } else {
245            end = safe_copy(dst, other);
246        }
247        if size > end {
248            set_zero(&mut dst[end..]);
249        }
250    }
251
252    /// Fill this buffer with zero
253    #[inline]
254    pub fn zero(&mut self) {
255        set_zero(self);
256    }
257
258    /// Fill specified region of buffer[offset..(offset+len)] with zero
259    #[inline]
260    pub fn set_zero(&mut self, offset: usize, len: usize) {
261        let _len = self.len();
262        let mut end = offset + len;
263        if end > _len {
264            end = _len;
265        }
266        let buf = self.as_mut();
267        if offset > 0 || end < _len {
268            set_zero(&mut buf[offset..end]);
269        } else {
270            set_zero(buf);
271        }
272    }
273}
274
275/// Allocates a new memory with the same size and clone the content.
276/// If original buffer is a c reference, will get a owned buffer after clone().
277impl Clone for Buffer {
278    fn clone(&self) -> Self {
279        let mut new_buf = if self.is_aligned() {
280            Self::aligned(self.capacity() as i32).unwrap()
281        } else {
282            Self::alloc(self.capacity() as i32).unwrap()
283        };
284        if self.len() != self.capacity() {
285            new_buf.set_len(self.len());
286        }
287        safe_copy(new_buf.as_mut(), self.as_ref());
288        new_buf
289    }
290}
291
292/// Automatically free on drop when buffer is owned
293impl Drop for Buffer {
294    fn drop(&mut self) {
295        if self.is_owned() {
296            unsafe {
297                libc::free(self.buf_ptr);
298            }
299        }
300    }
301}
302
303/// Convert a owned Buffer to `Vec<u8>`. Panic when buffer is a ref.
304impl Into<Vec<u8>> for Buffer {
305    fn into(mut self) -> Vec<u8> {
306        if !self.is_owned() {
307            panic!("buffer is c ref, not owned");
308        }
309        // Change to not owned, to prevent drop()
310        self.size &= MAX_BUFFER_SIZE as u32 - 1;
311        return unsafe {
312            Vec::<u8>::from_raw_parts(self.buf_ptr as *mut u8, self.len(), self.capacity())
313        };
314    }
315}
316
317/// Convert `Vec<u8>` to Buffer, inherit the size and cap of Vec.
318impl From<Vec<u8>> for Buffer {
319    fn from(buf: Vec<u8>) -> Self {
320        let size = buf.len();
321        let cap = buf.capacity();
322        assert!(size < MAX_BUFFER_SIZE, "size {} >= {} is not supported", size, MAX_BUFFER_SIZE);
323        assert!(cap < MAX_BUFFER_SIZE, "cap {} >= {} is not supported", cap, MAX_BUFFER_SIZE);
324        // owned == true
325        let _size = size as u32 | MAX_BUFFER_SIZE as u32;
326        // mutable == true
327        let _cap = cap as u32 | MAX_BUFFER_SIZE as u32;
328        Buffer { buf_ptr: buf.leak().as_mut_ptr() as *mut libc::c_void, size: _size, cap: _cap }
329    }
330}
331
332impl Deref for Buffer {
333    type Target = [u8];
334
335    #[inline]
336    fn deref(&self) -> &[u8] {
337        self.as_ref()
338    }
339}
340
341impl AsRef<[u8]> for Buffer {
342    #[inline]
343    fn as_ref(&self) -> &[u8] {
344        self.as_ref()
345    }
346}
347
348/// On debug mode, will panic if the Buffer is not owned [Buffer::from_c_ref_const()]
349///
350/// On release will skip the check for speed.
351impl AsMut<[u8]> for Buffer {
352    #[inline]
353    fn as_mut(&mut self) -> &mut [u8] {
354        self.as_mut()
355    }
356}
357
358impl DerefMut for Buffer {
359    #[inline]
360    fn deref_mut(&mut self) -> &mut [u8] {
361        self.as_mut()
362    }
363}