Skip to content

Commit dfdd7ae

Browse files
hawkwhawkw
committed
fix: remove not adding slots to free lists
This fixes an issue where, when a slot is freed immediately by a call to `remove` while uncontended, it was not added to the page's free list, and thus never reused. This resulted in a memory leak. I've also refactored some internals. Signed-off-by: Eliza Weisman <eliza@buoyant.io>
1 parent 6537cc8 commit dfdd7ae

File tree

4 files changed

+86
-47
lines changed

4 files changed

+86
-47
lines changed

‎src/lib.rs‎

Lines changed: 29 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ macro_rules! thread_local {
178178
macro_rules! test_println {
179179
($($arg:tt)*) => {
180180
if cfg!(test) || cfg!(slab_print) {
181-
eprintln!("{:?} {}", crate::Tid::<crate::DefaultConfig>::current(), format_args!($($arg)*))
181+
println!("{:?} {}", crate::Tid::<crate::DefaultConfig>::current(), format_args!($($arg)*))
182182
}
183183
}
184184
}
@@ -352,10 +352,12 @@ impl<T, C: cfg::Config> Slab<T, C> {
352352
let tid = C::unpack_tid(idx);
353353

354354
test_println!("rm_deferred {:?}", tid);
355-
self.shards
356-
.get(tid.as_usize())
357-
.map(|shard| shard.remove(idx))
358-
.unwrap_or(false)
355+
let shard = self.shards.get(tid.as_usize());
356+
if tid.is_current() {
357+
shard.map(|shard| shard.remove_local(idx)).unwrap_or(false)
358+
} else {
359+
shard.map(|shard| shard.remove_remote(idx)).unwrap_or(false)
360+
}
359361
}
360362

361363
/// Removes the value associated with the given key from the slab, returning
@@ -411,9 +413,9 @@ impl<T, C: cfg::Config> Slab<T, C> {
411413
test_println!("rm {:?}", tid);
412414
let shard = &self.shards[tid.as_usize()];
413415
if tid.is_current() {
414-
shard.remove_local(idx)
416+
shard.take_local(idx)
415417
} else {
416-
shard.remove_remote(idx)
418+
shard.take_remote(idx)
417419
}
418420
}
419421

@@ -545,41 +547,52 @@ impl<T, C: cfg::Config> Shard<T, C> {
545547
})
546548
}
547549

