GO实现一个令牌桶限流

package program

import (
    "sync"
    "testing"
    "time"
)

func TestProgram(t *testing.T) {

    limiter := CreateTokenLimiter(5, time.Second)
    for i := 0; i < 10; i++ {
        if limiter.GetToken() {
            t.Log("get token")
        } else {
            t.Log("no token")
        }
    }
    time.Sleep(time.Second * 3)
    for i := 0; i < 10; i++ {
        if limiter.GetToken() {
            t.Log("get token")
        } else {
            t.Log("no token")
        }
    }

}

type TokenLimiter struct {
    Tokens   int
    Capacity int
    Rate     time.Duration
    LastTime time.Time
    mu       sync.Mutex
}

func CreateTokenLimiter(capacity int, rate time.Duration) *TokenLimiter {
    return &TokenLimiter{
        Tokens:   capacity,
        Capacity: capacity,
        Rate:     rate,
        LastTime: time.Now(),
    }
}

func (t *TokenLimiter) GetToken() bool {
    t.mu.Lock()
    defer t.mu.Unlock()
    tokens := time.Since(t.LastTime) / t.Rate
    t.Tokens += int(tokens)
    if t.Tokens > t.Capacity {
        t.Tokens = t.Capacity
    }
    t.LastTime = time.Now()
    if t.Tokens > 0 {
        t.Tokens--
        return true
    }
    return false
}