Vendor Leaky Bucket Implementation (#11560)

* add changes

* fix tests

* change to minute

* remove dep

* remove

* fix tests

* add test for period

* improve

* linter

* build files

* ci

* make it stricter

* fix tests

* fix

* Update beacon-chain/sync/rate_limiter.go

Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

Co-authored-by: terencechain <terence@prysmaticlabs.com>
Co-authored-by: Preston Van Loon <preston@prysmaticlabs.com>

container/leaky-bucket/BUILD.bazel

@@ -0,0 +1,22 @@
+load("@prysm//tools/go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "collector.go",
+        "heap.go",
+        "leakybucket.go",
+    ],
+    importpath = "github.com/prysmaticlabs/prysm/v3/container/leaky-bucket",
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = [
+        "collector_test.go",
+        "heap_test.go",
+        "leakybucket_test.go",
+    ],
+    embed = [":go_default_library"],
+)

container/leaky-bucket/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 kevinms
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

container/leaky-bucket/collector.go

@@ -0,0 +1,202 @@
+package leakybucket
+
+import (
+	"container/heap"
+	"sync"
+	"time"
+)
+
+//TODO: Finer grained locking.
+
+type bucketMap map[string]*LeakyBucket
+
+// A Collector can keep track of multiple LeakyBucket's. The caller does not
+// directly interact with the buckets, but instead addresses them by a string
+// key (e.g. IP address, hostname, hash, etc.) that is passed to most Collector
+// methods.
+//
+// All Collector methods are goroutine safe.
+type Collector struct {
+	buckets  bucketMap
+	heap     priorityQueue
+	rate     float64
+	capacity int64
+	period   time.Duration
+	lock     sync.Mutex
+	quit     chan bool
+}
+
+// NewCollector creates a new Collector. When new buckets are created within
+// the Collector, they will be assigned the capacity and rate of the Collector.
+// A Collector does not provide a way to change the rate or capacity of
+// bucket's within it. If different rates or capacities are required, either
+// use multiple Collector's or manage your own LeakyBucket's.
+//
+// If deleteEmptyBuckets is true, a concurrent goroutine will be run that
+// watches for bucket's that become empty and automatically removes them,
+// freeing up memory resources.
+func NewCollector(rate float64, capacity int64, period time.Duration, deleteEmptyBuckets bool) *Collector {
+	c := &Collector{
+		buckets:  make(bucketMap),
+		heap:     make(priorityQueue, 0, 4096),
+		rate:     rate,
+		capacity: capacity,
+		period:   period,
+		quit:     make(chan bool),
+	}
+
+	if deleteEmptyBuckets {
+		c.PeriodicPrune()
+	}
+
+	return c
+}
+
+// Free releases the collector's resources. If the collector was created with
+// deleteEmptyBuckets = true, then the goroutine looking for empty buckets,
+// will be stopped.
+func (c *Collector) Free() {
+	c.Reset()
+	close(c.quit)
+}
+
+// Reset removes all internal buckets and resets the collector back to as if it
+// was just created.
+func (c *Collector) Reset() {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	// Let the garbage collector do all the work.
+	c.buckets = make(bucketMap)
+	c.heap = make(priorityQueue, 0, 4096)
+}
+
+// Capacity returns the collector's capacity.
+func (c *Collector) Capacity() int64 {
+	return c.capacity
+}
+
+// Rate returns the collector's rate.
+func (c *Collector) Rate() float64 {
+	return c.rate
+}
+
+// Remaining returns the remaining capacity of the internal bucket associated
+// with key.  If key is not associated with a bucket internally, it is treated
+// as being empty.
+func (c *Collector) Remaining(key string) int64 {
+	return c.capacity - c.Count(key)
+}
+
+// Count returns the count of the internal bucket associated with key. If key
+// is not associated with a bucket internally, it is treated as being empty.
+func (c *Collector) Count(key string) int64 {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	b, ok := c.buckets[key]
+	if !ok || b == nil {
+		return 0
+	}
+
+	return b.Count()
+}
+
+// TillEmpty returns how much time must pass until the internal bucket
+// associated with key is empty. If key is not associated with a bucket
+// internally, it is treated as being empty.
+func (c *Collector) TillEmpty(key string) time.Duration {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	b, ok := c.buckets[key]
+	if !ok || b == nil {
+		return 0
+	}
+
+	return b.TillEmpty()
+}
+
+// Remove deletes the internal bucket associated with key. If key is not
+// associated with a bucket internally, nothing is done.
+func (c *Collector) Remove(key string) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	b, ok := c.buckets[key]
+	if !ok || b == nil {
+		return
+	}
+
+	delete(c.buckets, b.key)
+	heap.Remove(&c.heap, b.index)
+}
+
+// Add 'amount' to the internal bucket associated with key, up to it's
+// capacity. Returns how much was added to the bucket. If the return is less
+// than 'amount', then the bucket's capacity was reached.
+//
+// If key is not associated with a bucket internally, a new bucket is created
+// and amount is added to it.
+func (c *Collector) Add(key string, amount int64) int64 {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
+	b, ok := c.buckets[key]
+
+	if !ok || b == nil {
+		// Create a new bucket.
+		b = &LeakyBucket{
+			key:      key,
+			capacity: c.capacity,
+			rate:     c.rate,
+			period:   c.period,
+			p:        now(),
+		}
+		c.heap.Push(b)
+		c.buckets[key] = b
+	}
+
+	n := b.Add(amount)
+
+	if n > 0 {
+		heap.Fix(&c.heap, b.index)
+	}
+
+	return n
+}
+
+// Prune removes all empty buckets in the collector.
+func (c *Collector) Prune() {
+	c.lock.Lock()
+	for c.heap.Peak() != nil {
+		b := c.heap.Peak()
+
+		if now().Before(b.p) {
+			// The bucket isn't empty.
+			break
+		}
+
+		// The bucket should be empty.
+		delete(c.buckets, b.key)
+		heap.Remove(&c.heap, b.index)
+	}
+	c.lock.Unlock()
+}
+
+// PeriodicPrune runs a concurrent goroutine that calls Prune() at the given
+// time interval.
+func (c *Collector) PeriodicPrune() {
+	go func() {
+		ticker := time.NewTicker(c.period)
+		for {
+			select {
+			case <-ticker.C:
+				c.Prune()
+			case <-c.quit:
+				ticker.Stop()
+				return
+			}
+		}
+	}()
+}

