From 4e70b7fa9a1720f48c56a51e95ce3e6e63bd76d6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82?= <pawel@uber.com>
Date: Tue, 2 Mar 2021 21:31:26 -0400
Subject: [PATCH] Allow configuring custom slack (#64)

Resolves #20

This is to allow people to configure custom slack - we previously
defaulted to 10, and only allowed to disable it.

The tests are a bit painful to read, I had a stab at generating
the same set of tests for both with, and without per, but that was
even more hard to read.
---
 ratelimit.go      |  7 +++++++
 ratelimit_test.go | 40 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/ratelimit.go b/ratelimit.go
index 89142bd..b5b16e5 100644
--- a/ratelimit.go
+++ b/ratelimit.go
@@ -100,6 +100,13 @@ func (o slackOption) apply(c *config) {
 // previously "unspent" requests for future bursts of traffic.
 var WithoutSlack Option = slackOption(0)
 
+// WithSlack configures custom slack.
+// Slack allows the limiter to accumulate "unspent" requests
+// for future bursts of traffic.
+func WithSlack(slack int) Option {
+	return slackOption(slack)
+}
+
 type perOption time.Duration
 
 func (p perOption) apply(c *config) {
diff --git a/ratelimit_test.go b/ratelimit_test.go
index a9e1cc8..2b7b56c 100644
--- a/ratelimit_test.go
+++ b/ratelimit_test.go
@@ -202,7 +202,6 @@ func TestSlack(t *testing.T) {
 	//   - faster limiter running for 1 second
 	//   - slack accumulated by the faster limiter during the two seconds.
 	//     it was blocked by slower limiter.
-
 	tests := []struct {
 		msg  string
 		opt  []Option
@@ -213,12 +212,50 @@ func TestSlack(t *testing.T) {
 			// 2*10 + 1*100 + 1*10 (slack)
 			want: 130,
 		},
+		{
+			msg: "slack of 10, like default",
+			opt: []Option{WithSlack(10)},
+			// 2*10 + 1*100 + 1*10 (slack)
+			want: 130,
+		},
+		{
+			msg: "slack of 20",
+			opt: []Option{WithSlack(20)},
+			// 2*10 + 1*100 + 1*20 (slack)
+			want: 140,
+		},
+		{
+			// Note this is bigger then the rate of the limiter.
+			msg: "slack of 150",
+			opt: []Option{WithSlack(150)},
+			// 2*10 + 1*100 + 1*150 (slack)
+			want: 270,
+		},
 		{
 			msg: "no option, defaults to 10, with per",
 			// 2*(10*2) + 1*(100*2) + 1*10 (slack)
 			opt:  []Option{Per(500 * time.Millisecond)},
 			want: 230,
 		},
+		{
+			msg: "slack of 10, like default, with per",
+			opt: []Option{WithSlack(10), Per(500 * time.Millisecond)},
+			// 2*(10*2) + 1*(100*2) + 1*10 (slack)
+			want: 230,
+		},
+		{
+			msg: "slack of 20, with per",
+			opt: []Option{WithSlack(20), Per(500 * time.Millisecond)},
+			// 2*(10*2) + 1*(100*2) + 1*20 (slack)
+			want: 240,
+		},
+		{
+			// Note this is bigger then the rate of the limiter.
+			msg: "slack of 150, with per",
+			opt: []Option{WithSlack(150), Per(500 * time.Millisecond)},
+			// 2*(10*2) + 1*(100*2) + 1*150 (slack)
+			want: 370,
+		},
 	}
 
 	for _, tt := range tests {
@@ -240,5 +277,4 @@ func TestSlack(t *testing.T) {
 			})
 		})
 	}
-
 }