package main

import (
    "fmt"
    "sync"
    "time"
)

type LeakyBucketLimiter struct {
    Rate          int        // 漏桶速率，每秒处理的请求数
    Capacity      int        // 漏桶容量，最多可存储请求数
    CurrentReqNum int        // 当前漏斗中待处理的请求数，表示当前漏桶中的请求数
    LastTime      time.Time  // 上次请求时间
    Lock          sync.Mutex // 请求锁
}

func (l *LeakyBucketLimiter) GetCurrentReqNum() int {
    return l.CurrentReqNum
}
func (l *LeakyBucketLimiter) GetRate() int {
    return l.Rate
}
func (l *LeakyBucketLimiter) GetCapacity() int {
    return l.Capacity
}

// NewLeakyBucketLimiter 初始化
func NewLeakyBucketLimiter(rate, capacity int) *LeakyBucketLimiter {
    return &LeakyBucketLimiter{
       Rate:     rate,
       Capacity: capacity,
       LastTime: time.Now(),
    }
}

func (l *LeakyBucketLimiter) Allow() bool {
    l.Lock.Lock()
    defer l.Lock.Unlock()

    // 计算两次请求间隔内应该要处理的请求数
    leakyReqCount := int(time.Since(l.LastTime).Seconds() * float64(l.Rate))

    if leakyReqCount > 0 {
       l.CurrentReqNum -= leakyReqCount
       l.LastTime = time.Now()
    }

    if l.CurrentReqNum < 0 {
       l.CurrentReqNum = 0
    }

    // 如果当前请求数在容量范围内，则请求通过且当前请求数+1(剩余待处理请求数+1)
    if l.CurrentReqNum < l.Capacity {
       l.CurrentReqNum++
       return true
    }
    return false
}

func MockRequest(n int, d time.Duration, l *LeakyBucketLimiter) {
    for i := 0; i < n; i++ {
       time.Sleep(d * time.Millisecond)
       if l.Allow() {
          fmt.Printf("第%d个请求通过\n", i+1)
       } else {
          fmt.Printf("第%d个请求被限流\n", i+1)
       }
    }
}

func main() {
    fmt.Println("=================漏桶算法=================")
    // 创建一个漏桶，速率为每秒处理4个请求，容量为5个请求
    limiter := NewLeakyBucketLimiter(4, 5)
     // 发送10个请求，每50ms发送一个
    MockRequest(10, 50, limiter)
    fmt.Println("------------------------------------------")
}

=================漏桶算法=================
第1个请求通过
第2个请求通过
第3个请求通过
第4个请求通过
第5个请求通过
第6个请求通过
第7个请求被限流
第8个请求被限流
第9个请求被限流
第10个请求通过
------------------------------------------

package main

import (
    "fmt"
    "sync"
    "time"
)

type TokenBucketLimiter struct {
    Rate            int        // 令牌产生的速度，即每秒生成多少个令牌
    Capacity        int        // 令牌桶容量，最多可存储令牌数
    CurrentTokenNum int        // 当前令牌数
    LastTime        time.Time  // 上次请求时间
    Lock            sync.Mutex // 请求锁
}

func (t *TokenBucketLimiter) GetCurrentTokenNum() int {
    return t.CurrentTokenNum
}
func (t *TokenBucketLimiter) GetRate() int {
    return t.Rate
}
func (t *TokenBucketLimiter) GetCapacity() int {
    return t.Capacity
}

// NewTokenBucketLimiter 初始化
func NewTokenBucketLimiter(rate, capacity int) *TokenBucketLimiter {
    return &TokenBucketLimiter{
       Rate:     rate,
       Capacity: capacity,
       LastTime: time.Now(),
       CurrentTokenNum: Capacity，
    }
}

func (l *TokenBucketLimiter) Allow() bool {
    l.Lock.Lock()
    defer l.Lock.Unlock()

    // 计算两次请求间隔内生成的令牌数
    tokenCount := int(time.Since(l.LastTime).Seconds() * float64(l.Rate))

    if tokenCount > 0 {
       l.CurrentTokenNum += tokenCount
       l.LastTime = time.Now()
    }

    // 令牌数不能超过令牌桶的容量
    if l.CurrentTokenNum > l.Capacity {
       l.CurrentTokenNum = l.Capacity
    }

    // 桶内还有令牌，获取一个令牌
    if l.CurrentTokenNum > 0 {
       l.CurrentTokenNum--
       return true
    }
    return false
}