container/leaky-bucket/collector_test.go

@@ -0,0 +1,210 @@
+package leakybucket
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestNewCollector(t *testing.T) {
+	rate := 1.0
+	capacity := int64(2)
+	c := NewCollector(rate, capacity, time.Second, true)
+
+	if c.buckets == nil {
+		t.Fatal("Didn't initialize priority?!")
+	}
+	if c.heap == nil {
+		t.Fatal("Didn't initialize priority?!")
+	}
+	if c.rate != rate || c.Rate() != rate {
+		t.Fatal("Wrong rate?!")
+	}
+	if c.capacity != capacity || c.Capacity() != capacity {
+		t.Fatal("Wrong capacity?!")
+	}
+
+	c.Free()
+}
+
+func TestNewCollector_LargerPeriod(t *testing.T) {
+	testNow := now
+	now = time.Now
+	defer func() {
+		now = testNow
+	}()
+	rate := 10.0
+	capacity := int64(20)
+	c := NewCollector(rate, capacity, 5*time.Second, true)
+
+	c.Add("test", 10)
+	c.Add("test", 10)
+
+	if c.Remaining("test") != 0 {
+		t.Errorf("Excess capacity exists of: %d", c.Remaining("test"))
+	}
+	time.Sleep(1 * time.Second)
+	if c.Remaining("test") >= 20 {
+		t.Errorf("Excess capacity exists in: %d", c.Remaining("test"))
+	}
+	time.Sleep(4 * time.Second)
+
+	if c.Add("test", 10) != 10 {
+		t.Errorf("Internal counter not refreshed: %d", c.Count("test"))
+	}
+	c.Free()
+}
+
+var collectorSimple = testSet{
+	capacity: int64(5),
+	rate:     1.0,
+	set: []actionSet{
+		{},
+		{1, "add", 1},
+		{1, "time-set", time.Nanosecond},
+		{1, "till", time.Second - time.Nanosecond},
+		{1, "time-set", time.Second - time.Nanosecond},
+		{1, "till", time.Nanosecond},
+		{0, "time-set", time.Second},
+		{0, "till", time.Duration(0)},
+		{1, "add", 1},
+		{1, "time-add", time.Second / 2},
+		{1, "till", time.Second / 2},
+		{2, "add", 1},
+		{2, "time-add", time.Second/2 - time.Nanosecond},
+		{0, "time-add", time.Second * time.Duration(5)},
+		{1, "add", 1},
+		{2, "add", 1},
+		{3, "add", 1},
+		{4, "add", 1},
+		{5, "add", 1},
+		{5, "add", 1},
+		{5, "till", time.Second * 5},
+	},
+}
+
+var collectorVaried = testSet{
+	capacity: 1000,
+	rate:     60.0,
+	set: []actionSet{
+		{},
+		{100, "add", 100},
+		{100, "time-set", time.Nanosecond},
+		{1000, "add", 1000},
+		{1000, "add", 1},
+		{940, "time-set", time.Second},
+	},
+}
+
+func runKey(t *testing.T, c *Collector, key string, test *testSet) {
+	setElapsed(0)
+	capacity := c.Capacity()
+
+	for i, v := range test.set {
+		switch v.action {
+		case "add":
+			count := c.Count(key)
+			remaining := test.capacity - count
+			amount := int64(v.value.(int))
+			n := c.Add(key, amount)
+			if n < amount {
+				// The bucket should be full now.
+				if n < remaining {
+					t.Fatalf("Test %d: Bucket should have been filled by this Add()?!", i)
+				}
+			}
+		case "time-set":
+			setElapsed(v.value.(time.Duration))
+		case "time-add":
+			addToElapsed(v.value.(time.Duration))
+		case "till":
+			dt := c.TillEmpty(key)
+			if dt != v.value.(time.Duration) {
+				t.Fatalf("%s -> Test %d: Bad TillEmpty(). Expected %v, got %v", key, i, v.value, dt)
+			}
+		}
+		count := c.Count(key)
+		if count != v.count {
+			t.Fatalf("%s -> Test %d: Bad count. Expected %d, got %d", key, i, v.count, count)
+		}
+		if count > capacity {
+			t.Fatalf("%s -> Test %d: Went over capacity?!", key, i)
+		}
+		if c.Remaining(key) != test.capacity-v.count {
+			t.Fatalf("Test %d: Expected remaining value %d, got %d",
+				i, test.capacity-v.count, c.Remaining(key))
+		}
+	}
+}
+
+func TestCollector(t *testing.T) {
+	setElapsed(0)
+
+	tests := []testSet{
+		collectorSimple,
+		collectorSimple,
+		collectorVaried,
+	}
+
+	for i, test := range tests {
+		fmt.Println("Running testSet:", i)
+
+		key := "127.0.0.1"
+
+		c := NewCollector(test.rate, test.capacity, time.Second, false)
+
+		// Run and test Remove()
+		runKey(t, c, key, &test)
+		c.Remove(key)
+		if c.Count(key) > 0 {
+			t.Fatal("Key still has a count after removal?!")
+		}
+
+		// Run again and test Prune()
+		runKey(t, c, "127.0.0.1", &test)
+		c.Prune()
+		setElapsed(time.Hour)
+		c.Prune()
+
+		// Run again and test Reset().
+		runKey(t, c, "127.0.0.1", &test)
+		c.Reset()
+		if c.Count(key) != 0 {
+			t.Fatal("Key still has a count after removal?!")
+		}
+		if c.TillEmpty(key) != 0 {
+			t.Fatal("Key still has time till empty?!")
+		}
+
+		// Try to remove a non-exist bucket.
+		c.Remove("fake")
+		if c.Count("fake") != 0 {
+			t.Fatal("Key still has a count after removal?!")
+		}
+	}
+}
+
+func TestPeriodicPrune(t *testing.T) {
+	setElapsed(0)
+	key := "localhost"
+	c := NewCollector(1e7, 8, time.Second, false)
+	c.PeriodicPrune()
+	n := c.Add(key, 100)
+	if n != 8 {
+		t.Fatal("Didn't fill bucket?!")
+	}
+
+	fmt.Printf("TillEmpty(): %v\n", c.TillEmpty(key))
+
+	// Wait for the periodic prune.
+	wait := time.Millisecond
+	time.Sleep(wait)
+	setElapsed(wait)
+
+	count := c.Count(key)
+	if count != 0 {
+		t.Fatalf("Key's bucket is not empty: %d?!", count)
+	}
+
+	c.Free()
+}

