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