func MockRequest(n int, d time.Duration, l *TokenBucketLimiter) {
    for i := 0; i < n; i++ {
       time.Sleep(d * time.Millisecond)
       if l.Allow() {
          fmt.Printf("第%d个请求通过\n", i+1)
       } else {
          fmt.Printf("第%d个请求被限流\n", i+1)
       }
    }
}

func main() {
    fmt.Println("=================令牌桶算法=================")
    // 创建一个令牌桶，令牌的生成速度为每秒4个，桶容量为5个请求
    limiter := NewTokenBucketLimiter(4, 5)
    // 发送10个请求，每50ms发送一个
    MockRequest(10, 50, limiter)
    fmt.Println("------------------------------------------")
}

=================令牌桶算法=================
第1个请求被限流
第2个请求被限流
第3个请求被限流
第4个请求被限流
第5个请求通过
第6个请求被限流
第7个请求被限流
第8个请求被限流
第9个请求被限流
第10个请求通过
------------------------------------------

package main

import (
    "context"
    "fmt"
    "time"

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

const luaScript = `
local key = KEYS[1]  -- hmap的key
local rate = tonumber(ARGV[1]) -- 令牌生成速率
local capacity = tonumber(ARGV[2])  -- 桶的容量
local now = tonumber(ARGV[3])   -- 当前时间（毫秒）
local requested = tonumber(ARGV[4]) -- 请求的令牌数

-- 获取上次更新时间和当前令牌数
local last_time = redis.call('HGET', key, 'last_time')
local current_tokens = redis.call('HGET', key, 'tokens')

-- 首次访问，将last_time初始化为当前时间，当前令牌数初始化为桶容量
if last_time == false then
    last_time = now
    current_tokens = capacity
else
    last_time = tonumber(last_time)   -- 非首次访问，从redis取得上次访问时间
    current_tokens = tonumber(current_tokens) -- 非首次访问，从redis取得当前桶的令牌数
end

-- 计算时间差
local time_diff = now - last_time   -- 单位毫秒
-- 计算新的令牌数
local new_tokens = math.min(capacity, current_tokens + (time_diff * rate / 1000))

local allowed = new_tokens >= requested

if allowed then
    new_tokens = new_tokens - requested
end

redis.call('HSET', key, 'last_time', now)
redis.call('HSET', key, 'tokens', new_tokens)

return allowed
`

var ctx = context.Background()

// Redis客户端
var rdb *redis.Client

// 初始化Redis客户端
func initRedis() {
    rdb = redis.NewClient(&redis.Options{
       Addr: "localhost:8089",
    })
}

// TokenBucketLimiter 令牌桶限流器
type TokenBucketLimiter struct {
    key       string
    rate      float64
    capacity  int
    luaScript string
}

// NewTokenBucketLimiter 创建新的令牌桶限流器
func NewTokenBucketLimiter(key string, rate float64, capacity int) *TokenBucketLimiter {
    return &TokenBucketLimiter{
       key:       key,
       rate:      rate,
       capacity:  capacity,
       luaScript: luaScript,
    }
}

// Allow 允许请求
func (limiter *TokenBucketLimiter) Allow(requested int) (bool, error) {
    now := time.Now().UnixNano() / 1e6
    _, err := rdb.Eval(ctx, limiter.luaScript, []string{limiter.key}, limiter.rate, limiter.capacity, now, requested).Result()
    if err != nil {
       if err.Error() != "redis: nil" {
          return false, err
       }
       return false, nil
    }
    return true, nil
}

func main() {
    initRedis()

    limiter := NewTokenBucketLimiter("my_limiter:token_bucket", 5, 5)

    for i := 0; i < 20; i++ {
       allowed, err := limiter.Allow(1) // 一个请求，只需要一个令牌，所以这里传1
       if err != nil {
          fmt.Println("Error:", err)
          continue
       }
       if allowed {
          fmt.Printf("第%d个请求通过\n", i+1)
       } else {
          fmt.Printf("第%d个请求被限流\n", i+1)
       }
       time.Sleep(100 * time.Millisecond)
    }
}

第1个请求通过
第2个请求通过
第3个请求通过
第4个请求通过
第5个请求通过
第6个请求通过
第7个请求通过
第8个请求通过
第9个请求通过
第10个请求被限流
第11个请求通过
第12个请求被限流
第13个请求通过
第14个请求被限流
第15个请求通过
第16个请求被限流
第17个请求通过
第18个请求被限流
第19个请求通过
第20个请求被限流