picodata_plugin/
util.rs

1use crate::error_code::ErrorCode;
2use abi_stable::StableAbi;
3use std::ptr::NonNull;
4use tarantool::error::BoxError;
5use tarantool::error::TarantoolErrorCode;
6use tarantool::ffi::tarantool as ffi;
7
8////////////////////////////////////////////////////////////////////////////////
9// FfiSafeBytes
10////////////////////////////////////////////////////////////////////////////////
11
12/// A helper struct for passing byte slices over the ABI boundary.
13#[repr(C)]
14#[derive(StableAbi, Clone, Copy, Debug)]
15pub struct FfiSafeBytes {
16    pointer: NonNull<u8>,
17    len: usize,
18}
19
20impl FfiSafeBytes {
21    #[inline(always)]
22    pub fn len(self) -> usize {
23        self.len
24    }
25
26    #[inline(always)]
27    pub fn is_empty(self) -> bool {
28        self.len == 0
29    }
30
31    /// # Safety
32    ///
33    /// `pointer` and `len` must be correct pointer and length
34    #[inline(always)]
35    pub unsafe fn from_raw_parts(pointer: NonNull<u8>, len: usize) -> Self {
36        Self { pointer, len }
37    }
38
39    #[inline(always)]
40    pub fn into_raw_parts(self) -> (*mut u8, usize) {
41        (self.pointer.as_ptr(), self.len)
42    }
43
44    /// Converts `self` back to a borrowed string `&[u8]`.
45    ///
46    /// # Safety
47    /// `FfiSafeBytes` can only be constructed from a valid rust byte slice,
48    /// so you only need to make sure that the origial `&[u8]` outlives the lifetime `'a`.
49    ///
50    /// This should generally be true when borrowing strings owned by the current
51    /// function and calling a function via FFI, but borrowing global data or
52    /// data stored within a `Rc` for example is probably unsafe.
53    pub unsafe fn as_bytes<'a>(self) -> &'a [u8] {
54        std::slice::from_raw_parts(self.pointer.as_ptr(), self.len)
55    }
56}
57
58impl Default for FfiSafeBytes {
59    #[inline(always)]
60    fn default() -> Self {
61        Self {
62            pointer: NonNull::dangling(),
63            len: 0,
64        }
65    }
66}
67
68impl<'a> From<&'a [u8]> for FfiSafeBytes {
69    #[inline(always)]
70    fn from(value: &'a [u8]) -> Self {
71        Self {
72            pointer: as_non_null_ptr(value),
73            len: value.len(),
74        }
75    }
76}
77
78impl<'a> From<&'a str> for FfiSafeBytes {
79    #[inline(always)]
80    fn from(value: &'a str) -> Self {
81        Self {
82            pointer: as_non_null_ptr(value.as_bytes()),
83            len: value.len(),
84        }
85    }
86}
87
88////////////////////////////////////////////////////////////////////////////////
89// FfiSafeStr
90////////////////////////////////////////////////////////////////////////////////
91
92/// A helper struct for passing rust strings over the ABI boundary.
93///
94/// This type can only be constructed from a valid rust string, so it's not
95/// necessary to validate the utf8 encoding when converting back to `&str`.
96#[repr(C)]
97#[derive(StableAbi, Clone, Copy, Debug)]
98pub struct FfiSafeStr {
99    pointer: NonNull<u8>,
100    len: usize,
101}
102
103impl FfiSafeStr {
104    #[inline(always)]
105    pub fn len(self) -> usize {
106        self.len
107    }
108
109    #[inline(always)]
110    pub fn is_empty(self) -> bool {
111        self.len == 0
112    }
113
114    /// # Safety
115    ///
116    /// `pointer` and `len` must be correct pointer and length
117    #[inline(always)]
118    pub unsafe fn from_raw_parts(pointer: NonNull<u8>, len: usize) -> Self {
119        Self { pointer, len }
120    }
121
122    /// # Safety
123    /// `bytes` must represent a valid utf8 string.
124    pub unsafe fn from_utf8_unchecked(bytes: &[u8]) -> Self {
125        let pointer = as_non_null_ptr(bytes);
126        let len = bytes.len();
127        Self { pointer, len }
128    }
129
130    #[inline(always)]
131    pub fn into_raw_parts(self) -> (*mut u8, usize) {
132        (self.pointer.as_ptr(), self.len)
133    }
134
135    /// Converts `self` back to a borrowed string `&str`.
136    ///
137    /// # Safety
138    /// `FfiSafeStr` can only be constructed from a valid rust `str`,
139    /// so you only need to make sure that the origial `str` outlives the lifetime `'a`.
140    ///
141    /// This should generally be true when borrowing strings owned by the current
142    /// function and calling a function via FFI, but borrowing global data or
143    /// data stored within a `Rc` for example is probably unsafe.
144    #[inline]
145    pub unsafe fn as_str<'a>(self) -> &'a str {
146        if cfg!(debug_assertions) {
147            std::str::from_utf8(self.as_bytes()).expect("should only be used with valid utf8")
148        } else {
149            std::str::from_utf8_unchecked(self.as_bytes())
150        }
151    }
152
153    /// Converts `self` back to a borrowed string `&[u8]`.
154    ///
155    /// # Safety
156    /// `FfiSafeStr` can only be constructed from a valid rust byte slice,
157    /// so you only need to make sure that the original `&[u8]` outlives the lifetime `'a`.
158    ///
159    /// This should generally be true when borrowing strings owned by the current
160    /// function and calling a function via FFI, but borrowing global data or
161    /// data stored within a `Rc` for example is probably unsafe.
162    #[inline(always)]
163    pub unsafe fn as_bytes<'a>(self) -> &'a [u8] {
164        std::slice::from_raw_parts(self.pointer.as_ptr(), self.len)
165    }
166}
167
168impl Default for FfiSafeStr {
169    #[inline(always)]
170    fn default() -> Self {
171        Self {
172            pointer: NonNull::dangling(),
173            len: 0,
174        }
175    }
176}
177
178impl<'a> From<&'a str> for FfiSafeStr {
179    #[inline(always)]
180    fn from(value: &'a str) -> Self {
181        Self {
182            pointer: as_non_null_ptr(value.as_bytes()),
183            len: value.len(),
184        }
185    }
186}
187
188////////////////////////////////////////////////////////////////////////////////
189// RegionGuard
190////////////////////////////////////////////////////////////////////////////////
191
192// TODO: move to tarantool-module https://git.picodata.io/picodata/picodata/tarantool-module/-/issues/210
193pub struct RegionGuard {
194    save_point: usize,
195}
196
197impl RegionGuard {
198    /// TODO
199    #[inline(always)]
200    #[allow(clippy::new_without_default)]
201    pub fn new() -> Self {
202        // This is safe as long as the function is called within an initialized
203        // fiber runtime
204        let save_point = unsafe { ffi::box_region_used() };
205        Self { save_point }
206    }
207
208    /// TODO
209    #[inline(always)]
210    pub fn used_at_creation(&self) -> usize {
211        self.save_point
212    }
213}
214
215impl Drop for RegionGuard {
216    fn drop(&mut self) {
217        // This is safe as long as the function is called within an initialized
218        // fiber runtime
219        unsafe { ffi::box_region_truncate(self.save_point) }
220    }
221}
222
223////////////////////////////////////////////////////////////////////////////////
224// region allocation
225////////////////////////////////////////////////////////////////////////////////
226
227// TODO: move to tarantool module https://git.picodata.io/picodata/picodata/tarantool-module/-/issues/210
228/// TODO: doc
229#[inline]
230fn allocate_on_region(size: usize) -> Result<&'static mut [u8], BoxError> {
231    // SAFETY: requires initialized fiber runtime
232    let pointer = unsafe { ffi::box_region_alloc(size).cast::<u8>() };
233    if pointer.is_null() {
234        return Err(BoxError::last());
235    }
236    // SAFETY: safe because pointer is not null
237    let region_slice = unsafe { std::slice::from_raw_parts_mut(pointer, size) };
238    Ok(region_slice)
239}
240
241// TODO: move to tarantool module https://git.picodata.io/picodata/picodata/tarantool-module/-/issues/210
242/// Copies the provided `data` to the current fiber's region allocator returning
243/// a reference to the new allocation.
244///
245/// Use this to return dynamically sized values over the ABI boundary, for
246/// example in RPC handlers.
247///
248/// Note that the returned slice's lifetime is not really `'static`, but is
249/// determined by the following call to `box_region_truncate`.
250#[inline]
251pub fn copy_to_region(data: &[u8]) -> Result<&'static [u8], BoxError> {
252    let region_slice = allocate_on_region(data.len())?;
253    region_slice.copy_from_slice(data);
254    Ok(region_slice)
255}
256
257////////////////////////////////////////////////////////////////////////////////
258// RegionBuffer
259////////////////////////////////////////////////////////////////////////////////
260
261// TODO: move to tarantool module https://git.picodata.io/picodata/picodata/tarantool-module/-/issues/210
262/// TODO
263pub struct RegionBuffer {
264    guard: RegionGuard,
265
266    start: *mut u8,
267    count: usize,
268}
269
270impl RegionBuffer {
271    #[inline(always)]
272    #[allow(clippy::new_without_default)]
273    pub fn new() -> Self {
274        Self {
275            guard: RegionGuard::new(),
276            start: NonNull::dangling().as_ptr(),
277            count: 0,
278        }
279    }
280
281    #[track_caller]
282    pub fn push(&mut self, data: &[u8]) -> Result<(), BoxError> {
283        let added_count = data.len();
284        let new_count = self.count + added_count;
285        unsafe {
286            let save_point = ffi::box_region_used();
287            let pointer: *mut u8 = ffi::box_region_alloc(added_count) as _;
288
289            if pointer.is_null() {
290                #[rustfmt::skip]
291                return Err(BoxError::new(TarantoolErrorCode::MemoryIssue, format!("failed to allocate {added_count} bytes on the region allocator")));
292            }
293
294            if self.start.is_null() || pointer == self.start.add(self.count) {
295                // New allocation is contiguous with the previous one
296                memcpy(pointer, data.as_ptr(), added_count);
297                self.count = new_count;
298                if self.start.is_null() {
299                    self.start = pointer;
300                }
301            } else {
302                // New allocation is in a different slab, need to reallocate
303                ffi::box_region_truncate(save_point);
304
305                let new_count = self.count + added_count;
306                let pointer: *mut u8 = ffi::box_region_alloc(new_count) as _;
307                memcpy(pointer, self.start, self.count);
308                memcpy(pointer.add(self.count), data.as_ptr(), added_count);
309                self.start = pointer;
310                self.count = new_count;
311            }
312        }
313
314        Ok(())
315    }
316
317    #[inline(always)]
318    pub fn get(&self) -> &[u8] {
319        if self.start.is_null() {
320            // Cannot construct a slice from a null pointer even if len is 0
321            &[]
322        } else {
323            unsafe { std::slice::from_raw_parts(self.start, self.count) }
324        }
325    }
326
327    #[inline]
328    pub fn into_raw_parts(self) -> (&'static [u8], usize) {
329        let save_point = self.guard.used_at_creation();
330        std::mem::forget(self.guard);
331        if self.start.is_null() {
332            // Cannot construct a slice from a null pointer even if len is 0
333            return (&[], save_point);
334        }
335        let slice = unsafe { std::slice::from_raw_parts(self.start, self.count) };
336        (slice, save_point)
337    }
338}
339
340impl std::io::Write for RegionBuffer {
341    #[inline(always)]
342    fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
343        if let Err(e) = self.push(data) {
344            #[rustfmt::skip]
345            return Err(std::io::Error::new(std::io::ErrorKind::OutOfMemory, e.message()));
346        }
347
348        Ok(data.len())
349    }
350
351    #[inline(always)]
352    fn flush(&mut self) -> std::io::Result<()> {
353        Ok(())
354    }
355}
356
357#[inline(always)]
358unsafe fn memcpy(destination: *mut u8, source: *const u8, count: usize) {
359    let to = std::slice::from_raw_parts_mut(destination, count);
360    let from = std::slice::from_raw_parts(source, count);
361    to.copy_from_slice(from)
362}
363
364////////////////////////////////////////////////////////////////////////////////
365// DisplayErrorLocation
366////////////////////////////////////////////////////////////////////////////////
367
368// TODO: move to taratool-module https://git.picodata.io/picodata/picodata/tarantool-module/-/issues/211
369pub struct DisplayErrorLocation<'a>(pub &'a BoxError);
370
371impl std::fmt::Display for DisplayErrorLocation<'_> {
372    #[inline]
373    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
374        if let Some((file, line)) = self.0.file().zip(self.0.line()) {
375            write!(f, "{file}:{line}: ")?;
376        }
377        Ok(())
378    }
379}
380
381////////////////////////////////////////////////////////////////////////////////
382// DisplayAsHexBytesLimitted
383////////////////////////////////////////////////////////////////////////////////
384
385// TODO: move to taratool-module https://git.picodata.io/picodata/picodata/tarantool-module/-/merge_requests/523
386pub struct DisplayAsHexBytesLimitted<'a>(pub &'a [u8]);
387
388impl std::fmt::Display for DisplayAsHexBytesLimitted<'_> {
389    #[inline]
390    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
391        if self.0.len() > 512 {
392            f.write_str("<too-big-to-display>")
393        } else {
394            tarantool::util::DisplayAsHexBytes(self.0).fmt(f)
395        }
396    }
397}
398
399////////////////////////////////////////////////////////////////////////////////
400// miscellaneous
401////////////////////////////////////////////////////////////////////////////////
402
403#[inline(always)]
404fn as_non_null_ptr<T>(data: &[T]) -> NonNull<T> {
405    let pointer = data.as_ptr();
406    // SAFETY: slice::as_ptr never returns `null`
407    // Also I have to cast to `* mut` here even though we're not going to
408    // mutate it, because there's no constructor that takes `* const`....
409    unsafe { NonNull::new_unchecked(pointer as *mut _) }
410}
411
412// TODO: this should be in tarantool module
413pub fn tarantool_error_to_box_error(e: tarantool::error::Error) -> BoxError {
414    match e {
415        tarantool::error::Error::Tarantool(e) => e,
416        other => BoxError::new(ErrorCode::Other, other.to_string()),
417    }
418}
419
420////////////////////////////////////////////////////////////////////////////////
421// test
422////////////////////////////////////////////////////////////////////////////////
423
424#[cfg(feature = "internal_test")]
425mod test {
426    use super::*;
427
428    #[tarantool::test]
429    fn region_buffer() {
430        #[derive(serde::Serialize, Debug)]
431        struct S {
432            name: String,
433            x: f32,
434            y: f32,
435            array: Vec<(i32, i32, bool)>,
436        }
437
438        let s = S {
439            name: "foo".into(),
440            x: 4.2,
441            y: 6.9,
442            array: vec![(1, 2, true), (3, 4, false)],
443        };
444
445        let vec = rmp_serde::to_vec(&s).unwrap();
446        let mut buffer = RegionBuffer::new();
447        rmp_serde::encode::write(&mut buffer, &s).unwrap();
448        assert_eq!(vec, buffer.get());
449    }
450}
picodata_plugin/util.rs

picodata_plugin/
util.rs
