Go 语言结合 Redis 实现固定窗口、滑动窗口、令牌桶和漏桶限流算法的示例代码

固定窗口算法

原理：将时间划分为固定大小的窗口，在每个窗口内对请求进行计数。如果请求数超过设定的阈值，则拒绝后续请求，直到进入下一个窗口。
代码：

Go 复制代码

package main

import (
    "fmt"
    "time"

    "github.com/go-redis/redis/v8"
)

// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{
    Addr:     "localhost:6379", // Redis 服务器地址和端口
    Password: "",               // Redis 密码，这里为空表示无密码
    DB:       0,                // 使用默认的数据库编号
})

// 定义固定窗口算法的常量
const (
    windowSize = 60 // 窗口大小，单位为秒
    limit      = 100 // 窗口内允许的最大请求数
)

// fixedWindowRateLimit 函数实现了固定窗口限流算法
func fixedWindowRateLimit(key string) bool {
    // 获取当前时间的 Unix 时间戳
    currentTime := time.Now().Unix()
    // 计算当前窗口的起始时间
    windowStart := currentTime - (currentTime % windowSize)
    // 构建当前窗口对应的 Redis key
    windowKey := fmt.Sprintf("%s:%d", key, windowStart)

    // 对 Redis 中的当前窗口计数器进行自增操作，并获取自增后的计数值
    count, err := rdb.Incr(rdb.Context(), windowKey).Result()
    if err != nil {
        // 如果自增操作出现错误，打印错误信息并返回 false
        fmt.Println("Error incrementing counter:", err)
        return false
    }
    // 如果计数值为 1，说明是该窗口的第一个请求，为该窗口设置过期时间
    if count == 1 {
        rdb.Expire(rdb.Context(), windowKey, time.Duration(windowSize)*time.Second)
    }
    // 如果计数值超过了允许的最大请求数，返回 false 表示请求被限流
    if count > limit {
        return false
    }
    // 否则，返回 true 表示请求通过
    return true
}

func main() {
    // 调用 fixedWindowRateLimit 函数进行限流判断
    if fixedWindowRateLimit("api_request") {
        fmt.Println("请求通过")
    } else {
        fmt.Println("请求被限流")
    }
}

优缺点：实现简单，但可能会出现临界问题，即在窗口切换的瞬间可能会出现流量高峰。

滑动窗口算法

原理：将固定窗口进一步细分，通过多个小窗口来统计请求数，随着时间的推移，窗口不断滑动，从而更精确地控制流量。
代码：

Go 复制代码

package main

import (
    "fmt"
    "time"

    "github.com/go-redis/redis/v8"
)

// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{
    Addr:     "localhost:6379", // Redis 服务器地址和端口
    Password: "",               // Redis 密码，这里为空表示无密码
    DB:       0,                // 使用默认的数据库编号
})

// 定义滑动窗口算法的常量
const (
    windowSize = 60 // 窗口大小，单位为秒
    limit      = 100 // 窗口内允许的最大请求数
)

// slidingWindowRateLimit 函数实现了滑动窗口限流算法
func slidingWindowRateLimit(key string) bool {
    // 获取当前时间的 Unix 时间戳
    currentTime := time.Now().Unix()
    // 移除 Redis 有序集合中时间戳小于当前窗口起始时间的元素
    rdb.ZRemRangeByScore(rdb.Context(), key, "0", fmt.Sprintf("%d", currentTime-windowSize))
    // 向 Redis 有序集合中添加当前时间戳，分数为当前时间戳
    rdb.ZAdd(rdb.Context(), key, &redis.Z{
        Score:  float64(currentTime),
        Member: currentTime,
    })
    // 获取 Redis 有序集合中当前窗口内的元素数量
    count, err := rdb.ZCard(rdb.Context(), key).Result()
    if err != nil {
        // 如果获取元素数量操作出现错误，打印错误信息并返回 false
        fmt.Println("Error getting count:", err)
        return false
    }
    // 如果元素数量超过了允许的最大请求数，返回 false 表示请求被限流
    if count > limit {
        return false
    }
    // 否则，返回 true 表示请求通过
    return true
}

