leakey bucket

github.com

go-redisを用いたレートリミットのうちleakey bucketを使った実装。rateLimit()でレートリミットを実現している。設定しているのはredis_rateでどれくらいの流量を通すかを決定している。

package main

import (
    "errors"
    "fmt"
    "log"
    "net/http"
    "strconv"
    "time"

    "github.com/go-redis/redis/v8"
    "github.com/go-redis/redis_rate/v9"

    "github.com/uptrace/bunrouter"
)

var ErrRateLimited = errors.New("rate limited")

var limiter *redis_rate.Limiter

func main() {
    rdb := redis.NewClient(&redis.Options{
        Addr: "localhost:6379",
    })
    limiter = redis_rate.NewLimiter(rdb)

    router := bunrouter.New(
        bunrouter.Use(errorHandler),
        bunrouter.Use(rateLimit),
    )

    router.GET("/", indexHandler)

    log.Println("listening on http://localhost:9999")
    log.Println(http.ListenAndServe(":9999", router))
}

func indexHandler(w http.ResponseWriter, req bunrouter.Request) error {
    fmt.Println("indexHandler")
    _, err := w.Write([]byte("hello world"))
    return err
}

func rateLimit(next bunrouter.HandlerFunc) bunrouter.HandlerFunc {
    return func(w http.ResponseWriter, req bunrouter.Request) error {
        res, err := limiter.Allow(req.Context(), "project:123", redis_rate.PerSecond(1))
        if err != nil {
            return err
        }

        h := w.Header()
        h.Set("RateLimit-Remaining", strconv.Itoa(res.Remaining))

        if res.Allowed == 0 {
            fmt.Println("rate limited")
            seconds := int(res.RetryAfter / time.Second)
            h.Set("RateLimit-RetryAfter", strconv.Itoa(seconds))
            http.Error(w, "Internal Server Error", http.StatusInternalServerError)
            return ErrRateLimited
        }
        return next(w, req)
    }
}

func errorHandler(next bunrouter.HandlerFunc) bunrouter.HandlerFunc {
    return func(w http.ResponseWriter, req bunrouter.Request) error {
        err := next(w, req)

        switch err {
        case nil:
        case ErrRateLimited:
            w.WriteHeader(http.StatusTooManyRequests)
            _ = bunrouter.JSON(w, bunrouter.H{
                "message": "you are rate limited",
            })
        default:
            w.WriteHeader(http.StatusInternalServerError)
            _ = bunrouter.JSON(w, bunrouter.H{
                "message": err.Error(),
            })
        }

        return err
    }
}

このほかにも golang.org/x/time/rateを用いて実装をすることもできる。こちらはトークンバケットを用いた実装となっている。こちらはRedisなどを用いないので実装は単純になる分システムワイドに動作するわけではない。

pkg.go.dev

こちらはleaker bucketかつローカルで動作するライブラリ。上記のライブラリも計算量を抑えて高速に動作するとのこと。

github.com