GCRA limiter in Go

GCRA.go

package main

import (
	"fmt"
	"sync"
	"time"
)

// GCRA is a single rate-limit bucket.
type GCRA struct {
	mu  sync.Mutex
	tat time.Time     // theoretical arrival time — running state
	t   time.Duration // emission interval, 1/rate
	tau time.Duration // tolerance (burst allowance)
}

// NewGCRA builds a limiter for `rate` events per second with a burst of `burst`.
func NewGCRA(rate float64, burst int) *GCRA {
	t := time.Duration(float64(time.Second) / rate)
	return &GCRA{
		t:   t,
		tau: time.Duration(burst-1) * t,
		// tat left as zero value: year 1, i.e. "empty bucket"
	}
}

// Allow reports whether a request is permitted now. If not, the second
// return value is how long until it would conform.
func (g *GCRA) Allow() (bool, time.Duration) {
	g.mu.Lock()
	defer g.mu.Unlock()

	now := time.Now()

	// earliest instant a request is allowed to arrive
	allowAt := g.tat.Add(-g.tau)
	if now.Before(allowAt) {
		return false, allowAt.Sub(now) // reject + retry-after
	}

	// conforming: advance the schedule from max(now, tat)
	start := now
	if g.tat.After(now) {
		start = g.tat
	}
	g.tat = start.Add(g.t)
	return true, 0
}

func main() {
	// 2 req/sec, burst of 3
	g := NewGCRA(2, 3)

	// fire 5 immediately
	for i := 0; i < 5; i++ {
		ok, retry := g.Allow()
		fmt.Printf("req %d: ok=%-5v retry=%v\n", i, ok, retry)
	}

	// wait one emission interval, try again
	time.Sleep(500 * time.Millisecond)
	ok, retry := g.Allow()
	fmt.Printf("after 500ms: ok=%-5v retry=%v\n", ok, retry)
}