go.uber.org/ratelimit 原始碼分析

木易小熙發表於2024-08-28
GoMIT原始碼

go.uber.org/ratelimit 原始碼分析

go 提供了一用來介面限流的包。其中"go.uber.org/ratelimit" 包正是基於漏桶演算法實現的。

使用方式：

  1. 透過 ratelimit.New 建立限流器物件，引數為每秒允許的請求數（RPS）。
  2. 使用 Take() 方法來獲取限流許可，該方法會阻塞請求知道滿足限速要求。

官方示例：

import (
	"fmt"
	"time"

	"go.uber.org/ratelimit"
)

func main() {
    rl := ratelimit.New(100) // 每秒多少次

    prev := time.Now()
    for i := 0; i < 10; i++ {
        now := rl.Take()	// 平均時間
        fmt.Println(i, now.Sub(prev))
        prev = now
    }

    // Output:
    // 0 0
    // 1 10ms
    // 2 10ms
    // 3 10ms
    // 4 10ms
    // 5 10ms
    // 6 10ms
    // 7 10ms
    // 8 10ms
    // 9 10ms
}

ratelimit.New()指的是每秒平均多少次，在執行程式後，並不會嚴格按照官方給的樣例輸出。

原始碼分析

不僅知其然，還要知其所以然。

最大鬆弛量

傳統的漏桶演算法每隔請求的間隔是固定的，然而在實際上的互連網應用中，流量經常是突發性的。對於這種情況，uber引入了最大鬆弛量的概念。

假如我們要求每秒限定100個請求，平均每個請求間隔 10ms。但是實際情況下，有些間隔比較長，有些間隔比較短。如下圖所示：

請求 1 完成後，15ms 後，請求 2 才到來，可以對請求 2 立即處理。請求 2 完成後，5ms 後，請求 3 到來，這個時候距離上次請求還不足 10ms，因此還需要等待 5ms。

但是，對於這種情況，實際上三個請求一共消耗了 25ms 才完成，並不是預期的 20ms。在 uber-go 實現的 ratelimit 中，可以把之前間隔比較長的請求的時間，勻給後面的使用，保證每秒請求數 (RPS) 即可。

瞭解完這個字首知識就可以檢視原始碼了。

New（）

ratelimit.New() 內部呼叫的是 newAtomicInt64Based 方法。

type atomicInt64Limiter struct {
	prepadding [64]byte // 填充位元組，確保state獨佔一個快取行
	state      int64    // 最後一次許可權傳送的納秒時間戳，用於控制請求的速度
	postpadding [56]byte // 填充位元組，確保state獨佔一個快取行

	perRequest time.Duration	// 限流器放行週期，用於計算下一個許可權傳送的state的值
	maxSlack   time.Duration	// 最大鬆弛量
	clock      Clock	// 指向當前時間獲取函式的指標
}

// newAtomicBased返回一個新的基於原子的限制器。
func newAtomicInt64Based(rate int, opts ...Option) *atomicInt64Limiter {
	config := buildConfig(opts) // 載入配置，config.per 預設為 1s，config.slack 預設為 10
	perRequest := config.per / time.Duration(rate)
	l := &atomicInt64Limiter{
		perRequest: perRequest,
		maxSlack:   time.Duration(config.slack) * perRequest,	// 預設maxSlack為perRequest 10倍
		clock:      config.clock,
	}
	atomic.StoreInt64(&l.state, 0)	
	return l
}

Take()

// Take blocks to ensure that the time spent between multiple
// Take calls is on average time.Second/rate.
func (t *atomicInt64Limiter) Take() time.Time {
   var (
      newTimeOfNextPermissionIssue int64	// 下一次允許請求的時間
      now                          int64	// 當前時間
   )
   for {
      now = t.clock.Now().UnixNano()	
      timeOfNextPermissionIssue := atomic.LoadInt64(&t.state) // 上一次允許請求時間

      switch {
      case timeOfNextPermissionIssue == 0 || (t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest)):
        // if this is our first call or t.maxSlack == 0 we need to shrink issue time to now
         newTimeOfNextPermissionIssue = now
      case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack)+int64(t.perRequest):
         // a lot of nanoseconds passed since the last Take call
         // we will limit max accumulated time to maxSlack
         newTimeOfNextPermissionIssue = now - int64(t.maxSlack)
      default:
         // calculate the time at which our permission was issued
         newTimeOfNextPermissionIssue = timeOfNextPermissionIssue + int64(t.perRequest)
      }

      if atomic.CompareAndSwapInt64(&t.state, timeOfNextPermissionIssue, newTimeOfNextPermissionIssue) {
         break
      }
   }

   sleepDuration := time.Duration(newTimeOfNextPermissionIssue - now)
   if sleepDuration > 0 {
      t.clock.Sleep(sleepDuration)
      return time.Unix(0, newTimeOfNextPermissionIssue)
   }
   // return now if we don't sleep as atomicLimiter does
   return time.Unix(0, now)
}

switch 這塊挺繞的，剛開始一直以為timeOfNextPermissionIssue 為下次放行的時間戳，這樣的話當t.maxSlack = 0時，只要 now-timeOfNextPermissionIssue > 0 就應該放行。無法解釋(t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest))。

讓我們對上面的三個 case 分析一下

case 1

case timeOfNextPermissionIssue == 0 || (t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest))

這個比較好理解，我們仍以每秒100個請求為例，平均間隔 10ms。當本次請求時間與上次放行時間 > 時間間隔時即可放行，並記錄本次訪問時間，如圖：

case 2

case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack)+int64(t.perRequest)

這塊比較巧妙，假如鬆弛量是3 ms，當我們在第二次請求時的時間戳 > 13 ms，此時 newTimeOfNextPermissionIssue= now - maxSlack = 12 ms。

當 maxSlack 較大且與上次請求相隔較長時，後續的大量請求會被直接放行，以彌補此次浪費的時間。

假設第一次請求時間為0， maxSlack 為 100 ms，perRequest為10 ms，在第二次請求時與第一次間隔為 111 ms ，newTimeOfNextPermissionIssue = 111 - 100 = 11 ms。而 now 為 111 ms，限流器在後面的10次take中都會經過default直接放行，直到 newTimeOfNextPermissionIssue > now。

case 3

對於其它的請求， newTimeOfNextPermissionIssue = timeOfNextPermissionIssue + int64(t.perRequest)。

假如maxSlack為 100ms，perRequest 為 10ms，當請求2在15ms訪問後，state 更新為 10ms，這樣在請求3在20ms訪問時，不會出現攔截的情況。

小結

uber 對基於漏桶實現的 ratelimit 進行了一些最佳化，讓其限流更加的平滑。主要體現在兩點：

  1. 本次請求時間距離上次放行時間 > 時間間隔 + 鬆弛量時，後面10次的請求會根據情況直接放行
  2. 時間間隔 + 鬆弛量 >= 本次請求時間距離上次放行時間 > 時間間隔 ，state = state + perRequest

相關文章