Go常见限流算法代码

可以在程序中设置一个变量 count，当过来一个请求我就将这个数 +1，同时记录请求时间。
当下一个请求来的时候判断count的计数值是否超过设定的频次，以及当前请求的时间和第一次请求时间是否在 1 分钟内。
如果在 1 分钟内并且超过设定的频次则证明请求过多，后面的请求就拒绝掉。
如果该请求与第一个请求的间隔时间大于计数周期，且 count 值还在限流范围内，就重置 count。

type CountLimiter struct {
    rate  int64         // 计数周期内运行的最大请求数
    begin time.Time     // 当前轮计数开始时间
    count int64         // 当前计数周期累计的请求数
    cycle time.Duration // 计数周期，如统计1秒内的总请求数，那么计数周期就是1秒
    lock  sync.Mutex    // 判断能否放行时需要加锁操作
}

func NewCountLimiter(rate int64, cycle time.Duration) *CountLimiter {
    return &CountLimiter{
        rate:  rate,
        begin: time.Now(),
        count: 0,
        cycle: cycle,
        lock:  sync.Mutex{},
    }
}

func (c *CountLimiter) Allow() bool {
    c.lock.Lock()
    defer c.lock.Unlock()
    if c.count == c.rate { // 达到最大限流时，判断时间是否也超过统计周期了，超了可以重置限流器了，没有超说明还在当前统计周期内，应该拦截
        if time.Now().Sub(c.begin) > c.cycle {
            c.Reset()
            return true
        } else {
            return false
        }
    } else { // 还没有达到最大限流数，那么不管时间周期是否超出统计计数周期，都可以放行
        c.count++
        return true
    }
}

func (c *CountLimiter) Reset() {
    c.begin = time.Now()
    c.count = 0
}

func main() {
    countLimiter := NewCountLimiter(3, time.Second) // 1s内不可超过3个请求
    var wg sync.WaitGroup
    wg.Add(10)
    for i := 0; i < 10; i++ {
        go func(i int) {
            defer wg.Done()
            if countLimiter.Allow() {
                fmt.Println(fmt.Sprintf("当前时间%s,请求%d通过了", time.Now().String(), i))
            } else {
                fmt.Println(fmt.Sprintf("当前时间%s,请求%d被限流了", time.Now().String(), i))
            }
        }(i)
        time.Sleep(200 * time.Millisecond)
    }
    wg.Wait()
}

令牌按固定的速率被放入令牌桶中
桶中最多存放 B 个令牌，当桶满时，新添加的令牌被丢弃或拒绝
如果桶中的令牌不足 1个，则不会删除令牌，且请求将被限流（丢弃或阻塞等待）

package main

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

type TokenBucketLimiter struct {
    lock     sync.Mutex
    rate     time.Duration // 多长时间放入一个令牌，即放入令牌的速率
    capacity int64         // 令牌桶的容量，控制最多放入多少令牌，也即突发最大并发量
    tokens   int64         // 当前桶中已有的令牌数量
    lastTime time.Time     // 上次放入令牌的时间，避免开启协程定时去放入令牌，而是请求到来时懒加载的方式(now - lastTime) / rate放入令牌
}

func NewTokenBucketLimiter(rate time.Duration, capacity int64) *TokenBucketLimiter {
    if capacity < 1 {
        panic("token bucket capacity must be large 1")
    }
    return &TokenBucketLimiter{
        lock:     sync.Mutex{},
        rate:     rate,
        capacity: capacity,
        tokens:   0,
        lastTime: time.Time{},
    }
}

func (tbl *TokenBucketLimiter) Allow() bool {
    tbl.lock.Lock() // 加锁避免并发错误
    defer tbl.lock.Unlock()

    // 如果 now 与上次请求的间隔超过了 token rate
    // 则增加令牌，更新lastTime
    now := time.Now()
    if now.Sub(tbl.lastTime) > tbl.rate {
        tbl.tokens += int64((now.Sub(tbl.lastTime)) / tbl.rate) // 放入令牌
        if tbl.tokens > tbl.capacity {
            tbl.tokens = tbl.capacity // 总令牌数不能大于桶的容量
        }
        tbl.lastTime = now // 更新上次往桶中放入令牌的时间
    }

    if tbl.tokens > 0 { // 令牌数是否充足
        tbl.tokens -= 1
        return true
    }

    return false // 令牌不足，拒绝请求
}

