@@ -37,21 +37,17 @@ struct CPInner<D> {
37
37
/// https://github.com/bitcoindevkit/bdk/issues/1634
38
38
impl < D > Drop for CPInner < D > {
39
39
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.
41
41
let mut current = self . prev . take ( ) ;
42
+ // Collect nodes to drop later so we avoid recursive drop calls while not leaking memory.
42
43
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 ( ) ,
55
51
None => break ,
56
52
}
57
53
}
@@ -318,3 +314,67 @@ impl<D> IntoIterator for CheckPoint<D> {
318
314
}
319
315
}
320
316
}
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