基于redis的限流

2020-08-07  本文已影响0人  风亡小窝
package main

import (
    "github.com/garyburd/redigo/redis"
    "log"
    "time"
)

const script1 = `
-- key
local key = KEYS[1]
-- 最大存储的令牌数
local burst = tonumber(KEYS[2])
-- 每秒钟产生的令牌数
local limit = tonumber(KEYS[3])
-- 请求的令牌数
local N = tonumber(ARGV[1])

-- 上一次成功获取令牌的时间
local last = tonumber(redis.call('hget', key, 'last') or 0)

-- 当前存储的令牌数
local tokens = tonumber(redis.call('hget', key, 'tokens') or 0)

-- 当前时间
local time = redis.call('time')
local now = tonumber(time[1]) * 1000000 + tonumber(time[2])

-- 添加令牌的时间间隔
local interval = 1000000 / limit

-- 距离上次获取流逝的时间
local max_elapsed = (burst - tokens)*interval

local elapsed = now - last
if (max_elapsed < elapsed) then
    elapsed = max_elapsed
end


-- 补充令牌
local new_tokens = elapsed / interval
tokens = math.min(burst, tokens + new_tokens)

-- 消耗令牌
local fresh_permits = math.max(N - tokens, 0);
local wait_micros = fresh_permits * interval

if (tokens >= N) then
    redis.call('hset', key, 'tokens', tokens - N)
    redis.call('hset', key, 'last', now)
end

-- redis.replicate_commands()

--重置过期时间
redis.call('expire', key, 1000)

-- 返回需要等待的时间长度
return wait_micros
`

func Check() {
    pool := redis.Pool{
        Dial: func() (conn redis.Conn, err error) {
            return redis.Dial("tcp", "localhost:6379")
        },
        MaxIdle:     10,
        MaxActive:   100,
        IdleTimeout: 10,
        Wait:        true,
    }

    conn := pool.Get()
    defer conn.Close()

    for {
        reply, err := conn.Do("EVAL", script1, 3, "test_rate_limit_key", 10, 1, 1)
        if err != nil {
            panic(err)
        }

        log.Println(reply)
        s, err := redis.Int64(reply, err)
        time.Sleep(time.Duration(s) * time.Microsecond)
    }

}

https://www.jianshu.com/p/b3a02dad5adb

上一篇下一篇

猜你喜欢

热点阅读