@@ -107,10 +107,10 @@ pub fn needs_quota_warning(curr_percentage: u64, warned_at_percentage: u64) -> b
107107impl Context {
108108 /// Returns whether the quota value needs an update. If so, `update_recent_quota()` should be
109109 /// called.
110- pub ( crate ) async fn quota_needs_update ( & self , ratelimit_secs : u64 ) -> bool {
110+ pub ( crate ) async fn quota_needs_update ( & self , transport_id : u32 , ratelimit_secs : u64 ) -> bool {
111111 let quota = self . quota . read ( ) . await ;
112112 quota
113- . as_ref ( )
113+ . get ( & transport_id )
114114 . filter ( |quota| time_elapsed ( & quota. modified ) < Duration :: from_secs ( ratelimit_secs) )
115115 . is_none ( )
116116 }
@@ -155,10 +155,13 @@ impl Context {
155155 }
156156 }
157157
158- * self . quota . write ( ) . await = Some ( QuotaInfo {
159- recent : quota,
160- modified : tools:: Time :: now ( ) ,
161- } ) ;
158+ self . quota . write ( ) . await . insert (
159+ session. transport_id ( ) ,
160+ QuotaInfo {
161+ recent : quota,
162+ modified : tools:: Time :: now ( ) ,
163+ } ,
164+ ) ;
162165
163166 self . emit_event ( EventType :: ConnectivityChanged ) ;
164167 Ok ( ( ) )
@@ -203,27 +206,42 @@ mod tests {
203206 let mut tcm = TestContextManager :: new ( ) ;
204207 let t = & tcm. unconfigured ( ) . await ;
205208 const TIMEOUT : u64 = 60 ;
206- assert ! ( t. quota_needs_update( TIMEOUT ) . await ) ;
207-
208- * t. quota . write ( ) . await = Some ( QuotaInfo {
209- recent : Ok ( Default :: default ( ) ) ,
210- modified : tools:: Time :: now ( ) - Duration :: from_secs ( TIMEOUT + 1 ) ,
211- } ) ;
212- assert ! ( t. quota_needs_update( TIMEOUT ) . await ) ;
213-
214- * t. quota . write ( ) . await = Some ( QuotaInfo {
215- recent : Ok ( Default :: default ( ) ) ,
216- modified : tools:: Time :: now ( ) ,
217- } ) ;
218- assert ! ( !t. quota_needs_update( TIMEOUT ) . await ) ;
209+ assert ! ( t. quota_needs_update( 0 , TIMEOUT ) . await ) ;
210+
211+ * t. quota . write ( ) . await = {
212+ let mut map = BTreeMap :: new ( ) ;
213+ map. insert (
214+ 0 ,
215+ QuotaInfo {
216+ recent : Ok ( Default :: default ( ) ) ,
217+ modified : tools:: Time :: now ( ) - Duration :: from_secs ( TIMEOUT + 1 ) ,
218+ } ,
219+ ) ;
220+ map
221+ } ;
222+ assert ! ( t. quota_needs_update( 0 , TIMEOUT ) . await ) ;
223+
224+ * t. quota . write ( ) . await = {
225+ let mut map = BTreeMap :: new ( ) ;
226+ map. insert (
227+ 0 ,
228+ QuotaInfo {
229+ recent : Ok ( Default :: default ( ) ) ,
230+ modified : tools:: Time :: now ( ) ,
231+ } ,
232+ ) ;
233+ map
234+ } ;
235+ assert ! ( !t. quota_needs_update( 0 , TIMEOUT ) . await ) ;
219236
220237 t. evtracker . clear_events ( ) ;
221238 t. set_primary_self_addr ( "new@addr" ) . await ?;
222- assert ! ( t. quota. read( ) . await . is_none ( ) ) ;
239+ assert ! ( t. quota. read( ) . await . is_empty ( ) ) ;
223240 t. evtracker
224241 . get_matching ( |evt| matches ! ( evt, EventType :: ConnectivityChanged ) )
225242 . await ;
226- assert ! ( t. quota_needs_update( TIMEOUT ) . await ) ;
243+ assert ! ( t. quota_needs_update( 0 , TIMEOUT ) . await ) ;
244+
227245 Ok ( ( ) )
228246 }
229247}
0 commit comments