@@ -349,16 +349,31 @@ func (c *throttlingTrafficShapingController) PerformChecking(arg interface{}, ba
349
349
}
350
350
intervalCostTime := int64 (math .Round (float64 (batchCount * c .durationInSec * 1000 / tokenCount )))
351
351
for {
352
- currentTimeInMs := int64 (util .CurrentTimeMillis ())
353
- lastPassTimePtr := timeCounter .AddIfAbsent (arg , & currentTimeInMs )
352
+ var (
353
+ expectedTime int64
354
+ currentTimeInMs int64
355
+ lastPassTime int64
356
+ lastPassTimePtr * int64
357
+ )
358
+
359
+ currentTimeInMs = int64 (util .CurrentTimeMillis ())
360
+ lastPassTimePtr = timeCounter .AddIfAbsent (arg , & currentTimeInMs )
354
361
if lastPassTimePtr == nil {
355
- // first access arg
356
- return nil
362
+ // initialize pointer for first access
363
+ lastPassTimePtr = & currentTimeInMs
357
364
}
358
365
// load the last pass time
359
- lastPassTime := atomic .LoadInt64 (lastPassTimePtr )
360
- // calculate the expected pass time
361
- expectedTime := lastPassTime + intervalCostTime
366
+ lastPassTime = atomic .LoadInt64 (lastPassTimePtr )
367
+ // calculate expected pass time based on two scenarios:
368
+ // 1. first access or expired statistics window
369
+ // 2. normal within-window access
370
+ if lastPassTimePtr == & currentTimeInMs || lastPassTime < currentTimeInMs - (c .durationInSec * 1000 ) {
371
+ // adjust the time of the previous window to one second ago, and at most TokenCount tokens can pass through
372
+ expectedTime = currentTimeInMs - (c .durationInSec * 1000 ) + intervalCostTime
373
+ } else {
374
+ // normal cumulative calculation
375
+ expectedTime = lastPassTime + intervalCostTime
376
+ }
362
377
363
378
if expectedTime <= currentTimeInMs || expectedTime - currentTimeInMs < c .maxQueueingTimeMs {
364
379
if atomic .CompareAndSwapInt64 (lastPassTimePtr , lastPassTime , currentTimeInMs ) {
0 commit comments