@@ -37,21 +37,17 @@ struct CPInner<D> {
3737/// https://github.com/bitcoindevkit/bdk/issues/1634
3838impl < 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