ratelimit.go

// The ratelimit package provides an efficient token bucket implementation

// that can be used to limit the rate of arbitrary things.

// See http://en.wikipedia.org/wiki/Token_bucket.

package ratelimit

import (

    "math"

    "strconv"

    "sync"

    "time"

)

// Bucket represents a token bucket that fills at a predetermined rate.

// Methods on Bucket may be called concurrently.

//令牌桶  结构体

type Bucket struct {

    startTime    time.Time  //开始时间

    capacity     int64    //令牌桶容量

    quantum      int64     //

    fillInterval time.Duration   //

    // The mutex guards the fields following it.

    mu sync.Mutex

    // avail holds the number of available tokens

    // in the bucket, as of availTick ticks from startTime.

    // It will be negative when there are consumers

    // waiting for tokens.

    avail     int64

    availTick int64

}

// NewBucket returns a new token bucket that fills at the

// rate of one token every fillInterval, up to the given

// maximum capacity. Both arguments must be

// positive. The bucket is initially full.

func NewBucket(fillInterval time.Duration, capacity int64) *Bucket {

    return NewBucketWithQuantum(fillInterval, capacity, 1)

}

// rateMargin specifes the allowed variance of actual

// rate from specified rate. 1% seems reasonable.

const rateMargin = 0.01

// NewBucketWithRate returns a token bucket that fills the bucket

// at the rate of rate tokens per second up to the given

// maximum capacity. Because of limited clock resolution,

// at high rates, the actual rate may be up to 1% different from the

// specified rate.

func NewBucketWithRate(rate float64, capacity int64) *Bucket {

    for quantum := int64(1); quantum < 1<<50; quantum = nextQuantum(quantum) {

        fillInterval := time.Duration(1e9 * float64(quantum) / rate)

        if fillInterval <= 0 {

            continue

        }

        tb := NewBucketWithQuantum(fillInterval, capacity, quantum)

        if diff := math.Abs(tb.Rate() - rate); diff/rate <= rateMargin {

            return tb

        }

    }

    panic("cannot find suitable quantum for " + strconv.FormatFloat(rate, 'g', -1, 64))

}

// nextQuantum returns the next quantum to try after q.

// We grow the quantum exponentially, but slowly, so we

// get a good fit in the lower numbers.

func nextQuantum(q int64) int64 {

    q1 := q * 11 / 10

    if q1 == q {

        q1++

    }

    return q1

}

// NewBucketWithQuantum is similar to NewBucket, but allows

// the specification of the quantum size - quantum tokens

// are added every fillInterval.

func NewBucketWithQuantum(fillInterval time.Duration, capacity, quantum int64) *Bucket {

    if fillInterval <= 0 {

        panic("token bucket fill interval is not > 0")

    }

    if capacity <= 0 {

        panic("token bucket capacity is not > 0")

    }

    if quantum <= 0 {

        panic("token bucket quantum is not > 0")

    }

    return &Bucket{

        startTime:    time.Now(),

        capacity:     capacity,

        quantum:      quantum,

        avail:        capacity,

        fillInterval: fillInterval,

    }

}

// Wait takes count tokens from the bucket, waiting until they are

// available.

func (tb *Bucket) Wait(count int64) {

    if d := tb.Take(count); d > 0 {

        time.Sleep(d)

    }

}

// WaitMaxDuration is like Wait except that it will

// only take tokens from the bucket if it needs to wait

// for no greater than maxWait. It reports whether

// any tokens have been removed from the bucket

// If no tokens have been removed, it returns immediately.

func (tb *Bucket) WaitMaxDuration(count int64, maxWait time.Duration) bool {

    d, ok := tb.TakeMaxDuration(count, maxWait)

    if d > 0 {

        time.Sleep(d)

    }

    return ok

}

const infinityDuration time.Duration = 0x7fffffffffffffff

// Take takes count tokens from the bucket without blocking. It returns