func main() {
    // 调用 slidingWindowRateLimit 函数进行限流判断
    if slidingWindowRateLimit("api_request") {
        fmt.Println("请求通过")
    } else {
        fmt.Println("请求被限流")
    }
}

优缺点：比固定窗口算法更精确地控制流量，但实现相对复杂，且需要更多的 Redis 操作。

令牌桶算法

原理：系统以固定的速率向令牌桶中添加令牌，每个请求需要从令牌桶中获取一个或多个令牌才能被处理。如果令牌桶中没有足够的令牌，则请求被拒绝。
代码：

Go 复制代码

package main

import (
    "fmt"
    "time"

    "github.com/go-redis/redis/v8"
)

// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{
    Addr:     "localhost:6379", // Redis 服务器地址和端口
    Password: "",               // Redis 密码，这里为空表示无密码
    DB:       0,                // 使用默认的数据库编号
})

// 定义令牌桶算法的常量
const (
    capacity        = 100        // 令牌桶的容量
    rate            = 1          // 令牌生成速率，每秒生成 1 个令牌
    requiredTokens  = 1          // 每个请求所需的令牌数
)

// tokenBucketScript 是用于执行令牌桶算法的 Lua 脚本
var tokenBucketScript = `
-- 获取传入的令牌桶 key
local tokens_key = KEYS[1]
-- 构建存储上次更新时间的 key
local last_time_key = tokens_key .. ":last_time"
-- 获取令牌桶容量
local capacity = tonumber(ARGV[2])
-- 获取令牌生成速率
local rate = tonumber(ARGV[3])
-- 获取当前时间
local now = tonumber(ARGV[1])
-- 获取每个请求所需的令牌数
local required_tokens = tonumber(ARGV[4])

-- 从 Redis 中获取上次更新时间
local last_time = tonumber(redis.call('get', last_time_key))
-- 从 Redis 中获取当前令牌数
local tokens = tonumber(redis.call('get', tokens_key))

-- 如果上次更新时间为空，说明是首次请求，初始化相关值
if last_time == nil then
    last_time = now
    tokens = capacity
end

-- 计算从上次更新到现在生成的令牌数
local generated_tokens = (now - last_time) * rate
-- 更新令牌数，确保不超过令牌桶容量
tokens = math.min(capacity, tokens + generated_tokens)

-- 如果令牌数小于请求所需的令牌数，说明令牌不足，拒绝请求
if tokens < required_tokens then
    redis.call('set', last_time_key, last_time)
    redis.call('set', tokens_key, tokens)
    return 0
else
    -- 否则，处理请求，扣除相应的令牌数
    tokens = tokens - required_tokens
    redis.call('set', last_time_key, now)
    redis.call('set', tokens_key, tokens)
    return 1
end
`

// tokenBucketRateLimit 函数调用 Lua 脚本来实现令牌桶限流算法
func tokenBucketRateLimit(key string) bool {
    // 获取当前时间的 Unix 时间戳
    currentTime := time.Now().Unix()
    // 执行 Lua 脚本，并获取执行结果
    result, err := rdb.Eval(rdb.Context(), tokenBucketScript, []string{key}, currentTime, capacity, rate, requiredTokens).Int64()
    if err != nil {
        // 如果执行脚本出现错误，打印错误信息并返回 false
        fmt.Println("Error executing script:", err)
        return false
    }
    // 如果结果为 1，表示请求通过；否则，表示请求被限流
    return result == 1
}

func main() {
    // 调用 tokenBucketRateLimit 函数进行限流判断
    if tokenBucketRateLimit("api_request") {
        fmt.Println("请求通过")
    } else {
        fmt.Println("请求被限流")
    }
}

