Rate limiting com token bucket ou leaky bucket

1. Introdução ao Rate Limiting e seus Algoritmos Clássicos

Rate limiting é uma técnica essencial para controlar a taxa de requisições em sistemas distribuídos. Sem ele, serviços podem ser sobrecarregados por picos de tráfego, ataques DDoS ou clientes mal comportados. Em Go, a concorrência nativa com goroutines torna o rate limiting ainda mais relevante, pois precisamos proteger recursos compartilhados em ambientes altamente paralelos.

Dois algoritmos clássicos dominam o cenário: Token Bucket e Leaky Bucket. O primeiro permite rajadas controladas enquanto mantém uma taxa média; o segundo suaviza o tráfego com vazão constante. Ambos podem ser implementados elegantemente em Go usando canais, timers e mutexes.

2. Algoritmo Token Bucket: Conceitos e Implementação Manual

O Token Bucket funciona como um balde que acumula tokens a uma taxa constante. Cada requisição consome um token; se o balde está vazio, a requisição é rejeitada ou bloqueada. A capacidade máxima do balde define o burst permitido.

package tokenbucket

import (
    "sync"
    "time"
)

type TokenBucket struct {
    capacity   int
    tokens     int
    rate       time.Duration // intervalo entre reposições
    mu         sync.Mutex
    lastRefill time.Time
}

func NewTokenBucket(capacity int, rate time.Duration) *TokenBucket {
    return &TokenBucket{
        capacity:   capacity,
        tokens:     capacity,
        rate:       rate,
        lastRefill: time.Now(),
    }
}

func (tb *TokenBucket) Allow() bool {
    tb.mu.Lock()
    defer tb.mu.Unlock()

    now := time.Now()
    elapsed := now.Sub(tb.lastRefill)
    tokensToAdd := int(elapsed / tb.rate)
    if tokensToAdd > 0 {
        tb.tokens = min(tb.capacity, tb.tokens+tokensToAdd)
        tb.lastRefill = now
    }

    if tb.tokens > 0 {
        tb.tokens--
        return true
    }
    return false
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

Esta implementação usa sync.Mutex para segurança em múltiplas goroutines. O Allow() retorna imediatamente, ideal para cenários onde perda de requisições é aceitável. Para esperar até um token disponível, basta adicionar um loop com time.Sleep.

3. Algoritmo Leaky Bucket: Conceitos e Implementação Manual

O Leaky Bucket modela um balde furado que vaza a uma taxa constante. Requisições entram em uma fila; se o balde transborda, requisições são descartadas. A implementação com canais em Go é natural:

package leakybucket

import (
    "sync"
    "time"
)

type LeakyBucket struct {
    capacity  int
    queue     chan struct{}
    leakRate  time.Duration
    wg        sync.WaitGroup
    closeOnce sync.Once
    done      chan struct{}
}

func NewLeakyBucket(capacity int, leakRate time.Duration) *LeakyBucket {
    lb := &LeakyBucket{
        capacity: capacity,
        queue:    make(chan struct{}, capacity),
        leakRate: leakRate,
        done:     make(chan struct{}),
    }
    lb.wg.Add(1)
    go lb.leak()
    return lb
}

func (lb *LeakyBucket) leak() {
    defer lb.wg.Done()
    ticker := time.NewTicker(lb.leakRate)
    defer ticker.Stop()
    for {
        select {
        case <-ticker.C:
            select {
            case <-lb.queue:
            default:
            }
        case <-lb.done:
            return
        }
    }
}

func (lb *LeakyBucket) Allow() bool {
    select {
    case lb.queue <- struct{}{}:
        return true
    default:
        return false
    }
}

func (lb *LeakyBucket) Close() {
    lb.closeOnce.Do(func() {
        close(lb.done)
        lb.wg.Wait()
    })
}

Aqui, o canal com buffer de tamanho capacity atua como a fila. Uma goroutine background remove elementos na taxa definida. Se o canal está cheio, Allow() retorna falso imediatamente.

4. Uso da Biblioteca golang.org/x/time/rate (Token Bucket Oficial)

O pacote oficial golang.org/x/time/rate implementa Token Bucket de forma otimizada e testada. Oferece três modos de operação:

  • Allow(): retorna bool (não bloqueante)
  • Reserve(): reserva um token futuro
  • Wait(): bloqueia até token disponível
package main

import (
    "fmt"
    "net/http"
    "time"
    "golang.org/x/time/rate"
)

func main() {
    limiter := rate.NewLimiter(rate.Limit(10), 20) // 10 req/s, burst 20

    http.HandleFunc("/api", func(w http.ResponseWriter, r *http.Request) {
        if !limiter.Allow() {
            w.Header().Set("Retry-After", "1")
            http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
            return
        }
        fmt.Fprintln(w, "OK")
    })

    http.ListenAndServe(":8080", nil)
}

O burst de 20 permite que picos curtos sejam absorvidos, enquanto a taxa média se mantém em 10 req/s. Para testes com concorrência:

func testConcurrency() {
    limiter := rate.NewLimiter(5, 10)
    var wg sync.WaitGroup
    for i := 0; i < 20; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            if limiter.Allow() {
                fmt.Printf("Goroutine %d: allowed\n", id)
            } else {
                fmt.Printf("Goroutine %d: rejected\n", id)
            }
        }(i)
    }
    wg.Wait()
}

