Пишем высокопроизводительный rate-limiter на golang

Недавно я познакомился с golang и в качестве практики решил реализовать алгоритмы rate limiting'а из статьи c arxiv. В этой статье мы рассмотрим написание rate-limiter'а в виде отдельно стоящего сервиса и в виде middleware, используя in-memory кэш или Redis, добавим поддержку REST и gRPC, подключим Prometheus и поймём как добиться атомарности операций с минимальными задержками.

Кому полезно

Новичкам в golang или же людям, которые по какой-то причине также решили использовать собственные реализации rate-limiter'ов заместо уже существующих решений.

Предисловие

Указанная ранее статья покрывает 3 алгоритма rate-limiting'а:

1. Token bucket

2. Fixed window

3. Rolling window

Суть данных алгоритмов вы можете прочитать самостоятельно, я же сфокусируюсь на том, как написать данные алгоритмы таким образом, чтобы rate-limiter работал быстро и надёжно.

Конфигурация

Первое, с чего стоит начать, это конфигурация вашего limiter'а. В своей реализации я ставил следующие требования:

• Возможность переключаться между алгоритмами

• Поддержка in-memory кэша и Redis

• Метрики в Prometheus

• Поддержка REST и gRPC

Из этих требований вытекает следующий файл конфигурации:

store: "redis"
use_algo: "token_bucket"
metrics:
  use_prometheus: true
  port: 9091
algo_settings:
  capacity: 10
  rate: 0.1
backends:
  in_memory:
    default_ttl: 30 # seconds
    eviction_time: 120 # seconds
    shards_count: 256
  redis:
    addr: "redis:6379"
    password: "m1re1di77sss"
    db: 0
    default_ttl: 1000

Рассмотрим его подробнее:

• store отвечает за используемое хранилище, может быть "redis" или же "in_memory"

• use_algo отвечает за используемый алгоритм

• metrics позволяет включать и отключать поддержку выгрузки и подсчёта метрик во время работы, а также за порт, на котором позже будет поднят мультиплексер для метрик

• algo_settings для каждого алгоритма разные, отвечает за настройки, которые будут использоваться алгоритмом

• backends - настройки каждого из хранилищ, в in_memory default_ttl отвечает за ttl ключей в хранилище, eviction_time отвечает за частоту очистки in-memory хранилища от протухших ключей, shard_count отвечает за то, на сколько частей in-memory кэш будет шардирован, число должно быть степенью 2, позже я объясню почему

Конфигурация для middleware выглядит похожим образом:

type InMemoryBackendSettings struct {
	DefaultTtl   int64  `yaml:"default_ttl"`
	EvictionTime int64  `yaml:"eviction_time"`
	ShardsCount  uint32 `yaml:"shards_count"`
}

type RedisBackendSettings struct {
	Addr       string `yaml:"addr"`
	Password   string `yaml:"password"`
	Db         int    `yaml:"db"`
	DefaultTtl int64  `yaml:"default_ttl"`
}

type AlgoSettings struct {
	Capacity    *float64 `yaml:"capacity"`
	Rate        *float64 `yaml:"rate"`
	WindowSize  *int64   `yaml:"window_size"`
	MaxRequests *int64   `yaml:"max_requests"`
}

Бэкенды для limiter'а

Думаю, ни для кого не секрет, что Redis умеет выполнять Lua скрипты атомарно, что является идеальной функциональной возможностью для надёжного rate-limiting'а. В таком случае наш бэкенд limiter'а также должен уметь выполнять операции атомарно, отсюда вытекает интерфейс бэкенда:

type InMemoryStoreInterface interface {
	Update(key string, newFn func() any, fn func(v any) error) error
	Close() error
}

Update принимает на вход ключ, по которому мы хотим принять решение, а также 2 колбэка:

1. newFn - функция создания нового объекта в хранилище, в случае если данный ключ в хранилище ещё не появлялся. Это позволяет избежать создания данного объекта в случае, если ключ уже содержится в хранилище. Конечно, можно было сначала проверять и только потом добавлять, но это добавляет лишней логики клиенту, использующему хранилище.

2. fn - собственно функция, содержащая логику limiting'а

In-memory кэш

Наше хранилище под капотом будет использовать мапу, во избежание race condition'ов мы будем использовать mutex. Но ставить mutex на всю мапу слишком дорого, так как другие клиенты всё ещё должны пользоваться хранилищем.
Отсюда вытекает необходимость в шардировании. Пока одна горутина выполняет логику по ключу x, мы всё ещё можем работать с ключом y:

type shard struct {
	mu      sync.Mutex
	entries map[string]*cacheEntry
}

type cacheEntry struct {
	entry      any
	lastAccess int64
}

type ShardedInMemoryStore struct {
	shards         []*shard
	shardsCount    uint32
	expirationTime int64
	ticker         *time.Ticker
	done           chan any
	once           sync.Once
}