container/leaky-bucket/heap.go

@@ -0,0 +1,47 @@
+package leakybucket
+
+import "fmt"
+
+// Based on the example implementation of priority queue found in the
+// container/heap package docs: https://golang.org/pkg/container/heap/
+type priorityQueue []*LeakyBucket
+
+func (pq priorityQueue) Len() int {
+	return len(pq)
+}
+
+func (pq priorityQueue) Peak() *LeakyBucket {
+	if len(pq) <= 0 {
+		return nil
+	}
+	return pq[0]
+}
+
+func (pq priorityQueue) Less(i, j int) bool {
+	return pq[i].p.Before(pq[j].p)
+}
+
+func (pq priorityQueue) Swap(i, j int) {
+	pq[i], pq[j] = pq[j], pq[i]
+	pq[i].index = i
+	pq[j].index = j
+}
+
+func (pq *priorityQueue) Push(x interface{}) {
+	n := len(*pq)
+	b, ok := x.(*LeakyBucket)
+	if !ok {
+		panic(fmt.Sprintf("%T", x))
+	}
+	b.index = n
+	*pq = append(*pq, b)
+}
+
+func (pq *priorityQueue) Pop() interface{} {
+	old := *pq
+	n := len(old)
+	b := old[n-1]
+	b.index = -1 // for safety
+	*pq = old[0 : n-1]
+	return b
+}