548-
fn remove(&self, idx: usize) -> bool {
550+
fn remove_local(&self, idx: usize) -> bool {
551+
debug_assert_eq!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
552+
let (addr, page_index) = Self::page_indices(idx);
553+
554+
if page_index > self.shared.len() {
555+
return false;
556+
}
557+
558+
self.shared[page_index].remove(addr, C::unpack_gen(idx), self.local(page_index))
559+
}
560+
561+
fn remove_remote(&self, idx: usize) -> bool {
549562
debug_assert_eq!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
550563
let (addr, page_index) = Self::page_indices(idx);
551564

552565
if page_index > self.shared.len() {
553566
return false;
554567
}
555568

556-
self.shared[page_index].remove(addr, C::unpack_gen(idx))
569+
let shared = &self.shared[page_index];
570+
shared.remove(addr, C::unpack_gen(idx), shared.free_list())
557571
}
558572

559573
/// Remove an item on the shard's local thread.
560-
fn remove_local(&self, idx: usize) -> Option<T> {
574+
fn take_local(&self, idx: usize) -> Option<T> {
561575
debug_assert_eq!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
562576
let (addr, page_index) = Self::page_indices(idx);
563577

564578
test_println!("-> remove_local {:?}", addr);
565579

566580
self.shared
567581
.get(page_index)?
568-
.remove_local(self.local(page_index), addr, C::unpack_gen(idx))
582+
.take(addr, C::unpack_gen(idx), self.local(page_index))
569583
}
570584

571585
/// Remove an item, while on a different thread from the shard's local thread.
572-
fn remove_remote(&self, idx: usize) -> Option<T> {
586+
fn take_remote(&self, idx: usize) -> Option<T> {
573587
debug_assert_eq!(Tid::<C>::from_packed(idx).as_usize(), self.tid);
574588
debug_assert!(Tid::<C>::current().as_usize() != self.tid);
575589

576590
let (addr, page_index) = Self::page_indices(idx);
577591

578-
test_println!("-> remove_remote {:?}; page {:?}", addr, page_index);
592+
test_println!("-> take_remote {:?}; page {:?}", addr, page_index);
579593

580-
self.shared
581-
.get(page_index)?
582-
.remove_remote(addr, C::unpack_gen(idx))
594+
let shared = self.shared.get(page_index)?;
595+
shared.take(addr, C::unpack_gen(idx), shared.free_list())
583596
}
584597

585598
#[inline(always)]

‎src/page/mod.rs‎

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@ impl<C: cfg::Config> Addr<C> {
3838
}
3939
}
4040

41+
pub(crate) trait FreeList<C> {
42+
fn push<T>(&self, new_head: usize, slot: &Slot<T, C>);
43+
}
44+
4145
impl<C: cfg::Config> Pack<C> for Addr<C> {
4246
const LEN: usize = C::MAX_PAGES + C::ADDR_INDEX_SHIFT;
4347

@@ -92,6 +96,13 @@ impl Local {
9296
}
9397
}
9498

99+
impl<C: cfg::Config> FreeList<C> for Local {
100+
fn push<T>(&self, new_head: usize, slot: &Slot<T, C>) {
101+
slot.set_next(self.head());
102+
self.set_head(new_head);
103+
}
104+
}
105+
95106
impl<T, C: cfg::Config> Shared<T, C> {
96107
const NULL: usize = Addr::<C>::NULL;
97108

@@ -191,59 +202,55 @@ impl<T, C: cfg::Config> Shared<T, C> {
191202
})
192203
}
193204

194-
pub(crate) fn remove(&self, addr: Addr<C>, gen: slot::Generation<C>) -> bool {
205+
pub(crate) fn remove<F: FreeList<C>>(
206+
&self,
207+
addr: Addr<C>,
208+
gen: slot::Generation<C>,
209+
free_list: &F,
210+
) -> bool {
195211
let offset = addr.offset() - self.prev_sz;
196212

197213
test_println!("-> offset {:?}", offset);
198214

199215
self.slab.with(|slab| {
200216
let slab = unsafe { &*slab }.as_ref();
201217
if let Some(slot) = slab.and_then(|slab| slab.get(offset)) {
202-
slot.remove(gen)
218+
slot.remove(gen, offset, free_list)
203219
} else {
204220
false
205221
}
206222
})
207223
}
208224

209-
pub(crate) fn remove_local(
225+
pub(crate) fn take<F>(
210226
&self,
211-
local: &Local,
212227
addr: Addr<C>,
213228
gen: slot::Generation<C>,
214-
) -> Option<T> {
215-
let offset = addr.offset() - self.prev_sz;
216-
217-
test_println!("-> offset {:?}", offset);
218-
219-
self.slab.with(|slab| {
220-
let slab = unsafe { &*slab }.as_ref()?;
221-
let slot = slab.get(offset)?;
222-
let val = slot.remove_value(gen)?;
223-
slot.set_next(local.head());
224-
local.set_head(offset);
225-
Some(val)
226-
})
227-
}
228-
229-
pub(crate) fn remove_remote(&self, addr: Addr<C>, gen: slot::Generation<C>) -> Option<T> {
229+
free_list: &F,
230+
) -> Option<T>
231+
where
232+
F: FreeList<C>,
233+
{
230234
let offset = addr.offset() - self.prev_sz;
231235

232-
test_println!("-> offset {:?}", offset);
236+
test_println!("-> take: offset {:?}", offset);
233237

234238
self.slab.with(|slab| {
235239
let slab = unsafe { &*slab }.as_ref()?;
236240
let slot = slab.get(offset)?;
237-
let val = slot.remove_value(gen)?;
238-
self.remote.push(offset, |next| slot.set_next(next));
239-
Some(val)
241+
slot.remove_value(gen, offset, free_list)
240242
})
241243
}
242244

243245
pub(crate) fn iter(&self) -> Option<Iter<'_, T, C>> {
244246
let slab = self.slab.with(|slab| unsafe { (&*slab).as_ref() });
245247
slab.map(|slab| slab.iter().filter_map(Slot::value as fn(_) -> _))
246248
}
249+
250+
#[inline(always)]
251+
pub(crate) fn free_list(&self) -> &impl FreeList<C> {
252+
&self.remote
253+
}
247254
}
248255

249256
impl fmt::Debug for Local {

‎src/page/slot.rs‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use super::FreeList;
12
use crate::sync::{
23
atomic::{self, AtomicUsize, Ordering},
34
CausalCell,
@@ -26,7 +27,6 @@ pub(crate) struct Generation<C = cfg::DefaultConfig> {
2627
value: usize,
2728
_cfg: PhantomData<fn(C)>,
2829
}
29-
3030
struct LifecycleGen<C>(Generation<C>);
3131

3232
#[repr(transparent)]
@@ -212,7 +212,12 @@ impl<T, C: cfg::Config> Slot<T, C> {
212212
}
213213

214214
#[inline]
215-
pub(super) fn remove(&self, gen: Generation<C>) -> bool {
215+
pub(super) fn remove<F: FreeList<C>>(
216+
&self,
217+
gen: Generation<C>,
218+
offset: usize,
219+
free: &F,
220+
) -> bool {
216221
let mut lifecycle = self.lifecycle.load(Ordering::Acquire);
217222
let mut curr_gen;
218223

@@ -265,11 +270,17 @@ impl<T, C: cfg::Config> Slot<T, C> {
265270

266271
// Otherwise, we can remove the slot now!
267272
test_println!("-> remove deferred; can remove now");
268-
self.remove_value(curr_gen).is_some()
273+
let removed = self.remove_value(curr_gen, offset, free).is_some();
274+
removed
269275
}
270276

271277
#[inline]
272-
pub(super) fn remove_value(&self, gen: Generation<C>) -> Option<T> {
278+
pub(super) fn remove_value<F: FreeList<C>>(
279+
&self,
280+
gen: Generation<C>,
281+
offset: usize,
282+
free: &F,
283+
) -> Option<T> {
273284
let mut lifecycle = self.lifecycle.load(Ordering::Acquire);
274285
let mut advanced = false;
275286
let next_gen = gen.advance();
@@ -305,7 +316,9 @@ impl<T, C: cfg::Config> Slot<T, C> {
305316
test_println!("-> advanced gen; lifecycle={:#x}; refs={:?};", actual, refs);
306317
if refs.value == 0 {
307318
test_println!("-> ok to remove!");
308-
return self.item.with_mut(|item| unsafe { (*item).take() });
319+
let item = self.item.with_mut(|item| unsafe { (*item).take() });
320+
free.push(offset, self);
321+
return item;
309322
}
310323

311324
// Otherwise, a reference must be dropped before we can

‎src/page/stack.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,15 @@ impl<C: cfg::Config> TransferStack<C> {
2525
}
2626
}
2727

28-
pub(super) fn push(&self, value: usize, before: impl Fn(usize)) {
28+
fn push(&self, new_head: usize, before: impl Fn(usize)) {
2929
let mut next = self.head.load(Ordering::Relaxed);
3030
loop {
3131
test_println!("-> next {:#x}", next);
3232
before(next);
3333

3434
match self
3535
.head
36-
.compare_exchange(next, value, Ordering::Release, Ordering::Relaxed)
36+
.compare_exchange(next, new_head, Ordering::Release, Ordering::Relaxed)
3737
{
3838
// lost the race!
3939
Err(actual) => {
@@ -49,6 +49,12 @@ impl<C: cfg::Config> TransferStack<C> {
4949
}
5050
}
5151

52+
impl<C: cfg::Config> super::FreeList<C> for TransferStack<C> {
53+
fn push<T>(&self, new_head: usize, slot: &super::Slot<T, C>) {
54+
self.push(new_head, |next| slot.set_next(next))
55+
}
56+
}
57+
5258
impl<C> fmt::Debug for TransferStack<C> {
5359
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
5460
f.debug_struct("TransferStack")

0 commit comments

Comments
 (0)