限流策略之令牌桶限流

一、限流模型

常见的限流模型有:

其他限流模型不再过多赘述，这里只聊一聊令牌桶限流。

二、令牌桶限流原理

令牌桶算法是一种常见的限流算法，用于控制在一个时间段内通过的请求速率。它的原理相对简单，基于一个桶（Bucket）来控制请求的流量。
下面是令牌桶算法的基本原理：

（图片来源于网络）

三、使用介绍

代码关键部分均做注释，不过多详解

桶结构和配置：

type UserTokenBucket struct {  
    Capacity int64 // 令牌桶的容量  
    Rate float64 // 令牌放入速率  
    Tokens float64 // 当前令牌数量  
    LastToken time.Time // 上一次放令牌的时间  
    Requests int // 当前用户的请求计数器  
    Interval time.Duration // 时间间隔  
    Lock sync.Mutex // 互斥锁  
    TimerStarted bool // 计时器是否已启动  
    Timer *time.Timer // 计时器  
}  
  
type LimitHandlerConfig struct {  
    MaxConn int64 // 最大连接数  
    Interval time.Duration // 时间间隔  
    MaxReq int // 最大请求次数  
}

实例桶结构和判断是否允许请求：

func GetUserTokenBucket(key string, lock *sync.Mutex, userBuckets map[string]*UserTokenBucket, config LimitHandlerConfig) *UserTokenBucket {  
    lock.Lock()  
    defer lock.Unlock()  
  
    // 检查用户是否已有令牌桶  
    tb, found := userBuckets[key]  
    if !found {  
        // 创建新的令牌桶  
        tb = &UserTokenBucket{  
            Capacity: config.MaxConn,  
            Rate: 1.0,  
            Tokens: 0,  
            LastToken: time.Now(),  
            Requests: 0,  
            Interval: config.Interval,  
            Lock: sync.Mutex{},  
        }  
        // 将令牌桶添加到用户令牌桶映射中  
        userBuckets[key] = tb  
  
        // 启动定时器以重置用户请求计数器  
        go func() {  
            time.Sleep(config.Interval)  
            lock.Lock()  
            delete(userBuckets, key)  
            lock.Unlock()  
        }()  
    }  
    return tb  
}  
  
func (tb *UserTokenBucket) Allow(maxRequests int, duration time.Duration) bool {  
    tb.Lock.Lock()  
    defer tb.Lock.Unlock()  
  
    now := time.Now()  
    // 计算需要放的令牌数量  
    tb.Tokens += tb.Rate * now.Sub(tb.LastToken).Seconds()  
    // 限制令牌数量不超过容量  
    if tb.Tokens > float64(tb.Capacity) {  
        tb.Tokens = float64(tb.Capacity)  
    }  
  
    // 判断是否允许请求  
    if tb.Requests >= maxRequests && duration > time.Since(tb.LastToken) {  
        return false  
    } else if duration < time.Since(tb.LastToken) {  
        // 用户超过请求限制，重置计数器和计时器  
        tb.Requests = 0  
        tb.LastToken = now  
        if tb.TimerStarted {  
            tb.Timer.Stop()  
        }  
        tb.TimerStarted = false  
    }  
    // 更新计数器和令牌桶状态  
    tb.Requests++  
    tb.Tokens -= 1  
    tb.LastToken = now  
  
    // 启动计时器以重置计数器  
    if !tb.TimerStarted {  
        tb.Timer = time.AfterFunc(tb.Interval, func() {  
            tb.Lock.Lock()  
            tb.Requests = 0  
            tb.TimerStarted = false  
            tb.Lock.Unlock()  
        })  
        tb.TimerStarted = true  
    }  
    return true  
}

代码逻辑：

限流中间件：

func NewLimitHandler(config tools.LimitHandlerConfig) gin.HandlerFunc {  
    userBuckets := make(map[string]*tools.UserTokenBucket)  
    lock := sync.Mutex{}  
    return func(c *gin.Context) { 
        //获取用户标识示例：
        //具体根据自己代码修改
        ip := c.ClientIP()  
        ua := c.Request.UserAgent()  
        Key := ip + "_" + ua  

        // 获取或创建用户的令牌桶  
        tb := tools.GetUserTokenBucket(key, &lock, userBuckets, config)  
        // 检查用户是否超过了时间间隔内的请求限制  
        if !tb.Allow(config.MaxReq, config.Interval) {  
            c.String(503, "Too many requests")  
            c.Abort()  
            return  
        }  
        // 允许请求通过  
        c.Next()  
    }  
}

路由方法：

//根据自己代码逻辑修改，可以将结构体实例写到配置文件中，这里不过多展开
var limitHandlerConfig = LimitHandlerConfig{  
    MaxConn: 10,  
    Interval: time.Second * 5,  
    MaxReq: 5,  
}

r.GET("/helloWorld", logic.NewLimitHandler(limitHandlerConfig), func(c *gin.Context) {  
    c.String(200, "Hello, World!")  
})

总结

令牌桶具有：

等优点。

总体而言，令牌桶算法是一种灵活、平滑和可控制的限流算法，适用于多种不同类型的应用场景。