func NewCache(
	expirationTime int64,
	evictionInterval time.Duration,
	shardCount uint32,
) *ShardedInMemoryStore {
	ticker := time.NewTicker(evictionInterval)
	cache := &ShardedInMemoryStore{
		shards:         make([]*shard, shardCount),
		shardsCount:    shardCount,
		expirationTime: expirationTime,
		done:           make(chan any),
		ticker:         ticker,
	}

	for i := range shardCount {
		cache.shards[i] = &shard{entries: make(map[string]*cacheEntry)}
	}
	go func() {
		defer ticker.Stop()
		for {
			select {
			case <-cache.done:
				return
			case <-ticker.C:
				cache.handleEviction()
			}
		}
	}()

	return cache
}

Также удаление протухших ключей теперь не будет лочить всю мапу, а только один конкретный шард:

func (cache *ShardedInMemoryStore) handleEviction() {
	currentTime := time.Now().Unix()
	for _, shard := range cache.shards {
		shard.mu.Lock()
		for key, value := range shard.entries {
			if currentTime-cache.expirationTime > value.lastAccess {
				delete(shard.entries, key)
			}
		}
		shard.mu.Unlock()
	}
}

А также ранее упомянутый Update:

func (cache *ShardedInMemoryStore) Update(key string, createFn func() any, fn func(v any) error) error {
	shard := cache.getShard(key)
	shard.mu.Lock()
	defer shard.mu.Unlock()
	now := time.Now().Unix()
	value, present := shard.entries[key]
	if !present {
		value = &cacheEntry{entry: createFn(), lastAccess: now}
		shard.entries[key] = value
	}
	value.lastAccess = now
	return fn(value.entry)
}

Именно здесь реализуется столь необходимая атомарность, мы лочим mutex на шарде и выполняем операцию, не блокируя другие шарды и не допуская race condition'ов.

Так почему количество шардов должно быть степенью двойки?

Для того, чтобы распределять клиентские ключи более-менее равномерно по шардам нужно использоват функцию хэширования, которая не только будет раскладывать ключи равномерно, но и будет быстрой. Под такую функцию подходит Fowler–Noll–Vo. А чтобы после этого быстро замапить хэш на шард можно воспользоваться трюком. Вкратце можно очень быстро искать остатки от деления на степень двойки просто используя формулу: hash & (shardsCount - 1):

func (cache *ShardedInMemoryStore) getShard(key string) *shard {
	h := uint32(2166136261)
	for i := 0; i < len(key); i++ {
		h ^= uint32(key[i])
		h *= 16777619
	}
	return cache.shards[h&(cache.shardsCount-1)]
}

Redis кэш

С Redis всё проще:

package cache

import (
	"context"

	"github.com/redis/go-redis/v9"
)

type RedisStoreInterface interface {
	RunLuaScript(ctx context.Context, script *redis.Script, keys []string, args ...any) (bool, int64, error)
	DefaultTtl() int64
}

Передаём контекст, скрипт и массив ключей, в данном случае массив ключей это требование библиотеки для работы с Redis в Golang. Массив всегда будет состоять из одного элемента.

DefaultTtl - это getter, который используется для token bucket алгоритма, чтобы удалением ключа при использовании Redis занимался сам Redis, а не приложение.

package cache

import (
	"context"
	"fmt"

	"github.com/redis/go-redis/v9"
)

type RedisStore struct {
	client     *redis.Client
	defaultTtl int64
}

func NewRedisStore(client *redis.Client, defaultTtl int64) *RedisStore {
	return &RedisStore{client: client, defaultTtl: defaultTtl}
}

// RunLuaScript runs a Lua script that returns {allowed(0|1), retryAfter}.
func (r *RedisStore) RunLuaScript(
	ctx context.Context,
	script *redis.Script,
	keys []string,
	args ...any,
) (bool, int64, error) {
	res, err := script.Run(ctx, r.client, keys, args...).Int64Slice()
	if err != nil {
		return false, 0, err
	}
	if len(res) != 2 {
		return false, 0, fmt.Errorf("unexpected length of script output")
	}
	return res[0] == 1, res[1], nil
}

func (r *RedisStore) DefaultTtl() int64 {
	return r.defaultTtl
}

Используем созданные хранилища

Теперь рассмотрим функцию запуска limiter'а:

package main

import (
	"context"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"qwrttqr-rate-limiter/core/config"
	"qwrttqr-rate-limiter/core/server/algos"
	"qwrttqr-rate-limiter/core/server/facade"
	"qwrttqr-rate-limiter/core/server/limiting"
	"qwrttqr-rate-limiter/core/server/metrics"
	ratelimiterv1 "qwrttqr-rate-limiter/gen/ratelimiter/v1"
	"strconv"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
)

