用 Go 語言做了一個分佈式限流器，看看實現的方法與步驟

項目中需要對 api 的接口進行限流，但是麻煩的是，api 可能有多個節點，傳統的本地限流無法處理這個問題。限流的算法有很多，比如計數器法，漏斗法，令牌桶法，等等。各有利弊，相關博文網上很多，這裏不再贅述。

項目的要求主要有以下幾點：

支持本地 / 分佈式限流，接口統一
支持多種限流算法的切換
方便配置，配置方式不確定

go 語言不是很支持 OOP，我在實現的時候是按 Java 的思路走的，所以看起來有點不倫不類，希望能拋磚引玉。

1. 接口定義

package ratelimit

import "time"

// 限流器接口
type Limiter interface {
    Acquire() error
    TryAcquire() bool
}

// 限流定義接口
type Limit interface {
    Name() string
    Key() string
    Period() time.Duration
    Count() int32
    LimitType() LimitType
}

// 支持 burst
type BurstLimit interface {
    Limit
    BurstCount() int32
}

// 分佈式定義的 burst
type DistLimit interface {
    Limit
    ClusterNum() int32
}

type LimitType int32
const (
    CUSTOM LimitType = iota
    IP
)

Limiter 接口參考了 Google 的 guava 包裏的 Limiter 實現。Acquire 接口是阻塞接口，其實還需要加上 context 來保證調用鏈安全，因爲實際項目中並沒有用到 Acquire 接口，所以沒有實現完善；同理，超時時間的支持也可以通過添加新接口繼承自 Limiter 接口來實現。TryAcquire 會立即返回。

Limit 抽象了一個限流定義，Key() 方法返回這個 Limit 的唯一標識，Name() 僅作輔助，Period() 表示週期，單位是秒，Count() 表示週期內的最大次數，LimitType()表示根據什麼來做區分，如 IP，默認是 CUSTOM.

BurstLimit 提供突發的能力，一般是配合令牌桶算法。DistLimit 新增 ClusterNum() 方法，因爲 mentor 要求分佈式遇到錯誤的時候，需要退化爲單機版本，退化的策略即是：2 節點總共 100QPS，如果出現分區，每個節點需要調整爲各 50QPS

2. LocalCounterLimiter

package ratelimit

import (
    "errors"
    "fmt"
    "math"
    "sync"
    "sync/atomic"
    "time"
)

// todo timer 需要 stop
type localCounterLimiter struct {
    limit Limit

    limitCount int32 // 內部使用，對 limit.count 做了 <0 時的轉換

    ticker *time.Ticker
    quit chan bool

    lock sync.Mutex
    newTerm *sync.Cond
    count int32
}

func (lim *localCounterLimiter) init() {
    lim.newTerm = sync.NewCond(&lim.lock)
    lim.limitCount = lim.limit.Count()

    if lim.limitCount < 0 {
        lim.limitCount = math.MaxInt32 // count 永遠不會大於 limitCount，後面的寫法保證溢出也沒問題
    } else if lim.limitCount == 0  {
        // 禁止訪問, 會無限阻塞
    } else {
        lim.ticker = time.NewTicker(lim.limit.Period())
        lim.quit = make(chan bool, 1)

        go func() {
            for {
                select {
                case <- lim.ticker.C:
                    fmt.Println("ticker .")
                    atomic.StoreInt32(&lim.count, 0)
                    lim.newTerm.Broadcast()

                    //lim.newTerm.L.Unlock()
                case <- lim.quit:
                    fmt.Println("work well .")
                    lim.ticker.Stop()
                    return
                }
            }
        }()
    }
}

// todo 需要機制來防止無限阻塞, 不超時也應該有個極限時間
func (lim *localCounterLimiter) Acquire() error {
    if lim.limitCount == 0 {
        return errors.New("rate limit is 0, infinity wait")
    }

    lim.newTerm.L.Lock()
    for lim.count >= lim.limitCount {
        // block instead of spinning
        lim.newTerm.Wait()
        //fmt.Println(count, lim.limitCount)
    }
    lim.count++
    lim.newTerm.L.Unlock()

    return nil
}

func (lim *localCounterLimiter) TryAcquire() bool {
    count := atomic.AddInt32(&lim.count, 1)
    if count > lim.limitCount {
        return false
    } else {
        return true
    }
}

代碼很簡單，就不多說了

3. LocalTokenBucketLimiter
golang 的官方庫裏提供了一個 ratelimiter，就是採用令牌桶的算法。所以這裏並沒有重複造輪子，直接代理了 ratelimiter。

package ratelimit

import (
    "context"
    "golang.org/x/time/rate"
    "math"
)

type localTokenBucketLimiter struct {
    limit Limit

    limiter *rate.Limiter // 直接複用令牌桶的
}

func (lim *localTokenBucketLimiter) init() {
    burstCount := lim.limit.Count()
    if burstLimit, ok := lim.limit.(BurstLimit); ok {
        burstCount = burstLimit.BurstCount()
    }

    count := lim.limit.Count()
    if count < 0 {
        count = math.MaxInt32
    }

    f := float64(count) / lim.limit.Period().Seconds()
    if f < 0 {
        f = float64(rate.Inf) // 無限
    } else if f == 0 {
        panic("爲 0 的時候，底層實現有問題")
    }

    lim.limiter = rate.NewLimiter(rate.Limit(f), int(burstCount))
}

func (lim *localTokenBucketLimiter) Acquire() error {
    err := lim.limiter.Wait(context.TODO())
    return err
}

