Rate limiting for frontends

vendor/github.com/vulcand/oxy/ratelimit/bucketset.go

@@ -0,0 +1,108 @@
+package ratelimit
+
+import (
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/mailgun/timetools"
+	"sort"
+)
+
+// TokenBucketSet represents a set of TokenBucket covering different time periods.
+type TokenBucketSet struct {
+	buckets   map[time.Duration]*tokenBucket
+	maxPeriod time.Duration
+	clock     timetools.TimeProvider
+}
+
+// newTokenBucketSet creates a `TokenBucketSet` from the specified `rates`.
+func NewTokenBucketSet(rates *RateSet, clock timetools.TimeProvider) *TokenBucketSet {
+	tbs := new(TokenBucketSet)
+	tbs.clock = clock
+	// In the majority of cases we will have only one bucket.
+	tbs.buckets = make(map[time.Duration]*tokenBucket, len(rates.m))
+	for _, rate := range rates.m {
+		newBucket := newTokenBucket(rate, clock)
+		tbs.buckets[rate.period] = newBucket
+		tbs.maxPeriod = maxDuration(tbs.maxPeriod, rate.period)
+	}
+	return tbs
+}
+
+// Update brings the buckets in the set in accordance with the provided `rates`.
+func (tbs *TokenBucketSet) Update(rates *RateSet) {
+	// Update existing buckets and delete those that have no corresponding spec.
+	for _, bucket := range tbs.buckets {
+		if rate, ok := rates.m[bucket.period]; ok {
+			bucket.update(rate)
+		} else {
+			delete(tbs.buckets, bucket.period)
+		}
+	}
+	// Add missing buckets.
+	for _, rate := range rates.m {
+		if _, ok := tbs.buckets[rate.period]; !ok {
+			newBucket := newTokenBucket(rate, tbs.clock)
+			tbs.buckets[rate.period] = newBucket
+		}
+	}
+	// Identify the maximum period in the set
+	tbs.maxPeriod = 0
+	for _, bucket := range tbs.buckets {
+		tbs.maxPeriod = maxDuration(tbs.maxPeriod, bucket.period)
+	}
+}
+
+func (tbs *TokenBucketSet) Consume(tokens int64) (time.Duration, error) {
+	var maxDelay time.Duration = UndefinedDelay
+	var firstErr error = nil
+	for _, tokenBucket := range tbs.buckets {
+		// We keep calling `Consume` even after a error is returned for one of
+		// buckets because that allows us to simplify the rollback procedure,
+		// that is to just call `Rollback` for all buckets.
+		delay, err := tokenBucket.consume(tokens)
+		if firstErr == nil {
+			if err != nil {
+				firstErr = err
+			} else {
+				maxDelay = maxDuration(maxDelay, delay)
+			}
+		}
+	}
+	// If we could not make ALL buckets consume tokens for whatever reason,
+	// then rollback consumption for all of them.
+	if firstErr != nil || maxDelay > 0 {
+		for _, tokenBucket := range tbs.buckets {
+			tokenBucket.rollback()
+		}
+	}
+	return maxDelay, firstErr
+}
+
+func (tbs *TokenBucketSet) GetMaxPeriod() time.Duration {
+	return tbs.maxPeriod
+}
+
+// debugState returns string that reflects the current state of all buckets in
+// this set. It is intended to be used for debugging and testing only.
+func (tbs *TokenBucketSet) debugState() string {
+	periods := sort.IntSlice(make([]int, 0, len(tbs.buckets)))
+	for period := range tbs.buckets {
+		periods = append(periods, int(period))
+	}
+	sort.Sort(periods)
+	bucketRepr := make([]string, 0, len(tbs.buckets))
+	for _, period := range periods {
+		bucket := tbs.buckets[time.Duration(period)]
+		bucketRepr = append(bucketRepr, fmt.Sprintf("{%v: %v}", bucket.period, bucket.availableTokens))
+	}
+	return strings.Join(bucketRepr, ", ")
+}
+
+func maxDuration(x time.Duration, y time.Duration) time.Duration {
+	if x > y {
+		return x
+	}
+	return y
+}