container/leaky-bucket/heap_test.go

@@ -0,0 +1,107 @@
+package leakybucket
+
+import (
+	"testing"
+	"time"
+)
+
+func TestLen(t *testing.T) {
+	q := make(priorityQueue, 0, 4096)
+
+	if q.Len() != 0 {
+		t.Fatal("Queue should be empty?!")
+	}
+
+	for i := 1; i <= 5; i++ {
+		b := NewLeakyBucket(1.0, 5, time.Second)
+		q.Push(b)
+
+		l := q.Len()
+		if l != i {
+			t.Fatalf("Expected length %d, got %d", i, l)
+		}
+	}
+	for i := 4; i >= 0; i-- {
+		q.Pop()
+
+		l := q.Len()
+		if l != i {
+			t.Fatalf("Expected length %d, got %d", i, l)
+		}
+	}
+}
+
+func TestPeak(t *testing.T) {
+	q := make(priorityQueue, 0, 4096)
+
+	for i := 0; i < 5; i++ {
+		b := NewLeakyBucket(1.0, 5, time.Second)
+		q.Push(b)
+	}
+}
+
+func TestLess(t *testing.T) {
+	q := make(priorityQueue, 0, 4096)
+
+	for i := 0; i < 5; i++ {
+		b := NewLeakyBucket(1.0, 5, time.Second)
+		b.p = now().Add(time.Duration(i))
+		q.Push(b)
+	}
+
+	for i, j := 0, 4; i < 5; i, j = i+1, j-1 {
+		if i < j && !q.Less(i, j) {
+			t.Fatal("Less is more?!")
+		}
+	}
+}
+
+func TestSwap(t *testing.T) {
+	q := make(priorityQueue, 0, 4096)
+
+	for i := 0; i < 5; i++ {
+		b := NewLeakyBucket(1.0, 5, time.Second)
+		q.Push(b)
+	}
+
+	i := 2
+	j := 4
+
+	bi := q[i]
+	bj := q[j]
+
+	q.Swap(i, j)
+
+	if bi != q[j] || bj != q[i] {
+		t.Fatal("Element weren't swapped?!")
+	}
+}
+
+func TestPush(t *testing.T) {
+	q := make(priorityQueue, 0, 4096)
+
+	for i := 0; i < 5; i++ {
+		b := NewLeakyBucket(1.0, 5, time.Second)
+		q.Push(b)
+
+		if b != q[len(q)-1] {
+			t.Fatal("Push should append to queue.")
+		}
+	}
+}
+
+func TestPop(t *testing.T) {
+	q := make(priorityQueue, 0, 4096)
+
+	for i := 1; i <= 5; i++ {
+		b := NewLeakyBucket(1.0, 5, time.Second)
+		q.Push(b)
+	}
+
+	for i := 1; i <= 5; i++ {
+		b := q[len(q)-1]
+		if b != q.Pop() {
+			t.Fatal("Pop should remove from end of queue.")
+		}
+	}
+}

