概要
Go と Echo フレームワークで Webサーバーを開発しており、 golang.org/x/time/rate パッケージを用いて Rate Limit を実装しましたのでご紹介します。
スケールアウトを考慮するのであれば、 OpenResty 用の Lua スクリプトを書き、Redis などのKVSでアクセス情報を保持するのが良いと思います。
しかし、今回は小規模利用かつシングルプロセスで起動することを前提としているので、Go のコードとインメモリでのデータ管理に留めることにしました。
ちなみに、 Rate Limit については、他にも uber-go/ratelimit というパッケージもあります。
リポジトリの README.md のFAQを読むと、複雑なユースケースであれば x/time/rate の方が良いとのことなので、こちらを使用します。
環境情報
この記事で扱っているバージョンは以下のとおりです。
- 言語
- Go (
go version go1.24.1 darwin/arm64)
- Go (
- パッケージ
Rate Limit の実装
Echo フレームワークのミドルウェアとして利用すること、IPアドレスではなくAPI Keyごとに制限することを念頭にコードを書いています。
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() などを実装していませんでした。
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 関数でロックがかけられています。）
つまり、以下のように rateLimiter に sync.Mutex をセットする必要はありません。
type rateLimiter struct {
limiter *rate.Limiter
lastAccess time.Time
mu sync.Mutex // 不要
}
しかし、 RateLimitByApiKey() と clean() はどちらも rlm に対して読み書きを行います。整合性を担保するために、どちらかで rlm が参照されているときにはロックをかける必要があります。
（そもそも、 Go の map は同時書き込みを許さないため排他制御が必要となります。）
Limit Struct の mu はこのような理由で存在しています。
ミドルウェアの呼び出し順
ミドルウェアの呼び出し順に注意する必要があります。具体的には以下の順番で処理を進めるのが良いです。
- API Key による認証
-
RateLimitByApiKey()の呼び出し
RateLimitByApiKey() では API Key の有効性を検証していません。
もし、順番を逆にしてミドルウェアを呼び出すと、不正な API Key をセットして大量にリクエストが送られてきたときに、メモリを大量に消費してしまう可能性があります。
sync.RWMutex の利用
このコードでは sync.Mutex を利用しています。
しかし、読み込みがメインで、書き込みがそこまで多くないのであれば sync.RWMutex を利用するのも良いでしょう。
ただ、そこはこの記事の主旨と乖離するため詳細は省きます。