func main() {
    tbl := NewTokenBucketLimiter(10, 5) // 每10ms放一个令牌，1s放100个，桶容量（最大突发流量）为5
    for i := 0; i < 10; i++ {
        fmt.Println(tbl.Allow())  // 模拟突发流量10个请求，超过桶容量5
    }
    time.Sleep(100 * time.Millisecond)
    fmt.Println(tbl.Allow())
}

漏桶具有固定容量，出水速率是固定常量（流出请求）
如果桶是空的，则不需流出水滴
可以以任意速率流入水滴到漏桶（流入请求）
如果流入水滴超出了桶的容量，则流入的水滴溢出（新请求被拒绝）
漏桶限制的是常量流出速率（即流出速率是一个固定常量值），所以最大的速率就是出水的速率，不能出现突发流量。

package main

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

type Result struct {
    Msg string // 根据实际情况定义返回结果需要哪些字段
}

type Handler func() Result // 处理函数的形式也应该根据具体需要而定

type Task struct {
    id      int64       // 任务id
    result  chan Result // 任务的执行结果，即请求的响应结果
    handler Handler     // 请求的执行函数
}

func NewTask(id int, handler Handler) Task {
    return Task{
        handler: handler,
        result:  make(chan Result),
        id:      int64(id),
    }
}

type LeakyBucketLimiter struct {
    bucketSize int64     // 桶的大小
    workerNum  int64     // 工作者数量，即最大并发数
    taskChan   chan Task // 用于存放请求
}

func NewLeakyBucketLimiter(bucketSize, workerNum int64) *LeakyBucketLimiter {
    if capacity < 1 {
        panic("capacity must be large 1")
    }
    if workerNum < 1 {
        panic("workerNum must be large 1")
    }
    return &LeakyBucketLimiter{
        bucketSize: bucketSize,
        workerNum:  workerNum,
        taskChan:   make(chan Task, bucketSize),
    }
}

func (lbl *LeakyBucketLimiter) AddTask(task Task) bool { // 类似其他限流算法的Allow方法
    // 如果木桶已经满了，或者任务执行失败或超时了，返回false
    select {
    case lbl.taskChan <- task:  // 利用了select的特性判断是否能往通道中添加任务
    default:
        fmt.Printf("请求%d被拒绝了\n", task.id)
        return false
    }

    // 如果成功入桶，调用者会等待Task的Handler执行结果
    // 由于Task的result是无缓冲的通道，不应该让其无限等待阻塞，否则出现问题时，不往该chan写，就会一直阻塞在这里了，泄漏
    // 因此设置一个超时时间
    //resp := <-task.result
    //fmt.Printf("请求%d运行成功，结果为：%v\n", task.id, resp)
    select {
    case resp := <-task.result:
        fmt.Printf("请求%d运行成功，结果为：%v\n", task.id, resp)
    case <-time.After(5 * time.Second): // 超时时间可以稍微设置长一点点，因为任务放入桶中后，可能需要排队一点时间才被拉取出来执行
        return false // 这里超时当被限流处理
    }

    return true
}

func (lbl *LeakyBucketLimiter) Start(ctx context.Context) {
    // 开启workerNum个协程从木桶拉取任务执行
    for i := 0; int64(i) < lbl.workerNum; i++ {
        go func(ctx context.Context) {
            defer func() { // 铁则：开启的子协程一定要捕获异常，否则一旦出现异常会依次上抛，上抛也一直不捕获，会导致程序退出
                if err := recover(); err != any(nil) {
                    fmt.Println("捕获到异常")
                }
            }()

            for { // 持续监听，拉取任务执行
                select {
                case <-ctx.Done():
                    fmt.Println("退出工作")
                    return
                default:
                    task := <-lbl.taskChan
                    result := task.handler()
                    task.result <- result // 处理结果写入对应Task的结果通道
                }
            }
        }(ctx)
    }
}

func main() {
    bucket := NewLeakyBucketLimiter(10, 4)
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    bucket.Start(ctx) // 开启消费者
    // 模拟20个并发请求
    var wg sync.WaitGroup
    wg.Add(20)
    for i := 0; i < 20; i++ {
        go func(id int) {
            defer wg.Done()
            task := NewTask(id, func() Result { // 这里的func应该根据实际需要定义为handler，写具体的业务逻辑和返回的result
                time.Sleep(300 * time.Millisecond)  // 模拟业务逻辑消耗的时间
                return Result{}
            })
            bucket.AddTask(task)  // 请求入桶
        }(i)
    }
    wg.Wait()
    time.Sleep(10 * time.Second)
}