11use parking_lot:: Mutex ;
22use rand:: prelude:: IndexedRandom ;
3- use std:: collections:: { btree_map:: Entry , BTreeMap , VecDeque } ;
3+ use std:: collections:: { btree_map:: Entry , BTreeMap , HashSet , VecDeque } ;
44use std:: time:: Instant ;
55
66use crate :: topology:: { Limits , TopologyManager } ;
@@ -25,6 +25,7 @@ pub(crate) struct ConnectionManager {
2525 pub pub_key : Arc < TransportPublicKey > ,
2626 courtesy_links : Arc < Mutex < VecDeque < CourtesyLink > > > ,
2727 max_courtesy_links : usize ,
28+ pending_courtesy : Arc < Mutex < HashSet < PeerId > > > ,
2829}
2930
3031#[ derive( Clone ) ]
@@ -123,9 +124,39 @@ impl ConnectionManager {
123124 pub_key : Arc :: new ( pub_key) ,
124125 courtesy_links : Arc :: new ( Mutex :: new ( VecDeque :: new ( ) ) ) ,
125126 max_courtesy_links : if is_gateway { MAX_COURTESY_LINKS } else { 0 } ,
127+ pending_courtesy : Arc :: new ( Mutex :: new ( HashSet :: new ( ) ) ) ,
126128 }
127129 }
128130
131+ fn remember_courtesy_intent ( & self , peer : & PeerId ) {
132+ if !self . is_gateway {
133+ return ;
134+ }
135+ let mut pending = self . pending_courtesy . lock ( ) ;
136+ pending. insert ( peer. clone ( ) ) ;
137+ tracing:: debug!(
138+ %peer,
139+ pending = pending. len( ) ,
140+ "remember_courtesy_intent: recorded pending courtesy join"
141+ ) ;
142+ }
143+
144+ fn take_pending_courtesy ( & self , peer : & PeerId ) -> bool {
145+ if !self . is_gateway {
146+ return false ;
147+ }
148+ let mut pending = self . pending_courtesy . lock ( ) ;
149+ let removed = pending. remove ( peer) ;
150+ if removed {
151+ tracing:: debug!(
152+ %peer,
153+ pending = pending. len( ) ,
154+ "take_pending_courtesy: consuming pending courtesy flag"
155+ ) ;
156+ }
157+ removed
158+ }
159+
129160 fn register_courtesy_connection ( & self , peer : & PeerId ) -> Option < PeerId > {
130161 if !self . is_gateway || self . max_courtesy_links == 0 {
131162 return None ;
@@ -295,6 +326,7 @@ impl ConnectionManager {
295326 }
296327
297328 if courtesy_join {
329+ self . remember_courtesy_intent ( peer_id) ;
298330 tracing:: debug!(
299331 %peer_id,
300332 open,
@@ -474,6 +506,8 @@ impl ConnectionManager {
474506 . fetch_add ( 1 , std:: sync:: atomic:: Ordering :: SeqCst ) ;
475507 std:: mem:: drop ( lop) ;
476508
509+ let courtesy = courtesy || self . take_pending_courtesy ( & peer) ;
510+
477511 if courtesy {
478512 self . register_courtesy_connection ( & peer)
479513 } else {
0 commit comments