实现限流器

题目描述

实现一个令牌桶限流器。

参考答案

type RateLimiter struct {
    rate     int           // 令牌生成速率（每秒）
    capacity int           // 桶容量
    tokens   int           // 当前令牌数
    lastTime time.Time     // 上次取令牌时间
    mu       sync.Mutex
}

func NewRateLimiter(rate, capacity int) *RateLimiter {
    return &RateLimiter{
        rate:     rate,
        capacity: capacity,
        tokens:   capacity,
        lastTime: time.Now(),
    }
}

func (rl *RateLimiter) Allow() bool {
    rl.mu.Lock()
    defer rl.mu.Unlock()
    
    now := time.Now()
    elapsed := now.Sub(rl.lastTime).Seconds()
    
    // 计算新增令牌
    newTokens := int(elapsed * float64(rl.rate))
    rl.tokens = min(rl.tokens+newTokens, rl.capacity)
    rl.lastTime = now
    
    if rl.tokens > 0 {
        rl.tokens--
        return true
    }
    return false
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

使用示例

limiter := NewRateLimiter(10, 100) // 每秒 10 个令牌，桶容量 100

for i := 0; i < 20; i++ {
    if limiter.Allow() {
        fmt.Println("Request allowed")
    } else {
        fmt.Println("Request denied")
    }
    time.Sleep(100 * time.Millisecond)
}

关键点

1. 令牌生成：根据时间差计算
2. 桶容量限制：不超过容量
3. 线程安全：使用互斥锁