1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
use std::cell::Cell;
use fastrand::Rng;
use crate::loom::{
sync::{
Mutex,
atomic::{AtomicU32, Ordering},
},
thread_local,
};
pub(crate) struct PageControl {
// Used to synchronize page allocations.
lock: Mutex<()>,
// Used to distribute `Idr::insert()` across existing pages.
// It improves performance by reducing contention.
allocated: AtomicU32,
}
impl Default for PageControl {
fn default() -> Self {
Self {
allocated: AtomicU32::new(0),
lock: Mutex::new(()),
}
}
}
impl PageControl {
pub(crate) fn get_or_lock<R>(
&self,
get: impl Fn() -> *const R,
alloc: impl FnOnce(),
) -> *const R {
let ptr = get();
// The fast path, the page is already allocated.
if !ptr.is_null() {
return ptr;
}
let _guard = self.lock.lock().expect("lock poisoned");
// Re-check if the page is allocated while acquiring the lock.
let ptr = get();
if !ptr.is_null() {
return ptr;
}
// Actually allocate the page.
alloc();
let ptr = get();
debug_assert!(!ptr.is_null());
// Use `Relaxed` ordering here because no need to synchronize with `choose()`,
// it's only for performance optimization and doesn't affect correctness.
self.allocated.fetch_add(1, Ordering::Relaxed);
ptr
}
pub(crate) fn choose<'a, P, R>(
&self,
pages: &'a [P],
f: impl Fn(&'a P) -> Option<R>,
) -> Option<R> {
// Use `Relaxed` ordering here because no need to synchronize with
// `get_or_lock()`, it's only for performance optimization and doesn't
// affect correctness either older or newer values are read.
let allocated = self.allocated.load(Ordering::Relaxed);
debug_assert!(allocated as usize <= pages.len());
// Randomly choose a page to start from.
// It helps to distribute the load more evenly and reduce contention.
if allocated > 0 {
let start_idx = gen_u32(allocated);
for page in &pages[start_idx as usize..allocated as usize] {
if let Some(ret) = f(page) {
return Some(ret);
}
}
}
// If we haven't found a page yet, try all pages.
// Either we will find a page or create a new one.
for page in pages {
if let Some(ret) = f(page) {
return Some(ret);
}
}
None
}
pub(crate) fn allocated(&self) -> u32 {
self.allocated.load(Ordering::Relaxed)
}
}
thread_local! {
static RNG: Cell<Rng> = Cell::new(Rng::with_seed(0xef6_f79e_d30b_a75a));
}
fn gen_u32(upper: u32) -> u32 {
RNG.with(|cell| {
let mut rng = cell.replace(Rng::with_seed(0));
let ret = rng.u32(0..upper);
cell.set(rng);
ret
})
}