优缺点：能够平滑地控制流量，允许一定程度的突发流量，但实现相对复杂。

漏桶限流算法

原理：请求就像水一样注入漏桶，漏桶以固定的速率处理请求。如果请求的注入速率超过漏桶的处理速率，多余的请求将被丢弃。
代码：

Go 复制代码

package main

import (
    "fmt"
    "time"

    "github.com/go-redis/redis/v8"
)

// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{
    Addr:     "localhost:6379", // Redis 服务器地址和端口
    Password: "",               // Redis 密码，这里为空表示无密码
    DB:       0,                // 使用默认的数据库编号
})

// 定义漏桶算法的常量
const (
    capacity     = 100        // 漏桶的容量
    rate         = 1          // 漏桶的流出速率，每秒流出 1 个单位
    requestSize  = 1          // 每个请求的大小
)

// leakyBucketScript 是用于执行漏桶算法的 Lua 脚本
var leakyBucketScript = `
-- 获取传入的漏桶 key
local bucket_key = KEYS[1]
-- 构建存储上次更新时间的 key
local last_time_key = bucket_key .. ":last_time"
-- 获取漏桶容量
local capacity = tonumber(ARGV[2])
-- 获取漏桶流出速率
local rate = tonumber(ARGV[3])
-- 获取当前时间
local now = tonumber(ARGV[1])
-- 获取每个请求的大小
local request_size = tonumber(ARGV[4])

-- 从 Redis 中获取上次更新时间
local last_time = tonumber(redis.call('get', last_time_key))
-- 从 Redis 中获取当前漏桶中的水量
local water = tonumber(redis.call('get', bucket_key))

-- 如果上次更新时间为空，说明是首次请求，初始化相关值
if last_time == nil then
    last_time = now
    water = 0
end

-- 计算从上次更新到现在流出的水量
local outflow = (now - last_time) * rate
-- 更新漏桶中的水量，确保不小于 0
water = math.max(0, water - outflow)

-- 如果漏桶中的水量加上当前请求的大小超过了漏桶容量，说明漏桶已满，拒绝请求
if water + request_size > capacity then
    redis.call('set', last_time_key, last_time)
    redis.call('set', bucket_key, water)
    return 0
else
    -- 否则，处理请求，增加漏桶中的水量
    water = water + request_size
    redis.call('set', last_time_key, now)
    redis.call('set', bucket_key, water)
    return 1
end
`

// leakyBucketRateLimit 函数调用 Lua 脚本来实现漏桶限流算法
func leakyBucketRateLimit(key string) bool {
    // 获取当前时间的 Unix 时间戳
    currentTime := time.Now().Unix()
    // 执行 Lua 脚本，并获取执行结果
    result, err := rdb.Eval(rdb.Context(), leakyBucketScript, []string{key}, currentTime, capacity, rate, requestSize).Int64()
    if err != nil {
        // 如果执行脚本出现错误，打印错误信息并返回 false
        fmt.Println("Error executing script:", err)
        return false
    }
    // 如果结果为 1，表示请求通过；否则，表示请求被限流
    return result == 1
}

func main() {
    // 调用 leakyBucketRateLimit 函数进行限流判断
    if leakyBucketRateLimit("api_request") {
        fmt.Println("请求通过")
    } else {
        fmt.Println("请求被限流")
    }
}

优缺点：可以平滑地处理请求，保证请求以固定的速率被处理，但无法应对突发流量。

总结

算法	实现复杂度	限流精度	突发流量处理	适用场景
固定窗口算法	简单	低	差	流量平稳、精度要求低的场景
滑动窗口算法	较复杂	高	一般	流量波动大、精度要求高的场景
令牌桶算法	较复杂	高	好	可容忍突发流量、需平滑限流的场景
漏桶算法	较复杂	高	差	对流量稳定性要求极高的场景