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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
use crate::allocator::{Allocator, AllocatorBase, FreeOnlyAllocator};
use core::ptr::NonNull;
/// Collects frees and publishes remote frees in batches grouped by owning worker.
///
/// Frees owned by the worker represented by the [`Allocator`] handle that created
/// the batch are reclaimed immediately. All other frees are linked into intrusive
/// chains and published when [`Self::flush`] is called or the batch is dropped.
pub struct RemoteFreeBatch<'a> {
source: RemoteFreeBatchSource<'a>,
chains: Vec<WorkerRemoteFreeChain>,
}
enum RemoteFreeBatchSource<'a> {
Allocator(&'a Allocator),
FreeOnly(&'a FreeOnlyAllocator),
}
struct WorkerRemoteFreeChain {
// INVARIANT: `head` and `tail` are valid, uniquely freed allocation offsets
// whose intrusive links form a private chain owned by the batch.
worker_index: u32,
head: usize,
tail: usize,
}
impl RemoteFreeBatchSource<'_> {
fn base(&self) -> &AllocatorBase {
match self {
Self::Allocator(allocator) => allocator.base(),
Self::FreeOnly(allocator) => allocator.base(),
}
}
/// Find the offset for an allocation pointer.
///
/// # Safety
/// - `ptr` must be a valid pointer into this allocator region.
unsafe fn offset(&self, ptr: NonNull<u8>) -> usize {
match self {
Self::Allocator(allocator) => unsafe { allocator.offset(ptr) },
Self::FreeOnly(allocator) => unsafe { allocator.offset(ptr) },
}
}
}
impl<'a> RemoteFreeBatch<'a> {
fn new(source: RemoteFreeBatchSource<'a>) -> Self {
Self {
source,
chains: Vec::new(),
}
}
/// Free a block of memory from this allocator region.
///
/// Locally owned allocations are reclaimed immediately. Remote frees remain
/// private to this batch until it is flushed or dropped.
///
/// # Safety
/// - `ptr` must point to a valid allocation in this allocator region.
/// - The `ptr` must not have been freed before or added to another free batch.
pub unsafe fn free(&mut self, ptr: NonNull<u8>) {
// SAFETY: The caller guarantees that the pointer refers to a valid allocation
// in this allocator region.
let offset = unsafe { self.source.offset(ptr) };
// SAFETY: The offset was derived from the caller-provided allocation pointer.
unsafe { self.free_offset(offset) };
}
/// Free a block of memory from this allocator region.
///
/// Locally owned allocations are reclaimed immediately. Remote frees remain
/// private to this batch until it is flushed or dropped.
///
/// # Safety
/// - `offset` must identify a valid allocation in this allocator region.
/// - The `offset` must not have been freed before or added to another free batch.
pub unsafe fn free_offset(&mut self, offset: usize) {
// SAFETY: The caller guarantees that `offset` refers to a valid allocation.
let Some((allocation_indexes, worker_index)) = (unsafe {
self.source
.base()
.allocation_indexes_and_assigned_worker(offset)
}) else {
return;
};
if let RemoteFreeBatchSource::Allocator(allocator) = &self.source {
if allocator.worker_index() == worker_index {
// SAFETY: The indexes came from the caller's valid offset and the
// ownership check above confirms its slab is local.
unsafe { allocator.free_local(allocation_indexes) };
return;
}
}
// SAFETY: The caller guarantees that `offset` refers to a valid, uniquely
// freed allocation and `worker_index` was read from its slab metadata.
unsafe { self.push_remote(worker_index, offset) };
}
/// Publish all remote frees currently held by this batch.
///
/// The batch may be reused after flushing. Capacity allocated for destination
/// workers is retained.
pub fn flush(&mut self) {
while let Some(chain) = self.chains.pop() {
// SAFETY: Chains can only be created from offsets accepted by the unsafe
// `free` methods and remain private until they are removed here.
unsafe {
self.source.base().publish_remote_free_chain(
chain.worker_index,
chain.head,
chain.tail,
)
};
}
}
/// Add an allocation to its worker's private remote-free chain.
///
/// # Safety
/// - `offset` must refer to a valid, uniquely freed allocation.
/// - `worker_index` must be the worker currently assigned to its slab.
unsafe fn push_remote(&mut self, worker_index: u32, offset: usize) {
if let Some(chain) = self
.chains
.iter_mut()
.find(|chain| chain.worker_index == worker_index)
{
// SAFETY: Guaranteed by the caller; `chain.head` is another valid
// offset in the same private chain.
unsafe { self.source.base().set_remote_free_next(offset, chain.head) };
chain.head = offset;
return;
}
// Register the destination before modifying the allocation. If allocation
// fails, dropping the batch can still publish every previously linked chain.
self.chains.push(WorkerRemoteFreeChain {
worker_index,
head: offset,
tail: offset,
});
}
}
impl Drop for RemoteFreeBatch<'_> {
fn drop(&mut self) {
self.flush();
}
}
impl Allocator {
/// Create a batch that groups remote frees by owning worker.
///
/// Locally owned allocations are reclaimed immediately. Remote frees are
/// published by [`RemoteFreeBatch::flush`] or when the batch is dropped.
pub fn remote_free_batch(&self) -> RemoteFreeBatch<'_> {
RemoteFreeBatch::new(RemoteFreeBatchSource::Allocator(self))
}
}
impl FreeOnlyAllocator {
/// Create a batch that groups remote frees by owning worker.
///
/// Remote frees are published by [`RemoteFreeBatch::flush`] or when the
/// batch is dropped.
pub fn remote_free_batch(&self) -> RemoteFreeBatch<'_> {
RemoteFreeBatch::new(RemoteFreeBatchSource::FreeOnly(self))
}
}