@@ -488,12 +488,18 @@ where
488488 return Ok ( ( ) ) ;
489489 }
490490
491- let result = self . regenerate_and_broadcast_spend_if_necessary_internal ( ) . await ;
492-
493- // Release the pending sweep flag again, regardless of result.
494- self . pending_sweep . store ( false , Ordering :: Release ) ;
491+ // Use an RAII guard so the flag is released even if this future is dropped mid-await
492+ // (e.g. cancelled by `tokio::time::timeout` or `select!`). A bare `store(false)` after
493+ // the await would never run on cancellation, leaving the sweeper permanently disabled.
494+ struct PendingSweepGuard < ' a > ( & ' a AtomicBool ) ;
495+ impl < ' a > Drop for PendingSweepGuard < ' a > {
496+ fn drop ( & mut self ) {
497+ self . 0 . store ( false , Ordering :: Release ) ;
498+ }
499+ }
500+ let _guard = PendingSweepGuard ( & self . pending_sweep ) ;
495501
496- result
502+ self . regenerate_and_broadcast_spend_if_necessary_internal ( ) . await
497503 }
498504
499505 /// Regenerates and broadcasts the spending transaction for any outputs that are pending
@@ -1176,3 +1182,151 @@ where
11761182 Ok ( ( best_block, OutputSweeperSync { sweeper } ) )
11771183 }
11781184}
1185+
1186+ #[ cfg( all( test, feature = "std" ) ) ]
1187+ mod tests {
1188+ use super :: * ;
1189+ use crate :: chain:: transaction:: OutPoint ;
1190+ use crate :: sign:: { ChangeDestinationSource , OutputSpender } ;
1191+ use crate :: util:: async_poll:: dummy_waker;
1192+ use crate :: util:: logger:: Record ;
1193+ use crate :: util:: native_async:: MaybeSend ;
1194+
1195+ use bitcoin:: hashes:: Hash as _;
1196+ use bitcoin:: secp256k1:: All ;
1197+ use bitcoin:: transaction:: Version ;
1198+ use bitcoin:: { Amount , BlockHash , ScriptBuf , Transaction , TxOut , Txid } ;
1199+
1200+ use core:: future as core_future;
1201+ use core:: pin:: pin;
1202+ use core:: sync:: atomic:: Ordering ;
1203+ use core:: task:: Poll ;
1204+
1205+ struct DummyBroadcaster ;
1206+ impl BroadcasterInterface for DummyBroadcaster {
1207+ fn broadcast_transactions ( & self , _: & [ ( & Transaction , TransactionType ) ] ) { }
1208+ }
1209+
1210+ struct DummyFeeEstimator ;
1211+ impl FeeEstimator for DummyFeeEstimator {
1212+ fn get_est_sat_per_1000_weight ( & self , _: ConfirmationTarget ) -> u32 {
1213+ 1000
1214+ }
1215+ }
1216+
1217+ struct DummyFilter ;
1218+ impl Filter for DummyFilter {
1219+ fn register_tx ( & self , _: & Txid , _: & bitcoin:: Script ) { }
1220+ fn register_output ( & self , _: WatchedOutput ) { }
1221+ }
1222+
1223+ struct DummyLogger ;
1224+ impl Logger for DummyLogger {
1225+ fn log ( & self , _: Record ) { }
1226+ }
1227+
1228+ struct DummyOutputSpender ;
1229+ impl OutputSpender for DummyOutputSpender {
1230+ fn spend_spendable_outputs (
1231+ & self , _: & [ & SpendableOutputDescriptor ] , _: Vec < TxOut > , _: ScriptBuf , _: u32 ,
1232+ _: Option < LockTime > , _: & Secp256k1 < All > ,
1233+ ) -> Result < Transaction , ( ) > {
1234+ Ok ( Transaction {
1235+ version : Version :: TWO ,
1236+ lock_time : LockTime :: ZERO ,
1237+ input : Vec :: new ( ) ,
1238+ output : Vec :: new ( ) ,
1239+ } )
1240+ }
1241+ }
1242+
1243+ struct DummyChangeDestSource ;
1244+ impl ChangeDestinationSource for DummyChangeDestSource {
1245+ fn get_change_destination_script < ' a > (
1246+ & ' a self ,
1247+ ) -> impl Future < Output = Result < ScriptBuf , ( ) > > + MaybeSend + ' a {
1248+ core_future:: ready ( Ok ( ScriptBuf :: new ( ) ) )
1249+ }
1250+ }
1251+
1252+ struct PendingKVStore ;
1253+ impl KVStore for PendingKVStore {
1254+ fn read (
1255+ & self , _: & str , _: & str , _: & str ,
1256+ ) -> impl Future < Output = Result < Vec < u8 > , io:: Error > > + ' static + MaybeSend {
1257+ core_future:: ready ( Err ( io:: Error :: new ( io:: ErrorKind :: NotFound , "" ) ) )
1258+ }
1259+ fn write (
1260+ & self , _: & str , _: & str , _: & str , _: Vec < u8 > ,
1261+ ) -> impl Future < Output = Result < ( ) , io:: Error > > + ' static + MaybeSend {
1262+ core_future:: pending ( )
1263+ }
1264+ fn remove (
1265+ & self , _: & str , _: & str , _: & str , _: bool ,
1266+ ) -> impl Future < Output = Result < ( ) , io:: Error > > + ' static + MaybeSend {
1267+ core_future:: ready ( Ok ( ( ) ) )
1268+ }
1269+ fn list (
1270+ & self , _: & str , _: & str ,
1271+ ) -> impl Future < Output = Result < Vec < String > , io:: Error > > + ' static + MaybeSend {
1272+ core_future:: ready ( Ok ( Vec :: new ( ) ) )
1273+ }
1274+ }
1275+
1276+ #[ test]
1277+ fn pending_sweep_flag_resets_after_future_drop ( ) {
1278+ let best_block = BlockLocator :: new ( BlockHash :: all_zeros ( ) , 1_000 ) ;
1279+
1280+ let sweeper: OutputSweeper <
1281+ DummyBroadcaster ,
1282+ Box < DummyChangeDestSource > ,
1283+ DummyFeeEstimator ,
1284+ DummyFilter ,
1285+ PendingKVStore ,
1286+ DummyLogger ,
1287+ DummyOutputSpender ,
1288+ > = OutputSweeper :: new (
1289+ best_block,
1290+ DummyBroadcaster ,
1291+ DummyFeeEstimator ,
1292+ None ,
1293+ DummyOutputSpender ,
1294+ Box :: new ( DummyChangeDestSource ) ,
1295+ PendingKVStore ,
1296+ DummyLogger ,
1297+ ) ;
1298+
1299+ // Inject a tracked output directly so the sweep loop has work to do.
1300+ let descriptor = SpendableOutputDescriptor :: StaticOutput {
1301+ outpoint : OutPoint { txid : Txid :: all_zeros ( ) , index : 0 } ,
1302+ output : TxOut { value : Amount :: from_sat ( 100_000 ) , script_pubkey : ScriptBuf :: new ( ) } ,
1303+ channel_keys_id : None ,
1304+ } ;
1305+ sweeper. sweeper_state . lock ( ) . unwrap ( ) . outputs . push ( TrackedSpendableOutput {
1306+ descriptor,
1307+ channel_id : None ,
1308+ counterparty_node_id : None ,
1309+ status : OutputSpendStatus :: PendingInitialBroadcast { delayed_until_height : None } ,
1310+ } ) ;
1311+
1312+ // Start a sweep, poll once (the persist step stays Pending because our KVStore's
1313+ // `write` future is `future::pending()`), then drop the future to mimic
1314+ // cancellation - the sort of thing a `tokio::time::timeout` wrapper produces.
1315+ {
1316+ let mut fut = pin ! ( sweeper. regenerate_and_broadcast_spend_if_necessary( ) ) ;
1317+ let waker = dummy_waker ( ) ;
1318+ let mut ctx = task:: Context :: from_waker ( & waker) ;
1319+ assert ! ( matches!( fut. as_mut( ) . poll( & mut ctx) , Poll :: Pending ) ) ;
1320+ }
1321+
1322+ // Once the future has been dropped, `pending_sweep` must be cleared. The bug
1323+ // is that the flag is only ever cleared after the inner future returns, so a
1324+ // dropped future leaves it stuck `true` and every subsequent call to
1325+ // `regenerate_and_broadcast_spend_if_necessary` short-circuits with `Ok(())`,
1326+ // permanently disabling the sweeper.
1327+ assert ! (
1328+ !sweeper. pending_sweep. load( Ordering :: Acquire ) ,
1329+ "pending_sweep flag was not reset when the future was dropped" ,
1330+ ) ;
1331+ }
1332+ }
0 commit comments