漏桶算法

近期在研究Jaeger，Jaeger中有一種採集策略是速率限制類型，內部使用的是漏桶算法，在這裏研究了下Jaeger漏桶算法的實現原理，自己仿照其實現了一個rateLimiter，並進行了相關測試，下面是主要實現。

• lck：lck是互斥鎖，主要用來防止併發情況下產生錯誤。
• rate：速率，即接口每秒限制多少個請求。在這裏也就是水滴從漏桶中流出的速度，同時也是餘量增加的速度。
• balance：漏桶的空閒餘量，會隨着漏桶滴水逐漸變大；如果將請求添加到漏桶中，會逐漸變小。當請求到來時，如果餘量不足1，那麼表明不能容下當前的請求，當前的請求會被拒絕。
• limit：漏桶的最大容量。
• lastTime：上次調用Check函數的時間。用於計算時間差dur，然後計算這段時間漏桶流出的水w，增加的餘量=流出的水量w=時間*速率=dur*rate。

rateLimiter實現代碼：

package ratelimiter

import (
	"sync"
	"time"
)

type rateLimiter struct {
	lck      *sync.Mutex
	rate     float64   //最大速率限制
	balance  float64   //漏桶的餘量
	limit    float64   //漏桶的最大容量限制
	lastTime time.Time //上次檢查的時間
}

func NewRateLimiter(limitPerSecond int, balance int) *rateLimiter {
	return &rateLimiter{
		lck:      new(sync.Mutex),
		rate:     float64(limitPerSecond),
		balance:  float64(balance),
		limit:    float64(balance),
		lastTime: time.Now(),
	}
}

func (r *rateLimiter) Check() bool {
	ok := false
	r.lck.Lock()
	now := time.Now()
	dur := now.Sub(r.lastTime).Seconds()
	r.lastTime = now
	water := dur * r.rate //計算這段時間內漏桶流出水的流量water
	r.balance += water    //漏桶流出water容量的水，自然漏桶的餘量多出water
	if r.balance > r.limit {
		r.balance = r.limit
	}
	if r.balance >= 1 { //漏桶餘量足夠容下當前的請求
		r.balance -= 1
		ok = true
	}
	r.lck.Unlock()
	return ok
}

單元測試代碼：

package ratelimiter

import (
	"fmt"
	"testing"
	"time"
)

func TestRateLimiter_Check(t *testing.T) {
	limiter := NewRateLimiter(10, 1)
	start := time.Now()
	count := 0
	for i := 0; i < 1e3; i++ {
		if limiter.Check() {
			fmt.Println(i)
			count++
		}
		time.Sleep(time.Millisecond)
	}
	fmt.Println("count:", count)
	fmt.Println(time.Now().Sub(start).Seconds())
}

測試結果：每秒限制10個，1.3秒接受了13個請求，加上一開始填滿桶的一個，共14個。