5. Implementação de Leaky Bucket com Pacotes da Comunidade

A Uber oferece go.uber.org/ratelimit, uma implementação de Leaky Bucket com foco em desempenho. Diferente da nossa implementação manual, ela usa um algoritmo baseado em tempo para evitar goroutines background:

package main

import (
    "fmt"
    "net/http"
    "time"
    "go.uber.org/ratelimit"
)

func main() {
    rl := ratelimit.New(10) // 10 requisições por segundo

    http.HandleFunc("/api", func(w http.ResponseWriter, r *http.Request) {
        now := time.Now()
        rl.Take() // bloqueia até poder processar
        fmt.Printf("Request processed at %v\n", now)
        fmt.Fprintln(w, "OK")
    })

    http.ListenAndServe(":8080", nil)
}

A biblioteca da Uber é mais eficiente que nossa implementação com canais, pois evita alocações de goroutines e usa apenas cálculos matemáticos. Ideal para sistemas de alta performance.

6. Rate Limiting em Middleware HTTP e gRPC

Criar middlewares reutilizáveis é uma prática recomendada. Exemplo para HTTP com rate.Limiter:

func RateLimitMiddleware(limiter *rate.Limiter) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            if !limiter.Allow() {
                w.Header().Set("Retry-After", "1")
                http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
                return
            }
            next.ServeHTTP(w, r)
        })
    }
}

Para gRPC, usamos interceptors:

import (
    "go.uber.org/ratelimit"
    "google.golang.org/grpc"
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/status"
)

func UnaryRateLimitInterceptor(rl ratelimit.Limiter) grpc.UnaryServerInterceptor {
    return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
        rl.Take()
        return handler(ctx, req)
    }
}

7. Boas Práticas, Testes e Cenários Avançados

Testar rate limiting requer controle do tempo. Use clock.Mock de pacotes como github.com/benbjohnson/clock:

func TestTokenBucket_Allow(t *testing.T) {
    mockClock := clock.NewMock()
    tb := &TokenBucket{
        capacity:   10,
        tokens:     10,
        rate:       100 * time.Millisecond,
        lastRefill: mockClock.Now(),
    }

    // Consome todos os tokens
    for i := 0; i < 10; i++ {
        assert.True(t, tb.Allow())
    }
    assert.False(t, tb.Allow()) // sem tokens

    mockClock.Add(200 * time.Millisecond) // 2 tokens adicionados
    assert.True(t, tb.Allow())
    assert.True(t, tb.Allow())
    assert.False(t, tb.Allow())
}

Para rate limiting por chave (ex.: IP), use um mapa sincronizado com limpeza periódica:

type RateLimiterStore struct {
    mu       sync.RWMutex
    limiters map[string]*rate.Limiter
    rate     rate.Limit
    burst    int
}

func (s *RateLimiterStore) GetLimiter(key string) *rate.Limiter {
    s.mu.RLock()
    limiter, ok := s.limiters[key]
    s.mu.RUnlock()
    if ok {
        return limiter
    }

    s.mu.Lock()
    defer s.mu.Unlock()
    limiter = rate.NewLimiter(s.rate, s.burst)
    s.limiters[key] = limiter
    return limiter
}

Periodicamente, remova limitadores inativos usando time.Ticker e verificando lastAccess.

8. Comparação Final: Token Bucket vs Leaky Bucket em Go

Característica Token Bucket Leaky Bucket
Rajadas (bursts) Permitidas (até capacidade) Não permitidas
Vazão Taxa média + burst Taxa constante
Comportamento Acumula tokens para uso futuro Processa em ritmo constante
Perda de requisições Quando sem tokens Quando fila cheia
Implementação em Go golang.org/x/time/rate go.uber.org/ratelimit
Melhor para APIs públicas, spikes de tráfego Filas de jobs, processamento estável

Escolha Token Bucket quando precisar absorver picos curtos sem rejeitar requisições legítimas — ideal para APIs REST públicas.

Escolha Leaky Bucket quando a taxa de saída precisa ser estritamente controlada — como em sistemas de processamento batch ou integrações com limites contratuais.

Para projetos Go reais, comece com golang.org/x/time/rate (Token Bucket oficial) e migre para go.uber.org/ratelimit apenas se precisar de vazão estritamente constante. Ambos são battle-tested e oferecem excelente desempenho em ambientes concorrentes.

Referências