From ae36b93bc275784b8dce6b7290341972515c2ac7 Mon Sep 17 00:00:00 2001
From: Pawel Krolikowski <pawel@uber.com>
Date: Sat, 9 Jul 2022 23:03:39 -0700
Subject: [PATCH 1/2] Add a test verifying initial startup sequence

See https://github.com/uber-go/ratelimit/pull/95#discussion_r915251700

From that discussion I wasn't sure whether the proposed the initial
startup sequence of the limiter - i.e. whether at startup we always
block, or always allow.

Since we didn't seem to have that codified (perhaps apart from the
`example_test.go`) this PR adds a test to verify this.

This is still slightly (2/1000) flaky, but I think that's good enough
to add this in - should be valuable anyway.
---
 ratelimit_test.go | 66 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/ratelimit_test.go b/ratelimit_test.go
index 16f0a19..6bc26bf 100644
--- a/ratelimit_test.go
+++ b/ratelimit_test.go
@@ -22,6 +22,8 @@ type testRunner interface {
 	// afterFunc executes a func at a given time.
 	// not using clock.AfterFunc because andres-erbsen/clock misses a nap there.
 	afterFunc(d time.Duration, fn func())
+	// some tests want raw access to the clock.
+	getClock() *clock.Mock
 }
 
 type runnerImpl struct {
@@ -89,6 +91,10 @@ func (r *runnerImpl) createLimiter(rate int, opts ...Option) Limiter {
 	return r.constructor(rate, opts...)
 }
 
+func (r *runnerImpl) getClock() *clock.Mock {
+	return r.clock
+}
+
 // startTaking tries to Take() on passed in limiters in a loop/goroutine.
 func (r *runnerImpl) startTaking(rls ...Limiter) {
 	r.goWait(func() {
@@ -198,6 +204,66 @@ func TestPer(t *testing.T) {
 	})
 }
 
+// TestInitial verifies that the initial sequence is scheduled as expected.
+func TestInitial(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		msg  string
+		opts []Option
+	}{
+		{
+			msg: "With Slack",
+		},
+		{
+			msg:  "Without Slack",
+			opts: []Option{WithoutSlack},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.msg, func(t *testing.T) {
+			runTest(t, func(r testRunner) {
+				rl := r.createLimiter(10, tt.opts...)
+
+				var (
+					clk            = r.getClock()
+					prev time.Time = clk.Now()
+
+					results         []time.Duration
+					startWg, doneWg sync.WaitGroup
+				)
+				startWg.Add(3)
+				doneWg.Add(3)
+
+				for i := 0; i < 3; i++ {
+					go func() {
+						startWg.Done()
+						ts := rl.Take()
+						// Below changes while "racy" should be synchronized by the limiter.
+						results = append(results, ts.Sub(prev))
+						prev = ts
+						doneWg.Done()
+					}()
+				}
+
+				startWg.Wait()
+				clk.Add(time.Second)
+				doneWg.Wait()
+
+				assert.Equal(t,
+					[]time.Duration{
+						0,
+						time.Millisecond * 100,
+						time.Millisecond * 100,
+					},
+					results,
+					"bad timestamps for inital takes",
+				)
+			})
+		})
+	}
+}
+
 func TestSlack(t *testing.T) {
 	// To simulate slack, we combine two limiters.
 	// - First, we start a single goroutine with both of them,

From a0a8d91cf8ea649f3f2dd23fda324afbf5723373 Mon Sep 17 00:00:00 2001
From: Pawel Krolikowski <pawel@uber.com>
Date: Sat, 9 Jul 2022 23:55:15 -0700
Subject: [PATCH 2/2] channels are great

---
 ratelimit_test.go | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/ratelimit_test.go b/ratelimit_test.go
index 6bc26bf..2ed72d6 100644
--- a/ratelimit_test.go
+++ b/ratelimit_test.go
@@ -226,29 +226,30 @@ func TestInitial(t *testing.T) {
 				rl := r.createLimiter(10, tt.opts...)
 
 				var (
-					clk            = r.getClock()
-					prev time.Time = clk.Now()
+					clk  = r.getClock()
+					prev = clk.Now()
 
-					results         []time.Duration
-					startWg, doneWg sync.WaitGroup
+					results = make(chan time.Time)
+					have    []time.Duration
+					startWg sync.WaitGroup
 				)
 				startWg.Add(3)
-				doneWg.Add(3)
 
 				for i := 0; i < 3; i++ {
 					go func() {
 						startWg.Done()
-						ts := rl.Take()
-						// Below changes while "racy" should be synchronized by the limiter.
-						results = append(results, ts.Sub(prev))
-						prev = ts
-						doneWg.Done()
+						results <- rl.Take()
 					}()
 				}
 
 				startWg.Wait()
 				clk.Add(time.Second)
-				doneWg.Wait()
+
+				for i := 0; i < 3; i++ {
+					ts := <-results
+					have = append(have, ts.Sub(prev))
+					prev = ts
+				}
 
 				assert.Equal(t,
 					[]time.Duration{
@@ -256,7 +257,7 @@ func TestInitial(t *testing.T) {
 						time.Millisecond * 100,
 						time.Millisecond * 100,
 					},
-					results,
+					have,
 					"bad timestamps for inital takes",
 				)
 			})