Initial implementation of key-value rate limits

ratelimits/gcra.go

@@ -0,0 +1,121 @@
+package ratelimits
+
+import (
+	"math"
+	"time"
+
+	"github.com/jmhodges/clock"
+)
+
+// divThenRound divides two int64s and rounds the result to the nearest integer.
+// This is used to calculate request intervals and costs in nanoseconds.
+func divThenRound(x, y int64) int64 {
+	return int64(math.Round(float64(x) / float64(y)))
+}
+
+// maybeSpend uses the GCRA algorithm to decide whether to allow a request. It
+// returns a Decision struct with the result of the decision and the updated
+// TAT. The cost must be 0 or greater and <= the burst capacity of the limit.
+func maybeSpend(clk clock.Clock, rl limit, tat time.Time, cost int64) *Decision {
+	if cost < 0 || cost > rl.Burst {
+		panic("invalid cost for maybeSpend")
+	}
+	nowUnix := clk.Now().UnixNano()
+	tatUnix := tat.UnixNano()
+
+	// If the TAT is in the future, use it as the starting point for the
+	// calculation. Otherwise, use the current time. This is to prevent the
+	// bucket from being filled with capacity from the past.
+	if nowUnix > tatUnix {
+		tatUnix = nowUnix
+	}
+
+	// Compute the cost increment.
+	emissionInterval := divThenRound(rl.Period.Nanoseconds(), rl.Count)
+	costIncrement := emissionInterval * cost
+
+	// Deduct the cost to find the new TAT and residual capacity.
+	newTAT := tatUnix + costIncrement
+	burstOffset := emissionInterval * rl.Burst
+	difference := nowUnix - (newTAT - burstOffset)
+	residual := divThenRound(difference, emissionInterval)
+
+	if costIncrement <= 0 && residual == 0 {
+		// Edge case: no cost to consume and no capacity to consume it from.
+		return &Decision{
+			Allowed:   false,
+			Remaining: 0,
+			RetryIn:   time.Duration(emissionInterval),
+			ResetIn:   time.Duration(tatUnix - nowUnix),
+			newTAT:    time.Unix(0, tatUnix).UTC(),
+		}
+	}
+
+	if residual < 0 {
+		// Too little capacity to satisfy the cost, deny the request.
+		remaining := divThenRound(nowUnix-(tatUnix-burstOffset), emissionInterval)
+		return &Decision{
+			Allowed:   false,
+			Remaining: int(remaining),
+			RetryIn:   -time.Duration(difference),
+			ResetIn:   time.Duration(tatUnix - nowUnix),
+			newTAT:    time.Unix(0, tatUnix).UTC(),
+		}
+	}
+
+	// There is enough capacity to satisfy the cost, allow the request.
+	var retryIn time.Duration
+	if residual == 0 {
+		// This request will empty the bucket.
+		retryIn = time.Duration(emissionInterval)
+	}
+	return &Decision{
+		Allowed:   true,
+		Remaining: int(residual),
+		RetryIn:   retryIn,
+		ResetIn:   time.Duration(newTAT - nowUnix),
+		newTAT:    time.Unix(0, newTAT).UTC(),
+	}
+}
+
+// maybeRefund uses the Generic Cell Rate Algorithm (GCRA) to attempt to refund
+// the cost of a request which was previously spent. The refund cost must be 0
+// or greater. A cost will only be refunded up to the burst capacity of the
+// limit. A partial refund is still considered successful.
+func maybeRefund(clk clock.Clock, rl limit, tat time.Time, cost int64) *Decision {
+	if cost <= 0 || cost > rl.Burst {
+		panic("invalid cost for maybeRefund")
+	}
+	nowUnix := clk.Now().UnixNano()
+	tatUnix := tat.UnixNano()
+
+	// If the TAT is in the past, use the current time as the starting point.
+	if nowUnix > tatUnix {
+		tatUnix = nowUnix
+	}
+
+	// Compute the refund increment.
+	emissionInterval := divThenRound(rl.Period.Nanoseconds(), rl.Count)
+	refundIncrement := emissionInterval * cost
+
+	// Subtract the refund increment from the TAT to find the new TAT.
+	newTAT := tatUnix - refundIncrement
+
+	// Ensure the new TAT is not earlier than now.
+	if newTAT < nowUnix {
+		newTAT = nowUnix
+	}
+
+	// Calculate the new capacity.
+	burstOffset := emissionInterval * rl.Burst
+	difference := nowUnix - (newTAT - burstOffset)
+	residual := divThenRound(difference, emissionInterval)
+
+	return &Decision{
+		Allowed:   (newTAT != tatUnix),
+		Remaining: int(residual),
+		RetryIn:   time.Duration(0),
+		ResetIn:   time.Duration(newTAT - nowUnix),
+		newTAT:    time.Unix(0, newTAT).UTC(),
+	}
+}

ratelimits/gcra_test.go