func main() {
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	cfg, err := config.ReadConfig()
	if err != nil {
		return fmt.Errorf("failed to read config: %w", err)
	}

	rl, err := limiting.NewRateLimiter(cfg)

	if err != nil {
		return fmt.Errorf("create rate limiter: %w", err)
	}
	defer rl.Close()

	var limiter algos.Limiter = rl

	errCh := make(chan error, 3)

	shutdownMetrics := func(ctx context.Context) {}

	if cfg.Metrics.UsePrometheus {
		reg := prometheus.NewRegistry()
		rec := metrics.NewPromRecorder(reg)
		limiter = metrics.Instrument(limiter, rec)

		metricsMux := http.NewServeMux()
		metricsMux.Handle("GET /metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))

		metricsSrv := &http.Server{Addr: ":" + strconv.Itoa(cfg.Metrics.Port), Handler: metricsMux, ReadHeaderTimeout: 5 * time.Second}
		go func() { errCh <- metricsSrv.ListenAndServe() }()
		log.Printf("metrics available on %s", metricsSrv.Addr)

		shutdownMetrics = func(ctx context.Context) {
			_ = metricsSrv.Shutdown(ctx)
		}
	}

	limitingMux := http.NewServeMux()
	limitingMux.HandleFunc("GET /health", facade.HealthHandler)
	limitingMux.HandleFunc("POST /limit", facade.LimitHandler(limiter))

	limitingSrv := &http.Server{Addr: ":8080", Handler: limitingMux, ReadHeaderTimeout: 5 * time.Second}
	grpcSrv := grpc.NewServer()
	ratelimiterv1.RegisterRateLimiterServer(grpcSrv, facade.NewGRPCServer(limiter))

	tcpLis, err := net.Listen("tcp", ":8081")
	if err != nil {
		return fmt.Errorf("listen grpc: %w", err)
	}

	go func() { errCh <- limitingSrv.ListenAndServe() }()
	log.Printf("listening http on %s", limitingSrv.Addr)
	go func() { errCh <- grpcSrv.Serve(tcpLis) }()
	log.Printf("listening grpc on %s", tcpLis.Addr())

	select {
	case err := <-errCh:
		return err
	case <-ctx.Done():
	}

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	grpcDone := make(chan struct{})
	go func() { grpcSrv.GracefulStop(); close(grpcDone) }()

	select {
	case <-grpcDone:
	case <-shutdownCtx.Done():
		grpcSrv.Stop()
	}

	shutdownMetrics(shutdownCtx)
	return limitingSrv.Shutdown(shutdownCtx)
}

Давайте по порядку:

ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()

Создаём shutdown context, который срабатывает на нажатие Ctrl-C либо при получении syscall.SIGTERM (сигнал, посылаемый родительским процессом на завершение нашего процесса, например когда тушим докер-контейнер). defer stop() нужен для того, чтобы прекратить обрабатывать данные сигналы, сразу перед возвратом из run.

cfg, err := config.ReadConfig()
if err != nil {
    return fmt.Errorf("failed to read config: %w", err)
}

rl, err := limiting.NewRateLimiter(cfg)

if err != nil {
    return fmt.Errorf("create rate limiter: %w", err)
}
defer rl.Close()

Далее читаем конфиг и создаём сам limiter, у limiter'а есть функция Close, нужная для закрытия соединения с Redis. в in-memory кэше она просто закрывает канал, что останавливает удаление ключей из кэша.

var limiter algos.Limiter = rl

errCh := make(chan error, 3)

Следующие строки позволяют создать указатель с типом интерфейса на limiter:

package algos

import (
	"context"
)

type Request struct {
	Key            string
	RequiredTokens float64
}

type Decision struct {
	Allowed bool
	// RetryAfter is seconds
	RetryAfter int64
}

type Limiter interface {
	Allow(ctx context.Context, req Request) (Decision, error)
}

Так как NewRateLimiter возвращает не интерфейс, а указатель на структуру, поддерживающую Close и Allow:

type RateLimiter struct {
	algos.Limiter
	io.Closer
}

func NewRateLimiter(cfg config.StandaloneServiceConfiguration) (_ *RateLimiter, err error) {
	var (
		mem    cache.InMemoryStoreInterface
		rdb    *redis.Client
		closer io.Closer
	)
	switch cfg.Store {
	case "in_memory":
		cacheInstance := cache.NewCache(
			cfg.Backends.InMemory.DefaultTtl,
			time.Duration(cfg.Backends.InMemory.EvictionTime)*time.Second,
			cfg.Backends.InMemory.ShardsCount,
		)
		mem, closer = cacheInstance, cacheInstance
	case "redis":
		redisClient := redis.NewClient(&redis.Options{
			Addr:     cfg.Backends.Redis.Addr,
			Password: cfg.Backends.Redis.Password,
			DB:       cfg.Backends.Redis.Db,
		})
		rdb, closer = redisClient, redisClient
	default:
		return nil, fmt.Errorf("unsupported store: %q", cfg.Store)
	}

	defer func() {
		if err != nil {
			_ = closer.Close()
		}
	}()
	var limiter algos.Limiter

	switch cfg.UsedAlgo {
	case "token_bucket":
		if err := algos.ValidateTokenBucketConfig(cfg.AlgoSettings); err != nil {
			return nil, err
		}
		store, err := algos.NewTokenBucketStore(cfg.Store, mem, rdb, cfg.Backends.Redis.DefaultTtl)
		if err != nil {
			return nil, err
		}
		limiter = &algos.TokenBucketLimiter{
			Capacity: *cfg.AlgoSettings.Capacity,
			Rate:     *cfg.AlgoSettings.Rate,
			Store:    store,
			Now:      realClock,
		}
	case "fixed_window":
		if err := algos.ValidateFixedWindowConfig(cfg.AlgoSettings); err != nil {
			return nil, err
		}
		store, err := algos.NewFixedWindowStore(cfg.Store, mem, rdb, cfg.Backends.Redis.DefaultTtl)
		if err != nil {
			return nil, err
		}
		limiter = &algos.FixedWindowLimiter{
			MaxRequests: *cfg.AlgoSettings.MaxRequests,
			WindowSize:  *cfg.AlgoSettings.WindowSize,
			Store:       store,
			Now:         realClock,
		}
	case "rolling_window":
		if err := algos.ValidateRollingWindowConfig(cfg.AlgoSettings); err != nil {
			return nil, err
		}
		store, err := algos.NewRollingWindowStore(cfg.Store, mem, rdb, cfg.Backends.Redis.DefaultTtl)
		if err != nil {
			return nil, err
		}
		limiter = &algos.RollingWindowLimiter{
			MaxRequests: *cfg.AlgoSettings.MaxRequests,
			WindowSize:  *cfg.AlgoSettings.WindowSize,
			Store:       store,
			Now:         realClock,
		}

	default:
		return nil, fmt.Errorf("unsupported rate limiting algorithm: %s", cfg.UsedAlgo)
	}

	return &RateLimiter{Limiter: limiter, Closer: closer}, nil
}

Возвращаясь обратно в run:

shutdownMetrics := func(ctx context.Context) {}

if cfg.Metrics.UsePrometheus {
    reg := prometheus.NewRegistry()
    rec := metrics.NewPromRecorder(reg)
    limiter = metrics.Instrument(limiter, rec)

    metricsMux := http.NewServeMux()
    metricsMux.Handle("GET /metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))

    metricsSrv := &http.Server{Addr: ":" + strconv.Itoa(cfg.Metrics.Port), Handler: metricsMux, ReadHeaderTimeout: 5 * time.Second}
    go func() { errCh <- metricsSrv.ListenAndServe() }()
    log.Printf("metrics available on %s", metricsSrv.Addr)

    shutdownMetrics = func(ctx context.Context) {
        _ = metricsSrv.Shutdown(ctx)
    }
}

Мы помечали rl, как структуру, реализующую интерфейс как раз ради поддержки метрик здесь. О том, как метрики работают я расскажу позже, сейчас достаточно понимания того, что мы создаём специальный HTTP сервер, с которого Prometheus будет забирать метрики. Для подсчёта данных метрик мы используем Registry, а в функции instrument мы оборачиваем limiter в мидлварь, которая будет обновлять счётчики.

Функция shutdownMetrics обновляется в зависимости от того, включаем ли у себя метрики или нет, это нужно, чтобы сделать graceful shutdown данного сервера.

limitingMux := http.NewServeMux()
	limitingMux.HandleFunc("GET /health", facade.HealthHandler)
	limitingMux.HandleFunc("POST /limit", facade.LimitHandler(limiter))

	limitingSrv := &http.Server{Addr: ":8080", Handler: limitingMux, ReadHeaderTimeout: 5 * time.Second}
	grpcSrv := grpc.NewServer()
	ratelimiterv1.RegisterRateLimiterServer(grpcSrv, facade.NewGRPCServer(limiter))

	tcpLis, err := net.Listen("tcp", ":8081")
	if err != nil {
		return fmt.Errorf("listen grpc: %w", err)
	}

	go func() { errCh <- limitingSrv.ListenAndServe() }()
	log.Printf("listening http on %s", limitingSrv.Addr)
	go func() { errCh <- grpcSrv.Serve(tcpLis) }()
	log.Printf("listening grpc on %s", tcpLis.Addr())

	select {
	case err := <-errCh:
		return err
	case <-ctx.Done():
	}

После этого создаём HTTP и gRPC сервера и ждём в select на ранее созданном NotifyContext, до момента пока не придёт сигнал на завершение процесса (os.Interrupt, syscall.SIGTERM) или ошибка.

shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()

grpcDone := make(chan struct{})
go func() { grpcSrv.GracefulStop(); close(grpcDone) }()

select {
case <-grpcDone:
case <-shutdownCtx.Done():
    grpcSrv.Stop()
}

shutdownMetrics(shutdownCtx)
return limitingSrv.Shutdown(shutdownCtx)

В случае, если к нам пришла ошибка или же сигнал на завершение процесса, создаём контекст с таймаутом и просим все наши сервера выключиться. В особенности это касается gRPC сервера, он самостоятельно поддерживает GracefulStop, но если он не успеет сделать это за 10 секунд, мы выключим его принудительно.

Используем ранее созданные бэкенды

Помните, что мы создали Redis и InMemory бэкенды для наших алгоритмов, теперь можно посмотреть, как мы пользуемся ими, рассмотрим пример с Rolling Window:

type RollingWindowStore interface {
	CheckAndIncrement(ctx context.Context, key string, windowSize, maxRequests, now int64) (bool, int64, error)
}

type RollingWindowLimiter struct {
	WindowSize  int64
	MaxRequests int64
	Store       RollingWindowStore
	Now         func() int64
}

func (rwl *RollingWindowLimiter) Configure() {
	if rwl.Now == nil {
		rwl.Now = func() int64 {
			return time.Now().Unix()
		}
	}
}

func ValidateRollingWindowConfig(cfg config.AlgoSettings) error {
	if *cfg.MaxRequests <= 0 {
		return fmt.Errorf("the max requests param should be strictly greater then 0")
	}
	return CheckRequiredFields(cfg, []string{"window_size", "max_requests"})
}

type InMemoryRollingWindowStore struct {
	Backend cache.InMemoryStoreInterface
}

type RedisRollingWindowStore struct {
	Backend cache.RedisStoreInterface
}

func NewRollingWindowStore(store string, cacheInstance cache.InMemoryStoreInterface, redisClient *redis.Client, redisDefaultKeyTTL int64) (RollingWindowStore, error) {
	switch store {
	case "in_memory":
		return &InMemoryRollingWindowStore{Backend: cacheInstance}, nil
	case "redis":
		return &RedisRollingWindowStore{Backend: cache.NewRedisStore(redisClient, redisDefaultKeyTTL)}, nil
	default:
		return nil, fmt.Errorf("unsupported store: %s", store)
	}
}

Ещё один интерфейс(RollingWindowStore), нужен для генерализации работы с бэкендом, таким образом, сам алгоритм будет просто вызывать функцию CheckAndIncrement, передавая параметры, а бэкенд будет выполнять её в атомарном варианте.

func (rwl *RollingWindowLimiter) Configure() {
	if rwl.Now == nil {
		rwl.Now = func() int64 {
			return time.Now().Unix()
		}
	}
}

Данная функция нужна для тестов, так как в тестах придётся мокать функцию получения текущего времени.

type RollingWindowState struct {
	Timestamps []int64
}

func (rws *RollingWindowState) check(
	windowSize,
	maxRequests,
	now int64,
) (bool, int64) {
	windowStart := now - windowSize
	i := 0
	for ; i < len(rws.Timestamps); i++ {
		if rws.Timestamps[i] > windowStart {
			break
		}
	}
	if i > 0 {
		copy(rws.Timestamps, rws.Timestamps[i:])
		rws.Timestamps = rws.Timestamps[:len(rws.Timestamps)-i]
	}
	if int64(len(rws.Timestamps)) < maxRequests {
		rws.Timestamps = append(rws.Timestamps, now)
		return true, 0
	}
	retryAfter := rws.Timestamps[0] - windowStart
	return false, retryAfter
}

Сам алгоритм устроен несложно, мы просто отслеживаем, какие временные метки уже вышли за временное окно, и если оставшихся меток меньше, чем лимит запросов, клиент может выполнить очередной запрос. Иначе нужно ответить, через какое количество секунд освободится хотя бы одно место.

Далее подвязываем данную функцию к бэкенду:

func (store *InMemoryRollingWindowStore) CheckAndIncrement(
	_ context.Context,
	key string,
	windowSize,
	maxRequests,
	now int64,
) (allowed bool, retryAfter int64, err error) {
	err = store.Backend.Update(
		key,
		func() any { return &RollingWindowState{Timestamps: make([]int64, 0, 8)} },
		func(v any) error {
			st, ok := v.(*RollingWindowState)
			if !ok {
				return fmt.Errorf("unexpected state type %T for key %q", v, key)
			}
			allowed, retryAfter = st.check(windowSize, maxRequests, now)
			return nil
		})
	return allowed, retryAfter, err
}

То же самое в Redis:

var rollingWindowScript = redis.NewScript(`
local requestId = ARGV[1]
local windowSize = tonumber(ARGV[2])
local maxRequests = tonumber(ARGV[3])
local now = tonumber(ARGV[4])

local windowStart = now - windowSize

redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", windowStart)
local count = redis.call("ZCARD", KEYS[1])

if count < maxRequests then
    redis.call("ZADD", KEYS[1], now, requestId)
    redis.call("EXPIRE", KEYS[1], windowSize)
    return {1, 0}
else
	local oldest = redis.call("ZRANGE", KEYS[1], 0, 0, "WITHSCORES")
	local retryAfter = 0
	if oldest and oldest[2] then
		retryAfter = tonumber(oldest[2]) - windowStart
	end
    return {0, retryAfter}
end
`)

func (s *RedisRollingWindowStore) CheckAndIncrement(ctx context.Context, key string, windowSize, maxRequests, now int64) (bool, int64, error) {
	requestId := uuid.New().String()
	return s.Backend.RunLuaScript(ctx, rollingWindowScript, []string{key}, requestId, windowSize, maxRequests, now)
}

В Redis реализации стоит обратить внимание на то, что мы генерируемrequestId, всё потому, что мы храним наши элементы в множестве, а в множестве могут храниться только уникальные элементы. Но как тогда быть в ситуации, когда 2 и более запросов прилетают в одну секунду? Для этого и создан requestId, он добавляет уникальность каждому запросу, поэтому если в одну секунду прилетит 1000 запросов, они все посчитаются как разные запросы.

А что с метриками?

Как я ранее говорил, rate-limiter поддерживает подсчёт метрик. Начнём с функции instrument, о которой говорилось ранее:

package metrics

import (
	"context"
	"qwrttqr-rate-limiter/core/server/algos"
	"time"
)

type instrumented struct {
	next algos.Limiter
	rec  Recorder
}

func Instrument(next algos.Limiter, rec Recorder) algos.Limiter {
	return &instrumented{next: next, rec: rec}
}

func (i *instrumented) Allow(ctx context.Context, req algos.Request) (algos.Decision, error) {
	start := time.Now()
	dec, err := i.next.Allow(ctx, req)
	elapsed := time.Since(start)

	switch {
	case err != nil:
		i.rec.ObserveDecision(Errored, elapsed)
	case dec.Allowed:
		i.rec.ObserveDecision(Allowed, elapsed)
	default:
		i.rec.ObserveDecision(Denied, elapsed)
	}
	return dec, err
}

Данная функция просто возвращает структуру, которая обладает методом Allow, таким образом мы подменяем наш оригинальный limiter обёрткой, которая поддерживает подсчёт метрик.

Recorder реализован так:

package metrics

import "time"

type Result uint8

const (
	Allowed Result = iota
	Denied
	Errored
	NumResults
)

type Recorder interface {
	ObserveDecision(r Result, elapsed time.Duration)
}

И конкретная реализация данного интерфейса под Prometheus выглядит так:

package metrics

import (
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

type PromRecorder struct {
	decisions    [NumResults]prometheus.Counter
	durationHist [NumResults]prometheus.Observer
}

func NewPromRecorder(reg prometheus.Registerer) *PromRecorder {
	counterVec := prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "ratelimiter",
		Name:      "decisions_total",
		Help:      "Each decision count by its type.",
	}, []string{"result"})

	reqDurHistVec := prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace: "ratelimiter",
		Name:      "decision_duration_seconds",
		Help:      "Execution time of one decision in seconds.",
		Buckets:   []float64{0.00005, 0.0001, 0.0005, 0.001, 0.1, 0.5, 1, 2},
	}, []string{"result"})
	reg.MustRegister(counterVec)
	reg.MustRegister(reqDurHistVec)
	p := &PromRecorder{}

	p.decisions[Allowed] = counterVec.WithLabelValues("allowed")
	p.decisions[Denied] = counterVec.WithLabelValues("denied")
	p.decisions[Errored] = counterVec.WithLabelValues("errored")
	p.durationHist[Allowed] = reqDurHistVec.WithLabelValues("allowed")
	p.durationHist[Denied] = reqDurHistVec.WithLabelValues("denied")
	p.durationHist[Errored] = reqDurHistVec.WithLabelValues("errored")

	return p
}
func (p *PromRecorder) ObserveDecision(r Result, dur time.Duration) {
	p.decisions[r].Inc()
	p.durationHist[r].Observe(dur.Seconds())
}

