限流算法

什么是限流

限流是一种应用系统常用的保护机制,用于防止服务因流量激增而崩溃或性能下降。

常见限流算法

令牌桶

令牌桶工作原理如下:

1. 系统中有一个容量为 R 的虚拟桶,每隔 T 时间单位会往桶里添加一个令牌。这样就限制了平均流量不能超过 R/T。
2. 当客户端发起请求时,会从桶中取出一个令牌。如果桶中没有令牌,则请求会被阻塞或拒绝。
3. 当桶中的令牌数达到最大值 B 时,新添加的令牌会被丢弃。这可以防止流量突增导致的短期高峰。
4. 桶中的令牌数可以动态变化。例如在高峰时段,可以降低令牌生成速率 R/T 来控制流量;在低峰时段,则可以提高令牌生成速率。

package main

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

type TokenBucket struct {
    capacity     int           // 桶的容量
    rate         int           // 令牌生成速率
    tokens       int           // 当前令牌数量
    lastRefill   time.Time     // 上次填充令牌的时间
    mu           sync.Mutex    // 互斥锁保护共享资源
}

func NewTokenBucket(capacity, rate int) *TokenBucket {
    return &TokenBucket{
        capacity:   capacity,
        rate:       rate,
        tokens:     capacity,
        lastRefill: time.Now(),
    }
}

func (tb *TokenBucket) refill() {
    now := time.Now()
    elapsed := now.Sub(tb.lastRefill).Seconds()
    tb.lastRefill = now

    // 根据经过的时间计算新令牌数量
    newTokens := int(elapsed * float64(tb.rate))
    if newTokens > 0 {
        tb.tokens += newTokens
        if tb.tokens > tb.capacity {
            tb.tokens = tb.capacity
        }
    }
}

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

    tb.refill()

    if tb.tokens > 0 {
        tb.tokens--
        return true
    }

    return false
}

func main() {
    bucket := NewTokenBucket(10, 1) // 容量为10，每秒生成1个令牌

    var wg sync.WaitGroup

    for i := 0; i < 20; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()

            if bucket.Allow() {
                fmt.Printf("请求 %d 通过\n", id)
            } else {
                fmt.Printf("请求 %d 被限流\n", id)
            }

            time.Sleep(500 * time.Millisecond) // 模拟请求处理时间
        }(i)
    }

    wg.Wait()
}

滑动窗口

工作原理如下:

1. 维护一个固定大小的时间窗口,通常是最近n秒内的请求数据。
2. 每次请求到来时,将其记录在当前窗口中。
3. 如果当前窗口中的请求数量超过了设定的阈值,则拒绝后续请求。
4. 随着时间推移,窗口会自动向前滑动,旧的请求数据会被"滑出"窗口。

package main

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

type SlidingWindow struct {
    interval    time.Duration // 窗口大小
    maxRequests int           // 最大请求数
    mu          sync.Mutex    // 互斥锁保护共享资源
    requests    []time.Time   // 请求时间戳
}

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

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

    now := time.Now()
    // 移除窗口之外的请求
    validTime := now.Add(-sw.interval)
    for len(sw.requests) > 0 && sw.requests[0].Before(validTime) {
        sw.requests = sw.requests[1:]
    }

    // 检查是否可以允许新的请求
    if len(sw.requests) < sw.maxRequests {
        sw.requests = append(sw.requests, now)
        return true
    }

    return false
}

func main() {
    window := NewSlidingWindow(5, 10*time.Second) // 10秒内最多5个请求

    var wg sync.WaitGroup

    for i := 0; i < 10; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()

            if window.Allow() {
                fmt.Printf("请求 %d 通过\n", id)
            } else {
                fmt.Printf("请求 %d 被限流\n", id)
            }

            time.Sleep(1 * time.Second) // 模拟请求间隔
        }(i)
    }

    wg.Wait()
}

漏斗算法

1. 将请求视为流入漏斗的水,漏斗代表系统的处理能力。
2. 漏斗有一个固定的容量,代表系统能够短暂承受的最大请求数。
3. 每个时间单位,漏斗会以一个固定的速率"漏出"一部分请求,代表系统的处理能力。
4. 如果新的请求流入时,漏斗已经满了,那么这些请求就会被丢弃或阻塞。

package main

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

type LeakyBucket struct {
    capacity     int           // 漏斗容量
    rate         int           // 出水速率（每秒处理请求数）
    tokens       int           // 当前令牌数量
    lastLeakTime time.Time     // 上次漏水时间
    mu           sync.Mutex    // 互斥锁保护共享资源
}

func NewLeakyBucket(capacity, rate int) *LeakyBucket {
    return &LeakyBucket{
        capacity:     capacity,
        rate:         rate,
        tokens:       0,
        lastLeakTime: time.Now(),
    }
}

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

    now := time.Now()
    elapsed := now.Sub(lb.lastLeakTime).Seconds()

    // 计算该时间段内漏掉的令牌数量
    leakedTokens := int(elapsed * float64(lb.rate))
    lb.tokens -= leakedTokens

    if lb.tokens < 0 {
        lb.tokens = 0
    }

    lb.lastLeakTime = now

    if lb.tokens < lb.capacity {
        lb.tokens++
        return true
    }

    return false
}

func main() {
    bucket := NewLeakyBucket(10, 1) // 容量为10，每秒处理1个请求

    var wg sync.WaitGroup

    for i := 0; i < 20; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()

            if bucket.Allow() {
                fmt.Printf("请求 %d 通过\n", id)
            } else {
                fmt.Printf("请求 %d 被限流\n", id)
            }

            time.Sleep(500 * time.Millisecond) // 模拟请求处理时间
        }(i)
    }

    wg.Wait()
}