Squashed 'src/github.com/tsenart/tb/' changes from 1fda845..9fbec38

9fbec38 Travis: download tests deps
aef8c07 Updated Travis config
0c24aae Move back to Travis
cf767b3 Fix doc typo
c86e63e Implement Wait method on Bucket
f38b46e TestBucket_Take_Throughput: GOMAXPROCS(2)

git-subtree-dir: src/github.com/tsenart/tb
git-subtree-split: 9fbec382fc274088cc45d8d3c0731d7da5d7e46c

Author: mindreframer

.travis.yml

@@ -0,0 +1,17 @@
+language: go
+
+go:
+  - 1.2
+  - 1.3
+  - tip
+
+install:
+  - go get -v code.google.com/p/go.tools/cmd/vet
+  - go get -v github.com/golang/lint/golint
+  - go get -d -t -v ./...
+  - go build -v ./...
+
+script:
+  - go vet ./...
+  - $HOME/gopath/bin/golint .
+  - go test -v -parallel=8 ./...

README.md

@@ -1,4 +1,4 @@
-# Token Bucket (tb) [![Build Status](https://drone.io/github.com/tsenart/tb/status.png)](https://drone.io/github.com/tsenart/tb/latest) [![GoDoc](https://godoc.org/github.com/tsenart/tb?status.png)](https://godoc.org/github.com/tsenart/tb)
+# Token Bucket (tb) [![Build Status](https://secure.travis-ci.org/tsenart/tb.png)](http://travis-ci.org/tsenart/tb) [![GoDoc](https://godoc.org/github.com/tsenart/tb?status.png)](https://godoc.org/github.com/tsenart/tb)
 
 This package provides a generic lock-free implementation of the "Token bucket"
 algorithm where handling of non-conformity is left to the user.

bucket.go

@@ -11,25 +11,29 @@ type Bucket struct {
 	inc      int64
 	tokens   int64
 	capacity int64
+	freq     time.Duration
 	closing  chan struct{}
 }
 
 // NewBucket returns a full Bucket with c capacity and starts a filling
 // go-routine which ticks every freq. The number of tokens added on each tick
 // is computed dynamically to be even across the duration of a second.
 //
-// If freq == -1 then it will be adjusted to 1/c seconds. Otherwise,
-// If freq < 1/c seconds, the filling go-routine won't be started.
+// If freq == -1 then the filling go-routine won't be started. Otherwise,
+// If freq < 1/c seconds, then it will be adjusted to 1/c seconds.
 func NewBucket(c int64, freq time.Duration) *Bucket {
 	b := &Bucket{tokens: c, capacity: c, closing: make(chan struct{})}
 
 	if freq == -1 {
-		freq = time.Duration(1e9 / c)
-	} else if freq.Seconds() < 1/float64(c) {
 		return b
+	} else if evenFreq := time.Duration(1e9 / c); freq < evenFreq {
+		freq = evenFreq
 	}
 
-	go b.fill(freq)
+	b.freq = freq
+	b.inc = int64(math.Floor(.5 + (float64(c) * freq.Seconds())))
+
+	go b.fill()
 
 	return b
 }
@@ -76,17 +80,43 @@ func (b *Bucket) Put(n int64) (added int64) {
 	}
 }
 
+// Wait waits for n amount of tokens to be available.
+// If n tokens are immediatelly available it doesn't sleep.
+// Otherwise, it sleeps the minimum amount of time required for the remaining
+// tokens to be available. It returns the wait duration.
+//
+// This method is thread-safe.
+func (b *Bucket) Wait(n int64) time.Duration {
+	var rem int64
+	if rem = n - b.Take(n); rem == 0 {
+		return 0
+	}
+
+	var wait time.Duration
+	for rem > 0 {
+		wait += b.wait(rem)
+		time.Sleep(wait)
+		rem -= b.Take(rem)
+	}
+	return wait
+}
+
 // Close stops the filling go-routine given it was started.
 func (b *Bucket) Close() error {
 	close(b.closing)
 	return nil
 }
 
