令牌桶算法的Golang实现

令牌桶算法常常用于实现限速器，并能够允许突发数据的发送。

这里尝试用Golang实现一下令牌桶算法。

第一个版本参考

package main

import (
    "fmt"
    "time"
)

func logs(args ...interface{}) {
    fmt.Println(args...)
}

func tokenBucket(limit int, rate int) chan struct{} {
    tb := make(chan struct{}, limit)
    ticker := time.NewTicker(time.Duration(1) * time.Second)
    for i := 0; i < limit; i++ {
        tb <- struct{}{}
    }

    go func() {
        for {
            for i := 0; i < rate; i++ {
                tb <- struct{}{}
            }
            <-ticker.C
        }
    }()
    return tb
}

func popToken(bucket chan struct{}, n int) {
    for i := 0; i < n; i++ {
        <-bucket
    }
}

func testTokenBucket() {
    rate := 10
    tb := tokenBucket(20, rate)

    dataLen := 100
    for i := 0; i <= dataLen; i += rate {
        popToken(tb, rate)
        logs(i)
    }
}

func main() {
    testTokenBucket()
}

第二个版本参考

type TokenBucket struct {
    rate            int
    limit           int
    currentAmount   int
    lastConsumeTime int64
}

func currentTime() int64 {
    return time.Now().Unix()
}

func (tb *TokenBucket) popToken(n int) {
    ticker := time.NewTicker(500 * time.Millisecond)
    for n > tb.currentAmount {
        pre := tb.currentAmount + int(currentTime()-tb.lastConsumeTime)*tb.rate
        if pre > tb.limit {
            tb.currentAmount = tb.limit
        } else {
            tb.currentAmount = pre
        }
        <-ticker.C
    }
    tb.lastConsumeTime = currentTime()
    tb.currentAmount -= n
}

func newTokenBucket(limit int, rate int) *TokenBucket {
    tb := TokenBucket{
        rate:            rate,
        limit:           limit,
        currentAmount:   0,
        lastConsumeTime: currentTime(),
    }
    return &tb
}