Здесь мы создаём счётчики и гистограммы под каждый результат rate-limiter'а и реализуем ObserveDecision, который их обновляет.

gRPC и REST

gRPC

Пишем .proto файл:

syntax = "proto3";
package ratelimiter.v1;
option go_package = "qwrttqr-rate-limiter/gen/ratelimiter/v1;ratelimiterv1";
import "google/protobuf/duration.proto";

service RateLimiter {
  rpc Allow(AllowRequest) returns (AllowResponse);
}

message AllowRequest {
  string key  = 1;
  double tokens = 2;
}

message AllowResponse {
  bool allowed = 1;
  google.protobuf.Duration retryAfter = 2;
}

Вызываем protoc и остаётся просто реализовать полученные интерфейсы:

protoc -I proto --go_out=. --go_opt=module=qwrttqr-rate-limiter --go-grpc_out=. --go-grpc_opt=module=qwrttqr-rate-limiter ratelimiter/v1/ratelimiter.proto

package facade

import (
	"context"
	"qwrttqr-rate-limiter/core/server/algos"
	ratelimiterv1 "qwrttqr-rate-limiter/gen/ratelimiter/v1"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/durationpb"
)

type GRPCServer struct {
	ratelimiterv1.UnimplementedRateLimiterServer
	l algos.Limiter
}

