golang 限流控制

限流

日常开发中，一般会遇到几种场景需要限流，比如有个api-server， 需要限制单个用户的调用频率，避免用户恶意刷接口或者突发大流量导致服务不可用等，这边记录几个常用的限流方法。

并发控制

简单的并发控制，用户的所有请求丢到一个channel里，再指定一个带缓冲区的channel为并发池，缓冲池的长度就是可以同时存在的协程数量，然后将执行完的任务根据需要直接返回或者丢到另外一个channel 里, 这样做的问题是如果任务太多，后面的任务会慢慢等待(因为channel阻塞机制)，用户体验不是太好。

Code:

package main

import (
	"fmt"
	"sync"
)


func setJob() <-chan int {
    var wg sync.WaitGroup
    jobChan := make(chan int, 50)
    wg.Add(1)
    go func(){
        for i := 0; i < 50; i++ {
            jobChan <- i
        }
        close(jobChan)
        wg.Done()
    }()
    wg.Wait()
    return jobChan
}

func main() {
	var wg sync.WaitGroup
    // 创建一个需要处理的数据来源, 假设是个channel
    jobChan := setJob()

    // 将结果存入一个channel
    resChan := make(chan int , 50)

    // 设置一个并发池
	buckets := make(chan bool, 10)
	for job := range jobChan {
		buckets <- true
		wg.Add(1)
		go func(job int) {
            res := 10 * job
            resChan <- res
			<-buckets
			wg.Done()
		}(job)
	}

    wg.Wait()
    close(resChan)
    tmpA := []int{}
	for r := range(resChan) {
		tmpA = append(tmpA, r)
	}
    fmt.Println(tmpA)
    
}

漏桶

简单的限流，指定一个大小固定的桶(bucket),以指定速率流入流出，如果桶满了，则拒绝后续请求，好处是简单，但是由于rate是固定的，导致了在无法在业务突发高峰时候(比如活动期间)有比较好的适配性。

Code:

package main

import (
    "fmt"
    "math"
    "time"
)

// 漏桶限流器
type BucketLimit struct {
    rate       float64 //漏桶中水的漏出速率, 即每秒流多少水
    bucketSize float64 //漏桶最多能装的水大小
    lastAccessTime   time.Time //上次访问时间
    curWater   float64 //当前桶里面的水
}

func NewBucketLimit(rate float64, bucketSize int64) *BucketLimit {
    return &BucketLimit{
        bucketSize: float64(bucketSize),
        rate:       rate,
        curWater:   0,
    }
}

func (b *BucketLimit) AllowControl() bool {
    now := time.Now()
    pastTime := now.Sub(b.lastAccessTime)

    // 当前剩余水量，当前水量减去距离上次访问的流出水量，如果流完了，即剩余水量为0
    b.curWater = math.Max(0, b.curWater - float64(pastTime) * b.rate)

    b.lastAccessTime = now

    // 当前水量必须小于桶的总量，不然则流出了
    if b.curWater < b.bucketSize {
        b.curWater = b.curWater + 1
        return true
    }
    return false
}

func main() {
    // 创建一个流出速率为1qps,桶的总量为2的限流器
    limit := NewBucketLimit(1, 2)

    // 在桶里放入1000滴水
    for i := 0; i < 1000; i++ {
        allow := limit.AllowControl()
        if allow {
            fmt.Printf("第%d滴水, 顺利流出\n", i)
            continue
        } else {
            fmt.Printf("第%d滴水, 溢出丢弃\n", i)
            time.Sleep(time.Millisecond * 100)
        }
    }
}

令牌桶

令牌桶相比漏桶可以对突发的流量做一定程度的处理，该方法意思是一定速率往一个桶里放令牌(token)，同时往桶里注水，每消耗一滴水，需要一定数量的令牌，这里代码假设需要2块令牌才能消耗1滴水，所以处理水滴的速度主要取决于令牌的数量。假设放入令牌的速度是4块/s，每秒消耗的水滴数量是1滴，即每秒消耗2块令牌，那么每秒多出来的2块令牌就会积累在桶里，直到桶满位置，然后因为某次活动大放水，每秒消耗的水滴数量突然变成了3滴，即每秒需要消耗6块另外，大于每秒放入的令牌数量，由于之前桶里有积累的令牌，所以即使放水量加大，依然可以在令牌桶消耗完之前快速处理。

Code:

package main
 
import (
    "fmt"
    "time"
    "math"
)
 
// 令牌桶限流器
type BucketLimit struct {
    rate       float64 //令牌桶放令牌的速率
    bucketSize float64 //令牌桶最多能装的令牌数量
    lastAccessTime   time.Time //上次访问时间
    curTokenNum  float64 //桶里当前的令牌数量
}

func NewBucketLimit(rate float64, bucketSize int64) *BucketLimit {
    return &BucketLimit{
        bucketSize: float64(bucketSize),
        rate:       rate,
        curTokenNum:   0,
    }
}

func (b *BucketLimit) AllowControl(tokenNeed float64) bool {
    now := time.Now()
    pastTime := now.Sub(b.lastAccessTime)
 
    // 在距离上次访问期间一共可发放了多少令牌
    newTokenNum := float64(pastTime) / float64(b.rate)

    // 剩余令牌数量不能超过桶的总空间
    b.curTokenNum = math.Min(b.bucketSize, b.curTokenNum + newTokenNum)

    b.lastAccessTime = now
    // tokenNeed 指处理一个请求需要的令牌数量
    if tokenNeed > b.curTokenNum {
        return false
    } else {
        b.curTokenNum = b.curTokenNum - tokenNeed
        return true
    }
}
 
func main() {
    // 创建一个放令牌速率为10qps,桶的总量为20的限流器
    limit := NewBucketLimit(10, 20)

    // 在桶里放入100滴水
    for i := 0; i < 100; i++ {
        // 处理1滴水需要消耗19块令牌
        if i == 50 {
            time.Sleep(3 * time.Second)
        }
        allow := limit.AllowControl(19)
        if allow {
            fmt.Printf("第%d滴水, 顺利流出\n", i)
            continue
        } else {
            fmt.Printf("第%d滴水, 溢出丢弃\n", i)
            time.Sleep(time.Millisecond * 100)
        }
    }
}

他山之石

如果不自己造轮子，其实已经有很多别人造好的轮子，简单记录下，按需求选用:

• “https://golang.org/x/net/netutil” 官方自带的http限流工具

伪代码

import(
    "fmt"
    "net"
    "golang.org/x/net/netutil"
)
 
func main() {
    l, err := net.Listen("tcp", "127.0.0.1:0")
    if err != nil {
        fmt.Fatalf("Listen: %v", err)
    }
    defer l.Close()
    // 开启限流， 其实就算是记录器，max 就算最大的并发数量
    l = LimitListener(l, max)
    http.Serve(l, http.HandlerFunc())
    ......
}

• “https://github.com/didip/tollbooth” 1.8K star的http限流中间件
• “https://golang.org/x/time/rate” 官方的令牌桶限流代码包
• “https://github.com/uber-go/ratelimit” uber的开源库，基于漏斗算法实现了一个限制器，也值得学习一下

后续代码新增和更新就放这里了:

• “https://github.com/zyxpaomian/rate_limiter”

个人公众号, 分享一些日常开发，运维工作中的日常以及一些学习感悟，欢迎大家互相学习，交流