From 6a1ccb5f74abf95116363fc93269a29619b64855 Mon Sep 17 00:00:00 2001
From: Sebastian Machuca <sebastian@hachiko.energy>
Date: Wed, 31 Dec 2025 14:36:18 +1100
Subject: [PATCH 1/2] Add concurrency support

---
 token/rate.go      |  5 +++++
 token/rate_test.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/token/rate.go b/token/rate.go
index d9fc558..23aa43c 100644
--- a/token/rate.go
+++ b/token/rate.go
@@ -2,6 +2,7 @@ package token
 
 import (
 	"errors"
+	"sync"
 	"time"
 )
 
@@ -18,6 +19,7 @@ func (c realClock) Now() time.Time {
 // Limiter implements a token bucket rate limiter. It allows a burst of
 // requests up to capacity, then refills tokens at the specified rate per second.
 type Limiter struct {
+	mu                     sync.Mutex
 	capacity, tokens, rate float64
 	lastRefillAt           time.Time
 	clock                  clock
@@ -54,6 +56,9 @@ func NewLimiterWithClock(capacity, rate float64, clock clock) (*Limiter, error)
 // available and returns true. If no tokens are available, it returns false
 // without blocking.
 func (lim *Limiter) Allow() bool {
+	lim.mu.Lock()
+	defer lim.mu.Unlock()
+
 	lim.refill()
 
 	if lim.tokens >= 1 {
diff --git a/token/rate_test.go b/token/rate_test.go
index b2bac05..abca8a4 100644
--- a/token/rate_test.go
+++ b/token/rate_test.go
@@ -1,6 +1,8 @@
 package token_test
 
 import (
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -106,3 +108,48 @@ func TestLimiter_Allow(t *testing.T) {
 		})
 	}
 }
+
+func TestLimiter_Allow_Concurrent(t *testing.T) {
+	lim, err := token.NewLimiter(100, 0)
+	require.NoError(t, err)
+
+	var allowed atomic.Int64
+	var wg sync.WaitGroup
+
+	// Launch 200 goroutines, but only 100 should be allowed
+	for range 200 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if lim.Allow() {
+				allowed.Add(1)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	// With capacity 100 and rate 0, exactly 100 should be allowed
+	require.Equal(t, int64(100), allowed.Load(), "expected exactly 100 requests to be allowed")
+}
+
+func TestLimiter_Allow_ConcurrentWithRefill(t *testing.T) {
+	lim, err := token.NewLimiter(10, 1000)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+
+	// Hammer the limiter from multiple goroutines
+	for range 100 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for range 100 {
+				lim.Allow()
+			}
+		}()
+	}
+
+	wg.Wait()
+	// If we get here without race detector complaints, the test passes
+}

From 7bbd142760fbf42db6dca67b89b88662248c2043 Mon Sep 17 00:00:00 2001
From: Sebastian Machuca <sebastian@hachiko.energy>
Date: Wed, 31 Dec 2025 14:38:09 +1100
Subject: [PATCH 2/2] lint

---
 token/rate_test.go | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/token/rate_test.go b/token/rate_test.go
index abca8a4..f551420 100644
--- a/token/rate_test.go
+++ b/token/rate_test.go
@@ -113,14 +113,19 @@ func TestLimiter_Allow_Concurrent(t *testing.T) {
 	lim, err := token.NewLimiter(100, 0)
 	require.NoError(t, err)
 
-	var allowed atomic.Int64
-	var wg sync.WaitGroup
+	var (
+		allowed atomic.Int64
+		wg      sync.WaitGroup
+	)
 
 	// Launch 200 goroutines, but only 100 should be allowed
+
 	for range 200 {
 		wg.Add(1)
+
 		go func() {
 			defer wg.Done()
+
 			if lim.Allow() {
 				allowed.Add(1)
 			}
@@ -142,8 +147,10 @@ func TestLimiter_Allow_ConcurrentWithRefill(t *testing.T) {
 	// Hammer the limiter from multiple goroutines
 	for range 100 {
 		wg.Add(1)
+
 		go func() {
 			defer wg.Done()
+
 			for range 100 {
 				lim.Allow()
 			}