func NewGRPCServer(l algos.Limiter) *GRPCServer {
	return &GRPCServer{l: l}
}

func (s *GRPCServer) Allow(ctx context.Context, in *ratelimiterv1.AllowRequest) (*ratelimiterv1.AllowResponse, error) {
	tokens := in.GetTokens()
	if tokens == 0.0 {
		tokens = 1.0
	}
	req, err := NewRequest(in.GetKey(), tokens)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, err.Error())
	}
	d, err := s.l.Allow(ctx, req)
	if err != nil {
		return nil, status.Error(codes.Unavailable, "rate limiter error")
	}

	return &ratelimiterv1.AllowResponse{
		Allowed:    d.Allowed,
		RetryAfter: durationpb.New(time.Duration(d.RetryAfter) * time.Second),
	}, nil
}

REST

Для REST всё совсем просто:

package facade

import (
	"net/http"
	"qwrttqr-rate-limiter/core/server/algos"
	"strconv"
)

func HealthHandler(w http.ResponseWriter, _ *http.Request) {
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("ok"))
}

func LimitHandler(l algos.Limiter) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		key := r.Header.Get("X-Client-Key")
		tokens := 1.0
		if s := r.Header.Get("X-Required-Tokens"); s != "" {
			n, err := strconv.ParseFloat(s, 64)
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			tokens = n
		}
		req, err := NewRequest(key, tokens)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		d, err := l.Allow(r.Context(), req)
		switch {
		case err != nil:
			http.Error(w, "rate limiter error", http.StatusInternalServerError) // store failure
		case !d.Allowed:
			w.Header().Set("Retry-After", strconv.FormatInt(d.RetryAfter, 10))
			http.Error(w, "rate limit exceeded", http.StatusTooManyRequests)
		}
	}
}

