@@ -117,7 +117,7 @@ impl AddressGasLimiterInner {
117117 // Only clean up stale buckets every `cleanup_interval` blocks
118118 if block_number. is_multiple_of ( self . config . cleanup_interval ) {
119119 self . address_buckets
120- . retain ( |_, bucket| bucket. available <= bucket. capacity ) ;
120+ . retain ( |_, bucket| bucket. available < bucket. capacity ) ;
121121 }
122122
123123 active_addresses - self . address_buckets . len ( )
@@ -217,7 +217,46 @@ mod tests {
217217
218218 // Everyone should get some gas back
219219 assert ! ( limiter. consume_gas( searcher1, 1_000_000 ) . is_ok( ) ) ; // Had 9.5M + 1M refill, now 9.5M
220- assert ! ( limiter. consume_gas( searcher2, 1_000_000 ) . is_ok( ) ) ; // Had 9.25M + 1M refill, now 9.25M
220+ assert ! ( limiter. consume_gas( searcher2, 1_000_000 ) . is_ok( ) ) ; // Had 9.25M + 1M refill, now 9.25M
221221 assert ! ( limiter. consume_gas( attacker, 1_000_000 ) . is_ok( ) ) ; // Had 5M + 1M refill, now 5M
222222 }
223+
224+ #[ test]
225+ fn test_bucket_cleanup ( ) {
226+ // Test that unused buckets get cleaned up properly
227+ let config = create_test_config ( 1000 , 1000 , 10 ) ;
228+ let limiter = AddressGasLimiter :: new ( config) ;
229+
230+ let addr1 = Address :: from ( [ 0x1 ; 20 ] ) ;
231+ let addr2 = Address :: from ( [ 0x2 ; 20 ] ) ;
232+
233+ // Create buckets for both
234+ assert ! ( limiter. consume_gas( addr1, 100 ) . is_ok( ) ) ;
235+ assert ! ( limiter. consume_gas( addr2, 100 ) . is_ok( ) ) ;
236+
237+ let inner = limiter. inner . as_ref ( ) . unwrap ( ) ;
238+ assert_eq ! ( inner. address_buckets. len( ) , 2 ) ;
239+
240+ // Refill for several blocks - addr1 stays at full capacity (unused)
241+ // but addr2 continues to be used
242+ for block in 1 ..=10 {
243+ limiter. refresh ( block) ;
244+
245+ if block > 1 {
246+ // addr1 is now full and unused
247+ // addr2 continues to use gas
248+ assert ! ( limiter. consume_gas( addr2, 100 ) . is_ok( ) ) ;
249+ }
250+ }
251+
252+ // After cleanup at block 10 (multiple of 5), addr1 should be removed
253+ // because it's at full capacity (unused), while addr2 remains
254+ assert_eq ! (
255+ inner. address_buckets. len( ) ,
256+ 1 ,
257+ "Unused bucket (addr1) should have been cleaned up"
258+ ) ;
259+ assert ! ( inner. address_buckets. contains_key( & addr2) ) ;
260+ assert ! ( !inner. address_buckets. contains_key( & addr1) ) ;
261+ }
223262}
0 commit comments