Explorando o pacote Rate e o algoritmo de Token Bucket

14/05/2025 892 palavras 5 minutos

Conteúdos

No post sobre Circuit Breaker, citei que atualmente é comum que a sua aplicação tenha que se comunicar com outras e, com isso, estratégias de controle de tráfego se tornam essenciais. Recentemente descobri o Token Bucket, uma estratégia baseada em tokens para controlar o tráfego.

Imagine que você tenha 5 ingressos para um brinquedo e que a cada hora você ganhe um novo ingresso, mas nunca podendo exceder o limite de 5. A cada ida ao brinquedo, um ticket é usado. Dessa forma, ao utilizar todos os seus tickets, você não pode ir ao brinquedo até ganhar um novo ticket. É um algoritmo bem interessante, utilizado, por exemplo, pelo PIX, para controlar a busca de chaves e evitando que um malfeitor pegue dados de usuários.

Show me the code!

Traduzindo para código, vamos propor um cenário em que 10 goroutines sejam inicializadas e vão executar uma ação qualquer. Caso você ainda não conheça sobre goroutines e concorrência, recomendo o meu post sobre o tema.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
package main

import	"golang.org/x/time/rate"

func main() {
  c := make(chan int)

  for i := range 10 {
    go doSomething(i, c)
  }

  for range 10 {
    <-c
  }
}

func doSomething(x int, c chan int) {
  fmt.Printf("goroutine %d did something\n", x)

  c <- x
}

Ao executar o código acima, pode-se notar que as goroutines foram executadas sem nenhum controle maior.

Pacote rate

O algoritmo de Token Bucket vem da necessidade de controlar essa execução e, como sempre, a própria linguagem e/ou a comunidade Go nos fornecem uma solução, neste caso o pacote rate. Seu uso é bem simples, primeiro é necessário inicializar um novo Limitter, que recebe a quantidade máxima de tokens disponível e um Limit, que define a taxa com que novos tokens são gerados por segundo.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
package main

import	"golang.org/x/time/rate"

func main() {
  c := make(chan int)
  r := rate.Limit(1.0) // refreshes 1 token per second
  l := rate.NewLimiter(r, 5) // allows 5 tokens max

  ... // do more stuff
}

Estratégia de permissão

O legal do pacote rate é que ele possui três estratégias diferentes. A primeira delas é a estratégia de allow.

1
2
3
4
5
6
7
func doSomethingWithAllow(l *rate.Limiter, x int, c chan int) {
  if l.Allow() {
    fmt.Printf("Allowing %d to run\n", x)
  }

  c <- x
}

Nesta estratégia, a execução será permitida se existir um token para ser consumido naquele momento e, caso não exista, nada será executado.

Estratégia de espera

A segunda estratégia é chamada de wait e é provavelmente a estratégia mais comum de ser utilizada, segundo o autor do pacote.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
func doSomethingWithWait(ctx context.Context, l *rate.Limiter, x int, c chan int) {
  err := l.Wait(ctx)
  if err != nil {
    fmt.Printf("Error waiting for %d: %v\n", x, err)
    c <- x
    return
  }

  fmt.Printf("Allowing %d to run\n", x)
  c <- x
}

Nesta estratégia, a goroutine vai esperar até existir um token disponível para ser usado.

Estratégia de reserva

Por fim, temos a estratégia de reserve. Como o nome diz, você reserva o próximo ticket disponível e aguarda até o momento dele ser emitido.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
func doSomethingWithReserve(l *rate.Limiter, x int, c chan int) {
  r := l.Reserve()
  if !r.OK() {
    return
  }

  fmt.Printf("Reserving %d to run\n", x)
  d := r.Delay()
  time.Sleep(d)
  fmt.Printf("Allowing %d to run\n", x)

  c <- x
}

Esta estratégia é bem parecida com a de wait, contudo nos permite controlar com mais detalhes e saber quando é possível executar novamente.

Extras

Além disso, todas essas estratégias permitem você queimar mais de um ticket, caso seja necessário. Por exemplo, a função Allow() vira AllowN(t time.Time, n int), onde t é o tempo para acontecer n eventos.

Outra funcionalidade que achei bem legal, foi a possiblidade de se fazer um Circuit Breaker simples utilizando esse pacote e o tipo Sometimes. Nele é possível configurar os seguintes parâmetros:

First: as primeiras N chamadas serão executadas.
Every: A cada N chamadas executadas, uma será bloqueada.
Inverval: o limite de chamadas no tempo.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
package main

import (
  "fmt"

  "golang.org/x/time/rate"
)

func main() {
  s := rate.Sometimes{Every: 2}

  for i := range 10 {
    s.Do(func() { fmt.Printf("Allowing %d to run!\n", i) })
  }
}

Como podemos ver, a cada duas chamadas, uma é bloqueada.

Conclusão

O token bucket é um algoritmo de controle de tráfego extremamente interessante, já que ele permite barrar execuções excedentes e agendar execuções futuras. A implementação em Go é simples de se utilizar, além de ser segura ao se utilizar em múltiplas goroutines.

O pacote rate ainda provê uma simples implementação de circuit breaker, evitando assim a adição de mais bibliotecas de terceiros em seu código, caso você deseje utilizar ambos os métodos em sua aplicação.

Contudo, vale a pena citar que o pacote ainda é experimental, podendo mudar o contrato entre as versões, ser descontinuado ou até mesmo incorporado a stdlib.

Para ver todos os exemplos, acesse este repositório.