Middleware

Для реализации поддержки middleware можно использовать такую схему:

1. Приложение создаёт себе бэкенд, используя Redis или оставаясь на in-memory кэше

2. Используя созданный бэкенд создаём себе rate-limiter, с нужным алгоритмом

package middleware

import (
	"qwrttqr-rate-limiter/core/config"
	"qwrttqr-rate-limiter/core/server/algos"
	"qwrttqr-rate-limiter/core/server/cache"
	"time"

	"github.com/redis/go-redis/v9"
)

var realClock = func() int64 { return time.Now().Unix() }

type Backend interface {
	TokenBucket(cfg config.AlgoSettings) (algos.Limiter, error)
	FixedWindow(cfg config.AlgoSettings) (algos.Limiter, error)
	RollingWindow(cfg config.AlgoSettings) (algos.Limiter, error)
}

type memoryBackend struct {
	mem cache.InMemoryStoreInterface
}

func NewInMemoryBackend(cfg config.InMemoryBackendSettings) Backend {
	return &memoryBackend{mem: cache.NewCache(
		cfg.DefaultTtl,
		time.Duration(cfg.EvictionTime)*time.Second,
		cfg.ShardsCount),
	}
}

type redisBackend struct {
	redis cache.RedisStoreInterface
}

func NewRedisBackend(client *redis.Client, redisDefaultKeyTTL int64) Backend {
	return &redisBackend{redis: cache.NewRedisStore(client, redisDefaultKeyTTL)}
}

