@@ -36,6 +36,7 @@ use crate::util::ser::Writeable;
36
36
37
37
use core:: fmt;
38
38
use core:: ops:: Deref ;
39
+ use core:: sync:: atomic:: { AtomicBool , Ordering } ;
39
40
use crate :: io;
40
41
use crate :: sync:: Mutex ;
41
42
use crate :: prelude:: * ;
@@ -263,6 +264,7 @@ pub struct OnionMessenger<
263
264
intercept_messages_for_offline_peers : bool ,
264
265
pending_intercepted_msgs_events : Mutex < Vec < Event > > ,
265
266
pending_peer_connected_events : Mutex < Vec < Event > > ,
267
+ pending_events_processor : AtomicBool ,
266
268
}
267
269
268
270
/// [`OnionMessage`]s buffered to be sent.
@@ -1017,6 +1019,28 @@ where
1017
1019
}
1018
1020
}
1019
1021
1022
+ macro_rules! drop_handled_events_and_abort { ( $self: expr, $res: expr, $offset: expr, $event_queue: expr) => {
1023
+ // We want ot make sure to cleanly abort upon event handling failure. To this end, we drop all
1024
+ // successfully handled events from the given queue, reset the events processing flag, and
1025
+ // return, to have the events eventually replayed upon next invocation.
1026
+ {
1027
+ let mut queue_lock = $event_queue. lock( ) . unwrap( ) ;
1028
+
1029
+ // We skip `$offset` result entries to reach the ones relevant for the given `$event_queue`.
1030
+ let mut res_iter = $res. iter( ) . skip( $offset) ;
1031
+
1032
+ // Keep all events which previously error'd *or* any that have been added since we dropped
1033
+ // the Mutex before.
1034
+ queue_lock. retain( |_| res_iter. next( ) . map_or( true , |r| r. is_err( ) ) ) ;
1035
+
1036
+ if $res. iter( ) . any( |r| r. is_err( ) ) {
1037
+ // We failed handling some events. Return to have them eventually replayed.
1038
+ $self. pending_events_processor. store( false , Ordering :: Release ) ;
1039
+ return ;
1040
+ }
1041
+ }
1042
+ } }
1043
+
1020
1044
impl < ES : Deref , NS : Deref , L : Deref , NL : Deref , MR : Deref , OMH : Deref , APH : Deref , CMH : Deref >
1021
1045
OnionMessenger < ES , NS , L , NL , MR , OMH , APH , CMH >
1022
1046
where
@@ -1093,6 +1117,7 @@ where
1093
1117
intercept_messages_for_offline_peers,
1094
1118
pending_intercepted_msgs_events : Mutex :: new ( Vec :: new ( ) ) ,
1095
1119
pending_peer_connected_events : Mutex :: new ( Vec :: new ( ) ) ,
1120
+ pending_events_processor : AtomicBool :: new ( false ) ,
1096
1121
}
1097
1122
}
1098
1123
@@ -1331,42 +1356,57 @@ where
1331
1356
pub async fn process_pending_events_async < Future : core:: future:: Future < Output = Result < ( ) , ReplayEvent > > + core:: marker:: Unpin , H : Fn ( Event ) -> Future > (
1332
1357
& self , handler : H
1333
1358
) {
1334
- let mut intercepted_msgs = Vec :: new ( ) ;
1335
- let mut peer_connecteds = Vec :: new ( ) ;
1336
- {
1337
- let mut pending_intercepted_msgs_events =
1338
- self . pending_intercepted_msgs_events . lock ( ) . unwrap ( ) ;
1339
- let mut pending_peer_connected_events =
1340
- self . pending_peer_connected_events . lock ( ) . unwrap ( ) ;
1341
- core:: mem:: swap ( & mut * pending_intercepted_msgs_events, & mut intercepted_msgs) ;
1342
- core:: mem:: swap ( & mut * pending_peer_connected_events, & mut peer_connecteds) ;
1359
+ if self . pending_events_processor . compare_exchange ( false , true , Ordering :: Acquire , Ordering :: Relaxed ) . is_err ( ) {
1360
+ return ;
1343
1361
}
1344
1362
1345
- let mut futures = Vec :: with_capacity ( intercepted_msgs. len ( ) ) ;
1346
- for ( node_id, recipient) in self . message_recipients . lock ( ) . unwrap ( ) . iter_mut ( ) {
1347
- if let OnionMessageRecipient :: PendingConnection ( _, addresses, _) = recipient {
1348
- if let Some ( addresses) = addresses. take ( ) {
1349
- futures. push ( handler ( Event :: ConnectionNeeded { node_id : * node_id, addresses } ) ) ;
1363
+ {
1364
+ let intercepted_msgs = self . pending_intercepted_msgs_events . lock ( ) . unwrap ( ) . clone ( ) ;
1365
+ let mut futures = Vec :: with_capacity ( intercepted_msgs. len ( ) ) ;
1366
+ for ( node_id, recipient) in self . message_recipients . lock ( ) . unwrap ( ) . iter_mut ( ) {
1367
+ if let OnionMessageRecipient :: PendingConnection ( _, addresses, _) = recipient {
1368
+ if let Some ( addresses) = addresses. take ( ) {
1369
+ futures. push ( handler ( Event :: ConnectionNeeded { node_id : * node_id, addresses } ) ) ;
1370
+ }
1350
1371
}
1351
1372
}
1352
- }
1353
1373
1354
- for ev in intercepted_msgs {
1355
- if let Event :: OnionMessageIntercepted { .. } = ev { } else { debug_assert ! ( false ) ; }
1356
- futures. push ( handler ( ev) ) ;
1374
+ // The offset in the `futures` vec at which `intercepted_msgs` start. We don't bother
1375
+ // replaying `ConnectionNeeded` events.
1376
+ let intercepted_msgs_offset = futures. len ( ) ;
1377
+
1378
+ for ev in intercepted_msgs {
1379
+ if let Event :: OnionMessageIntercepted { .. } = ev { } else { debug_assert ! ( false ) ; }
1380
+ futures. push ( handler ( ev) ) ;
1381
+ }
1382
+ // Let the `OnionMessageIntercepted` events finish before moving on to peer_connecteds
1383
+ let res = crate :: util:: async_poll:: MultiResultFuturePoller :: new ( futures) . await ;
1384
+ drop_handled_events_and_abort ! ( self , res, intercepted_msgs_offset, self . pending_intercepted_msgs_events) ;
1357
1385
}
1358
- // Let the `OnionMessageIntercepted` events finish before moving on to peer_connecteds
1359
- crate :: util:: async_poll:: MultiResultFuturePoller :: new ( futures) . await ;
1360
1386
1361
- if peer_connecteds. len ( ) <= 1 {
1362
- for event in peer_connecteds { handler ( event) . await ; }
1363
- } else {
1364
- let mut futures = Vec :: new ( ) ;
1365
- for event in peer_connecteds {
1366
- futures. push ( handler ( event) ) ;
1387
+ {
1388
+ let peer_connecteds = self . pending_peer_connected_events . lock ( ) . unwrap ( ) . clone ( ) ;
1389
+ let num_peer_connecteds = peer_connecteds. len ( ) ;
1390
+ if num_peer_connecteds <= 1 {
1391
+ for event in peer_connecteds {
1392
+ if handler ( event) . await . is_ok ( ) {
1393
+ self . pending_peer_connected_events . lock ( ) . unwrap ( ) . drain ( ..num_peer_connecteds) ;
1394
+ } else {
1395
+ // We failed handling the event. Return to have it eventually replayed.
1396
+ self . pending_events_processor . store ( false , Ordering :: Release ) ;
1397
+ return ;
1398
+ }
1399
+ }
1400
+ } else {
1401
+ let mut futures = Vec :: new ( ) ;
1402
+ for event in peer_connecteds {
1403
+ futures. push ( handler ( event) ) ;
1404
+ }
1405
+ let res = crate :: util:: async_poll:: MultiResultFuturePoller :: new ( futures) . await ;
1406
+ drop_handled_events_and_abort ! ( self , res, 0 , self . pending_peer_connected_events) ;
1367
1407
}
1368
- crate :: util:: async_poll:: MultiResultFuturePoller :: new ( futures) . await ;
1369
1408
}
1409
+ self . pending_events_processor . store ( false , Ordering :: Release ) ;
1370
1410
}
1371
1411
}
1372
1412
@@ -1406,17 +1446,24 @@ where
1406
1446
CMH :: Target : CustomOnionMessageHandler ,
1407
1447
{
1408
1448
fn process_pending_events < H : Deref > ( & self , handler : H ) where H :: Target : EventHandler {
1449
+ if self . pending_events_processor . compare_exchange ( false , true , Ordering :: Acquire , Ordering :: Relaxed ) . is_err ( ) {
1450
+ return ;
1451
+ }
1452
+
1409
1453
for ( node_id, recipient) in self . message_recipients . lock ( ) . unwrap ( ) . iter_mut ( ) {
1410
1454
if let OnionMessageRecipient :: PendingConnection ( _, addresses, _) = recipient {
1411
1455
if let Some ( addresses) = addresses. take ( ) {
1412
1456
let _ = handler. handle_event ( Event :: ConnectionNeeded { node_id : * node_id, addresses } ) ;
1413
1457
}
1414
1458
}
1415
1459
}
1416
- let mut events = Vec :: new ( ) ;
1460
+ let intercepted_msgs;
1461
+ let peer_connecteds;
1417
1462
{
1418
- let mut pending_intercepted_msgs_events = self . pending_intercepted_msgs_events . lock ( ) . unwrap ( ) ;
1463
+ let pending_intercepted_msgs_events = self . pending_intercepted_msgs_events . lock ( ) . unwrap ( ) ;
1464
+ intercepted_msgs = pending_intercepted_msgs_events. clone ( ) ;
1419
1465
let mut pending_peer_connected_events = self . pending_peer_connected_events . lock ( ) . unwrap ( ) ;
1466
+ peer_connecteds = pending_peer_connected_events. clone ( ) ;
1420
1467
#[ cfg( debug_assertions) ] {
1421
1468
for ev in pending_intercepted_msgs_events. iter ( ) {
1422
1469
if let Event :: OnionMessageIntercepted { .. } = ev { } else { panic ! ( ) ; }
@@ -1425,13 +1472,16 @@ where
1425
1472
if let Event :: OnionMessagePeerConnected { .. } = ev { } else { panic ! ( ) ; }
1426
1473
}
1427
1474
}
1428
- core:: mem:: swap ( & mut * pending_intercepted_msgs_events, & mut events) ;
1429
- events. append ( & mut pending_peer_connected_events) ;
1430
1475
pending_peer_connected_events. shrink_to ( 10 ) ; // Limit total heap usage
1431
1476
}
1432
- for ev in events {
1433
- handler. handle_event ( ev) ;
1434
- }
1477
+
1478
+ let res = intercepted_msgs. into_iter ( ) . map ( |ev| handler. handle_event ( ev) ) . collect :: < Vec < _ > > ( ) ;
1479
+ drop_handled_events_and_abort ! ( self , res, 0 , self . pending_intercepted_msgs_events) ;
1480
+
1481
+ let res = peer_connecteds. into_iter ( ) . map ( |ev| handler. handle_event ( ev) ) . collect :: < Vec < _ > > ( ) ;
1482
+ drop_handled_events_and_abort ! ( self , res, 0 , self . pending_peer_connected_events) ;
1483
+
1484
+ self . pending_events_processor . store ( false , Ordering :: Release ) ;
1435
1485
}
1436
1486
}
1437
1487
0 commit comments