@@ -28,21 +28,24 @@ struct CPInner {
2828/// https://github.com/bitcoindevkit/bdk/issues/1634
2929impl 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