#!/usr/bin/env bash
set -euo pipefail
mkdir -p internal/ratelimit
cat > go.mod <<'EOF'
module gate
go 1.22
EOF
cat > internal/ratelimit/bucket.go <<'EOF'
// Package ratelimit implements a token bucket.
package ratelimit
import (
"sync"
"time"
)
// Bucket allows burst requests up to capacity, refilling at rate per second.
type Bucket struct {
mu sync.Mutex
capacity float64
rate float64
tokens float64
last time.Time
now func() time.Time
}
// New builds a bucket that starts full.
func New(capacity, ratePerSec float64) *Bucket {
return &Bucket{
capacity: capacity,
rate: ratePerSec,
tokens: capacity,
last: time.Now(),
now: time.Now,
}
}
// Allow reports whether one request may proceed, consuming a token if so.
func (b *Bucket) Allow() bool {
b.mu.Lock()
defer b.mu.Unlock()
now := b.now()
elapsed := now.Sub(b.last).Seconds()
b.tokens += elapsed * b.rate
if b.tokens > b.capacity {
b.tokens = b.capacity
}
if b.tokens < 1 {
return false
}
b.tokens--
return true
}
EOF
cat > internal/ratelimit/bucket_test.go <<'EOF'
package ratelimit
import (
"testing"
"time"
)
func TestBurstThenRefill(t *testing.T) {
now := time.Now()
b := New(3, 1)
b.now = func() time.Time { return now }
b.last = now
for i := 0; i < 3; i++ {
if !b.Allow() {
t.Fatalf("request %d should have been allowed", i+1)
}
}
if b.Allow() {
t.Fatal("the bucket should be empty")
}
now = now.Add(2 * time.Second)
if !b.Allow() {
t.Fatal("two seconds should have refilled two tokens")
}
}
EOF
git init -q . && git add -A && git -c user.email=e@x -c user.name=n commit -qm base