golang+redis 实现分布式限流

实现分布式滑动窗口限流

复制代码
package main

import (
	"context"
	"fmt"
	"time"

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

// Redis 客户端
var ctx = context.Background()
var rdb = redis.NewClient(&redis.Options{
	Addr:     "localhost:6379", // Redis 服务器地址
	Password: "",               // 无密码
	DB:       0,                // 默认数据库
})

// 限流配置
const (
	windowSize  = 10               // 时间窗口(秒)
	maxRequests = 5                // 允许的最大请求数
	requestKey  = "sliding_window" // Redis Key
)

// 限流检查
func isAllowed() bool {
	now := time.Now().Unix()           // 当前时间戳
	minTime := now - int64(windowSize) // 窗口开始时间

	// 使用 Lua 脚本保证操作的原子性
	luaScript := `
        redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", ARGV[1]) -- 清除过期请求
        local reqCount = redis.call("ZCOUNT", KEYS[1], "-inf", "+inf") -- 获取当前窗口内请求数
        if reqCount < tonumber(ARGV[2]) then
            redis.call("ZADD", KEYS[1], ARGV[3], ARGV[3]) -- 添加当前请求时间戳
            redis.call("EXPIRE", KEYS[1], ARGV[4]) -- 设置过期时间,防止 key 长期存在
            return 1
        else
            return 0
        end
    `

	result, err := rdb.Eval(ctx, luaScript, []string{requestKey}, minTime, maxRequests, now, windowSize).Int()
	if err != nil {
		fmt.Println("Redis error:", err)
		return false
	}
	return result == 1
}

func main() {
	for i := 0; i < 10; i++ {
		if isAllowed() {
			fmt.Println("Request allowed")
		} else {
			fmt.Println("Rate limit exceeded")
		}
		time.Sleep(1 * time.Second) // 模拟请求间隔
	}
}

令牌桶限流

复制代码
package main

import (
	"context"
	"fmt"
	"time"

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

// Redis 客户端
var ctx = context.Background()
var rdb = redis.NewClient(&redis.Options{
	Addr:     "localhost:6379", // Redis 地址
	Password: "",               // Redis 认证密码(无则留空)
	DB:       0,                // 使用默认数据库
})

// 令牌桶配置
const (
	bucketKey  = "token_bucket" // Redis 令牌桶 Key
	bucketSize = 10             // 令牌桶容量
	refillRate = 2              // 每秒补充的令牌数量
)

// 初始化令牌桶
func initBucket() {
	rdb.Set(ctx, bucketKey, bucketSize, 0) // 初始化令牌桶,存满令牌
}

// 获取令牌
func getToken() bool {
	luaScript := `
        local tokens = redis.call("GET", KEYS[1])
        if tokens == false then
            redis.call("SET", KEYS[1], ARGV[2])
            tokens = ARGV[2]
        end
        tokens = tonumber(tokens)
        if tokens > 0 then
            redis.call("DECR", KEYS[1])
            return 1
        else
            return 0
        end
    `
	result, err := rdb.Eval(ctx, luaScript, []string{bucketKey}, 1, bucketSize).Int()
	if err != nil {
		fmt.Println("Redis Error:", err)
		return false
	}
	return result == 1
}

// 定期补充令牌
func refillTokens() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for range ticker.C {
		rdb.IncrBy(ctx, bucketKey, refillRate)
		currentTokens, _ := rdb.Get(ctx, bucketKey).Int()
		if currentTokens > bucketSize {
			rdb.Set(ctx, bucketKey, bucketSize, 0)
		}
		fmt.Printf("Tokens refilled: %d\n", currentTokens)
	}
}

func main() {
	initBucket()
	go refillTokens() // 启动令牌补充协程

	for i := 0; i < 20; i++ {
		if getToken() {
			fmt.Println("Request allowed")
		} else {
			fmt.Println("Rate limit exceeded")
		}
		time.Sleep(50 * time.Millisecond) // 模拟请求间隔
	}
}
相关推荐
aashuii3 小时前
go客户端ssh交换机
开发语言·golang·ssh
是紫焅呢3 小时前
E结构体基础.go
开发语言·后端·golang·学习方法·visual studio code
clt1233213 小时前
golang excel导出时需要显示刷新
开发语言·后端·golang
mxpan3 小时前
Golang 与 C/C++ 交互实践
c语言·c++·golang
Silverdew*3 小时前
vs code配置go开发环境以及问题解决 could not import cannot find package in GOROOT or GOPATH
开发语言·后端·golang
拍客圈4 小时前
单服务器部署多个Discuz! X3.5站点并独立Redis配置方案
运维·服务器·redis
FmZero4 小时前
Redis使用规范
java·redis·mybatis
乐世东方客8 小时前
Kafka使用Elasticsearch Service Sink Connector直接传输topic数据到Elasticsearch
分布式·elasticsearch·kafka
火龙谷10 小时前
【hadoop】搭建考试环境(单机)
大数据·hadoop·分布式
计算机毕设定制辅导-无忧学长10 小时前
Kafka 可靠性保障:消息确认与事务机制(二)
分布式·kafka·linq