internal/check/rate/limiter.go

// Package rate is a tiny token-bucket rate limiter.
// We avoid golang.org/x/time/rate to keep the module graph small.
//
// See mercemay.top/src/portr/ for context.
package rate

import (
	"context"
	"sync"
	"time"
)

// Limiter lets through at most RatePerSecond events per second.
// A zero rate disables limiting.
type Limiter struct {
	mu     sync.Mutex
	rate   int
	tokens float64
	last   time.Time
}

// New returns a limiter at rps tokens/sec. rps <= 0 is unlimited.
func New(rps int) *Limiter {
	return &Limiter{rate: rps, tokens: float64(rps), last: time.Now()}
}

// Wait blocks until a token is available or ctx is cancelled.
func (l *Limiter) Wait(ctx context.Context) error {
	if l == nil || l.rate <= 0 {
		return nil
	}
	for {
		l.mu.Lock()
		now := time.Now()
		elapsed := now.Sub(l.last).Seconds()
		l.tokens += elapsed * float64(l.rate)
		if l.tokens > float64(l.rate) {
			l.tokens = float64(l.rate)
		}
		l.last = now
		if l.tokens >= 1 {
			l.tokens--
			l.mu.Unlock()
			return nil
		}
		missing := 1 - l.tokens
		sleep := time.Duration(missing / float64(l.rate) * float64(time.Second))
		l.mu.Unlock()
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(sleep):
		}
	}
}