Vectrum
Sign In
86 lines · 1.6 KB · Shell
#!/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