@@ -28,21 +28,24 @@ struct CPInner {
28
28
/// https://github.com/bitcoindevkit/bdk/issues/1634
29
29
impl Drop for CPInner {
30
30
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.
32
32
let mut current = self . prev . take ( ) ;
33
+ // Collect nodes to drop later so we avoid recursive drop calls while not leaking memory.
33
34
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.
36
37
//
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 ( ) ,
46
49
None => break ,
47
50
}
48
51
}
@@ -262,3 +265,77 @@ impl IntoIterator for CheckPoint {
262
265
}
263
266
}
264
267
}
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