限流算法

计数器

在固定时间间隔内，处理请求有上限，请求超出部分丢弃。

package main

import (
	"sync"
	"time"

	klog "k8s.io/klog/v2"
)

type counterRateLimiter struct {
	m             sync.Mutex
	startPartTime int64
	endPartTime   int64
	maxCount      int
	currCount     int
}

func (c *counterRateLimiter) isLimit() bool {
	c.m.Lock()
	defer c.m.Unlock()

	currentTime := time.Now().UnixMilli()
	// 区间考虑前闭后开
	if currentTime >= c.endPartTime {
		c.startPartTime += (currentTime - c.startPartTime) / 1000 * 1000
		c.endPartTime = c.startPartTime + 1000
		c.currCount = 0
	}

	if c.currCount < c.maxCount {
		c.currCount++
		return true
	}
	return false
}

func main() {
	crl := counterRateLimiter{
		m:             sync.Mutex{},
		startPartTime: time.Now().UnixMilli(),
		endPartTime:   time.Now().UnixMilli() + 1000,
        // 1s处理5个请求
		maxCount:      5,
		currCount:     0,
	}

	for i := 0; i < 6; i++ {
		klog.Infof("isLimit is %v", crl.isLimit())
	}
	time.Sleep(500 * time.Millisecond)
	for i := 0; i < 6; i++ {
		klog.Infof("isLimit is %v", crl.isLimit())
	}
	time.Sleep(time.Second)
	for i := 0; i < 6; i++ {
		klog.Infof("isLimit is %v", crl.isLimit())
	}
}

在单秒内没有请求时，可以应对突发流量。

漏桶

漏桶（装请求）容量固定，请求下发速率固定，请求超出漏桶容量部分丢弃。

package main

import (
	"sync"
	"time"

	klog "k8s.io/klog/v2"
)

type waterBucketRateLimiter struct {
	m             sync.Mutex
	interval      int
	bucketCap     int
	countInBucket int
	ch            chan int
}

func initWaterBucketRateLimiter(maxCount, bucketCap int) *waterBucketRateLimiter {
	klog.Infof("init water bucket rate limiter")
	w := &waterBucketRateLimiter{
		m:             sync.Mutex{},
		interval:      1000 / maxCount,
		bucketCap:     bucketCap,
		countInBucket: 0,
		ch:            make(chan int, maxCount),
	}

	go w.handle()
	return w
}

// 新增请求
func (w *waterBucketRateLimiter) addWater() bool {
	w.m.Lock()
	defer w.m.Unlock()

	if w.countInBucket < w.bucketCap {
		w.countInBucket++
		return true
	}
	return false
}

func (w *waterBucketRateLimiter) handle() {
	t := time.NewTicker(time.Duration(w.interval) * time.Millisecond)
	defer t.Stop()

	for {
		select {
		case <-t.C:
			w.m.Lock()

			if w.countInBucket == 0 {
				continue
			}
			w.countInBucket--
			w.ch <- 1

			w.m.Unlock()
		}
	}
}

func main() {
	wbrl := initWaterBucketRateLimiter(5, 10)

	for i := 0; i < 10; i++ {
		wbrl.addWater()
	}

	for _ = range wbrl.ch {
		klog.Info("get item from rate limiter")
	}
}

固定间隔200ms/个下发请求，平滑处理流量。

令牌桶

令牌桶（装令牌，不装请求）容量固定，令牌产生速度固定，取到令牌后处理请求，消耗令牌速度不固定，请求没有取到令牌丢弃。

package main

import (
	"sync"
	"time"

	klog "k8s.io/klog/v2"
)

type tokenBucketRateLimiter struct {
	m                  sync.Mutex
	interval           int
	bucketCap          int
	tokenCountInBucket int
}

func initTokenBucketRateLimiter(maxCount, bucketCap int, wg *sync.WaitGroup, stopCh chan bool) *tokenBucketRateLimiter {
	w := &tokenBucketRateLimiter{
		m:                  sync.Mutex{},
		interval:           1000 / maxCount,
		bucketCap:          bucketCap,
		tokenCountInBucket: bucketCap,
	}

	go w.putTokenIntoBucket(wg, stopCh)
	return w
}

func (tbrl *tokenBucketRateLimiter) putTokenIntoBucket(wg *sync.WaitGroup, stopCh chan bool) {
	t := time.NewTicker(time.Duration(tbrl.interval) * time.Millisecond)
	defer t.Stop()

	wg.Add(1)

	for {
		select {
		case <-t.C:
			tbrl.m.Lock()
			if tbrl.tokenCountInBucket < tbrl.bucketCap {
				tbrl.tokenCountInBucket++
			}
			tbrl.m.Unlock()
		case <-stopCh:
			wg.Done()
			return
		}
	}
}

func (tbrl *tokenBucketRateLimiter) isLimit() bool {
	tbrl.m.Lock()
	defer tbrl.m.Unlock()

	if tbrl.tokenCountInBucket > 0 {
		tbrl.tokenCountInBucket--
		return true
	}
	return false
}

func main() {
	wg := sync.WaitGroup{}
	stopCh := make(chan bool)
	tbrl := initTokenBucketRateLimiter(5, 10, &wg, stopCh)

	for i := 0; i < 20; i++ {
		klog.Infof("get token result is %v", tbrl.isLimit())
	}
	time.Sleep(time.Second)
	for i := 0; i < 6; i++ {
		klog.Infof("get token result is %v", tbrl.isLimit())
	}
	stopCh <- true

	wg.Wait()
}

与漏桶相比，令牌桶不是固定间隔下发请求，可能因没有请求且令牌桶容量大，令牌桶中存了大量令牌，应对突发流量时可以快速下发一部分。