lockfree_object_pool/
page.rs

Help
1use std::{
2    cell::UnsafeCell,
3    sync::atomic::{AtomicU32, Ordering},
4};
5
6pub struct Page<T> {
7    data: [UnsafeCell<T>; 32],
8    free: AtomicU32,
9}
10
11pub type PageId = u8;
12
13impl<T> Page<T> {
14    #[inline]
15    pub fn new<I>(init: I) -> Self
16    where
17        I: Fn() -> T,
18    {
19        Self {
20            data: [
21                UnsafeCell::new(init()),
22                UnsafeCell::new(init()),
23                UnsafeCell::new(init()),
24                UnsafeCell::new(init()),
25                UnsafeCell::new(init()),
26                UnsafeCell::new(init()),
27                UnsafeCell::new(init()),
28                UnsafeCell::new(init()),
29                UnsafeCell::new(init()),
30                UnsafeCell::new(init()),
31                UnsafeCell::new(init()),
32                UnsafeCell::new(init()),
33                UnsafeCell::new(init()),
34                UnsafeCell::new(init()),
35                UnsafeCell::new(init()),
36                UnsafeCell::new(init()),
37                UnsafeCell::new(init()),
38                UnsafeCell::new(init()),
39                UnsafeCell::new(init()),
40                UnsafeCell::new(init()),
41                UnsafeCell::new(init()),
42                UnsafeCell::new(init()),
43                UnsafeCell::new(init()),
44                UnsafeCell::new(init()),
45                UnsafeCell::new(init()),
46                UnsafeCell::new(init()),
47                UnsafeCell::new(init()),
48                UnsafeCell::new(init()),
49                UnsafeCell::new(init()),
50                UnsafeCell::new(init()),
51                UnsafeCell::new(init()),
52                UnsafeCell::new(init()),
53            ],
54            free: AtomicU32::new(u32::MAX),
55        }
56    }
57
58    #[cfg(test)]
59    pub(crate) fn is_full(&self) -> bool {
60        self.free.load(Ordering::Relaxed) == 0
61    }
62
63    #[cfg(test)]
64    pub(crate) fn get_mask(&self) -> u32 {
65        self.free.load(Ordering::Relaxed)
66    }
67
68    #[inline]
69    pub fn alloc(&self) -> Option<PageId> {
70        self.free
71            .fetch_update(Ordering::SeqCst, Ordering::Relaxed, |free| {
72                if free == 0 {
73                    None
74                } else {
75                    Some(free & (free - 1))
76                }
77            })
78            .ok()
79            .map(|free| free.trailing_zeros() as u8)
80    }
81
82    #[inline]
83    pub fn free(&self, id: &PageId) {
84        let mask: u32 = 1 << id;
85        self.free.fetch_or(mask, Ordering::SeqCst);
86    }
87
88    #[inline]
89    pub unsafe fn get(&self, id: &PageId) -> &T {
90        &*self.data[*id as usize].get()
91    }
92
93    #[inline]
94    #[allow(clippy::mut_from_ref)] // the function is marked as unsafe for a reason
95    pub unsafe fn get_mut(&self, id: &PageId) -> &mut T {
96        &mut *self.data[*id as usize].get()
97    }
98}
99
100unsafe impl<T: Send> Send for Page<T> {} // normal rules apply
101unsafe impl<T: Sync> Sync for Page<T> {} // normal rules apply
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    #[test]
107    fn test_page_01() {
108        let page = Page::<u32>::new(|| 0);
109        assert!(!page.is_full());
110        assert_eq!(page.get_mask(), u32::MAX);
111    }
112
113    #[test]
114    fn test_page_02() {
115        let page = Page::<u32>::new(|| 0);
116
117        let item1 = page.alloc();
118        assert!(item1.is_some());
119        assert_eq!(item1.unwrap(), 0);
120        assert_eq!(
121            format!("{:b}", page.get_mask()),
122            "11111111111111111111111111111110"
123        );
124
125        let item2 = page.alloc();
126        assert!(item2.is_some());
127        assert_eq!(item2.unwrap(), 1);
128        assert_eq!(
129            format!("{:b}", page.get_mask()),
130            "11111111111111111111111111111100"
131        );
132
133        let item3 = page.alloc();
134        assert!(item3.is_some());
135        assert_eq!(item3.unwrap(), 2);
136        assert_eq!(
137            format!("{:b}", page.get_mask()),
138            "11111111111111111111111111111000"
139        );
140
141        page.free(&item2.unwrap());
142        assert_eq!(
143            format!("{:b}", page.get_mask()),
144            "11111111111111111111111111111010"
145        );
146
147        page.free(&item1.unwrap());
148        assert_eq!(
149            format!("{:b}", page.get_mask()),
150            "11111111111111111111111111111011"
151        );
152
153        page.free(&item3.unwrap());
154        assert_eq!(
155            format!("{:b}", page.get_mask()),
156            "11111111111111111111111111111111"
157        );
158    }
159
160    #[test]
161    fn test_page_03() {
162        let page = Page::<u32>::new(|| 0);
163        for i in 0..32 {
164            assert!(!page.is_full());
165
166            let item = page.alloc();
167            assert!(item.is_some());
168            assert_eq!(item.unwrap(), i);
169        }
170        assert!(page.is_full());
171        let item = page.alloc();
172        assert!(item.is_none());
173    }
174}
lockfree_object_pool/page.rs

lockfree_object_pool/
page.rs
