RedisとLuaでSliding window logを実装してGolangから呼ぶ
この記事は GENDA Advent Calendar 2023 20日目の記事です。
本当は
「俺の全プログラミングスキルを集結し、麻雀の上がり判定と点数計算を1Sec未満で行う(前編)」
という記事を書き上げたかったのですが、師走に敗北しました。
株式会社GENDA テックリードにしがみです
TL;DR
Sliding window log は Rate limiting のアルゴリズム。
Redis と Lua を使うと Sliding window log が簡単に実装できる。
Lua で実装することで Lua サポートされているRedisライブラリがある言語なら、
どの言語でも簡単に利用できる。
Sliding window logとは？
Rate limiting アルゴリズムの一つ。
Wikipedia によると他にはこのようなアルゴリズムがあるようです。
- Token bucket
- Leaky bucket
- Fixed window counter
- Sliding window counter
Sliding window log は他に比べるとこのような特徴があります。
- n回 / m分 を他と比べ正確に実現できる
- Sliding window counter や Fixed window counter と比べ、期間分のタイムスタンプを保持管理するため、実装コスト、ストレージ(メモリ)コストがかかる。
前置き
今回は理由があって Redis を選定したわけではなく、
Redis を使うことが先に決まっていて、Rate limiting をやってみようという順番でした。
Sliging window log に決めた理由は、実装的に勉強になると思ったからです。
結果として Redis を使うと Sliding window log を楽に実現できたためとても満足しています。
もちろん RDBMS でも NoSQL でも他の In memory store でも実現することはできると思います。
Luaの有無について
Lua の使用有無にはメリットデメリットあったのですが今回は Lua を使いました
メリット
Lua を使うと Atomic 性が保証される(Pipelining や Transaction の手間が省ける)
呼び出し言語側がロジックを持たずに済む
デメリット
Lua の学習コスト
テスト、デバッグ容易性の低さ
最終的には勉強になるし Lua でやってみようと決めました。
すごくシンプルな Lua スクリプトで実装できたので Lua の学習コストというデメリットは結果的に少なかったです。
GoのRedis ライブラリについて
次の二つが Go で有名な Redis ライブラリかなと思います。
redigo の方は使ってないのでわからないですが、go-redis の方は Redis 公式のライブラリだけあって、Lua サポートがしっかりされていると感じました。
Redis には Lua の script を cache する機能もあるみたいですが、go-redis がよしなにやってくれるので全く気にせずに済みました。
コード
package main
import (
"context"
"errors"
"fmt"
"time"
"github.com/redis/go-redis/v9"
)
func connectRedis() *redis.Client {
return redis.NewClient(&redis.Options{
Addr: ":8080",
Password: "",
DB: 0,
})
}
func main() {
rdb := connectRedis()
ratelimiter := NewSlidingWindowStaticKey(rdb, 4, 2, "your-key") // 4秒に2回まで
ctx := context.Background()
loop := 10 * time.Second
count := 0
execCount := 0
fmt.Printf("設定：%d 回 / %.0f 秒\n", ratelimiter.allowedMaxCount, ratelimiter.window.Seconds())
for begin := time.Now(); time.Since(begin) < loop; {
count++
ok := callProcWithRateLimit(ctx, ratelimiter, "your_key")
if ok {
execCount++
}
}
fmt.Printf("実行 / 呼び出し = %d / %d\n", execCount, count)
}
func callProcWithRateLimit(ctx context.Context, ratelimit *SlidingWindow, key string) bool {
allowed, err := ratelimit.Allow(ctx, time.Now())
if err != nil {
panic(err)
}
if !allowed {
return false
}
proc()
return true
}
func proc() {
fmt.Printf("%s: heavy calc\n", time.Now().Format(time.TimeOnly))
}
var ErrEmptyKey = errors.New("key is empty")
type (
SlidingWindow struct {
window time.Duration
allowedMaxCount int64
keyGen KeyGenFunc
rdb redis.UniversalClient
}
)
type KeyGenFunc func() string
func NewSlidingWindow(rdb redis.UniversalClient, windowSeconds, allowedMaxCount int64, keyGen KeyGenFunc) *SlidingWindow {
return &SlidingWindow{
rdb: rdb,
window: time.Duration(windowSeconds) * time.Second,
allowedMaxCount: allowedMaxCount,
keyGen: keyGen,
}
}
func NewSlidingWindowStaticKey(rdb redis.UniversalClient, windowSeconds, allowedMaxCount int64, key string) *SlidingWindow {
return NewSlidingWindow(rdb, windowSeconds, allowedMaxCount, func() string {
return key
})
}
func (sw SlidingWindow) Allow(ctx context.Context, now time.Time) (bool, error) {
return allow(ctx, sw.rdb, sw.keyGen(), sw.window, sw.allowedMaxCount, now)
}
func allow(ctx context.Context, rdb redis.UniversalClient, key string,
window time.Duration, limit int64, now time.Time,
) (bool, error) {
if key == "" {
return false, ErrEmptyKey
}
sec := window.Seconds()
if sec <= 0 {
panic("SlidingWindow: second must be positive")
}
if limit <= 0 {
panic("SlidingWindow: limit must be positive")
}
ts := now.UnixMicro()
if ts <= 0 {
panic("SlidingWindow: now must be positive")
}
return slidingWindow.Run(ctx, rdb, []string{key}, sec, limit, ts).Bool() // nolint:wrapcheck
}
// (now-window, now].
var slidingWindow = redis.NewScript(`
local windowMicro = tonumber(ARGV[1]) * 1000000
local current_time = redis.call('TIME')
local ts = tonumber(current_time[1]) * 1000000 + tonumber(current_time[2])
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ts - windowMicro)
if redis.call('ZCARD', KEYS[1]) < tonumber(ARGV[2]) then
redis.call('ZADD', KEYS[1], ts, current_time[1] .. current_time[2])
redis.call('EXPIRE', KEYS[1], ARGV[1])
return 1
end
return 0
`)
実行
go mod init example.com
go mod tidy
// Redis起動
go run main.go
// 以下実行結果のstdout
設定：2 回 / 4 秒
11:01:11: heavy calc
11:01:11: heavy calc
11:01:15: heavy calc
11:01:15: heavy calc
11:01:19: heavy calc
11:01:19: heavy calc
実行 / 呼び出し = 6 / 7229
ちゃんと4秒に2回まで許可されていますね
Discussion