// the time that the caller should wait until the tokens are actually

// available.

//

// Note that if the request is irrevocable - there is no way to return

// tokens to the bucket once this method commits us to taking them.

func (tb *Bucket) Take(count int64) time.Duration {

    d, _ := tb.take(time.Now(), count, infinityDuration)

    return d

}

// TakeMaxDuration is like Take, except that

// it will only take tokens from the bucket if the wait

// time for the tokens is no greater than maxWait.

//

// If it would take longer than maxWait for the tokens

// to become available, it does nothing and reports false,

// otherwise it returns the time that the caller should

// wait until the tokens are actually available, and reports

// true.

func (tb *Bucket) TakeMaxDuration(count int64, maxWait time.Duration) (time.Duration, bool) {

    return tb.take(time.Now(), count, maxWait)

}

// TakeAvailable takes up to count immediately available tokens from the

// bucket. It returns the number of tokens removed, or zero if there are

// no available tokens. It does not block.

func (tb *Bucket) TakeAvailable(count int64) int64 {

    return tb.takeAvailable(time.Now(), count)

}

// takeAvailable is the internal version of TakeAvailable - it takes the

// current time as an argument to enable easy testing.

func (tb *Bucket) takeAvailable(now time.Time, count int64) int64 {

    if count <= 0 {

        return 0

    }

    tb.mu.Lock()

    defer tb.mu.Unlock()

    tb.adjust(now)

    if tb.avail <= 0 {

        return 0

    }

    if count > tb.avail {

        count = tb.avail

    }

    tb.avail -= count

    return count

}

// Available returns the number of available tokens. It will be negative

// when there are consumers waiting for tokens. Note that if this

// returns greater than zero, it does not guarantee that calls that take

// tokens from the buffer will succeed, as the number of available

// tokens could have changed in the meantime. This method is intended

// primarily for metrics reporting and debugging.

func (tb *Bucket) Available() int64 {

    return tb.available(time.Now())

}

// available is the internal version of available - it takes the current time as

// an argument to enable easy testing.

func (tb *Bucket) available(now time.Time) int64 {

    tb.mu.Lock()

    defer tb.mu.Unlock()

    tb.adjust(now)

    return tb.avail

}

// Capacity returns the capacity that the bucket was created with.

func (tb *Bucket) Capacity() int64 {

    return tb.capacity

}

// Rate returns the fill rate of the bucket, in tokens per second.

func (tb *Bucket) Rate() float64 {

    return 1e9 * float64(tb.quantum) / float64(tb.fillInterval)

}

// take is the internal version of Take - it takes the current time as

// an argument to enable easy testing.

func (tb *Bucket) take(now time.Time, count int64, maxWait time.Duration) (time.Duration, bool) {

    if count <= 0 {

        return 0, true

    }

    tb.mu.Lock()

    defer tb.mu.Unlock()

    currentTick := tb.adjust(now)

    avail := tb.avail - count

    if avail >= 0 {

        tb.avail = avail

        return 0, true

    }

    // Round up the missing tokens to the nearest multiple

    // of quantum - the tokens won't be available until

    // that tick.

    endTick := currentTick + (-avail+tb.quantum-1)/tb.quantum

    endTime := tb.startTime.Add(time.Duration(endTick) * tb.fillInterval)

    waitTime := endTime.Sub(now)

    if waitTime > maxWait {

        return 0, false

    }

    tb.avail = avail

    return waitTime, true

}

// adjust adjusts the current bucket capacity based on the current time.

// It returns the current tick.

func (tb *Bucket) adjust(now time.Time) (currentTick int64) {

    currentTick = int64(now.Sub(tb.startTime) / tb.fillInterval)

    if tb.avail >= tb.capacity {

        return

    }

    tb.avail += (currentTick - tb.availTick) * tb.quantum

    if tb.avail > tb.capacity {

        tb.avail = tb.capacity

    }

    tb.availTick = currentTick

    return

}