container/leaky-bucket/leakybucket.go

@@ -0,0 +1,155 @@
+/*
+Package leakybucket implements a scalable leaky bucket algorithm.
+
+There are at least two different definitions of the leaky bucket algorithm.
+This package implements the leaky bucket as a meter. For more details see:
+
+https://en.wikipedia.org/wiki/Leaky_bucket#As_a_meter
+
+This means it is the exact mirror of a token bucket.
+
+	// New LeakyBucket that leaks at the rate of 0.5/sec and a total capacity of 10.
+	b := NewLeakyBucket(0.5, 10)
+
+	b.Add(5)
+	b.Add(5)
+	// Bucket is now full!
+
+	n := b.Add(1)
+	// n == 0
+
+
+A Collector is a convenient way to keep track of multiple LeakyBucket's.
+Buckets are associated with string keys for fast lookup. It can dynamically
+add new buckets and automatically remove them as they become empty, freeing
+up resources.
+
+	// New Collector that leaks at 1 MiB/sec, a total capacity of 10 MiB and
+	// automatic removal of bucket's when they become empty.
+	const megabyte = 1<<20
+	c := NewCollector(megabyte, megabyte*10, true)
+
+	// Attempt to add 100 MiB to a bucket associated with an IP.
+	n := c.Add("192.168.0.42", megabyte*100)
+
+	// 100 MiB is over the capacity, so only 10 MiB is actually added.
+	// n equals 10 MiB.
+*/
+package leakybucket
+
+import (
+	"math"
+	"time"
+)
+
+// Makes it easy to test time based things.
+var now = time.Now
+
+// LeakyBucket represents a bucket that leaks at a constant rate.
+type LeakyBucket struct {
+	// The identifying key, used for map lookups.
+	key string
+
+	// How large the bucket is.
+	capacity int64
+
+	// Amount the bucket leaks per time duration.
+	rate float64
+
+	// The priority of the bucket in a min-heap priority queue, where p is the
+	// exact time the bucket will have leaked enough to be empty. Buckets that
+	// are empty or will be the soonest are at the top of the heap. This allows
+	// for quick pruning of empty buckets that scales very well. p is adjusted
+	// any time an amount is added to the Queue().
+	p time.Time
+
+	// The time duration through which the leaky bucket is
+	// assessed.
+	period time.Duration
+
+	// The index is maintained by the heap.Interface methods.
+	index int
+}
+
+// NewLeakyBucket creates a new LeakyBucket with the give rate and capacity.
+func NewLeakyBucket(rate float64, capacity int64, period time.Duration) *LeakyBucket {
+	return &LeakyBucket{
+		rate:     rate,
+		capacity: capacity,
+		period:   period,
+		p:        now(),
+	}
+}
+
+// Count returns the bucket's current count.
+func (b *LeakyBucket) Count() int64 {
+	if !now().Before(b.p) {
+		return 0
+	}
+
+	nsRemaining := float64(b.p.Sub(now()))
+	nsPerDrip := float64(b.period) / b.rate
+	count := int64(math.Ceil(nsRemaining / nsPerDrip))
+
+	return count
+}
+
+// Rate returns the amount the bucket leaks per second.
+func (b *LeakyBucket) Rate() float64 {
+	return b.rate
+}
+
+// Capacity returns the bucket's capacity.
+func (b *LeakyBucket) Capacity() int64 {
+	return b.capacity
+}
+
+// Remaining returns the bucket's remaining capacity.
+func (b *LeakyBucket) Remaining() int64 {
+	return b.capacity - b.Count()
+}
+
+// ChangeCapacity changes the bucket's capacity.
+//
+// If the bucket's current count is greater than the new capacity, the count
+// will be decreased to match the new capacity.
+func (b *LeakyBucket) ChangeCapacity(capacity int64) {
+	diff := float64(capacity - b.capacity)
+
+	if diff < 0 && b.Count() > capacity {
+		// We are shrinking the capacity and the new bucket size can't hold all
+		// the current contents. Dump the extra and adjust the time till empty.
+		nsPerDrip := float64(b.period) / b.rate
+		b.p = now().Add(time.Duration(nsPerDrip * float64(capacity)))
+	}
+	b.capacity = capacity
+}
+
+// TillEmpty returns how much time must pass until the bucket is empty.
+func (b *LeakyBucket) TillEmpty() time.Duration {
+	return b.p.Sub(now())
+}
+
+// Add 'amount' to the bucket's count, up to it's capacity. Returns how much
+// was added to the bucket. If the return is less than 'amount', then the
+// bucket's capacity was reached.
+func (b *LeakyBucket) Add(amount int64) int64 {
+	count := b.Count()
+	if count >= b.capacity {
+		// The bucket is full.
+		return 0
+	}
+
+	if !now().Before(b.p) {
+		// The bucket needs to be reset.
+		b.p = now()
+	}
+	remaining := b.capacity - count
+	if amount > remaining {
+		amount = remaining
+	}
+	t := time.Duration(float64(b.period) * (float64(amount) / b.rate))
+	b.p = b.p.Add(t)
+
+	return amount
+}

