Skip to content

Commit f7632fb

Browse files
committed
fix(core): Memory leak bugs in CheckPoint::drop impl
Fix memory leak bug in `CheckPoint::drop` by using `Arc::into_inner` if it is available (>= 1.70). 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 23dfd6e commit f7632fb

File tree

3 files changed

+100
-12
lines changed

3 files changed

+100
-12
lines changed

crates/core/Cargo.toml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,3 +23,9 @@ serde = ["dep:serde", "bitcoin/serde", "hashbrown?/serde"]
2323
[dev-dependencies]
2424
bdk_chain = { path = "../chain" }
2525
bdk_testenv = { path = "../testenv", default-features = false }
26+
27+
[build-dependencies]
28+
version_check = "0.9"
29+
30+
[lints.rust]
31+
unexpected_cfgs = { level = "warn", check-cfg = ['cfg(has_arc_into_inner)'] }

crates/core/build.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
fn main() {
2+
if version_check::is_min_version("1.70.0").unwrap_or(false) {
3+
println!("cargo:rustc-cfg=has_arc_into_inner");
4+
}
5+
}

crates/core/src/checkpoint.rs

Lines changed: 89 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,21 +28,24 @@ struct CPInner {
2828
/// https://github.com/bitcoindevkit/bdk/issues/1634
2929
impl Drop for CPInner {
3030
fn drop(&mut self) {
31-
// Take out `prev` so its `drop` won't be called when this drop is finished
31+
// Take out `prev` so its `drop` won't be called when this drop is finished.
3232
let mut current = self.prev.take();
33+
// Collect nodes to drop later so we avoid recursive drop calls while not leaking memory.
3334
while let Some(arc_node) = current {
34-
// Get rid of the Arc around `prev` if we're the only one holding a ref
35-
// So the `drop` on it won't be called when the `Arc` is dropped.
35+
// Get rid of the `Arc` around `prev` if we're the only one holding a reference so the
36+
// `drop` on it won't be called when the `Arc` is dropped.
3637
//
37-
// FIXME: When MSRV > 1.70.0 this should use Arc::into_inner which actually guarantees
38-
// that no recursive drop calls can happen even with multiple threads.
39-
match Arc::try_unwrap(arc_node).ok() {
40-
Some(mut node) => {
41-
// Keep going backwards
42-
current = node.prev.take();
43-
// Don't call `drop` on `CPInner` since that risks it becoming recursive.
44-
core::mem::forget(node);
45-
}
38+
// Ideally we would only use `Arc::into_inner` which guarantees no memory leaks.
39+
// However, this is only available since 1.70, whereas 1.63 is our current MSRV.
40+
#[cfg(has_arc_into_inner)]
41+
#[allow(clippy::incompatible_msrv)]
42+
let arc_inner = Arc::into_inner(arc_node);
43+
#[cfg(not(has_arc_into_inner))]
44+
let arc_inner = Arc::try_unwrap(arc_node).ok();
45+
46+
match arc_inner {
47+
// Keep going backwards.
48+
Some(mut node) => current = node.prev.take(),
4649
None => break,
4750
}
4851
}
@@ -262,3 +265,77 @@ impl IntoIterator for CheckPoint {
262265
}
263266
}
264267
}
268+
269+
#[cfg(test)]
270+
mod tests {
271+
use super::*;
272+
273+
/// Make sure that dropping checkpoints does not result in recursion and stack overflow.
274+
#[test]
275+
fn checkpoint_drop_is_not_recursive() {
276+
use bitcoin::hashes::Hash;
277+
278+
let run = || {
279+
let mut cp = CheckPoint::new(BlockId {
280+
height: 0,
281+
hash: Hash::hash(b"genesis"),
282+
});
283+
284+
for height in 1u32..=(1024 * 10) {
285+
let hash = Hash::hash(height.to_be_bytes().as_slice());
286+
let block = BlockId { height, hash };
287+
cp = cp.push(block).unwrap();
288+
}
289+
290+
// `cp` would be dropped here.
291+
};
292+
std::thread::Builder::new()
293+
// Restrict stack size.
294+
.stack_size(32 * 1024)
295+
.spawn(run)
296+
.unwrap()
297+
.join()
298+
.unwrap();
299+
}
300+
301+
#[test]
302+
fn checkpoint_does_not_leak() {
303+
let mut cp = CheckPoint::new(BlockId {
304+
height: 0,
305+
hash: bitcoin::hashes::Hash::hash(b"genesis"),
306+
});
307+
308+
for height in 1u32..=1000 {
309+
let hash = bitcoin::hashes::Hash::hash(height.to_be_bytes().as_slice());
310+
let block = BlockId { height, hash };
311+
cp = cp.push(block).unwrap();
312+
}
313+
314+
let genesis = cp.get(0).expect("genesis exists");
315+
let weak = Arc::downgrade(&genesis.0);
316+
317+
// At this point there should be exactly two strong references to the
318+
// genesis checkpoint: the variable `genesis` and the chain `cp`.
319+
assert_eq!(
320+
Arc::strong_count(&genesis.0),
321+
2,
322+
"`cp` and `genesis` should be the only strong references",
323+
);
324+
325+
// Dropping the chain should remove one strong reference.
326+
drop(cp);
327+
assert_eq!(
328+
Arc::strong_count(&genesis.0),
329+
1,
330+
"`genesis` should be the last strong reference after `cp` is dropped",
331+
);
332+
333+
// Dropping the final reference should deallocate the node, so the weak
334+
// reference cannot be upgraded.
335+
drop(genesis);
336+
assert!(
337+
weak.upgrade().is_none(),
338+
"the checkpoint node should be freed when all strong references are dropped",
339+
);
340+
}
341+
}

0 commit comments

Comments
 (0)