【Go Web 開發】API 限流

如果您正在構建公共使用的 API,那麼很可能需要實現某種形式的限流,以防止客戶端過快地發出大量請求,從而給服務器帶來很大的壓力。

接下來的內容我們將創建中間件來實現這一點。

本質上,這個中間件檢查在過去 N 秒鐘內服務器接收到多少請求,如果請求太多就向客戶端發送 "429 Too Many Requests" 響應。我們需要將這個中間件放在業務處理程序之前,對請求被處理之前進行攔截,避免對請求進行 JSON 序列化或數據庫查詢等操作。

你將學習到:

全侷限流

我們先爲應用程序創建一個全侷限流器,然後慢慢深入。全侷限流將考慮 API 接收到的所有請求 (而不是爲每個客戶端設置單獨的限速)。

可以利用 x/time/rate 包來幫助我們,而不是從頭開始編寫限流邏輯,因其非常複雜和耗時。x/time/rate 包提供了經過測試的令牌桶限流器。先下載依賴包:

$ go get golang.org/x/time/rate@latest
go: downloading golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba 
go: added golang.org/x/time v0.0.0-20210220033141-f8bda1e9f3ba

在開始寫代碼之前,我們介紹下令牌桶限流是如何工作的。x/time/rate 包的官方文檔描述:

限流器能控制事件發生的頻率。限流器實現了一個大小爲 b 的 “令牌桶”,最初桶裝滿 token 並以每秒 r 個令牌的速率重新填充。

將這些內容放在 API 上下文中可以理解爲:

在實踐中,這意味着我們的應用程序允許最大 “突發”b 個連續的 HTTP 請求,但隨着時間的推移,它將允許平均每秒 r 個請求。

根據 x/time/rate 包,爲了創建一個令牌桶限流器,我們需要使用 NewLimiter() 函數,其包含如下參數:

 // Limit類型是float64的別名
func NewLimiter(r Limit, b int) *Limiter

因此,如果我們要創建一個每秒允許 2 個請求,支持 4 個突發請求的限流器,我們可以使用如下代碼實現:

//每秒允許2個請求,一次最多4個請求
limiter := rate.NewLimiter(2, 4)

執行全侷限流

前面我們從宏觀角度解釋限流,下面讓我們進入代碼,看看它在實踐中是如何工作的。

使用中間件模式的一個優點是,它可以包含 “初始化” 代碼,當我們用中間件包裝一些東西時,它只運行一次,就可以處理所有的請求。

func (app *application) exampleMiddleware(next http.Handler) http.Handler {
    //當使用中間件封裝後,這裏代碼只運行一次
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.request) {
        //這裏代碼在每個請求到來都會執行
        next.ServerHTTP(w, r)
    })
}

在我們的例子中,將創建一個 rateLimit()中間件方法,它將創建一個新的限流器作爲 “初始化” 代碼的一部分,然後對隨後處理的每個請求使用這個限流器。打開 cmd/api/middleware.go 文件並創建以下中間件:

package main

...

func (app *application)rateLimit(next http.Handler) http.Handler {
    //初始化限流器
    limit := rate.NewLimiter(2, 4)
    //運行的是一個閉包
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        //調用limit.Allow()函數檢查請求是否允許,不允許就返回429
        if !limit.Allow(){
            app.rateLimitExceededResponse(w, r)
            return
        }
        next.ServeHTTP(w, r)
    })
}

在這段代碼中,每當調用限流器上的 Allow() 方法時,只會從桶中消耗一個令牌。如果桶中沒有令牌,那麼 Allow() 將返回 false,以此觸發向客戶端返回 429 Too Many Requests 響應。

要注意的是 Allow() 方法實現使用了互斥鎖保護,併發使用是安全的。

下面到 cmd/api/errors.go 文件創建 rateLimitExeededRespose() 幫助函數:

package main

...

func (app *application) rateLimitExceededResponse(w http.ResponseWriter, r *http.Request) {
    message := "rate limit exceeded"
    app.errorResponse(w, r, http.StatusTooManyRequests, message)
}

最後,在 cmd/api/routes.go 文件中需要將 rateLimit()中間件添加到中間件服務鏈中。它應該在我們的 panicRecovery 中間件之前運行 (以便恢復 rateLimit() 中有任何 panic 發生),但是應該儘早使用它,以防止服務器進行不必要的工作。

File: cmd/api/routes.go

package main

...

func (app *application) routes() http.Handler {
    router := httprouter.New()

    router.NotFound = http.HandlerFunc(app.notFoundResponse)
    router.MethodNotAllowed = http.HandlerFunc(app.methodNotAllowedResponse)

    router.HandlerFunc(http.MethodGet, "/v1/healthcheck", app.healthcheckHandler)

    router.HandlerFunc(http.MethodGet, "/v1/movies", app.listMoviesHandler)
    router.HandlerFunc(http.MethodPost, "/v1/movies", app.createMovieHandler)
    router.HandlerFunc(http.MethodGet, "/v1/movies/:id", app.showMovieHandler)
    router.HandlerFunc(http.MethodPatch, "/v1/movies/:id", app.updateMovieHandler)
    router.HandlerFunc(http.MethodDelete, "/v1/movies/:id", app.deleteMovieHandler)
        //用ratelimit()中間件處理router
    return app.recoverPanic(app.rateLimit(router))
}

現在我們應該準備好進行測試了!重啓服務,然後在另一個終端窗口執行以下命令批量發送請求到 GET /v1/healthcheck 接口。你會看到如下響應:

$ for i in {1..6}; do curl http://localhost:4000/v1/healthcheck; done
{
        "status": "available",
        "system_info": {
                "environment": "development",
                "version": "1.0.0"
        }
}
{
        "status": "available",
        "system_info": {
                "environment": "development",
                "version": "1.0.0"
        }
}
{
        "status": "available",
        "system_info": {
                "environment": "development",
                "version": "1.0.0"
        }
}
{
        "status": "available",
        "system_info": {
                "environment": "development",
                "version": "1.0.0"
        }
}
{
        "error": "rate limit exceeded"
}
{
        "error": "rate limit exceeded"
}

從這裏我們可以看到,前 4 個請求成功了,因爲我們的限流器設置爲允許接收 “突發”4 個請求。一旦這 4 個請求被用完,桶中的令牌已經用完,API 將返回“rate limit exceeded” 錯誤響應。

如果您等待一秒鐘並重新運行此命令,您應該會發現第二批中的一些請求再次成功,這是因爲令牌桶以每秒兩個令牌的速度重新填充到桶中。

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/mo61fGQ7pnhnGLz3ctPzrw