container/leaky-bucket/leakybucket_test.go

@@ -0,0 +1,181 @@
+package leakybucket
+
+import (
+	"fmt"
+	"os"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+// Arbitrary start time.
+var start = time.Date(1990, 1, 2, 0, 0, 0, 0, time.UTC).Round(0)
+var elapsed int64
+
+// We provide atomic access to elapsed to avoid data races between multiple
+// concurrent goroutines during the tests.
+func getElapsed() time.Duration {
+	return time.Duration(atomic.LoadInt64(&elapsed))
+}
+func setElapsed(v time.Duration) {
+	atomic.StoreInt64(&elapsed, int64(v))
+}
+func addToElapsed(v time.Duration) {
+	atomic.AddInt64(&elapsed, int64(v))
+}
+
+func reset(t *testing.T, c *Collector) {
+	c.Reset()
+	setElapsed(0)
+}
+
+func TestNewLeakyBucket(t *testing.T) {
+	rate := 1.0
+	capacity := int64(5)
+	b := NewLeakyBucket(rate, capacity, time.Second)
+
+	if b.p != now() {
+		t.Fatal("Didn't initialize priority?!")
+	}
+	if b.rate != rate || b.Rate() != rate {
+		t.Fatal("Wrong rate?!")
+	}
+	if b.capacity != capacity || b.Capacity() != capacity {
+		t.Fatal("Wrong capacity?!")
+	}
+}
+
+type actionSet struct {
+	count  int64
+	action string
+	value  interface{}
+}
+
+type testSet struct {
+	capacity int64
+	rate     float64
+	set      []actionSet
+}
+
+var oneAtaTime = testSet{
+	capacity: 5,
+	rate:     1.0,
+	set: []actionSet{
+		{},
+		{1, "add", 1},
+		{1, "time-set", time.Nanosecond},
+		{1, "till", time.Second - time.Nanosecond},
+		{1, "time-set", time.Second - time.Nanosecond},
+		{1, "till", time.Nanosecond},
+		{0, "time-set", time.Second},
+		{0, "till", time.Duration(0)},
+
+		// Add a couple.
+		{1, "add", 1},
+		{1, "time-add", time.Second / 2},
+		{1, "till", time.Second / 2},
+		{2, "add", 1},
+		{2, "time-add", time.Second/2 - time.Nanosecond},
+
+		// Monkey with the capacity and make sure Count()/TillEmpty() are
+		// adjusted as needed.
+		{2, "cap", 5 + 1},
+		{2, "till", time.Second + time.Nanosecond},
+		{2, "cap", 5 - 1},
+		{2, "till", time.Second + time.Nanosecond},
+		{1, "cap", 1},
+		{1, "till", time.Second},
+		{1, "cap", 4},
+		{1, "till", time.Second},
+
+		// Test the full cases.
+		{0, "time-add", time.Second * time.Duration(5)},
+		{1, "add", 1},
+		{2, "add", 1},
+		{3, "add", 1},
+		{4, "add", 1},
+		{4, "add", 1},
+		{4, "till", time.Second * 4},
+	},
+}
+
+var varied = testSet{
+	capacity: 1000,
+	rate:     60.0,
+	set: []actionSet{
+		{},
+		{100, "add", 100},
+		{100, "time-set", time.Nanosecond},
+		{1000, "add", 1000},
+		{1000, "add", 1},
+		{940, "time-set", time.Second},
+	},
+}
+
+func runTest(t *testing.T, test *testSet) {
+	setElapsed(0)
+	b := NewLeakyBucket(test.rate, test.capacity, time.Second)
+
+	for i, v := range test.set {
+		switch v.action {
+		case "add":
+			count := b.Count()
+			remaining := test.capacity - count
+			amount := int64(v.value.(int))
+			n := b.Add(amount)
+			if n < amount {
+				// The bucket should be full now.
+				if n < remaining {
+					t.Fatalf("Test %d: Bucket should have been filled by this Add()?!", i)
+				}
+			}
+		case "time-set":
+			setElapsed(v.value.(time.Duration))
+		case "cap":
+			b.ChangeCapacity(int64(v.value.(int)))
+			test.capacity = b.Capacity()
+		case "time-add":
+			addToElapsed(v.value.(time.Duration))
+		case "till":
+			dt := b.TillEmpty()
+			if dt != v.value.(time.Duration) {
+				t.Fatalf("Test %d: Bad TillEmpty(). Expected %v, got %v", i, v.value, dt)
+			}
+		case "debug":
+			fmt.Println("elapsed:", getElapsed())
+			fmt.Println("tillEmpty:", b.TillEmpty())
+			fmt.Println("count:", b.Count())
+		}
+		c := b.Count()
+		if c != v.count {
+			t.Fatalf("Test %d: Bad count. Expected %d, got %d", i, v.count, c)
+		}
+		if c > test.capacity {
+			t.Fatalf("Test %d: Went over capacity?!", i)
+		}
+		if b.Remaining() != test.capacity-v.count {
+			t.Fatalf("Test %d: Expected remaining value %d, got %d",
+				i, test.capacity-v.count, b.Remaining())
+		}
+	}
+}
+
+func TestLeakyBucket(t *testing.T) {
+	tests := []testSet{
+		oneAtaTime,
+		varied,
+	}
+
+	for i, test := range tests {
+		fmt.Println("Running testSet:", i)
+		runTest(t, &test)
+	}
+}
+
+func TestMain(m *testing.M) {
+	// Override what now() function the leakybucket library uses.
+	// This greatly increases testability!
+	now = func() time.Time { return start.Add(getElapsed()) }
+
+	os.Exit(m.Run())
+}