From a2f4776db19d7a416b7b590ee1bfbf2f8b269d2d Mon Sep 17 00:00:00 2001
From: Sebastian Machuca <sebastian@hachiko.energy>
Date: Wed, 31 Dec 2025 14:58:46 +1100
Subject: [PATCH] Adding a registry

---
 token/registry.go      |  47 +++++++++++++++++
 token/registry_test.go | 117 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 164 insertions(+)
 create mode 100644 token/registry.go
 create mode 100644 token/registry_test.go

diff --git a/token/registry.go b/token/registry.go
new file mode 100644
index 0000000..8cbc071
--- /dev/null
+++ b/token/registry.go
@@ -0,0 +1,47 @@
+package token
+
+import (
+	"fmt"
+	"sync"
+)
+
+type (
+	Identifier string
+	Registry   struct {
+		mu             sync.Mutex
+		limiters       map[Identifier]*Limiter
+		capacity, rate float64
+	}
+)
+
+func NewRegistry(capacity, rate float64, users ...Identifier) (*Registry, error) {
+	limiters := make(map[Identifier]*Limiter)
+
+	for _, user := range users {
+		limiter, err := NewLimiter(capacity, rate)
+		if err != nil {
+			return nil, fmt.Errorf("fail to create a new limiter %w", err)
+		}
+
+		limiters[user] = limiter
+	}
+
+	return &Registry{
+		limiters: limiters,
+		capacity: capacity,
+		rate:     rate,
+	}, nil
+}
+
+func (r *Registry) Allow(key Identifier) bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	lim, ok := r.limiters[key]
+	if !ok {
+		lim, _ = NewLimiter(r.capacity, r.rate)
+		r.limiters[key] = lim
+	}
+
+	return lim.Allow()
+}
diff --git a/token/registry_test.go b/token/registry_test.go
new file mode 100644
index 0000000..336e225
--- /dev/null
+++ b/token/registry_test.go
@@ -0,0 +1,117 @@
+package token_test
+
+import (
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/serroba/rate/token"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewRegistry(t *testing.T) {
+	reg, err := token.NewRegistry(10, 2)
+	require.NoError(t, err)
+	require.NotNil(t, reg)
+}
+
+func TestNewRegistry_WithUsers(t *testing.T) {
+	reg, err := token.NewRegistry(10, 2, "alice", "bob")
+	require.NoError(t, err)
+	require.NotNil(t, reg)
+}
+
+func TestNewRegistry_InvalidCapacity(t *testing.T) {
+	_, err := token.NewRegistry(-1, 2, "alice")
+	require.Error(t, err)
+}
+
+func TestNewRegistry_InvalidRate(t *testing.T) {
+	_, err := token.NewRegistry(10, -1, "alice")
+	require.Error(t, err)
+}
+
+func TestRegistry_Allow_ExistingUser(t *testing.T) {
+	reg, err := token.NewRegistry(2, 0, "alice")
+	require.NoError(t, err)
+
+	require.True(t, reg.Allow("alice"))
+	require.True(t, reg.Allow("alice"))
+	require.False(t, reg.Allow("alice"))
+}
+
+func TestRegistry_Allow_NewUser(t *testing.T) {
+	reg, err := token.NewRegistry(2, 0)
+	require.NoError(t, err)
+
+	// First call for a new user should create limiter and allow
+	require.True(t, reg.Allow("alice"))
+	require.True(t, reg.Allow("alice"))
+	require.False(t, reg.Allow("alice"))
+}
+
+func TestRegistry_Allow_IndependentUsers(t *testing.T) {
+	reg, err := token.NewRegistry(1, 0)
+	require.NoError(t, err)
+
+	// Each user has their own bucket
+	require.True(t, reg.Allow("alice"))
+	require.True(t, reg.Allow("bob"))
+
+	// Both exhausted now
+	require.False(t, reg.Allow("alice"))
+	require.False(t, reg.Allow("bob"))
+}
+
+func TestRegistry_Allow_Concurrent(t *testing.T) {
+	reg, err := token.NewRegistry(100, 0)
+	require.NoError(t, err)
+
+	var (
+		allowed atomic.Int64
+		wg      sync.WaitGroup
+	)
+
+	// 50 goroutines per user, 4 users = 200 goroutines
+	users := []token.Identifier{"alice", "bob", "charlie", "diana"}
+	for _, user := range users {
+		for range 50 {
+			wg.Add(1)
+
+			go func(u token.Identifier) {
+				defer wg.Done()
+
+				if reg.Allow(u) {
+					allowed.Add(1)
+				}
+			}(user)
+		}
+	}
+
+	wg.Wait()
+
+	// Each user has capacity 100, only 50 requests each, so all should be allowed
+	require.Equal(t, int64(200), allowed.Load())
+}
+
+func TestRegistry_Allow_ConcurrentNewUsers(t *testing.T) {
+	reg, err := token.NewRegistry(5, 0)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+
+	// Create 100 different users concurrently
+	for i := range 100 {
+		wg.Add(1)
+
+		go func(id int) {
+			defer wg.Done()
+
+			user := token.Identifier(string(rune('a' + id%26)))
+			reg.Allow(user)
+		}(i)
+	}
+
+	wg.Wait()
+	// If we get here without panic or race, the test passes
+}