简单的高并发限流场景案例

ByWhat'sUs

简单的高并发限流场景案例

利用Go+Redis+Lua脚本实现当并发过高超过服务承受上限时，只允许部分流量进入服务的限流场景，简单案例

redis.go

package redis

import "github.com/go-redis/redis"

type InitRedis struct {
    RDB *redis.Client
}

func NewInitRedis() *InitRedis {
    return &InitRedis{
	RDB: redis.NewClient(&redis.Options{
	    Addr:     "127.0.0.1:6379",
	    Password: "",
	    DB:       0,
	}),
    }
}

func (r InitRedis) Nosql() *redis.Client {
    return r.RDB
}

lua.go

package lua

import (
    "fmt"
    "github.com/go-redis/redis"
)

var luaScript = `
	local key         = KEYS[1]; -- 键 key
	local now_time    = ARGV[1]; -- 值 value
	local before_time = ARGV[2]; -- 间隔时间之前，纳秒
	local period      = ARGV[3]; -- 时间间隔，多少秒内，设置过期时间
	local requests    = ARGV[4]; -- 时间间隔内的请求次数
 
	redis.pcall("zadd", key, now_time, now_time); -- zset结构设置一个key，zadd(key,value,scores) 
	redis.pcall("zremrangebyscore", key, 0, before_time); -- 移除scores范围在0到bofore_time之间的值，即移除时间窗口之前的行为记录，剩下的都是时间窗口内的
	local count = redis.pcall("zcard", key); -- 获取窗口内的请求数量
	redis.pcall("expire", key, period); -- 设置 zset 过期时间，避免不活跃用户持续占用内存
	
	if tonumber(count) > tonumber(requests) then
		return 2;
	end
	return 1;`

func RunLua(client *redis.Client) (string, error) {
    evalSha, err := client.ScriptLoad(luaScript).Result()
    if err != nil {
	fmt.Println(err)
    }

    return evalSha, nil
}

logic.go

package logic

import (
    "current-limit/lua"
    "fmt"
    "github.com/go-redis/redis"
    "time"
)

// 成功的次数
var successSort = 0

func DoRequest(client *redis.Client) {
    success := isPermitted("nothing", "add/Cart", 3, 10, client)
    if success {
	successSort++
	fmt.Println("成功", successSort) // 处理业务逻辑
    } else {
	fmt.Println("失败")
    }
}

func isPermitted(uid string, action string, period, maxCount int, client *redis.Client) bool {
    key := fmt.Sprintf("%v_%v", uid, action)
    now := time.Now().UnixNano() // 纳秒
    beforeTime := now - int64(period*1000000000)
    // lua脚本
    evalSha, _ := lua.RunLua(client)
    res, err := client.EvalSha(evalSha, []string{key}, now, beforeTime, period, maxCount).Result()
    if err != nil {
	panic(err)
    }
    if res.(int64) == int64(2) {
	return false
    }

    return true
}

limit.go

package main

import (
    "current-limit/logic"
    "current-limit/redis"
    "github.com/gin-gonic/gin"
    "net/http"
)

func main() {
    rDb := redis.NewInitRedis()
    // 1.创建路由
    r := gin.Default()
    // 2.绑定路由规则，执行的函数
    // gin.Context，封装了request和response
    r.GET("/count", func(c *gin.Context) {
	logic.DoRequest(rDb.Nosql())
	c.String(http.StatusOK, "ok")
    })

    // 3.监听端口，默认在8080
    // Run("里面不指定端口号默认为8080")
    r.Run(":8000")
}