限流算法实战：三种经典算法的Go实现与选型指南

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

package ratelimit

import (
	"sync"
	"time"
)

type TokenBucket struct {
	mu         sync.Mutex
	rate       float64   // 每秒生成令牌数
	burst      float64   // 桶容量
	tokens     float64   // 当前令牌数
	lastRefill time.Time // 上次补充时间
}

func NewTokenBucket(rate, burst float64) *TokenBucket {
	return &TokenBucket{
		rate:       rate,
		burst:      burst,
		tokens:     burst, // 初始满桶
		lastRefill: time.Now(),
	}
}

func (tb *TokenBucket) Allow() bool {
	tb.mu.Lock()
	defer tb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(tb.lastRefill).Seconds()
	// 按时间比例补充令牌
	tb.tokens += elapsed * tb.rate
	if tb.tokens > tb.burst {
		tb.tokens = tb.burst
	}
	tb.lastRefill = now

	if tb.tokens >= 1 {
		tb.tokens--
		return true
	}
	return false
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

type LeakyBucket struct {
	mu       sync.Mutex
	rate     float64   // 漏出速率（每秒处理请求数）
	capacity float64   // 桶容量
	water    float64   // 当前水量
	lastLeak time.Time // 上次漏水时间
}

func NewLeakyBucket(rate, capacity float64) *LeakyBucket {
	return &LeakyBucket{
		rate:     rate,
		capacity: capacity,
		water:    0,
		lastLeak: time.Now(),
	}
}

func (lb *LeakyBucket) Allow() bool {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(lb.lastLeak).Seconds()
	// 漏水
	lb.water -= elapsed * lb.rate
	if lb.water < 0 {
		lb.water = 0
	}
	lb.lastLeak = now

	if lb.water < lb.capacity {
		lb.water++
		return true
	}
	// 水满溢出
	return false
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

type SlidingWindow struct {
	mu       sync.Mutex
	rate     int           // 窗口内允许请求数
	window   time.Duration // 窗口大小
	requests []time.Time   // 请求时间戳
}

func NewSlidingWindow(rate int, window time.Duration) *SlidingWindow {
	return &SlidingWindow{
		rate:     rate,
		window:   window,
		requests: make([]time.Time, 0, rate),
	}
}

func (sw *SlidingWindow) Allow() bool {
	sw.mu.Lock()
	defer sw.mu.Unlock()

	now := time.Now()
	border := now.Add(-sw.window)

	// 移除过期请求
	valid := sw.requests[:0]
	for _, t := range sw.requests {
		if t.After(border) {
			valid = append(valid, t)
		}
	}
	sw.requests = valid

	if len(sw.requests) < sw.rate {
		sw.requests = append(sw.requests, now)
		return true
	}
	return false
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

// Redis令牌桶Lua脚本
const tokenBucketScript = `
local key = KEYS[1]
local rate = tonumber(ARGV[1])
local burst = tonumber(ARGV[2])
local now = tonumber(ARGV[3])

local data = redis.call('HMGET', key, 'tokens', 'last_refill')
local tokens = tonumber(data[1]) or burst
local lastRefill = tonumber(data[2]) or now

local elapsed = math.max(0, now - lastRefill)
tokens = math.min(burst, tokens + elapsed * rate)

if tokens >= 1 then
    tokens = tokens - 1
    redis.call('HMSET', key, 'tokens', tokens, 'last_refill', now)
    redis.call('EXPIRE', key, 60)
    return 1
else
    redis.call('HMSET', key, 'tokens', tokens, 'last_refill', now)
    redis.call('EXPIRE', key, 60)
    return 0
end
`

// 使用示例
func AllowRedis(rdb *redis.Client, key string, rate, burst float64) (bool, error) {
	result, err := rdb.Eval(context.Background(), tokenBucketScript,
		[]string{key},
		rate, burst, float64(time.Now().UnixNano())/1e9,
	).Int()
	if err != nil {
		return false, err
	}
	return result == 1, nil
}