Skip to content

Commit

Permalink
Merge pull request #31 from kingsamchen/fix/adjustAvailableToken
Browse files Browse the repository at this point in the history
Fix bonus tokens by function adjustavailableTokens (#30)
  • Loading branch information
manadart authored Oct 2, 2019
2 parents 6070dec + 4b7d3dd commit f60b320
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 2 deletions.
5 changes: 3 additions & 2 deletions ratelimit.go
Original file line number Diff line number Diff line change
Expand Up @@ -310,14 +310,15 @@ func (tb *Bucket) currentTick(now time.Time) int64 {
// available in the bucket at the given time, which must
// be in the future (positive) with respect to tb.latestTick.
func (tb *Bucket) adjustavailableTokens(tick int64) {
lastTick := tb.latestTick
tb.latestTick = tick
if tb.availableTokens >= tb.capacity {
return
}
tb.availableTokens += (tick - tb.latestTick) * tb.quantum
tb.availableTokens += (tick - lastTick) * tb.quantum
if tb.availableTokens > tb.capacity {
tb.availableTokens = tb.capacity
}
tb.latestTick = tick
return
}

Expand Down
25 changes: 25 additions & 0 deletions ratelimit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -382,6 +382,31 @@ func TestAvailable(t *testing.T) {

}

func TestNoBonusTokenAfterBucketIsFull(t *testing.T) {
tb := NewBucketWithQuantum(time.Second*1, 100, 20)
curAvail := tb.Available()
if curAvail != 100 {
t.Fatalf("initially: actual available = %d, expected = %d", curAvail, 100)
}

time.Sleep(time.Second * 5)

curAvail = tb.Available()
if curAvail != 100 {
t.Fatalf("after pause: actual available = %d, expected = %d", curAvail, 100)
}

cnt := tb.TakeAvailable(100)
if cnt != 100 {
t.Fatalf("taking: actual taken count = %d, expected = %d", cnt, 100)
}

curAvail = tb.Available()
if curAvail != 0 {
t.Fatalf("after taken: actual available = %d, expected = %d", curAvail, 0)
}
}

func BenchmarkWait(b *testing.B) {
tb := NewBucket(1, 16*1024)
for i := b.N - 1; i >= 0; i-- {
Expand Down

0 comments on commit f60b320

Please sign in to comment.