Replacing signed int with agnostic uint (#6)

ratelimit.go

@@ -8,16 +8,16 @@ import (
 
 // Limiter allows a burst of request during the defined duration
 type Limiter struct {
-	maxCount int64
-	count    int64
+	maxCount uint
+	count    uint
 	ticker   *time.Ticker
 	tokens   chan struct{}
 	ctx      context.Context
 }
 
 func (limiter *Limiter) run() {
 	for {
-		if limiter.count <= 0 {
+		if limiter.count == 0 {
 			<-limiter.ticker.C
 			limiter.count = limiter.maxCount
 		}
@@ -40,10 +40,10 @@ func (rateLimiter *Limiter) Take() {
 }
 
 // New creates a new limiter instance with the tokens amount and the interval
-func New(ctx context.Context, max int64, duration time.Duration) *Limiter {
+func New(ctx context.Context, max uint, duration time.Duration) *Limiter {
 	limiter := &Limiter{
-		maxCount: max,
-		count:    max,
+		maxCount: uint(max),
+		count:    uint(max),
 		ticker:   time.NewTicker(duration),
 		tokens:   make(chan struct{}),
 		ctx:      ctx,
@@ -56,8 +56,8 @@ func New(ctx context.Context, max int64, duration time.Duration) *Limiter {
 // NewUnlimited create a bucket with approximated unlimited tokens
 func NewUnlimited(ctx context.Context) *Limiter {
 	limiter := &Limiter{
-		maxCount: math.MaxInt64,
-		count:    math.MaxInt64,
+		maxCount: math.MaxUint,
+		count:    math.MaxUint,
 		ticker:   time.NewTicker(time.Millisecond),
 		tokens:   make(chan struct{}),
 		ctx:      ctx,

ratelimit_test.go

@@ -2,6 +2,8 @@ package ratelimit
 
 import (
 	"context"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -41,4 +43,33 @@ func TestRateLimit(t *testing.T) {
 		require.Equal(t, count, 1000)
 		require.True(t, took < time.Duration(1*time.Second))
 	})
+
+	t.Run("Concurrent Rate Limit Use", func(t *testing.T) {
+		limiter := New(context.Background(), 20, time.Duration(10*time.Second))
+		require.NotNil(t, limiter)
+		var count int32
+		expected := 40
+		start := time.Now()
+		var wg sync.WaitGroup
+		// expected
+		// - First burst of 20 go routines at 0 seconds
+		// - Second burst of 20 go routines at 10 seconds
+		// - Everything should complete right after 10 seconds
+		for i := 0; i < expected; i++ {
+			wg.Add(1)
+
+			go func() {
+				defer wg.Done()
+
+				limiter.Take()
+				atomic.AddInt32(&count, 1)
+			}()
+		}
+
+		wg.Wait()
+		took := time.Since(start)
+		require.Equal(t, expected, int(count))
+		require.True(t, took >= time.Duration(10*time.Second))
+		require.True(t, took <= time.Duration(12*time.Second))
+	})
 }