390 lines
9 KiB
Go
390 lines
9 KiB
Go
|
// Copyright 2014 Canonical Ltd.
|
||
|
|
// Licensed under the LGPLv3 with static-linking exception.
|
||
|
|
// See LICENCE file for details.
|
||
|
|
|
||
|
|
package ratelimit
|
||
|
|
|
||
|
|
import (
|
||
|
|
"math"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
gc "gopkg.in/check.v1"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestPackage(t *testing.T) {
|
||
|
|
gc.TestingT(t)
|
||
|
|
}
|
||
|
|
|
||
|
|
type rateLimitSuite struct{}
|
||
|
|
|
||
|
|
var _ = gc.Suite(rateLimitSuite{})
|
||
|
|
|
||
|
|
type takeReq struct {
|
||
|
|
time time.Duration
|
||
|
|
count int64
|
||
|
|
expectWait time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
var takeTests = []struct {
|
||
|
|
about string
|
||
|
|
fillInterval time.Duration
|
||
|
|
capacity int64
|
||
|
|
reqs []takeReq
|
||
|
|
}{{
|
||
|
|
about: "serial requests",
|
||
|
|
fillInterval: 250 * time.Millisecond,
|
||
|
|
capacity: 10,
|
||
|
|
reqs: []takeReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 0,
|
||
|
|
expectWait: 0,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 10,
|
||
|
|
expectWait: 0,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 1,
|
||
|
|
expectWait: 250 * time.Millisecond,
|
||
|
|
}, {
|
||
|
|
time: 250 * time.Millisecond,
|
||
|
|
count: 1,
|
||
|
|
expectWait: 250 * time.Millisecond,
|
||
|
|
}},
|
||
|
|
}, {
|
||
|
|
about: "concurrent requests",
|
||
|
|
fillInterval: 250 * time.Millisecond,
|
||
|
|
capacity: 10,
|
||
|
|
reqs: []takeReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 10,
|
||
|
|
expectWait: 0,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 2,
|
||
|
|
expectWait: 500 * time.Millisecond,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 2,
|
||
|
|
expectWait: 1000 * time.Millisecond,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 1,
|
||
|
|
expectWait: 1250 * time.Millisecond,
|
||
|
|
}},
|
||
|
|
}, {
|
||
|
|
about: "more than capacity",
|
||
|
|
fillInterval: 1 * time.Millisecond,
|
||
|
|
capacity: 10,
|
||
|
|
reqs: []takeReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 10,
|
||
|
|
expectWait: 0,
|
||
|
|
}, {
|
||
|
|
time: 20 * time.Millisecond,
|
||
|
|
count: 15,
|
||
|
|
expectWait: 5 * time.Millisecond,
|
||
|
|
}},
|
||
|
|
}, {
|
||
|
|
about: "sub-quantum time",
|
||
|
|
fillInterval: 10 * time.Millisecond,
|
||
|
|
capacity: 10,
|
||
|
|
reqs: []takeReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 10,
|
||
|
|
expectWait: 0,
|
||
|
|
}, {
|
||
|
|
time: 7 * time.Millisecond,
|
||
|
|
count: 1,
|
||
|
|
expectWait: 3 * time.Millisecond,
|
||
|
|
}, {
|
||
|
|
time: 8 * time.Millisecond,
|
||
|
|
count: 1,
|
||
|
|
expectWait: 12 * time.Millisecond,
|
||
|
|
}},
|
||
|
|
}, {
|
||
|
|
about: "within capacity",
|
||
|
|
fillInterval: 10 * time.Millisecond,
|
||
|
|
capacity: 5,
|
||
|
|
reqs: []takeReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 5,
|
||
|
|
expectWait: 0,
|
||
|
|
}, {
|
||
|
|
time: 60 * time.Millisecond,
|
||
|
|
count: 5,
|
||
|
|
expectWait: 0,
|
||
|
|
}, {
|
||
|
|
time: 60 * time.Millisecond,
|
||
|
|
count: 1,
|
||
|
|
expectWait: 10 * time.Millisecond,
|
||
|
|
}, {
|
||
|
|
time: 80 * time.Millisecond,
|
||
|
|
count: 2,
|
||
|
|
expectWait: 10 * time.Millisecond,
|
||
|
|
}},
|
||
|
|
}}
|
||
|
|
|
||
|
|
var availTests = []struct {
|
||
|
|
about string
|
||
|
|
capacity int64
|
||
|
|
fillInterval time.Duration
|
||
|
|
take int64
|
||
|
|
sleep time.Duration
|
||
|
|
|
||
|
|
expectCountAfterTake int64
|
||
|
|
expectCountAfterSleep int64
|
||
|
|
}{{
|
||
|
|
about: "should fill tokens after interval",
|
||
|
|
capacity: 5,
|
||
|
|
fillInterval: time.Second,
|
||
|
|
take: 5,
|
||
|
|
sleep: time.Second,
|
||
|
|
expectCountAfterTake: 0,
|
||
|
|
expectCountAfterSleep: 1,
|
||
|
|
}, {
|
||
|
|
about: "should fill tokens plus existing count",
|
||
|
|
capacity: 2,
|
||
|
|
fillInterval: time.Second,
|
||
|
|
take: 1,
|
||
|
|
sleep: time.Second,
|
||
|
|
expectCountAfterTake: 1,
|
||
|
|
expectCountAfterSleep: 2,
|
||
|
|
}, {
|
||
|
|
about: "shouldn't fill before interval",
|
||
|
|
capacity: 2,
|
||
|
|
fillInterval: 2 * time.Second,
|
||
|
|
take: 1,
|
||
|
|
sleep: time.Second,
|
||
|
|
expectCountAfterTake: 1,
|
||
|
|
expectCountAfterSleep: 1,
|
||
|
|
}, {
|
||
|
|
about: "should fill only once after 1*interval before 2*interval",
|
||
|
|
capacity: 2,
|
||
|
|
fillInterval: 2 * time.Second,
|
||
|
|
take: 1,
|
||
|
|
sleep: 3 * time.Second,
|
||
|
|
expectCountAfterTake: 1,
|
||
|
|
expectCountAfterSleep: 2,
|
||
|
|
}}
|
||
|
|
|
||
|
|
func (rateLimitSuite) TestTake(c *gc.C) {
|
||
|
|
for i, test := range takeTests {
|
||
|
|
tb := NewBucket(test.fillInterval, test.capacity)
|
||
|
|
for j, req := range test.reqs {
|
||
|
|
d, ok := tb.take(tb.startTime.Add(req.time), req.count, infinityDuration)
|
||
|
|
c.Assert(ok, gc.Equals, true)
|
||
|
|
if d != req.expectWait {
|
||
|
|
c.Fatalf("test %d.%d, %s, got %v want %v", i, j, test.about, d, req.expectWait)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rateLimitSuite) TestTakeMaxDuration(c *gc.C) {
|
||
|
|
for i, test := range takeTests {
|
||
|
|
tb := NewBucket(test.fillInterval, test.capacity)
|
||
|
|
for j, req := range test.reqs {
|
||
|
|
if req.expectWait > 0 {
|
||
|
|
d, ok := tb.take(tb.startTime.Add(req.time), req.count, req.expectWait-1)
|
||
|
|
c.Assert(ok, gc.Equals, false)
|
||
|
|
c.Assert(d, gc.Equals, time.Duration(0))
|
||
|
|
}
|
||
|
|
d, ok := tb.take(tb.startTime.Add(req.time), req.count, req.expectWait)
|
||
|
|
c.Assert(ok, gc.Equals, true)
|
||
|
|
if d != req.expectWait {
|
||
|
|
c.Fatalf("test %d.%d, %s, got %v want %v", i, j, test.about, d, req.expectWait)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
type takeAvailableReq struct {
|
||
|
|
time time.Duration
|
||
|
|
count int64
|
||
|
|
expect int64
|
||
|
|
}
|
||
|
|
|
||
|
|
var takeAvailableTests = []struct {
|
||
|
|
about string
|
||
|
|
fillInterval time.Duration
|
||
|
|
capacity int64
|
||
|
|
reqs []takeAvailableReq
|
||
|
|
}{{
|
||
|
|
about: "serial requests",
|
||
|
|
fillInterval: 250 * time.Millisecond,
|
||
|
|
capacity: 10,
|
||
|
|
reqs: []takeAvailableReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 0,
|
||
|
|
expect: 0,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 10,
|
||
|
|
expect: 10,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 1,
|
||
|
|
expect: 0,
|
||
|
|
}, {
|
||
|
|
time: 250 * time.Millisecond,
|
||
|
|
count: 1,
|
||
|
|
expect: 1,
|
||
|
|
}},
|
||
|
|
}, {
|
||
|
|
about: "concurrent requests",
|
||
|
|
fillInterval: 250 * time.Millisecond,
|
||
|
|
capacity: 10,
|
||
|
|
reqs: []takeAvailableReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 5,
|
||
|
|
expect: 5,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 2,
|
||
|
|
expect: 2,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 5,
|
||
|
|
expect: 3,
|
||
|
|
}, {
|
||
|
|
time: 0,
|
||
|
|
count: 1,
|
||
|
|
expect: 0,
|
||
|
|
}},
|
||
|
|
}, {
|
||
|
|
about: "more than capacity",
|
||
|
|
fillInterval: 1 * time.Millisecond,
|
||
|
|
capacity: 10,
|
||
|
|
reqs: []takeAvailableReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 10,
|
||
|
|
expect: 10,
|
||
|
|
}, {
|
||
|
|
time: 20 * time.Millisecond,
|
||
|
|
count: 15,
|
||
|
|
expect: 10,
|
||
|
|
}},
|
||
|
|
}, {
|
||
|
|
about: "within capacity",
|
||
|
|
fillInterval: 10 * time.Millisecond,
|
||
|
|
capacity: 5,
|
||
|
|
reqs: []takeAvailableReq{{
|
||
|
|
time: 0,
|
||
|
|
count: 5,
|
||
|
|
expect: 5,
|
||
|
|
}, {
|
||
|
|
time: 60 * time.Millisecond,
|
||
|
|
count: 5,
|
||
|
|
expect: 5,
|
||
|
|
}, {
|
||
|
|
time: 70 * time.Millisecond,
|
||
|
|
count: 1,
|
||
|
|
expect: 1,
|
||
|
|
}},
|
||
|
|
}}
|
||
|
|
|
||
|
|
func (rateLimitSuite) TestTakeAvailable(c *gc.C) {
|
||
|
|
for i, test := range takeAvailableTests {
|
||
|
|
tb := NewBucket(test.fillInterval, test.capacity)
|
||
|
|
for j, req := range test.reqs {
|
||
|
|
d := tb.takeAvailable(tb.startTime.Add(req.time), req.count)
|
||
|
|
if d != req.expect {
|
||
|
|
c.Fatalf("test %d.%d, %s, got %v want %v", i, j, test.about, d, req.expect)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rateLimitSuite) TestPanics(c *gc.C) {
|
||
|
|
c.Assert(func() { NewBucket(0, 1) }, gc.PanicMatches, "token bucket fill interval is not > 0")
|
||
|
|
c.Assert(func() { NewBucket(-2, 1) }, gc.PanicMatches, "token bucket fill interval is not > 0")
|
||
|
|
c.Assert(func() { NewBucket(1, 0) }, gc.PanicMatches, "token bucket capacity is not > 0")
|
||
|
|
c.Assert(func() { NewBucket(1, -2) }, gc.PanicMatches, "token bucket capacity is not > 0")
|
||
|
|
}
|
||
|
|
|
||
|
|
func isCloseTo(x, y, tolerance float64) bool {
|
||
|
|
return math.Abs(x-y)/y < tolerance
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rateLimitSuite) TestRate(c *gc.C) {
|
||
|
|
tb := NewBucket(1, 1)
|
||
|
|
if !isCloseTo(tb.Rate(), 1e9, 0.00001) {
|
||
|
|
c.Fatalf("got %v want 1e9", tb.Rate())
|
||
|
|
}
|
||
|
|
tb = NewBucket(2*time.Second, 1)
|
||
|
|
if !isCloseTo(tb.Rate(), 0.5, 0.00001) {
|
||
|
|
c.Fatalf("got %v want 0.5", tb.Rate())
|
||
|
|
}
|
||
|
|
tb = NewBucketWithQuantum(100*time.Millisecond, 1, 5)
|
||
|
|
if !isCloseTo(tb.Rate(), 50, 0.00001) {
|
||
|
|
c.Fatalf("got %v want 50", tb.Rate())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func checkRate(c *gc.C, rate float64) {
|
||
|
|
tb := NewBucketWithRate(rate, 1<<62)
|
||
|
|
if !isCloseTo(tb.Rate(), rate, rateMargin) {
|
||
|
|
c.Fatalf("got %g want %v", tb.Rate(), rate)
|
||
|
|
}
|
||
|
|
d, ok := tb.take(tb.startTime, 1<<62, infinityDuration)
|
||
|
|
c.Assert(ok, gc.Equals, true)
|
||
|
|
c.Assert(d, gc.Equals, time.Duration(0))
|
||
|
|
|
||
|
|
// Check that the actual rate is as expected by
|
||
|
|
// asking for a not-quite multiple of the bucket's
|
||
|
|
// quantum and checking that the wait time
|
||
|
|
// correct.
|
||
|
|
d, ok = tb.take(tb.startTime, tb.quantum*2-tb.quantum/2, infinityDuration)
|
||
|
|
c.Assert(ok, gc.Equals, true)
|
||
|
|
expectTime := 1e9 * float64(tb.quantum) * 2 / rate
|
||
|
|
if !isCloseTo(float64(d), expectTime, rateMargin) {
|
||
|
|
c.Fatalf("rate %g: got %g want %v", rate, float64(d), expectTime)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rateLimitSuite) TestNewWithRate(c *gc.C) {
|
||
|
|
for rate := float64(1); rate < 1e6; rate += 7 {
|
||
|
|
checkRate(c, rate)
|
||
|
|
}
|
||
|
|
for _, rate := range []float64{
|
||
|
|
1024 * 1024 * 1024,
|
||
|
|
1e-5,
|
||
|
|
0.9e-5,
|
||
|
|
0.5,
|
||
|
|
0.9,
|
||
|
|
0.9e8,
|
||
|
|
3e12,
|
||
|
|
4e18,
|
||
|
|
} {
|
||
|
|
checkRate(c, rate)
|
||
|
|
checkRate(c, rate/3)
|
||
|
|
checkRate(c, rate*1.3)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestAvailable(t *testing.T) {
|
||
|
|
for i, tt := range availTests {
|
||
|
|
tb := NewBucket(tt.fillInterval, tt.capacity)
|
||
|
|
if c := tb.takeAvailable(tb.startTime, tt.take); c != tt.take {
|
||
|
|
t.Fatalf("#%d: %s, take = %d, want = %d", i, tt.about, c, tt.take)
|
||
|
|
}
|
||
|
|
if c := tb.available(tb.startTime); c != tt.expectCountAfterTake {
|
||
|
|
t.Fatalf("#%d: %s, after take, available = %d, want = %d", i, tt.about, c, tt.expectCountAfterTake)
|
||
|
|
}
|
||
|
|
if c := tb.available(tb.startTime.Add(tt.sleep)); c != tt.expectCountAfterSleep {
|
||
|
|
t.Fatalf("#%d: %s, after some time it should fill in new tokens, available = %d, want = %d",
|
||
|
|
i, tt.about, c, tt.expectCountAfterSleep)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
func BenchmarkWait(b *testing.B) {
|
||
|
|
tb := NewBucket(1, 16*1024)
|
||
|
|
for i := b.N - 1; i >= 0; i-- {
|
||
|
|
tb.Wait(1)
|
||
|
|
}
|
||
|
|
}
|