Skip to main content

mace/utils/
mod.rs

1use std::{
2    cmp::Ordering,
3    fmt::Debug,
4    ops::{Deref, DerefMut, Range},
5    path::{Path, PathBuf},
6    ptr::null_mut,
7    sync::atomic::{
8        AtomicI64, AtomicU32,
9        Ordering::{AcqRel, Relaxed},
10    },
11};
12
13pub(crate) mod addr_trace;
14pub(crate) mod bitmap;
15pub(crate) mod block;
16pub(crate) mod countblock;
17pub(crate) mod data;
18#[cfg(feature = "failpoints")]
19pub(crate) mod failpoint;
20pub(crate) mod imtree;
21pub(crate) mod interval;
22pub(crate) mod lru;
23pub mod observe;
24pub(crate) mod options;
25pub(crate) mod spooky;
26pub(crate) mod varint;
27
28type Comparator<T> = fn(&T, &T) -> Ordering;
29
30/// Error codes for storage operations.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32#[repr(u8)]
33pub enum OpCode {
34    /// The key was not found.
35    NotFound,
36    /// Data corruption detected.
37    Corruption,
38    /// I/O error occurred.
39    IoError,
40    /// Version mismatch.
41    BadVersion,
42    /// Value or key too large.
43    TooLarge,
44    /// Storage is full.
45    NoSpace,
46    /// Transaction aborted.
47    AbortTx,
48    /// Try again later.
49    Again,
50    /// Invalid operation.
51    Invalid,
52    /// already exists
53    Exist,
54}
55
56impl std::fmt::Display for OpCode {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        write!(f, "{:?}", self)
59    }
60}
61
62impl std::error::Error for OpCode {}
63
64impl From<std::io::Error> for OpCode {
65    fn from(_: std::io::Error) -> Self {
66        OpCode::IoError
67    }
68}
69
70impl From<btree_store::Error> for OpCode {
71    fn from(err: btree_store::Error) -> Self {
72        match err {
73            btree_store::Error::NotFound => OpCode::NotFound,
74            btree_store::Error::Corruption => OpCode::Corruption,
75            btree_store::Error::TooLarge => OpCode::TooLarge,
76            btree_store::Error::Internal => OpCode::Invalid,
77            btree_store::Error::NoSpace => OpCode::NoSpace,
78            btree_store::Error::IoError => OpCode::IoError,
79            btree_store::Error::Invalid => OpCode::Invalid,
80            btree_store::Error::Duplicate => OpCode::Exist,
81            btree_store::Error::Conflict => OpCode::Again,
82        }
83    }
84}
85
86pub(crate) const NULL_PID: u64 = 0;
87/// it's same to DEFAULT_BUCKET_ID
88pub(crate) const ROOT_PID: u64 = 1;
89pub(crate) const NULL_ADDR: u64 = 0;
90/// starts from 1 make sure remapped addr will always greater than NULL_ADDR
91pub(crate) const INIT_ADDR: u64 = 1;
92pub const ADDR_LEN: usize = size_of::<u64>();
93pub(crate) const INIT_ID: u64 = 0;
94pub(crate) const INIT_CMD: u32 = 1;
95pub(crate) const NULL_CMD: u32 = u32::MAX;
96pub(crate) const INIT_WMK: u64 = 0;
97/// NOTE: must larger than wmk_oldest (which is 0 by default)
98pub(crate) const INIT_ORACLE: u64 = 1;
99pub(crate) const NULL_ORACLE: u64 = u64::MAX;
100
101pub(crate) const fn align_up(n: usize, align: usize) -> usize {
102    (n + (align - 1)) & !(align - 1)
103}
104
105#[allow(unused)]
106pub(crate) const fn align_down(n: usize, align: usize) -> usize {
107    n & !(align - 1)
108}
109
110pub(crate) fn raw_ptr_to_ref<'a, T>(x: *mut T) -> &'a T {
111    unsafe { &*x }
112}
113
114pub(crate) fn raw_ptr_to_ref_mut<'a, T>(x: *mut T) -> &'a mut T {
115    unsafe { &mut *x }
116}
117
118#[allow(unused)]
119pub(crate) const fn to_str(x: &[u8]) -> &str {
120    unsafe { std::str::from_utf8_unchecked(x) }
121}
122
123#[macro_export]
124macro_rules! static_assert {
125    ($cond:expr, $msg:expr) => {
126        const _: () = assert!($cond, $msg);
127    };
128
129    ($cond:expr) => {
130        const _: () = assert!($cond);
131    };
132}
133
134#[macro_export]
135macro_rules! number_to_slice {
136    ($num: expr, $slice:expr) => {
137        $slice.copy_from_slice(&$num.to_be_bytes());
138    };
139}
140#[macro_export]
141macro_rules! slice_to_number {
142    ($slice:expr, $num:ty) => {{ <$num>::from_be_bytes($slice.try_into().unwrap()) }};
143}
144
145pub fn rand_range(range: Range<usize>) -> usize {
146    rand::random_range(range)
147}
148
149static_assert!(size_of::<usize>() == size_of::<u64>());
150static_assert!(size_of::<u8>() == size_of::<bool>());
151
152/// A utility for managing temporary paths that are deleted on drop.
153pub struct RandomPath {
154    path: PathBuf,
155    del: bool,
156}
157
158impl Default for RandomPath {
159    fn default() -> Self {
160        Self::new()
161    }
162}
163
164impl RandomPath {
165    const PREFIX: &'static str = "mace_tmp_";
166
167    fn gen_path(root: &PathBuf) -> PathBuf {
168        static TID: AtomicI64 = AtomicI64::new(0);
169        let path = Path::new(&root);
170        loop {
171            let r = rand_range(1000..1000000);
172            let p = path.join(format!(
173                "{}{}{}{}",
174                Self::PREFIX,
175                std::process::id(),
176                TID.fetch_add(1, Relaxed),
177                r
178            ));
179            if !p.exists() {
180                return p;
181            }
182        }
183    }
184
185    /// Creates a new temporary path that will be deleted on drop.
186    pub fn tmp() -> Self {
187        Self {
188            path: Self::gen_path(&std::env::temp_dir()),
189            del: true,
190        }
191    }
192
193    /// Creates a new temporary path that will NOT be deleted on drop by default.
194    pub fn new() -> Self {
195        Self {
196            path: Self::gen_path(&std::env::temp_dir()),
197            del: false,
198        }
199    }
200
201    /// Creates a new temporary path under the given root.
202    pub fn from_root<P: AsRef<Path>>(root: P) -> Self {
203        Self {
204            path: Self::gen_path(&root.as_ref().to_path_buf()),
205            del: false,
206        }
207    }
208
209    /// Explicitly deletes the file or directory at this path.
210    pub fn unlink(&self) {
211        if self.path.exists() {
212            let _ = if self.path.is_file() {
213                std::fs::remove_file(&self.path)
214            } else {
215                std::fs::remove_dir_all(&self.path)
216            };
217        }
218    }
219}
220
221impl Deref for RandomPath {
222    type Target = PathBuf;
223
224    fn deref(&self) -> &Self::Target {
225        &self.path
226    }
227}
228
229impl Drop for RandomPath {
230    fn drop(&mut self) {
231        if self.del {
232            self.unlink();
233        }
234    }
235}
236
237struct MutRefInner<T> {
238    raw: T,
239    refcnt: AtomicU32,
240}
241
242impl<T> MutRefInner<T> {
243    fn new(x: T) -> Self {
244        Self {
245            raw: x,
246            refcnt: AtomicU32::new(1),
247        }
248    }
249}
250
251impl<T> Drop for MutRefInner<T> {
252    fn drop(&mut self) {}
253}
254
255pub struct MutRef<T> {
256    inner: *mut MutRefInner<T>,
257}
258
259unsafe impl<T> Send for MutRef<T> {}
260unsafe impl<T> Sync for MutRef<T> {}
261
262impl<T> MutRef<T> {
263    pub fn new(x: T) -> Self {
264        Self {
265            inner: Box::into_raw(Box::new(MutRefInner::new(x))),
266        }
267    }
268
269    #[allow(clippy::mut_from_ref)]
270    pub fn raw_ref(&self) -> &mut T {
271        unsafe { &mut (*self.inner).raw }
272    }
273
274    pub fn is_null(&self) -> bool {
275        self.inner.is_null()
276    }
277
278    pub fn reset(&mut self, x: T) {
279        let mut new_ref = MutRef::new(x);
280        std::mem::swap(self, &mut new_ref);
281    }
282
283    fn inc(&self) {
284        unsafe { (*self.inner).refcnt.fetch_add(1, Relaxed) };
285    }
286
287    fn dec(&self) -> u32 {
288        unsafe { (*self.inner).refcnt.fetch_sub(1, AcqRel) }
289    }
290}
291
292impl<T> Clone for MutRef<T> {
293    fn clone(&self) -> Self {
294        self.inc();
295        Self { inner: self.inner }
296    }
297}
298
299impl<T> Drop for MutRef<T> {
300    fn drop(&mut self) {
301        unsafe {
302            if !self.is_null() {
303                let old = self.dec();
304                if old == 1 {
305                    drop(Box::from_raw(self.inner));
306                }
307            }
308        }
309    }
310}
311
312impl<T> Default for MutRef<T> {
313    fn default() -> Self {
314        Self { inner: null_mut() }
315    }
316}
317
318impl<T> Deref for MutRef<T> {
319    type Target = T;
320
321    fn deref(&self) -> &Self::Target {
322        unsafe { &(*self.inner).raw }
323    }
324}
325
326impl<T> DerefMut for MutRef<T> {
327    fn deref_mut(&mut self) -> &mut Self::Target {
328        unsafe { &mut (*self.inner).raw }
329    }
330}
331
332#[derive(Debug)]
333pub(crate) struct Handle<T: Sized> {
334    raw: *mut T,
335}
336
337impl<T> Handle<T> {
338    pub(crate) fn new(x: T) -> Self {
339        Self {
340            raw: Box::into_raw(Box::new(x)),
341        }
342    }
343
344    pub(crate) fn inner(&self) -> *mut T {
345        self.raw
346    }
347
348    pub(crate) fn reclaim(&self) {
349        debug_assert!(
350            !self.raw.is_null(),
351            "Double reclaim detected or reclaiming null handle"
352        );
353        if !self.raw.is_null() {
354            unsafe { drop(Box::from_raw(self.raw)) };
355        }
356    }
357}
358
359impl<T> Deref for Handle<T> {
360    type Target = T;
361
362    fn deref(&self) -> &Self::Target {
363        raw_ptr_to_ref(self.raw)
364    }
365}
366
367impl<T> DerefMut for Handle<T> {
368    fn deref_mut(&mut self) -> &mut Self::Target {
369        raw_ptr_to_ref_mut(self.raw)
370    }
371}
372
373impl<T> From<*mut T> for Handle<T> {
374    fn from(value: *mut T) -> Self {
375        Self { raw: value }
376    }
377}
378
379// NOTE: have to manually impl Clone/Copy instead of derive
380impl<T> Copy for Handle<T> {}
381impl<T> Clone for Handle<T> {
382    fn clone(&self) -> Self {
383        *self
384    }
385}
386
387unsafe impl<T> Send for Handle<T> {}
388unsafe impl<T> Sync for Handle<T> {}
389
390#[repr(align(64))]
391pub(crate) struct CachePad<T> {
392    data: T,
393}
394
395impl<T> Deref for CachePad<T> {
396    type Target = T;
397
398    fn deref(&self) -> &Self::Target {
399        &self.data
400    }
401}
402
403impl<T> DerefMut for CachePad<T> {
404    fn deref_mut(&mut self) -> &mut Self::Target {
405        &mut self.data
406    }
407}
408
409impl<T> From<T> for CachePad<T> {
410    fn from(value: T) -> Self {
411        CachePad { data: value }
412    }
413}
414
415impl<T> Default for CachePad<T>
416where
417    T: Default,
418{
419    fn default() -> Self {
420        Self { data: T::default() }
421    }
422}
423
424impl<T> Clone for CachePad<T>
425where
426    T: Clone,
427{
428    fn clone(&self) -> Self {
429        Self {
430            data: self.data.clone(),
431        }
432    }
433}
434
435impl<T> Copy for CachePad<T> where T: Copy {}
436
437impl<T> Debug for CachePad<T>
438where
439    T: Debug,
440{
441    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
442        f.write_fmt(format_args!("{:?}", self.data))
443    }
444}
445
446impl<T> Ord for CachePad<T>
447where
448    T: Ord,
449{
450    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
451        self.data.cmp(&other.data)
452    }
453}
454
455impl<T> PartialOrd for CachePad<T>
456where
457    T: PartialOrd,
458{
459    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
460        self.data.partial_cmp(&other.data)
461    }
462}
463
464impl<T> PartialEq for CachePad<T>
465where
466    T: PartialEq,
467{
468    fn eq(&self, other: &Self) -> bool {
469        self.data.eq(&other.data)
470    }
471}
472
473impl<T> Eq for CachePad<T> where T: Eq {}
474
475#[cfg(test)]
476mod test {
477    use crate::utils::align_up;
478
479    #[test]
480    fn test_free_functions() {
481        assert_eq!(align_up(4, 8), 8);
482        assert_eq!(align_up(16, 8), 16);
483        assert_eq!(align_up(23, 8), 24);
484        assert_eq!(align_up(56, 8), 56);
485
486        static_assert!(true);
487        static_assert!(true, "damn");
488
489        let mut num = 233u64;
490        let mut s = &mut num.to_be_bytes()[0..size_of::<u64>()];
491        let new_num = slice_to_number!(s, u64);
492        assert_eq!(new_num, num);
493
494        s[0] = 1;
495        num = 114514;
496        number_to_slice!(num, &mut s);
497        let new_num = u64::from_be_bytes(s.try_into().unwrap());
498
499        assert_eq!(new_num, num);
500    }
501}