@@ -0,0 +1,174 @@
+package ratelimits
+
+import (
+	"testing"
+	"time"
+
+	"github.com/jmhodges/clock"
+	"github.com/letsencrypt/boulder/config"
+	"github.com/letsencrypt/boulder/test"
+)
+
+func Test_decide(t *testing.T) {
+	clk := clock.NewFake()
+	limit := limit{Burst: 10, Count: 1, Period: config.Duration{Duration: time.Second}}
+
+	// Begin by using 1 of our 10 requests.
+	d := maybeSpend(clk, limit, clk.Now(), 1)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 9)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Second)
+
+	// Immediately use another 9 of our remaining requests.
+	d = maybeSpend(clk, limit, d.newTAT, 9)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+
+	// Our new TAT should be 10 seconds (limit.Burst) in the future.
+	test.AssertEquals(t, d.newTAT, clk.Now().Add(time.Second*10))
+
+	// Let's try using just 1 more request without waiting.
+	d = maybeSpend(clk, limit, d.newTAT, 1)
+	test.Assert(t, !d.Allowed, "should not be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+
+	// Let's try being exactly as patient as we're told to be.
+	clk.Add(d.RetryIn)
+
+	// We are 1 second in the future, we should have 1 new request.
+	d = maybeSpend(clk, limit, d.newTAT, 1)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+
+	// Let's try waiting (10 seconds) for our whole bucket to refill.
+	clk.Add(d.ResetIn)
+
+	// We should have 10 new requests. If we use 1 we should have 9 remaining.
+	d = maybeSpend(clk, limit, d.newTAT, 1)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 9)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Second)
+
+	// Have you ever tried spending 0, like, just to see what happens?
+	d = maybeSpend(clk, limit, d.newTAT, 0)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 9)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Second)
+
+	// Spending 0 simply informed us that we still have 9 remaining, let's see
+	// what we have after waiting 20 hours.
+	clk.Add(20 * time.Hour)
+
+	// C'mon, big money, no whammies, no whammies, STOP!
+	d = maybeSpend(clk, limit, d.newTAT, 0)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 10)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Duration(0))
+
+	// Turns out that the most we can accrue is 10 (limit.Burst). Let's empty
+	// this bucket out so we can try something else.
+	d = maybeSpend(clk, limit, d.newTAT, 10)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+
+	// If you spend 0 while you have 0 you should get 0.
+	d = maybeSpend(clk, limit, d.newTAT, 0)
+	test.Assert(t, !d.Allowed, "should not be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+
+	// We don't play by the rules, we spend 1 when we have 0.
+	d = maybeSpend(clk, limit, d.newTAT, 1)
+	test.Assert(t, !d.Allowed, "should not be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+
+	// Okay, maybe we should play by the rules if we want to get anywhere.
+	clk.Add(d.RetryIn)
+
+	// Our patience pays off, we should have 1 new request. Let's use it.
+	d = maybeSpend(clk, limit, d.newTAT, 1)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+}
+
+func Test_maybeRefund(t *testing.T) {
+	clk := clock.NewFake()
+	limit := limit{Burst: 10, Count: 1, Period: config.Duration{Duration: time.Second}}
+
+	// Begin by using 1 of our 10 requests.
+	d := maybeSpend(clk, limit, clk.Now(), 1)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 9)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Second)
+
+	// Refund back to 10.
+	d = maybeRefund(clk, limit, d.newTAT, 1)
+	test.AssertEquals(t, d.Remaining, 10)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Duration(0))
+
+	// Spend 1 more of our 10 requests.
+	d = maybeSpend(clk, limit, d.newTAT, 1)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 9)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Second)
+
+	// Wait for our bucket to refill.
+	clk.Add(d.ResetIn)
+
+	// Attempt to refund from 10 to 11.
+	d = maybeRefund(clk, limit, d.newTAT, 1)
+	test.Assert(t, !d.Allowed, "should not be allowed")
+	test.AssertEquals(t, d.Remaining, 10)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Duration(0))
+
+	// Spend 10 all 10 of our requests.
+	d = maybeSpend(clk, limit, d.newTAT, 10)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 0)
+	test.AssertEquals(t, d.RetryIn, time.Second)
+	test.AssertEquals(t, d.ResetIn, time.Second*10)
+
+	// Attempt a refund of 10.
+	d = maybeRefund(clk, limit, d.newTAT, 10)
+	test.AssertEquals(t, d.Remaining, 10)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Duration(0))
+
+	// Wait 11 seconds to catching up to TAT.
+	clk.Add(11 * time.Second)
+
+	// Attempt to refund to 11, then ensure it's still 10.
+	d = maybeRefund(clk, limit, d.newTAT, 1)
+	test.Assert(t, !d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 10)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+	test.AssertEquals(t, d.ResetIn, time.Duration(0))
+
+	// Spend 5 of our 10 requests, then refund 1.
+	d = maybeSpend(clk, limit, d.newTAT, 5)
+	d = maybeRefund(clk, limit, d.newTAT, 1)
+	test.Assert(t, d.Allowed, "should be allowed")
+	test.AssertEquals(t, d.Remaining, 6)
+	test.AssertEquals(t, d.RetryIn, time.Duration(0))
+}