Go語言使用漏桶算法和令牌桶算法來實現(xiàn)API限流

更新時間：2024年11月21日 10:13:10 作者：左詩右碼

為防止服務(wù)器被過多的請求壓垮,限流是一個至關(guān)重要的技術(shù)手段,下面我們就來看看如何使用漏桶算法和令牌桶算法來實現(xiàn) API 的限流吧

限流的意義

限流的主要目的是保護系統(tǒng)資源，防止因請求量過大導(dǎo)致服務(wù)器崩潰。同時，它也能防止惡意用戶對系統(tǒng)的攻擊，確保服務(wù)的穩(wěn)定性和可用性。

兩種常見的限流算法

1.漏桶算法（Leaky Bucket）

漏桶算法將請求視為水滴，水滴先進入桶中，然后以固定的速率從桶中流出。如果請求的速率超過了桶的流出速率，多余的請求將會被丟棄。

這個算法的優(yōu)點很明顯，就是讓請求非常穩(wěn)定，但是缺點也很明顯，因為請求非常穩(wěn)定，就不適于一些秒殺等一些可能在某一段時間會有洪峰流量的場景。不太好適情況控制流量的進入。

2.令牌桶算法（Token Bucket）

令牌桶算法中，系統(tǒng)會以固定的速率向桶中加入令牌，每個請求需要獲取一個令牌才能執(zhí)行。如果桶中沒有足夠的令牌，請求將被拒絕。

代碼實現(xiàn)

在這個示例中，我們將展示如何在 Gin 框架中應(yīng)用這兩種算法來實現(xiàn) API 的限流。

package main

import (
	"fmt"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	ratelimit2 "github.com/juju/ratelimit" // 令牌桶算法
	ratelimit1 "go.uber.org/ratelimit"     // 漏桶算法
)

func pingHandler(c *gin.Context) {
	c.JSON(200, gin.H{
		"message": "pong",
	})
}

func pingHandler2(c *gin.Context) {
	c.JSON(200, gin.H{
		"message": "pong2",
	})
}

// rateLimit1 使用漏桶算法來限制請求速率
func rateLimit1() func(ctx *gin.Context) {
	// 漏桶算法，第一個參數(shù)為兩滴水滴之間的時間間隔。
	// 此時表示兩滴水之間的時間間隔是 100 納秒
	rl := ratelimit1.New(100)

	return func(ctx *gin.Context) {
		// 嘗試取出水滴
		if waitTime := rl.Take().Sub(time.Now()); waitTime > 0 {
			fmt.Printf("需要等待 %v 秒，下一滴水才會滴下來\n", waitTime)
			// 這里我們可以讓程序繼續(xù)等待，也可以直接拒絕掉
			// time.Sleep(waitTime)
			ctx.String(http.StatusOK, "rate limit, try again later")
			ctx.Abort()
			return
		}
		// 證明可以繼續(xù)執(zhí)行
		ctx.Next()
	}
}

// rateLimit2 使用令牌桶算法來限制請求速率
func rateLimit2() func(ctx *gin.Context) {
	// 令牌桶算法：第一個參數(shù)為每秒填充令牌的速率為多少
	// 第二個參數(shù)為令牌桶的容量
	// 這里表示每秒填充 10 個令牌
	rl := ratelimit2.NewBucket(time.Second, 10)

	return func(ctx *gin.Context) {
		// 嘗試取出令牌
		var num int64 = 1
		// 這里表示需要 num 個令牌和已經(jīng)取出的令牌數(shù)是否相等
		// 不相等，則表示超過了限流
                // 比如，假設(shè)每一個請求過來消耗2個令牌，但是從桶中取出的令牌個數(shù)為 1 ，那么則認為超過了限流（一般而言是一個請求消耗一個令牌，這里僅為舉例）
		if rl.TakeAvailable(num) != num {
			// 此次沒有取到令牌，說明超過了限流
			ctx.String(http.StatusOK, "rate limit, try again later")
			ctx.Abort()
			return
		}
		// 證明可以繼續(xù)執(zhí)行
		ctx.Next()
	}
}

func main() {
	r := gin.Default()

	// 漏桶算法限流
	r.GET("/ping", rateLimit1(), pingHandler)

	// 令牌桶算法限流
	r.GET("/ping2", rateLimit2(), pingHandler2)

	r.Run()
}