77
88use std:: collections:: { HashMap , HashSet } ;
99use std:: sync:: atomic:: { AtomicBool , Ordering } ;
10- use std:: sync:: { Arc , Mutex , RwLock } ;
10+ use std:: sync:: { Arc , Mutex , RwLock , Weak } ;
1111use std:: time:: { Duration , Instant , SystemTime , UNIX_EPOCH } ;
1212
1313use bdk_chain:: bdk_core:: spk_client:: {
@@ -95,7 +95,17 @@ impl ElectrumChainSource {
9595 }
9696
9797 pub ( super ) fn stop ( & self ) {
98- self . electrum_runtime_status . write ( ) . expect ( "lock" ) . stop ( ) ;
98+ let client = self . electrum_runtime_status . write ( ) . expect ( "lock" ) . stop ( ) ;
99+ if let Some ( client) = client {
100+ client. begin_shutdown ( ) ;
101+ }
102+ }
103+
104+ pub ( super ) fn begin_shutdown ( & self ) {
105+ let client = self . electrum_runtime_status . read ( ) . expect ( "lock" ) . client ( ) ;
106+ if let Some ( client) = client {
107+ client. begin_shutdown ( ) ;
108+ }
99109 }
100110
101111 pub ( crate ) async fn sync_onchain_wallet (
@@ -243,7 +253,7 @@ impl ElectrumChainSource {
243253 let sync_cman = Arc :: clone ( & channel_manager) ;
244254 let sync_cmon = Arc :: clone ( & chain_monitor) ;
245255 let sync_sweeper = Arc :: clone ( & output_sweeper) ;
246- let confirmables = vec ! [
256+ let confirmables: Vec < Arc < dyn Confirm + Sync + Send > > = vec ! [
247257 sync_cman as Arc <dyn Confirm + Sync + Send >,
248258 sync_cmon as Arc <dyn Confirm + Sync + Send >,
249259 sync_sweeper as Arc <dyn Confirm + Sync + Send >,
@@ -261,7 +271,8 @@ impl ElectrumChainSource {
261271 return Err ( Error :: TxSyncFailed ) ;
262272 } ;
263273
264- let res = electrum_client. sync_confirmables ( confirmables) . await ;
274+ let confirmable = electrum_client. wrap_confirmables ( & confirmables) ;
275+ let res = electrum_client. sync_confirmables ( vec ! [ confirmable] ) . await ;
265276
266277 if let Ok ( _) = res {
267278 let unix_time_secs_opt =
@@ -436,10 +447,10 @@ impl ElectrumRuntimeStatus {
436447 Ok ( ( ) )
437448 }
438449
439- pub ( super ) fn stop ( & mut self ) {
450+ pub ( super ) fn stop ( & mut self ) -> Option < Arc < ElectrumRuntimeClient > > {
440451 // Drop the client, but retain the registration inventory so we can replay it if we're
441452 // started again.
442- self . client = None ;
453+ self . client . take ( )
443454 }
444455
445456 fn client ( & self ) -> Option < Arc < ElectrumRuntimeClient > > {
@@ -471,6 +482,7 @@ struct ElectrumRuntimeClient {
471482 runtime : Arc < Runtime > ,
472483 config : Arc < Config > ,
473484 logger : Arc < Logger > ,
485+ confirm_gate : Arc < ConfirmGate > ,
474486}
475487
476488impl ElectrumRuntimeClient {
@@ -507,9 +519,23 @@ impl ElectrumRuntimeClient {
507519 runtime,
508520 config,
509521 logger,
522+ confirm_gate : Arc :: new ( ConfirmGate :: new ( ) ) ,
510523 } )
511524 }
512525
526+ fn begin_shutdown ( & self ) {
527+ self . confirm_gate . deactivate ( ) ;
528+ }
529+
530+ fn wrap_confirmables (
531+ & self , confirmables : & [ Arc < dyn Confirm + Sync + Send > ] ,
532+ ) -> Arc < dyn Confirm + Sync + Send > {
533+ Arc :: new ( ShutdownAwareConfirm :: new (
534+ Arc :: downgrade ( & self . confirm_gate ) ,
535+ confirmables. iter ( ) . map ( Arc :: downgrade) . collect ( ) ,
536+ ) )
537+ }
538+
513539 async fn sync_confirmables (
514540 & self , confirmables : Vec < Arc < dyn Confirm + Sync + Send > > ,
515541 ) -> Result < ( ) , Error > {
@@ -539,6 +565,10 @@ impl ElectrumRuntimeClient {
539565 Error :: TxSyncFailed
540566 } ) ?;
541567
568+ if !self . confirm_gate . is_active ( ) {
569+ return Err ( Error :: TxSyncFailed ) ;
570+ }
571+
542572 log_debug ! (
543573 self . logger,
544574 "Sync of Lightning wallet finished in {}ms." ,
@@ -795,6 +825,89 @@ impl ElectrumRuntimeClient {
795825 }
796826}
797827
828+ struct ConfirmGate {
829+ active : Mutex < bool > ,
830+ }
831+
832+ impl ConfirmGate {
833+ fn new ( ) -> Self {
834+ Self { active : Mutex :: new ( true ) }
835+ }
836+
837+ fn deactivate ( & self ) {
838+ * self . active . lock ( ) . expect ( "lock" ) = false ;
839+ }
840+
841+ fn is_active ( & self ) -> bool {
842+ * self . active . lock ( ) . expect ( "lock" )
843+ }
844+ }
845+
846+ struct ShutdownAwareConfirm {
847+ gate : Weak < ConfirmGate > ,
848+ confirmables : Vec < Weak < dyn Confirm + Sync + Send > > ,
849+ }
850+
851+ impl ShutdownAwareConfirm {
852+ fn new ( gate : Weak < ConfirmGate > , confirmables : Vec < Weak < dyn Confirm + Sync + Send > > ) -> Self {
853+ Self { gate, confirmables }
854+ }
855+
856+ fn with_confirmables < T > (
857+ & self , inactive_result : T , f : impl FnOnce ( & [ Arc < dyn Confirm + Sync + Send > ] ) -> T ,
858+ ) -> T {
859+ let Some ( gate) = self . gate . upgrade ( ) else {
860+ return inactive_result;
861+ } ;
862+ let active = gate. active . lock ( ) . expect ( "lock" ) ;
863+ if !* active {
864+ return inactive_result;
865+ }
866+
867+ let Some ( confirmables) : Option < Vec < Arc < dyn Confirm + Sync + Send > > > =
868+ self . confirmables . iter ( ) . map ( Weak :: upgrade) . collect ( )
869+ else {
870+ return inactive_result;
871+ } ;
872+ f ( & confirmables)
873+ }
874+ }
875+
876+ impl Confirm for ShutdownAwareConfirm {
877+ fn transactions_confirmed (
878+ & self , header : & bitcoin:: block:: Header ,
879+ txdata : & lightning:: chain:: transaction:: TransactionData < ' _ > , height : u32 ,
880+ ) {
881+ self . with_confirmables ( ( ) , |confirmables| {
882+ for confirmable in confirmables {
883+ confirmable. transactions_confirmed ( header, txdata, height) ;
884+ }
885+ } )
886+ }
887+
888+ fn transaction_unconfirmed ( & self , txid : & Txid ) {
889+ self . with_confirmables ( ( ) , |confirmables| {
890+ for confirmable in confirmables {
891+ confirmable. transaction_unconfirmed ( txid) ;
892+ }
893+ } )
894+ }
895+
896+ fn best_block_updated ( & self , header : & bitcoin:: block:: Header , height : u32 ) {
897+ self . with_confirmables ( ( ) , |confirmables| {
898+ for confirmable in confirmables {
899+ confirmable. best_block_updated ( header, height) ;
900+ }
901+ } )
902+ }
903+
904+ fn get_relevant_txids ( & self ) -> Vec < ( Txid , u32 , Option < bitcoin:: BlockHash > ) > {
905+ self . with_confirmables ( Vec :: new ( ) , |confirmables| {
906+ confirmables. iter ( ) . flat_map ( |confirmable| confirmable. get_relevant_txids ( ) ) . collect ( )
907+ } )
908+ }
909+ }
910+
798911impl Filter for ElectrumRuntimeClient {
799912 fn register_tx ( & self , txid : & Txid , script_pubkey : & Script ) {
800913 self . tx_sync . register_tx ( txid, script_pubkey)
@@ -803,3 +916,141 @@ impl Filter for ElectrumRuntimeClient {
803916 self . tx_sync . register_output ( output)
804917 }
805918}
919+
920+ #[ cfg( test) ]
921+ mod tests {
922+ use std:: sync:: atomic:: { AtomicUsize , Ordering } ;
923+ use std:: sync:: mpsc;
924+ use std:: thread;
925+
926+ use bitcoin:: blockdata:: constants:: genesis_block;
927+
928+ use super :: * ;
929+
930+ struct RecordingConfirm {
931+ calls : AtomicUsize ,
932+ relevant_txid : Txid ,
933+ }
934+
935+ impl RecordingConfirm {
936+ fn new ( relevant_txid : Txid ) -> Self {
937+ Self { calls : AtomicUsize :: new ( 0 ) , relevant_txid }
938+ }
939+ }
940+
941+ impl Confirm for RecordingConfirm {
942+ fn transactions_confirmed (
943+ & self , _header : & bitcoin:: block:: Header ,
944+ _txdata : & lightning:: chain:: transaction:: TransactionData < ' _ > , _height : u32 ,
945+ ) {
946+ self . calls . fetch_add ( 1 , Ordering :: AcqRel ) ;
947+ }
948+
949+ fn transaction_unconfirmed ( & self , _txid : & Txid ) {
950+ self . calls . fetch_add ( 1 , Ordering :: AcqRel ) ;
951+ }
952+
953+ fn best_block_updated ( & self , _header : & bitcoin:: block:: Header , _height : u32 ) {
954+ self . calls . fetch_add ( 1 , Ordering :: AcqRel ) ;
955+ }
956+
957+ fn get_relevant_txids ( & self ) -> Vec < ( Txid , u32 , Option < bitcoin:: BlockHash > ) > {
958+ vec ! [ ( self . relevant_txid, 0 , None ) ]
959+ }
960+ }
961+
962+ struct BlockingConfirm {
963+ calls : AtomicUsize ,
964+ started : Mutex < Option < mpsc:: SyncSender < ( ) > > > ,
965+ release : Mutex < mpsc:: Receiver < ( ) > > ,
966+ }
967+
968+ impl Confirm for BlockingConfirm {
969+ fn transactions_confirmed (
970+ & self , _header : & bitcoin:: block:: Header ,
971+ _txdata : & lightning:: chain:: transaction:: TransactionData < ' _ > , _height : u32 ,
972+ ) {
973+ }
974+
975+ fn transaction_unconfirmed ( & self , _txid : & Txid ) { }
976+
977+ fn best_block_updated ( & self , _header : & bitcoin:: block:: Header , _height : u32 ) {
978+ self . calls . fetch_add ( 1 , Ordering :: AcqRel ) ;
979+ if let Some ( started) = self . started . lock ( ) . expect ( "lock" ) . take ( ) {
980+ started. send ( ( ) ) . expect ( "test should still be waiting" ) ;
981+ }
982+ self . release . lock ( ) . expect ( "lock" ) . recv ( ) . expect ( "test should release callback" ) ;
983+ }
984+
985+ fn get_relevant_txids ( & self ) -> Vec < ( Txid , u32 , Option < bitcoin:: BlockHash > ) > {
986+ Vec :: new ( )
987+ }
988+ }
989+
990+ #[ test]
991+ fn confirm_callbacks_are_ignored_after_shutdown ( ) {
992+ let block = genesis_block ( Network :: Regtest ) ;
993+ let txid = block. txdata [ 0 ] . compute_txid ( ) ;
994+ let delegate = Arc :: new ( RecordingConfirm :: new ( txid) ) ;
995+ let delegate_dyn: Arc < dyn Confirm + Sync + Send > = delegate. clone ( ) ;
996+ let gate = Arc :: new ( ConfirmGate :: new ( ) ) ;
997+ let confirm =
998+ ShutdownAwareConfirm :: new ( Arc :: downgrade ( & gate) , vec ! [ Arc :: downgrade( & delegate_dyn) ] ) ;
999+
1000+ confirm. best_block_updated ( & block. header , 0 ) ;
1001+ assert_eq ! ( delegate. calls. load( Ordering :: Acquire ) , 1 ) ;
1002+ assert_eq ! ( confirm. get_relevant_txids( ) , vec![ ( txid, 0 , None ) ] ) ;
1003+
1004+ gate. deactivate ( ) ;
1005+ confirm. transactions_confirmed ( & block. header , & [ ] , 0 ) ;
1006+ confirm. transaction_unconfirmed ( & txid) ;
1007+ confirm. best_block_updated ( & block. header , 0 ) ;
1008+ assert_eq ! ( delegate. calls. load( Ordering :: Acquire ) , 1 ) ;
1009+ assert ! ( confirm. get_relevant_txids( ) . is_empty( ) ) ;
1010+ }
1011+
1012+ #[ test]
1013+ fn shutdown_waits_for_the_whole_confirm_callback ( ) {
1014+ let block = genesis_block ( Network :: Regtest ) ;
1015+ let txid = block. txdata [ 0 ] . compute_txid ( ) ;
1016+ let ( started_sender, started_receiver) = mpsc:: sync_channel ( 1 ) ;
1017+ let ( release_sender, release_receiver) = mpsc:: sync_channel ( 1 ) ;
1018+ let blocking = Arc :: new ( BlockingConfirm {
1019+ calls : AtomicUsize :: new ( 0 ) ,
1020+ started : Mutex :: new ( Some ( started_sender) ) ,
1021+ release : Mutex :: new ( release_receiver) ,
1022+ } ) ;
1023+ let trailing = Arc :: new ( RecordingConfirm :: new ( txid) ) ;
1024+ let blocking_dyn: Arc < dyn Confirm + Sync + Send > = blocking. clone ( ) ;
1025+ let trailing_dyn: Arc < dyn Confirm + Sync + Send > = trailing. clone ( ) ;
1026+ let gate = Arc :: new ( ConfirmGate :: new ( ) ) ;
1027+ let confirm = Arc :: new ( ShutdownAwareConfirm :: new (
1028+ Arc :: downgrade ( & gate) ,
1029+ vec ! [ Arc :: downgrade( & blocking_dyn) , Arc :: downgrade( & trailing_dyn) ] ,
1030+ ) ) ;
1031+
1032+ let callback = {
1033+ let confirm = Arc :: clone ( & confirm) ;
1034+ thread:: spawn ( move || confirm. best_block_updated ( & block. header , 0 ) )
1035+ } ;
1036+ started_receiver. recv ( ) . expect ( "callback should start" ) ;
1037+
1038+ let ( shutdown_done_sender, shutdown_done_receiver) = mpsc:: sync_channel ( 1 ) ;
1039+ let shutdown = thread:: spawn ( move || {
1040+ gate. deactivate ( ) ;
1041+ shutdown_done_sender. send ( ( ) ) . expect ( "test should still be waiting" ) ;
1042+ } ) ;
1043+ assert ! ( shutdown_done_receiver. recv_timeout( Duration :: from_millis( 50 ) ) . is_err( ) ) ;
1044+
1045+ release_sender. send ( ( ) ) . expect ( "callback should still be running" ) ;
1046+ callback. join ( ) . expect ( "callback should finish" ) ;
1047+ shutdown_done_receiver. recv ( ) . expect ( "shutdown should finish" ) ;
1048+ shutdown. join ( ) . expect ( "shutdown should not panic" ) ;
1049+
1050+ assert_eq ! ( blocking. calls. load( Ordering :: Acquire ) , 1 ) ;
1051+ assert_eq ! ( trailing. calls. load( Ordering :: Acquire ) , 1 ) ;
1052+ confirm. best_block_updated ( & block. header , 0 ) ;
1053+ assert_eq ! ( blocking. calls. load( Ordering :: Acquire ) , 1 ) ;
1054+ assert_eq ! ( trailing. calls. load( Ordering :: Acquire ) , 1 ) ;
1055+ }
1056+ }
0 commit comments