From 9634526a67b0aa553b5a1275987a8ae1ba8751b2 Mon Sep 17 00:00:00 2001
From: Adam Harrison <harrisonadamw@gmail.com>
Date: Thu, 2 Jul 2015 11:58:40 +0100
Subject: [PATCH] Add token bucket implementation

---
 router/token_bucket.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)
 create mode 100644 router/token_bucket.go

diff --git a/router/token_bucket.go b/router/token_bucket.go
new file mode 100644
index 00000000..254de6eb
--- /dev/null
+++ b/router/token_bucket.go
@@ -0,0 +1,42 @@
+package router
+
+import (
+	"time"
+)
+
+type TokenBucket struct {
+	capacity             int64         // Maximum capacity of bucket
+	tokenInterval        time.Duration // Token replenishment rate
+	refillDuration       time.Duration // Time to refill from empty
+	earliestUnspentToken time.Time
+}
+
+func NewTokenBucket(capacity int64, tokenInterval time.Duration) *TokenBucket {
+	tb := TokenBucket{
+		capacity:       capacity,
+		tokenInterval:  tokenInterval,
+		refillDuration: tokenInterval * time.Duration(capacity)}
+
+	tb.earliestUnspentToken = tb.capacityToken()
+
+	return &tb
+}
+
+func (tb *TokenBucket) Wait() {
+	// If earliest unspent token is in the future, sleep until then
+	time.Sleep(tb.earliestUnspentToken.Sub(time.Now()))
+
+	// Alternatively, enforce bucket capacity if necessary
+	capacityToken := tb.capacityToken()
+	if tb.earliestUnspentToken.Before(capacityToken) {
+		tb.earliestUnspentToken = capacityToken
+	}
+
+	// 'Remove' a token from the bucket
+	tb.earliestUnspentToken = tb.earliestUnspentToken.Add(tb.tokenInterval)
+}
+
+// Determine the historic token timestamp representing a full bucket
+func (tb *TokenBucket) capacityToken() time.Time {
+	return time.Now().Add(-tb.refillDuration).Truncate(tb.tokenInterval)
+}