关于golang:gozero源码阅读限流器第四期

go-zero 给咱们提供了两种限流器，而且都是基于 redis 实现的可分布式的

限流器外围文件带正文代码如下，大家能够参阅

• 计数器限流器 https://github.com/TTSimple/g…
• 令牌桶限流器 https://github.com/TTSimple/g…

咱们通过最小化代码来看看限流器的外围思路

繁难计数器算法

// 繁难计数器算法
type Counter struct {
    rate  int           // 计数周期内最多容许的申请数
    begin time.Time     // 计数开始工夫
    cycle time.Duration // 计数周期
    count int           // 计数周期内累计收到的申请数
    lock  sync.Mutex
}

func (l *Counter) Allow() bool {l.lock.Lock()
    defer l.lock.Unlock()

    if l.count == l.rate-1 {now := time.Now()
        if now.Sub(l.begin) >= l.cycle {
            // 速度容许范畴内，重置计数器
            l.Reset(now)
            return true
        } else {return false}
    } else {
        // 没有达到速率限度，计数加 1
        l.count++
        return true
    }
}

func (l *Counter) Set(r int, cycle time.Duration) {
    l.rate = r
    l.begin = time.Now()
    l.cycle = cycle
    l.count = 0
}

func (l *Counter) Reset(t time.Time) {
    l.begin = t
    l.count = 0
}

func Test_Counter(t *testing.T) {c := Counter{}
    c.Set(20, time.Second)
    reqTime := 2 * time.Second                     // 总申请工夫
    reqNum := 200                                  // 总申请次数
    reqInterval := reqTime / time.Duration(reqNum) // 每次申请距离
    var trueCount, falseCount int
    for i := 0; i < reqNum; i++ {go func() {if c.Allow() {trueCount++} else {falseCount++}
        }()
        time.Sleep(reqInterval)
    }
    fmt.Println("true count:", trueCount)
    fmt.Println("false count:", falseCount)
}

最终输入

// === RUN   Test_Counter
// true count:  44
// false count:  156
// --- PASS: Test_Counter (2.07s)

繁难令牌桶算法

// 繁难令牌桶算法
type TokenBucket struct {
    rate         int64 // 固定的 token 放入速率, r/s
    capacity     int64 // 桶的容量
    tokens       int64 // 桶中以后 token 数量
    lastTokenSec int64 // 桶上次放 token 的工夫戳 s

    lock sync.Mutex
}

// 判断是否可通过
func (l *TokenBucket) Allow() bool {l.lock.Lock()
    defer l.lock.Unlock()

    now := time.Now().Unix()
    // 先增加初始令牌
    l.tokens = l.tokens + (now-l.lastTokenSec)*l.rate
    if l.tokens > l.capacity {l.tokens = l.capacity}
    l.lastTokenSec = now
    if l.tokens > 0 {
        // 还有令牌，支付令牌
        l.tokens--
        return true
    }
    // 没有令牌, 则回绝
    return false
}

// 动静设置参数
// r rate
// c capacity
func (l *TokenBucket) Set(r, c int64) {
    l.rate = r
    l.capacity = c
    l.tokens = r
    l.lastTokenSec = time.Now().Unix()
}

func Test_TokenBucket(t *testing.T) {lb := &TokenBucket{}
    lb.Set(20, 20)
    requestTime := 2 * time.Second                             // 总申请工夫
    requestNum := 200                                          // 总申请次数
    requestInterval := requestTime / time.Duration(requestNum) // 每次申请距离
    var trueCount, falseCount int
    for i := 0; i < requestNum; i++ {go func() {if lb.Allow() {trueCount++} else {falseCount++}
        }()
        time.Sleep(requestInterval)
    }
    fmt.Println("true count:", trueCount)
    fmt.Println("false count:", falseCount)
}

最终输入

=== RUN   Test_TokenBucket
true count:  60
false count:  140
--- PASS: Test_TokenBucket (2.07s)

繁难漏桶算法

漏桶算法的分布式版本 go-zero 没有给咱们实现，咱们看看其外围算法，而后参照外围算法来实现分布式版本，给大家安排个作业 :)

// 繁难漏桶算法
type LeakyBucket struct {
    rate       float64 // 固定每秒出水速率
    capacity   float64 // 桶的容量
    water      float64 // 桶中以后水量
    lastLeakMs int64   // 桶上次漏水工夫戳 ms

    lock sync.Mutex
}

// 判断是否可通过
func (l *LeakyBucket) Allow() bool {l.lock.Lock()
    defer l.lock.Unlock()

    now := time.Now().UnixNano() / 1e6
    eclipse := float64((now - l.lastLeakMs)) * l.rate / 1000 // 先执行漏水
    l.water = l.water - eclipse                              // 计算残余水量
    l.water = math.Max(0, l.water)                           // 桶干了
    l.lastLeakMs = now
    if (l.water + 1) < l.capacity {
        // 尝试加水, 并且水还未满
        l.water++
        return true
    } else {
        // 水满，回绝加水
        return false
    }
}

// 动静设置参数
// r rate
// c capacity
func (l *LeakyBucket) Set(r, c float64) {
    l.rate = r
    l.capacity = c
    l.water = 0
    l.lastLeakMs = time.Now().UnixNano() / 1e6
}

func Test_LeakyBucket(t *testing.T) {lb := &LeakyBucket{}
    lb.Set(20, 20)
    reqTime := 2 * time.Second                     // 总申请工夫
    reqNum := 200                                  // 总申请次数
    reqInterval := reqTime / time.Duration(reqNum) // 每次申请距离
    var trueCount, falseCount int
    for i := 0; i < reqNum; i++ {go func() {if lb.Allow() {trueCount++} else {falseCount++}
        }()
        time.Sleep(reqInterval)
    }
    fmt.Println("true count:", trueCount)
    fmt.Println("false count:", falseCount)
}

最终输入

// === RUN   Test_LeakyBucket
// true count:  60
// false count:  140
// --- PASS: Test_LeakyBucket (2.06s)