Skip to content

Commit a3bbba8

Browse files
evanlinjinoleonardolima
authored andcommitted
fix(core): Memory leak bugs in CheckPoint::drop impl
Fix memory leak bug in `CheckPoint::drop` by using `Arc::into_inner`. Fix `CPInner::drop` logic so that if `CPInner::block` becomes generic and is of a type that required `drop`, it does not leak memory. Add tests for memory leak + stack overflow when dropping `CheckPoint`.
1 parent 0f5ac52 commit a3bbba8

File tree

1 file changed

+73
-13
lines changed

1 file changed

+73
-13
lines changed

crates/core/src/checkpoint.rs

Lines changed: 73 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -37,21 +37,17 @@ struct CPInner<D> {
3737
/// https://github.com/bitcoindevkit/bdk/issues/1634
3838
impl<D> Drop for CPInner<D> {
3939
fn drop(&mut self) {
40-
// Take out `prev` so its `drop` won't be called when this drop is finished
40+
// Take out `prev` so its `drop` won't be called when this drop is finished.
4141
let mut current = self.prev.take();
42+
// Collect nodes to drop later so we avoid recursive drop calls while not leaking memory.
4243
while let Some(arc_node) = current {
43-
// Get rid of the Arc around `prev` if we're the only one holding a ref
44-
// So the `drop` on it won't be called when the `Arc` is dropped.
45-
//
46-
// FIXME: When MSRV > 1.70.0 this should use Arc::into_inner which actually guarantees
47-
// that no recursive drop calls can happen even with multiple threads.
48-
match Arc::try_unwrap(arc_node).ok() {
49-
Some(mut node) => {
50-
// Keep going backwards
51-
current = node.prev.take();
52-
// Don't call `drop` on `CPInner` since that risks it becoming recursive.
53-
core::mem::forget(node);
54-
}
44+
// Get rid of the `Arc` around `prev` if we're the only one holding a reference so the
45+
// `drop` on it won't be called when the `Arc` is dropped.
46+
let arc_inner = Arc::into_inner(arc_node);
47+
48+
match arc_inner {
49+
// Keep going backwards.
50+
Some(mut node) => current = node.prev.take(),
5551
None => break,
5652
}
5753
}
@@ -318,3 +314,67 @@ impl<D> IntoIterator for CheckPoint<D> {
318314
}
319315
}
320316
}
317+
318+
#[cfg(test)]
319+
mod tests {
320+
use super::*;
321+
322+
/// Make sure that dropping checkpoints does not result in recursion and stack overflow.
323+
#[test]
324+
fn checkpoint_drop_is_not_recursive() {
325+
let run = || {
326+
let mut cp = CheckPoint::new(0, bitcoin::hashes::Hash::hash(b"genesis"));
327+
328+
for height in 1u32..=(1024 * 10) {
329+
let hash: BlockHash = bitcoin::hashes::Hash::hash(height.to_be_bytes().as_slice());
330+
cp = cp.push(height, hash).unwrap();
331+
}
332+
333+
// `cp` would be dropped here.
334+
};
335+
std::thread::Builder::new()
336+
// Restrict stack size.
337+
.stack_size(32 * 1024)
338+
.spawn(run)
339+
.unwrap()
340+
.join()
341+
.unwrap();
342+
}
343+
344+
#[test]
345+
fn checkpoint_does_not_leak() {
346+
let mut cp = CheckPoint::new(0, bitcoin::hashes::Hash::hash(b"genesis"));
347+
348+
for height in 1u32..=1000 {
349+
let hash: BlockHash = bitcoin::hashes::Hash::hash(height.to_be_bytes().as_slice());
350+
cp = cp.push(height, hash).unwrap();
351+
}
352+
353+
let genesis = cp.get(0).expect("genesis exists");
354+
let weak = Arc::downgrade(&genesis.0);
355+
356+
// At this point there should be exactly two strong references to the
357+
// genesis checkpoint: the variable `genesis` and the chain `cp`.
358+
assert_eq!(
359+
Arc::strong_count(&genesis.0),
360+
2,
361+
"`cp` and `genesis` should be the only strong references",
362+
);
363+
364+
// Dropping the chain should remove one strong reference.
365+
drop(cp);
366+
assert_eq!(
367+
Arc::strong_count(&genesis.0),
368+
1,
369+
"`genesis` should be the last strong reference after `cp` is dropped",
370+
);
371+
372+
// Dropping the final reference should deallocate the node, so the weak
373+
// reference cannot be upgraded.
374+
drop(genesis);
375+
assert!(
376+
weak.upgrade().is_none(),
377+
"the checkpoint node should be freed when all strong references are dropped",
378+
);
379+
}
380+
}

0 commit comments

Comments
 (0)