initial commit

README.md

@@ -0,0 +1,44 @@
+# ratelimit
+--
+    import "github.com/juju/ratelimit"
+
+The ratelimit package provides an efficient token bucket implementation. See
+http://en.wikipedia.org/wiki/Token_bucket.
+
+## Usage
+
+#### type TokenBucket
+
+```go
+type TokenBucket struct {
+}
+```
+
+TokenBucket represents a token bucket that fills at a predetermined rate.
+Methods on TokenBucket may be called concurrently.
+
+#### func  New
+
+```go
+func New(fillInterval time.Duration, capacity int64) *TokenBucket
+```
+New returns a new token bucket that fills at the rate of one token every
+fillInterval, up to the given maximum capacity. Both arguments must be positive.
+
+#### func (*TokenBucket) Get
+
+```go
+func (tb *TokenBucket) Get(count int64)
+```
+Get gets count tokens from the bucket, waiting until the tokens are available.
+
+#### func (*TokenBucket) GetNB
+
+```go
+func (tb *TokenBucket) GetNB(count int64) time.Duration
+```
+GetNB gets count tokens from the bucket without blocking. It returns the time to
+wait until the tokens are actually available.
+
+Note that if the request is irrevocable - there is no way to return tokens to
+the bucket once this method commits us to taking them.

ratelimit.go

@@ -0,0 +1,93 @@
+// The ratelimit package provides an efficient token bucket implementation.
+// See http://en.wikipedia.org/wiki/Token_bucket.
+package ratelimit
+
+import (
+	"sync"
+	"time"
+)
+
+// TODO what about aborting requests?
+
+// TokenBucket represents a token bucket
+// that fills at a predetermined rate.
+// Methods on TokenBucket may be called
+// concurrently.
+type TokenBucket struct {
+	mu           sync.Mutex
+	startTime    time.Time
+	capacity     int64
+	fillInterval time.Duration
+	availTick    int64
+	avail        int64
+}
+
+// New returns a new token bucket that fills at the
+// rate of one token every fillInterval, up to the given
+// maximum capacity. Both arguments must be
+// positive.
+func New(fillInterval time.Duration, capacity int64) *TokenBucket {
+	if fillInterval <= 0 {
+		panic("token bucket fill interval is not > 0")
+	}
+	if capacity <= 0 {
+		panic("token bucket capacity is not > 0")
+	}
+	return &TokenBucket{
+		startTime:    time.Now(),
+		capacity:     capacity,
+		fillInterval: fillInterval,
+	}
+}
+
+// Get gets count tokens from the bucket, waiting
+// until the tokens are available.
+func (tb *TokenBucket) Get(count int64) {
+	if d := tb.GetNB(count); d > 0 {
+		time.Sleep(d)
+	}
+}
+
+// GetNB gets count tokens from the bucket without
+// blocking. It returns the time to wait until the
+// tokens are actually available.
+//
+// Note that if the request is irrevocable - there
+// is no way to return tokens to the bucket once
+// this method commits us to taking them.
+func (tb *TokenBucket) GetNB(count int64) time.Duration {
+	return tb.getNB(time.Now(), count)
+}
+
+// getNB is the internal version of GetNB - it takes
+// the current time as an argument to enable easy testing.
+func (tb *TokenBucket) getNB(now time.Time, count int64) time.Duration {
+	if count <= 0 {
+		return 0
+	}
+	tb.mu.Lock()
+	defer tb.mu.Unlock()
+	currentTick := int64(now.Sub(tb.startTime) / tb.fillInterval)
+	tb.adjust(currentTick)
+
+	tb.avail -= count
+	if tb.avail >= 0 {
+		return 0
+	}
+	endTick := currentTick-tb.avail
+	endTime := tb.startTime.Add(time.Duration(endTick) * tb.fillInterval)
+	return endTime.Sub(now)
+}
+
+// adjust adjusts the current bucket capacity based
+// on the current tick.
+func (tb *TokenBucket) adjust(currentTick int64) {
+	if tb.avail >= tb.capacity {
+		return
+	}
+	tb.avail += currentTick - tb.availTick
+	if tb.avail > tb.capacity {
+		tb.avail = tb.capacity
+	}
+	tb.availTick = currentTick
+}

ratelimit_test.go

@@ -0,0 +1,120 @@
+package ratelimit
+import (
+	gc "launchpad.net/gocheck"
+
+	"testing"
+	"time"
+)
+
+func TestPackage(t *testing.T) {
+	gc.TestingT(t)
+}
+
+type rateLimitSuite struct {}
+
+var _ = gc.Suite(rateLimitSuite{})
+
+type req struct {
+	time time.Duration
+	count int64
+	expectWait time.Duration
+}
+
+var rateLimitTests = []struct {
+	about string
+	fillInterval time.Duration
+	capacity int64
+	reqs []req
+}{{
+	about: "serial requests",
+	fillInterval: 250 * time.Millisecond,
+	capacity: 10,
+	reqs: []req{{
+		time: 0,
+		count: 0,
+		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: []req{{
+		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: []req{{
+		time: 20 * time.Millisecond,
+		count: 15,
+		expectWait: 5 * time.Millisecond,
+	}},
+}, {
+	about: "sub-quantum time",
+	fillInterval: 10 * time.Millisecond,
+	capacity: 10,
+	reqs: []req{{
+		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: []req{{
+		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,
+	}},
+}}
+
+func (rateLimitSuite) TestRateLimit(c *gc.C) {
+	for i, test := range rateLimitTests {
+		tb := New(test.fillInterval, test.capacity)
+		for j, req := range test.reqs {
+			d := tb.getNB(tb.startTime.Add(req.time), req.count)
+			if d != req.expectWait {
+				c.Fatalf("test %d.%d, %s, got %v want %v", i, j, test.about, d, req.expectWait)
+			}
+		}
+	}
+}
+
+func (rateLimitSuite) TestPanics(c *gc.C) {
+	c.Assert(func() {New(0, 1)}, gc.PanicMatches, "token bucket fill interval is not > 0")
+	c.Assert(func() {New(-2, 1)}, gc.PanicMatches, "token bucket fill interval is not > 0")
+	c.Assert(func() {New(1, 0)}, gc.PanicMatches, "token bucket capacity is not > 0")
+	c.Assert(func() {New(1, -2)}, gc.PanicMatches, "token bucket capacity is not > 0")
+}