65aaa607ab
Change made by: - running "gvt fetch" on each of the packages mentioned in Godeps/Godeps.json - `rm -rf Godeps` - tweaking the build scripts to not mention Godeps - tweaking the build scripts to test `./lib/...`, `./cmd/...` explicitly (to avoid testing vendor) - tweaking the build scripts to not juggle GOPATH for Godeps and instead set GO15VENDOREXPERIMENT. This also results in some updated packages at the same time I bet. Building with Go 1.3 and 1.4 still *works* but won't use our vendored dependencies - the user needs to have the actual packages in their GOPATH then, which they'll get with a normal "go get". Building with Go 1.6+ will get our vendored dependencies by default even when not using our build script, which is nice. By doing this we gain some freedom in that we can pick and choose manually what to include in vendor, as it's not based on just dependency analysis of our own code. This is also a risk as we might pick up dependencies we are unaware of, as the build may work locally with those packages present in GOPATH. On the other hand the build server will detect this as it has no packages in it's GOPATH beyond what is included in the repo. Recommended tool to manage dependencies is github.com/FiloSottile/gvt.
390 lines
9.0 KiB
Go
390 lines
9.0 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)
|
|
}
|
|
}
|