func (lim *localTokenBucketLimiter) TryAcquire() bool {
    return lim.limiter.Allow()
}

4. RedisCounterLimiter

package ratelimit

import (
    "math"
    "sync"
    "xg-go/log"
    "xg-go/xg/common"
)

type redisCounterLimiter struct {
    limit      DistLimit
    limitCount int32 // 內部使用，對 limit.count 做了 <0 時的轉換

    redisClient *common.RedisClient

    once sync.Once // 退化爲本地計數器的時候使用
    localLim Limiter

    //script string
}

func (lim *redisCounterLimiter) init() {
    lim.limitCount = lim.limit.Count()
    if lim.limitCount < 0 {
        lim.limitCount = math.MaxInt32
    }

    //lim.script = buildScript()
}

//func buildScript() string {
//  sb := strings.Builder{}
//
//  sb.WriteString("local c")
//  sb.WriteString("\nc = redis.call('get',KEYS[1])")
//  // 調用不超過最大值，則直接返回
//  sb.WriteString("\nif c and tonumber(c) > tonumber(ARGV[1]) then")
//  sb.WriteString("\nreturn c;")
//  sb.WriteString("\nend")
//  // 執行計算器自加
//  sb.WriteString("\nc = redis.call('incr',KEYS[1])")
//  sb.WriteString("\nif tonumber(c) == 1 then")
//  sb.WriteString("\nredis.call('expire',KEYS[1],ARGV[2])")
//  sb.WriteString("\nend")
//  sb.WriteString("\nif tonumber(c) == 1 then")
//  sb.WriteString("\nreturn c;")
//
//  return sb.String()
//}

func (lim *redisCounterLimiter) Acquire() error {
    panic("implement me")
}

func (lim *redisCounterLimiter) TryAcquire() (success bool) {
    defer func() {
        // 一般是 redis 連接斷了，會觸發空指針
        if err := recover(); err != nil {
            //log.Errorw("TryAcquire err", common.ERR, err)
            //success = lim.degradeTryAcquire()
            //return
            success = true
        }

        // 沒有錯誤，判斷是否開啓了 local 如果開啓了，把它停掉
        //if lim.localLim != nil {
        //  // stop 線程安全
        //  lim.localLim.Stop()
        //}
    }()

    count, err := lim.redisClient.IncrBy(lim.limit.Key(), 1)
    //panic("模擬 redis 出錯")
    if err != nil {
        log.Errorw("TryAcquire err", common.ERR, err)
        panic(err)
    }

    // *2 是爲了保留久一點，便於觀察
    err = lim.redisClient.Expire(lim.limit.Key(), int(2 * lim.limit.Period().Seconds()))
    if err != nil {
        log.Errorw("TryAcquire error", common.ERR, err)
        panic(err)
    }

    // 業務正確的情況下 確認超限
    if int32(count) > lim.limitCount {
        return false
    }

    return true

    //keys := []string{lim.limit.Key()}
    //
    //log.Errorw("TryAcquire ", keys, lim.limit.Count(), lim.limit.Period().Seconds())
    //count, err := lim.redisClient.Eval(lim.script, keys, lim.limit.Count(), lim.limit.Period().Seconds())
    //if err != nil {
    //  log.Errorw("TryAcquire error", common.ERR, err)
    //  return false
    //}
    //
    //
    //typeName := reflect.TypeOf(count).Name()
    //log.Errorw(typeName)
    //
    //if count != nil && count.(int32) <= lim.limitCount {
    //
    //  return true
    //}
    //return false
}

func (lim *redisCounterLimiter) Stop() {
    // 判斷是否開啓了 local 如果開啓了，把它停掉
    if lim.localLim != nil {
        // stop 線程安全
        lim.localLim.Stop()
    }
}

func (lim *redisCounterLimiter) degradeTryAcquire() bool {
    lim.once.Do(func() {
        count := lim.limit.Count() / lim.limit.ClusterNum()
        limit := LocalLimit {
            name: lim.limit.Name(),
            key: lim.limit.Key(),
            count: count,
            period: lim.limit.Period(),
            limitType: lim.limit.LimitType(),
        }

        lim.localLim = NewLimiter(&limit)
    })

    return lim.localLim.TryAcquire()
}

代碼裏回退的部分註釋了，因爲線上爲了穩定，實習生的代碼畢竟，所以先不跑。

本來原有的思路是直接用 lua 腳本在 redis上保證原子操作，但是底層封裝的庫對於直接調 eval 跑的時候，會拋錯，而且 source 是 go-redis 裏面，趕 ddl 沒有時間去 debug，所以只能用 incrBy + expire 分開來。

5. RedisTokenBucketLimiter
令牌桶的狀態變量得放在一個線程安全 / 一致的地方，redis 是不二人選。但是令牌桶的算法核心是個延遲計算得到令牌數量，這個是一個很長的臨界區，所以要麼用分佈式鎖，要麼直接利用 redis 的單線程以原子方式跑。一般業界是後者，即 lua 腳本維護令牌桶的狀態變量、計算令牌。代碼類似這種

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
local intval = tonumber(ARGV[5])

local fill_time = capacity/rate
local ttl = math.floor(fill_time*2) * intval

local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end

local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
if allowed then
  new_tokens = filled_tokens - requested
end

redis.call("setex", tokens_key, ttl, new_tokens)
redis.call("setex", timestamp_key, ttl, now)

return { allowed, new_tokens }

文章首發：https://www.jb51.net/article/234478.htm

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/3mzyHdgnfcP1EUxoqz1MyA

文章首發：https://www.jb51.net/article/234478.htm

猜你喜歡