-func (b *Bucket) fill(freq time.Duration) {
-	ticker := time.NewTicker(freq)
-	defer ticker.Stop()
+// wait returns the minimum amount of time required for n tokens to be available.
+// if n > capacity, n will be adjusted to capacity
+func (b *Bucket) wait(n int64) time.Duration {
+	return time.Duration(int64(math.Ceil(math.Min(float64(n), float64(b.capacity))/float64(b.inc))) *
+		b.freq.Nanoseconds())
+}
 
-	b.inc = int64(math.Floor(.5 + (float64(b.capacity) * freq.Seconds())))
+func (b *Bucket) fill() {
+	ticker := time.NewTicker(b.freq)
+	defer ticker.Stop()
 
 	for _ = range ticker.C {
 		select {

bucket_test.go

@@ -1,25 +1,27 @@
 package tb
 
 import (
+	"fmt"
+	"runtime"
 	"testing"
 	"time"
 )
 
 func TestNewBucket(t *testing.T) {
 	t.Parallel()
 
-	b := NewBucket(10, 0)
+	b := NewBucket(10, -1)
 	b.Take(10)
 	time.Sleep(100 * time.Millisecond)
 	if w, g := int64(0), b.Take(1); w != g {
-		t.Fatal("Expected no filling when freq < 1/c")
+		t.Fatal("Expected no filling when freq == -1")
 	}
 }
 
 func TestBucket_Take_single(t *testing.T) {
 	t.Parallel()
 
-	b := NewBucket(10, -1)
+	b := NewBucket(10, 0)
 	defer b.Close()
 
 	ex := [...]int64{5, 5, 1, 1, 5, 4, 1, 0}
@@ -49,7 +51,7 @@ func TestBucket_Put_single(t *testing.T) {
 func TestBucket_Take_multi(t *testing.T) {
 	t.Parallel()
 
-	b := NewBucket(10, -1)
+	b := NewBucket(10, 0)
 	defer b.Close()
 
 	exs := [2][]int64{{4, 4, 2, 2}, {2, 2, 1, 1}}
@@ -86,11 +88,14 @@ func TestBucket_Put_multi(t *testing.T) {
 
 func TestBucket_Take_throughput(t *testing.T) {
 	t.Parallel()
+
 	if testing.Short() {
 		t.Skip("Skipping test in short mode")
 	}
 
-	b := NewBucket(1000, -1)
+	runtime.GOMAXPROCS(2)
+
+	b := NewBucket(1000, 0)
 	defer b.Close()
 
 	b.Take(1000)
@@ -110,7 +115,7 @@ func TestBucket_Take_throughput(t *testing.T) {
 }
 
 func BenchmarkBucket_Take_sequential(b *testing.B) {
-	bucket := NewBucket(int64(b.N), -1)
+	bucket := NewBucket(int64(b.N), 0)
 	defer bucket.Close()
 
 	b.ResetTimer()
@@ -130,8 +135,42 @@ func BenchmarkBucket_Put_sequential(b *testing.B) {
 		bucket.Put(8)
 	}
 }
+
+func TestBucket_Wait(t *testing.T) {
+	t.Parallel()
+
+	cases := map[*Bucket]time.Duration{
+		NewBucket(1e3, 500*time.Millisecond): 1 * time.Second,
+		NewBucket(1e3, 20*time.Millisecond):  1 * time.Second,
+		NewBucket(1e3, 1*time.Millisecond):   1 * time.Second,
+		NewBucket(1e3, 0):                    1 * time.Second,
+		NewBucket(2e3, 0):                    0,
+		NewBucket(3e3, 0):                    0,
+	}
+	errors := make(chan error, len(cases))
+
+	for bucket, wait := range cases {
+		go func(bucket *Bucket, wait time.Duration) {
+			defer bucket.Close()
+
+			if got := bucket.Wait(2000); int(wait.Seconds()) != int(got.Seconds()) {
+				errors <- fmt.Errorf("bucket.Wait(2000) with cap=%d, freq=%s: Want: %s, Got %s",
+					bucket.capacity, bucket.freq, wait, got)
+			} else {
+				errors <- nil
+			}
+		}(bucket, wait)
+	}
+
+	for i := 0; i < cap(errors); i++ {
+		if err := <-errors; err != nil {
+			t.Error(err)
+		}
+	}
+}
+
 func TestBucket_Close(t *testing.T) {
-	b := NewBucket(10000, -1)
+	b := NewBucket(10000, 0)
 	b.Close()
 	b.Take(10000)
 	time.Sleep(10 * time.Millisecond)

throttler.go

@@ -47,8 +47,9 @@ func (t *Throttler) Bucket(key string, rate int64) *Bucket {
 	t.mu.RUnlock()
 
 	if !ok {
-		b = NewBucket(rate, 0)
+		b = NewBucket(rate, -1)
 		b.inc = int64(math.Floor(.5 + (float64(b.capacity) * t.freq.Seconds())))
+		b.freq = t.freq
 		t.mu.Lock()
 		t.buckets[key] = b
 		t.mu.Unlock()
@@ -57,32 +58,18 @@ func (t *Throttler) Bucket(key string, rate int64) *Bucket {
 	return b
 }
 
-// Wait waits for n amount of tokens to be available, sleeping freq between each
-// take. It returns the wait duration and whether it had to wait or not.
+// Wait waits for n amount of tokens to be available.
+// If n tokens are immediatelly available it doesn't sleep. Otherwise, it sleeps
+// the minimum amount of time required for the remaining tokens to be available.
+// It returns the wait duration.
 //
 // If a Bucket (key, rate) doesn't exist yet, it is created.
 // If freq < 1/rate seconds, the effective wait rate won't be correct.
 //
 // You must call Close when you're done with the Throttler in order to not leak
 // a go-routine and a system-timer.
-func (t *Throttler) Wait(key string, n, rate int64) (time.Duration, bool) {
-	var (
-		got   int64
-		began = time.Now()
-	)
-
-	b := t.Bucket(key, rate)
-
-	if got = b.Take(n); got == n {
-		return time.Since(began), false
-	}
-
-	for got < n {
-		got += b.Take(n - got)
-		time.Sleep(t.freq)
-	}
-
-	return time.Since(began), true
+func (t *Throttler) Wait(key string, n, rate int64) time.Duration {
+	return t.Bucket(key, rate).Wait(n)
 }
 
 // Halt returns a bool indicating if the Bucket identified by key and rate has

throttler_test.go

@@ -56,16 +56,12 @@ func TestThrottler_Wait(t *testing.T) {
 	th := NewThrottler(1 * time.Millisecond)
 	defer th.Close()
 
-	if _, waited := th.Wait("a", 1000, 1000); waited {
+	if wait := th.Wait("a", 1000, 1000); wait > 0 {
 		t.Fatal("Didn't expect wait")
 	}
 
-	if took, waited := th.Wait("a", 1000, 1000); !waited || int(took.Seconds()) != 1 {
-		t.Fatalf("Expected wait of 1s. Got: %s", took)
-	}
-
-	if _, waited := th.Wait("b", 1000, 1000); waited {
-		t.Fatal("Didn't expect wait")
+	if wait := th.Wait("a", 2000, 1000); int(wait.Seconds()) != 2 {
+		t.Fatalf("Expected wait of 2s. Got: %s", wait)
 	}
 }