Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions token/registry.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
package token

import (
"fmt"
"sync"
)

Copy link

Copilot AI Dec 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Identifier type lacks documentation. Public types should have documentation comments explaining their purpose and usage.

Suggested change
// Identifier represents a unique key used to look up or create a Limiter in a Registry.

Copilot uses AI. Check for mistakes.
type (
Identifier string
Registry struct {
mu sync.Mutex
limiters map[Identifier]*Limiter
capacity, rate float64
}
)

Copy link

Copilot AI Dec 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The NewRegistry function lacks documentation. It should have a comment describing what it does, explaining its parameters (capacity, rate, and optional users), and documenting when it returns an error.

Suggested change
// NewRegistry creates a Registry that manages token limiters configured with the
// provided capacity and rate. For each optional user identifier supplied in
// users, a corresponding Limiter is created and stored in the registry.
//
// capacity specifies the maximum number of tokens that each limiter can hold,
// and rate specifies the rate at which tokens are replenished. An error is
// returned if creating any underlying Limiter fails.

Copilot uses AI. Check for mistakes.
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
}

Copy link

Copilot AI Dec 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Allow method lacks documentation. It should have a comment explaining its behavior, particularly that it creates a new limiter for unknown identifiers and returns whether the request is allowed.

Suggested change
// Allow reports whether a request for the given identifier is permitted,
// creating a new limiter with the registry's capacity and rate if none exists.

Copilot uses AI. Check for mistakes.
func (r *Registry) Allow(key Identifier) bool {
r.mu.Lock()
defer r.mu.Unlock()

lim, ok := r.limiters[key]
if !ok {
Copy link

Copilot AI Dec 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The error returned by NewLimiter is silently ignored. While this should never fail given the capacity and rate have been validated in NewRegistry, this could hide bugs if the Limiter implementation changes. Consider handling the error or at least adding a comment explaining why it's safe to ignore.

Suggested change
if !ok {
if !ok {
// NewLimiter is expected not to fail here because r.capacity and r.rate
// were validated when the Registry was created in NewRegistry.
// If NewLimiter's contract changes to allow errors for these values,
// this call site should be updated to handle them explicitly.

Copilot uses AI. Check for mistakes.
lim, _ = NewLimiter(r.capacity, r.rate)
r.limiters[key] = lim
}

return lim.Allow()
}
117 changes: 117 additions & 0 deletions token/registry_test.go
Original file line number Diff line number Diff line change
@@ -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
Copy link

Copilot AI Dec 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment is misleading - it states "Create 100 different users concurrently" but the actual implementation only creates 26 unique users due to the modulo operation on line 108. Update the comment to accurately reflect the test behavior.

Suggested change
// Create 100 different users concurrently
// Issue 100 concurrent requests across up to 26 different users

Copilot uses AI. Check for mistakes.
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
}