Go の Webサーバーで Rate Limit を実装する


# 概要

Go と [Echo](https://echo.labstack.com/) フレームワークで Webサーバーを開発しており、 [`golang.org/x/time/rate`](https://pkg.go.dev/golang.org/x/time/rate) パッケージを用いて Rate Limit を実装しましたのでご紹介します。

スケールアウトを考慮するのであれば、 OpenResty 用の Lua スクリプトを書き、Redis などのKVSでアクセス情報を保持するのが良いと思います。
しかし、今回は小規模利用かつシングルプロセスで起動することを前提としているので、Go のコードとインメモリでのデータ管理に留めることにしました。

ちなみに、 Rate Limit については、他にも [`uber-go/ratelimit`](https://github.com/uber-go/ratelimit) というパッケージもあります。
リポジトリの `README.md` のFAQを読むと、複雑なユースケースであれば `x/time/rate` の方が良いとのことなので、こちらを使用します。

# 環境情報

この記事で扱っているバージョンは以下のとおりです。

* 言語
  * Go (`go version go1.24.1 darwin/arm64`)
* パッケージ
  * [`github.com/labstack/echo/v4 v4.13.3`](https://pkg.go.dev/github.com/labstack/echo/v4@v4.13.3)
  * [`golang.org/x/time v0.11.0`](https://pkg.go.dev/golang.org/x/time@v0.11.0/rate)

# Rate Limit の実装

Echo フレームワークのミドルウェアとして利用すること、IPアドレスではなくAPI Keyごとに制限することを念頭にコードを書いています。

```go
package middlewares

import (
  "fmt"
  "sync"
  "time"

  "github.com/labstack/echo/v4"
  "golang.org/x/time/rate"
)

// NOTE:
// maxRequestPerSeconds < burstPerSecond < limiterExpiration

const (
  maxRequestPerSecond = 1
  burstPerSecond      = 3
  limiterExpiration   = 2 * time.Minute
  cleaningInterval    = 30 * time.Second
)

type rateLimiter struct {
  limiter    *rate.Limiter
  lastAccess time.Time
}

type LimitManager struct {
  rlm map[string]*rateLimiter // Rate Limit Map: it maps API keys to their respective rate limiters
  mu  sync.Mutex
}

func NewLimitManager() *LimitManager {
  lm := &LimitManager{
    rlm: make(map[string]*rateLimiter),
  }
  go lm.startCleaning()
  return lm
}

func (lm *LimitManager) withMuLock(f func()) {
  lm.mu.Lock()
  defer lm.mu.Unlock()
  f()
}

// RateLimitByApiKey is a middleware that limits the number of requests.
// Keep in mind it doesn't limit accesses if the API key is not provided in the `X-API-Key` header.
func (lm *LimitManager) RateLimitByApiKey(next echo.HandlerFunc) echo.HandlerFunc {
  return func(c echo.Context) error {
    apiKey := c.Request().Header.Get("X-API-Key")
    if apiKey == "" {
      return next(c)
    }

    var rl *rateLimiter
    lm.withMuLock(func() {
      now := time.Now()

      var exists bool
      if rl, exists = lm.rlm[apiKey]; exists {
        rl.lastAccess = now
      } else {
        rl = &rateLimiter{
          limiter:    rate.NewLimiter(rate.Limit(maxRequestPerSecond), burstPerSecond),
          lastAccess: now,
        }
        lm.rlm[apiKey] = rl
      }
    })
    if !rl.limiter.Allow() {
      return echo.ErrTooManyRequests
    }
    return next(c)
  }
}

func (lm *LimitManager) clean() {
  now := time.Now()
  lm.withMuLock(func() {
    for k, v := range lm.rlm {
      if now.Sub(v.lastAccess) > limiterExpiration {
        delete(lm.rlm, k)
        fmt.Printf("[Limiter] Deleted rate limiter for API key: %s\n", k)
      }
    }
  })
}

func (lm *LimitManager) startCleaning() {
  t := time.NewTicker(cleaningInterval)
  defer t.Stop()

  for range t.C {
    lm.clean()
  }
}
```

# 実装上の注意点と工夫

## メモリリークの対策

当初は `LimitManager` Struct を以下のように定義し、 `startCleaning()` などを実装していませんでした。

```go
type LimitManager struct {
  rlm map[string]*rate.Limiter  // x/time/rate で定義
  mu  sync.Mutex
}
```

しかし、このコードだと、一度でもアクセスのあった API Key に対する `*rate.Limiter` を保持し続けることになるのでメモリを大量に使ってしまうかもしれません。
そのため、最後のアクセス時刻を元に、一定期間アクセスがないリミッターを削除するクリーンアップ処理（`clean()` 関数）を、 `startCleaning()` によって定期的に実行しています。

## スレッドセーフティ

`golang.org/x/time/rate` の `*rate.Limiter` 自体はスレッドセーフなため、 `Allow()` の呼び出しでは競合が発生しません。（[ソースコードの `reserveN` 関数](https://cs.opensource.google/go/x/time/+/refs/tags/v0.11.0:rate/rate.go;l=338-339)でロックがかけられています。）
つまり、以下のように `rateLimiter` に `sync.Mutex` をセットする必要はありません。

```go
type rateLimiter struct {
  limiter    *rate.Limiter
  lastAccess time.Time
  mu         sync.Mutex  // 不要
}
```

しかし、 `RateLimitByApiKey()` と `clean()` はどちらも `rlm` に対して読み書きを行います。整合性を担保するために、どちらかで `rlm` が参照されているときにはロックをかける必要があります。
（そもそも、 Go の map は同時書き込みを許さないため排他制御が必要となります。）
`Limit` Struct の `mu` はこのような理由で存在しています。

## ミドルウェアの呼び出し順

ミドルウェアの呼び出し順に注意する必要があります。具体的には以下の順番で処理を進めるのが良いです。

1. API Key による認証
2. `RateLimitByApiKey()` の呼び出し

`RateLimitByApiKey()` では API Key の有効性を検証していません。

もし、順番を逆にしてミドルウェアを呼び出すと、不正な API Key をセットして大量にリクエストが送られてきたときに、メモリを大量に消費してしまう可能性があります。

## `sync.RWMutex` の利用

このコードでは `sync.Mutex` を利用しています。

しかし、読み込みがメインで、書き込みがそこまで多くないのであれば `sync.RWMutex` を利用するのも良いでしょう。
ただ、そこはこの記事の主旨と乖離するため詳細は省きます。