func newTokenBucket(cfg config.AlgoSettings, store algos.TokenBucketStore) (algos.Limiter, error) {
	if err := algos.ValidateTokenBucketConfig(cfg); err != nil {
		return nil, err
	}
	return &algos.TokenBucketLimiter{
		Capacity: *cfg.Capacity, Rate: *cfg.Rate, Store: store, Now: realClock,
	}, nil
}

func newFixedWindow(cfg config.AlgoSettings, store algos.FixedWindowStore) (algos.Limiter, error) {
	if err := algos.ValidateFixedWindowConfig(cfg); err != nil {
		return nil, err
	}
	return &algos.FixedWindowLimiter{
		WindowSize: *cfg.WindowSize, MaxRequests: *cfg.MaxRequests, Store: store, Now: realClock,
	}, nil
}

func newRollingWindow(cfg config.AlgoSettings, store algos.RollingWindowStore) (algos.Limiter, error) {
	if err := algos.ValidateRollingWindowConfig(cfg); err != nil {
		return nil, err
	}
	return &algos.RollingWindowLimiter{
		WindowSize: *cfg.WindowSize, MaxRequests: *cfg.MaxRequests, Store: store, Now: realClock,
	}, nil
}

func (b *memoryBackend) TokenBucket(cfg config.AlgoSettings) (algos.Limiter, error) {
	return newTokenBucket(cfg, &algos.InMemoryTokenBucketStore{Backend: b.mem})
}

func (b *memoryBackend) FixedWindow(cfg config.AlgoSettings) (algos.Limiter, error) {
	return newFixedWindow(cfg, &algos.InMemoryFixedWindowStore{Backend: b.mem})
}

func (b *memoryBackend) RollingWindow(cfg config.AlgoSettings) (algos.Limiter, error) {
	return newRollingWindow(cfg, &algos.InMemoryRollingWindowStore{Backend: b.mem})
}

