go gin 基于令牌桶实现的接口限流-526互联

令牌桶实现接口限流

简单有效.

但是不一定对得上我现在需要使用的业务.

新建 main.go

package main

import (
	"github.com/gin-gonic/gin"
	"sync"
	"time"
)

// 令牌桶demo

func main() {
	r := gin.Default()
	// 在路由中使用中间件
	r.Use(LimitHandler(1))
	r.GET("/", func(c *gin.Context) {
		c.String(200, "Hello, World!")
	})
	r.Run(":8080")
}

// 定义令牌桶
type TokenBucket struct {
	capacity  int64      // 桶的容量
	rate      float64    // 令牌放入速率
	tokens    float64    // 当前令牌数量
	lastToken time.Time  // 上一次放令牌的时间
	mtx       sync.Mutex // 互斥锁
}

// 根据令牌桶的状态来判断服务是否可用
func (tb *TokenBucket) Allow() bool {
	tb.mtx.Lock()
	defer tb.mtx.Unlock()
	now := time.Now()
	// 计算需要放的令牌数量
	tb.tokens = tb.tokens + tb.rate*now.Sub(tb.lastToken).Seconds()
	if tb.tokens > float64(tb.capacity) {
		tb.tokens = float64(tb.capacity)
	}
	// 判断是否允许请求
	if tb.tokens >= 1 {
		tb.tokens--
		tb.lastToken = now
		return true
	} else {
		return false
	}
}

// 使用中间件限流: 根据令牌桶算法
func LimitHandler(maxConn int) gin.HandlerFunc {
	tb := &TokenBucket{
		capacity:  int64(maxConn),
		rate:      0.2,
		tokens:    1,
		lastToken: time.Now(),
	}
	return func(c *gin.Context) {
		if !tb.Allow() {
			c.String(503, "Too many request")
			c.Abort()
			return
		}
		c.Next()
	}
}