@@ -10,8 +10,8 @@ use std::{
1010use async_trait:: async_trait;
1111use bytes:: Bytes ;
1212use futures:: future:: { self , AbortHandle } ;
13+ use lfu_cache:: TimedLfuCache ;
1314use log:: { debug, error, trace, warn} ;
14- use lru_time_cache:: { Entry , LruCache } ;
1515use shadowsocks:: {
1616 lookup_then,
1717 net:: UdpSocket as ShadowUdpSocket ,
@@ -42,14 +42,17 @@ pub trait UdpInboundWrite {
4242 async fn send_to ( & self , peer_addr : SocketAddr , remote_addr : & Address , data : & [ u8 ] ) -> io:: Result < ( ) > ;
4343}
4444
45+ type AssociationMap < W > = TimedLfuCache < SocketAddr , UdpAssociation < W > > ;
46+ type SharedAssociationMap < W > = Arc < Mutex < AssociationMap < W > > > ;
47+
4548/// UDP association manager
4649pub struct UdpAssociationManager < W >
4750where
4851 W : UdpInboundWrite + Clone + Send + Sync + Unpin + ' static ,
4952{
5053 respond_writer : W ,
5154 context : Arc < ServiceContext > ,
52- assoc_map : Arc < Mutex < LruCache < SocketAddr , UdpAssociation < W > > > > ,
55+ assoc_map : SharedAssociationMap < W > ,
5356 cleanup_abortable : AbortHandle ,
5457 balancer : PingBalancer ,
5558}
7780 ) -> UdpAssociationManager < W > {
7881 let time_to_live = time_to_live. unwrap_or ( crate :: DEFAULT_UDP_EXPIRY_DURATION ) ;
7982 let assoc_map = Arc :: new ( Mutex :: new ( match capacity {
80- Some ( capacity) => LruCache :: with_expiry_duration_and_capacity ( time_to_live , capacity ) ,
81- None => LruCache :: with_expiry_duration ( time_to_live) ,
83+ Some ( capacity) => TimedLfuCache :: with_capacity_and_expiration ( capacity , time_to_live ) ,
84+ None => TimedLfuCache :: with_expiration ( time_to_live) ,
8285 } ) ) ;
8386
8487 let cleanup_abortable = {
8790 loop {
8891 time:: sleep ( time_to_live) . await ;
8992
90- // iter() will trigger a cleanup of expired associations
91- let _ = assoc_map. lock ( ) . await . iter ( ) ;
93+ // cleanup expired associations
94+ let _ = assoc_map. lock ( ) . await . evict_expired ( ) ;
9295 }
9396 } ) ;
9497 tokio:: spawn ( cleanup_task) ;
@@ -107,23 +110,27 @@ where
107110 /// Sends `data` from `peer_addr` to `target_addr`
108111 pub async fn send_to ( & self , peer_addr : SocketAddr , target_addr : Address , data : & [ u8 ] ) -> io:: Result < ( ) > {
109112 // Check or (re)create an association
110- match self . assoc_map . lock ( ) . await . entry ( peer_addr) {
111- Entry :: Occupied ( occ) => {
112- let assoc = occ. into_mut ( ) ;
113- assoc. try_send ( ( target_addr, Bytes :: copy_from_slice ( data) ) )
114- }
115- Entry :: Vacant ( vac) => {
116- let assoc = vac. insert ( UdpAssociation :: new (
117- self . context . clone ( ) ,
118- peer_addr,
119- self . assoc_map . clone ( ) ,
120- self . balancer . clone ( ) ,
121- self . respond_writer . clone ( ) ,
122- ) ) ;
123- trace ! ( "created udp association for {}" , peer_addr) ;
124- assoc. try_send ( ( target_addr, Bytes :: copy_from_slice ( data) ) )
125- }
113+
114+ let mut assoc_map = self . assoc_map . lock ( ) . await ;
115+
116+ if let Some ( assoc) = assoc_map. get ( & peer_addr) {
117+ return assoc. try_send ( ( target_addr, Bytes :: copy_from_slice ( data) ) ) ;
126118 }
119+
120+ let assoc = UdpAssociation :: new (
121+ self . context . clone ( ) ,
122+ peer_addr,
123+ self . assoc_map . clone ( ) ,
124+ self . balancer . clone ( ) ,
125+ self . respond_writer . clone ( ) ,
126+ ) ;
127+
128+ trace ! ( "created udp association for {}" , peer_addr) ;
129+
130+ assoc. try_send ( ( target_addr, Bytes :: copy_from_slice ( data) ) ) ?;
131+ assoc_map. insert ( peer_addr, assoc) ;
132+
133+ Ok ( ( ) )
127134 }
128135}
129136
@@ -153,7 +160,7 @@ where
153160 fn new (
154161 context : Arc < ServiceContext > ,
155162 peer_addr : SocketAddr ,
156- assoc_map : Arc < Mutex < LruCache < SocketAddr , UdpAssociation < W > > > > ,
163+ assoc_map : SharedAssociationMap < W > ,
157164 balancer : PingBalancer ,
158165 respond_writer : W ,
159166 ) -> UdpAssociation < W > {
@@ -245,7 +252,7 @@ where
245252 bypassed_ipv4_socket : SpinMutex < UdpAssociationBypassState > ,
246253 bypassed_ipv6_socket : SpinMutex < UdpAssociationBypassState > ,
247254 proxied_socket : SpinMutex < UdpAssociationProxyState > ,
248- assoc_map : Arc < Mutex < LruCache < SocketAddr , UdpAssociation < W > > > > ,
255+ assoc_map : SharedAssociationMap < W > ,
249256 balancer : PingBalancer ,
250257 respond_writer : W ,
251258}
@@ -266,7 +273,7 @@ where
266273 fn new (
267274 context : Arc < ServiceContext > ,
268275 peer_addr : SocketAddr ,
269- assoc_map : Arc < Mutex < LruCache < SocketAddr , UdpAssociation < W > > > > ,
276+ assoc_map : SharedAssociationMap < W > ,
270277 balancer : PingBalancer ,
271278 respond_writer : W ,
272279 ) -> ( Arc < UdpAssociationContext < W > > , mpsc:: Sender < ( Address , Bytes ) > ) {
0 commit comments