func (b *redisBackend) TokenBucket(cfg config.AlgoSettings) (algos.Limiter, error) {
	return newTokenBucket(cfg, &algos.RedisTokenBucketStore{Backend: b.redis})
}

func (b *redisBackend) FixedWindow(cfg config.AlgoSettings) (algos.Limiter, error) {
	return newFixedWindow(cfg, &algos.RedisFixedWindowStore{Store: b.redis})
}

func (b *redisBackend) RollingWindow(cfg config.AlgoSettings) (algos.Limiter, error) {
	return newRollingWindow(cfg, &algos.RedisRollingWindowStore{Backend: b.redis})
}

Далее остаётся написать функцию, создающую непосредственно саму middleware, а именно функцию возвращающую тип, поддерживающий serveHTTP метод:

package middleware

import (
	"log"
	"net/http"
	"qwrttqr-rate-limiter/core/server/algos"
	"qwrttqr-rate-limiter/core/server/facade"
	"strconv"
)

func NewHTTPMiddleware(
	algo algos.Limiter,
	keyDer HTTPKeyDerivationFunction,
	costDer HTTPCostDerivationFunction,
	failOpen bool,
) func(http.Handler) http.Handler {
	if algo == nil || keyDer == nil {
		panic("middleware: limiter and key function are required")
	}
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			fail := func(stage string, err error) {
				log.Printf("limiter error on %s: %v", stage, err)
				if failOpen {
					next.ServeHTTP(w, r)
					return
				}
				http.Error(w, "service unavailable", http.StatusServiceUnavailable)
			}

			tokens := 1.0
			if costDer != nil {
				t, err := costDer(r)
				if err != nil {
					fail("cost derivation", err)
					return
				}
				tokens = t
			}

			key, err := keyDer(r)
			if err != nil {
				fail("key derivation", err)
				return
			}

			req, err := facade.NewRequest(key, tokens)
			if err != nil {
				fail("request creation", err)
				return
			}

			decision, err := algo.Allow(r.Context(), req)
			if err != nil {
				fail("allow", err)
				return
			}
			if !decision.Allowed {
				w.Header().Set("Retry-After", strconv.FormatInt(decision.RetryAfter, 10))
				http.Error(w, "rate limit exceeded", http.StatusTooManyRequests)
				return
			}

			next.ServeHTTP(w, r)
		})
	}
}

С gRPC необходимо реализовать 2 вида interceptor'а - унарный и потоковый:

func NewGRPCUnaryInterceptor(
	algo algos.Limiter,
	keyDer GRPCKeyDerivationFunction,
	costDer GRPCCostDerivationFunction,
	failOpen bool,
) grpc.UnaryServerInterceptor {
	check := newGRPCCheck(algo, keyDer, costDer, failOpen)
	return func(
		ctx context.Context,
		req any,
		info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler,
	) (any, error) {
		if err := check(ctx); err != nil {
			return nil, err
		}
		return handler(ctx, req)
	}
}

func NewGRPCStreamInterceptor(
	algo algos.Limiter,
	keyDer GRPCKeyDerivationFunction,
	costDer GRPCCostDerivationFunction,
	failOpen bool,
) grpc.StreamServerInterceptor {
	check := newGRPCCheck(algo, keyDer, costDer, failOpen)
	return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if err := check(ss.Context()); err != nil {
			return err
		}
		return handler(srv, ss)
	}
}

И check:

func newGRPCCheck(
	algo algos.Limiter,
	keyDer GRPCKeyDerivationFunction,
	costDer GRPCCostDerivationFunction,
	failOpen bool,
) func(ctx context.Context) error {
	if algo == nil || keyDer == nil {
		panic("middleware: limiter and key function are required")
	}
	return func(ctx context.Context) error {
		fail := func(stage string, err error) error {
			log.Printf("limiter error on %s: %v", stage, err)
			if failOpen {
				return nil
			}
			return status.Error(codes.Unavailable, "service unavailable")
		}

		tokens := 1.0
		if costDer != nil {
			t, err := costDer(ctx)
			if err != nil {
				return fail("cost derivation", err)
			}
			tokens = t
		}

		key, err := keyDer(ctx)
		if err != nil {
			return fail("key derivation", err)
		}

		req, err := facade.NewRequest(key, tokens)
		if err != nil {
			return fail("request creation", err)
		}

		decision, err := algo.Allow(ctx, req)
		if err != nil {
			return fail("allow", err)
		}
		if !decision.Allowed {
			md := metadata.Pairs("retry-after", strconv.FormatInt(decision.RetryAfter, 10))
			_ = grpc.SetTrailer(ctx, md)
			return status.Error(codes.ResourceExhausted, "rate limit exceeded")
		}
		return nil
	}
}

Как вы могли заметить, везде можно передать функции извлечение ключа и стоимости запроса. Это нужно так как в режиме middleware вы чаще всего будете писать сервисы, которые смотрят наружу, а значит будут принимать запросы от недоверенных клиентов.

Итог

Статья получилась достаточно объёмной, надеюсь она будет полезна ребятам, которые только начинают учить Golang)