github.com/lrita/ratelimit@v0.0.0-20190723030019-81504bd89bc5/ratelimit_test.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+int64(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+int64(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+int64(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+int64(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")
	c.Assert(func() { NewBucketWithQuantum(1, 1, 2) }, gc.PanicMatches, "token capacity need large than quantum")
}

func (rateLimitSuite) TestAvailable(c *gc.C) {
	tb := NewBucketWithQuantum(time.Hour, 2, 2)
	c.Assert(tb.Available(), gc.Equals, int64(2))
	tb.Take(1)
	c.Assert(tb.Available(), gc.Equals, int64(1))
	c.Assert(tb.TakeAvailable(1), gc.Equals, int64(1))
	c.Assert(tb.Available(), gc.Equals, int64(0))
	c.Assert(tb.WaitMaxDuration(1, time.Millisecond), gc.Equals, false)
}

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, 5, 5)
	if !isCloseTo(tb.Rate(), 50, 0.00001) {
		c.Fatalf("got %v want 50", tb.Rate())
	}
	tb.ResetRate(500.0, 50)
	if !isCloseTo(tb.Rate(), 500.0, 0.00001) {
		c.Fatalf("got %v want 500", tb.Rate())
	}
}

func (rateLimitSuite) TestResetRate(c *gc.C) {
	tb := NewBucketWithRate(500.0, 50)
	go func() {
		time.Sleep(500 * time.Millisecond)
		tb.ResetRate(1000.0, 100)
	}()
	count := 0
	begin := time.Now()
	for time.Since(begin) < time.Second {
		if tb.WaitMaxDuration(1, time.Millisecond) {
			count++
		}
	}
	if count < 600 {
		c.Fatalf("got %v want > 550", count)
	}
}

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) TestNewBucketWithRate(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,
		float64(1<<63 - 1),
	} {
		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 tb.Capacity() != tt.capacity {
			t.Fatalf("#%d: %s, take = %d, want = %d",
				i, tt.about, tb.Capacity(), 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 + int64(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)
	}
}

func BenchmarkNewBucket(b *testing.B) {
	for i := b.N - 1; i >= 0; i-- {
		NewBucketWithRate(4e18, 1<<62)
	}
}

func BenchmarkTakeParallel(b *testing.B) {
	tb := NewBucket(1, 16*1024)
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			tb.Take(1)
		}
	})
}

func BenchmarkTakeAvailableParallel(b *testing.B) {
	tb := NewBucket(1, 16*1024)
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			tb.TakeAvailable(1)
		}
	})
}

func BenchmarkResetRate(b *testing.B) {
	tb := NewBucketWithRate(500.0, 500)
	for i := 0; i < b.N; i++ {
		tb.ResetRate(float64(i+1), int64(i+1))
	}
}