go语言中的限流漏桶和令牌桶库怎么使用

令牌桶

令牌桶是一种流量控制算法，其基本思想是固定速率地往桶中放入令牌，任务需要从桶中取出令牌才能执行，桶中最多存放一定数量的令牌，多余的令牌将会被丢弃。当桶中没有令牌可取时，无法执行任务。这一算法常被用于限制系统的并发量，可以保证系统在高负荷情况下仍能稳定运行。

下面是一个使用golang实现令牌桶算法的代码示例：

import (
    "time"
)

type TokenBucket struct {
    capacity    int64           // 桶的容量（最多存放 Token 个令牌）
    rate        float64         // 令牌放置速度（每秒钟放置 rate 个令牌）
    tokens      chan struct{}   // 当前桶中 Token 的数量
    stop        chan struct{}   // 用于关闭 tokenChan
}

func NewTokenBucket(capacity int64, rate float64) *TokenBucket {
    return &TokenBucket{
        capacity: capacity,
        rate: rate,
        tokens: make(chan struct{}, capacity),
        stop: make(chan struct{}),
    }
}

func (tb *TokenBucket) Run() {
    interval := time.Duration(float64(time.Second) / tb.rate)
    ticker := time.NewTicker(interval)
    for {
        select {
            case <-ticker.C:
                select {
                    case tb.tokens <- struct{}{}:
                    default:
                }
            case <-tb.stop:
                close(tb.tokens)
                return
        }
    }
}

func (tb *TokenBucket) TryTake() bool {
    select {
        case _, ok := <-tb.tokens:
            return ok
        default:
            return false
    }
}

func (tb *TokenBucket) Close() {
    close(tb.stop)
}

上述代码实现了一个token桶结构体，其中capacity和rate分别表示令牌桶的最大容量和放置速度，tokens是一个channel，表示当前桶中的令牌数量，当桶中没有令牌时，无法执行任务。

Run()方法通过time包中的Ticker实现对令牌的定期放置，TryTake()方法则是获取桶中的令牌并返回其是否可用的方法，Close()方法用于关闭tokenChan。

限流漏桶

限流漏桶是另一种流量控制算法，其基本思想是通过一个漏嘴不断地漏水，如果水来的太快，水位就会涨上去，如果水来得太慢，那么有些水会闲置，不能得到利用。当水位达到设定的阈值时，流量控制开启，拒绝新的请求。

以下是一个使用golang实现限流漏桶算法的代码示例：

import (
    "time"
    "sync/atomic"
)

type LeakyBucket struct {
    capacity        int64           // 桶的容量（漏嘴流水的速度）
    rate            float64         // 请求到达的速度
    water           int64           // 当前漏桶里的水量
    stop            chan struct{}   // 关闭 bucketChan
}

func NewLeakyBucket(capacity int64, rate float64) *LeakyBucket {
    return &LeakyBucket{
        capacity: capacity,
        rate: rate,
        water: 0,
        stop: make(chan struct{}),
    }
}

func (lb *LeakyBucket) Run() {
    interval := time.Duration(float64(time.Second) / lb.rate)
    ticker := time.NewTicker(interval)
    for {
        select {
            case <-ticker.C:
                // 漏嘴流水的速度限制了水量增加的速度
                // 如果水太多了，就会溢出
                if atomic.LoadInt64(&lb.water) < lb.capacity {
                    atomic.AddInt64(&lb.water, 1)
                }
            case <-lb.stop:
                return
        }
    }
}

func (lb *LeakyBucket) TryTake() bool {
    if atomic.LoadInt64(&lb.water) > 0 {
        atomic.AddInt64(&lb.water, -1)
        return true
    }
    return false
}

func (lb *LeakyBucket) Close() {
    close(lb.stop)
}

上述代码实现了一个leaky bucket结构体，其中capacity和rate分别表示漏桶的最大容量和水流速度，water表示桶中的水量。Run()方法通过time包中的Ticker实现对漏嘴的周期性流水，TryTake()方法则是获取桶中的水并返回其是否可用的方法，Close()方法用于关闭bucketChan。

使用限流桶

限流桶和令牌桶常被用于限制系统的并发量，下面是一个使用限流桶限制每秒钟的HTTP请求数量的代码示例：

import (
    "net/http"
    "log"
    "github.com/RussellLuo/rate-limiters"
)

func main() {
    lb := rate_limiters.NewLeakyBucket(100, 200) // 容量是100，速度是每秒200个请求
    go lb.Run() // 启动漏桶流水
    
    handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // 尝试获取令牌，如果没有令牌则不执行任务
        if !lb.TryTake() {
            http.Error(w, "too many requests", http.StatusTooManyRequests)
            return
        }
        // 执行任务
        // ...
    })
    
    http.HandleFunc("/api", handler)
    log.Fatal(http.ListenAndServe(":8080", nil))
}

上述代码使用RussellLuo开发的rate-limiters库实现了一个限流桶，对每个请求尝试获取一个令牌，如果令牌不够用则返回HTTP 429 “too many requests”错误。

总结

限流算法是一种常用的流量控制算法，令牌桶和漏桶是两种基本的限流算法。无论是使用令牌桶还是漏桶，都需要定期对桶进行填充（令牌）或漏水（水），这个需要根据实际情况设置速度。在使用时，可以根据需求来选择合适的限流算法来控制系统的并发量，保证系统的稳定性。