91 lines
2.4 KiB
Go
91 lines
2.4 KiB
Go
// Copyright (c) Tailscale Inc & AUTHORS
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
// This is a modified, simplified version of code from golang.org/x/time/rate.
|
|
|
|
// Copyright 2015 The Go Authors. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// Package rate provides a rate limiter.
|
|
package rate
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"tailscale.com/tstime/mono"
|
|
)
|
|
|
|
// Limit defines the maximum frequency of some events.
|
|
// Limit is represented as number of events per second.
|
|
// A zero Limit is invalid.
|
|
type Limit float64
|
|
|
|
// Every converts a minimum time interval between events to a Limit.
|
|
func Every(interval time.Duration) Limit {
|
|
if interval <= 0 {
|
|
panic("invalid interval")
|
|
}
|
|
return 1 / Limit(interval.Seconds())
|
|
}
|
|
|
|
// A Limiter controls how frequently events are allowed to happen.
|
|
// It implements a [token bucket] of a particular size b,
|
|
// initially full and refilled at rate r tokens per second.
|
|
// Informally, in any large enough time interval,
|
|
// the Limiter limits the rate to r tokens per second,
|
|
// with a maximum burst size of b events.
|
|
// Use NewLimiter to create non-zero Limiters.
|
|
//
|
|
// [token bucket]: https://en.wikipedia.org/wiki/Token_bucket
|
|
type Limiter struct {
|
|
limit Limit
|
|
burst float64
|
|
mu sync.Mutex // protects following fields
|
|
tokens float64 // number of tokens currently in bucket
|
|
last mono.Time // the last time the limiter's tokens field was updated
|
|
}
|
|
|
|
// NewLimiter returns a new Limiter that allows events up to rate r and permits
|
|
// bursts of at most b tokens.
|
|
func NewLimiter(r Limit, b int) *Limiter {
|
|
if b < 1 {
|
|
panic("bad burst, must be at least 1")
|
|
}
|
|
return &Limiter{limit: r, burst: float64(b)}
|
|
}
|
|
|
|
// Allow reports whether an event may happen now.
|
|
func (lim *Limiter) Allow() bool {
|
|
return lim.allow(mono.Now())
|
|
}
|
|
|
|
func (lim *Limiter) allow(now mono.Time) bool {
|
|
lim.mu.Lock()
|
|
defer lim.mu.Unlock()
|
|
|
|
// If time has moved backwards, look around awkwardly and pretend nothing happened.
|
|
if now.Before(lim.last) {
|
|
lim.last = now
|
|
}
|
|
|
|
// Calculate the new number of tokens available due to the passage of time.
|
|
elapsed := now.Sub(lim.last)
|
|
tokens := lim.tokens + float64(lim.limit)*elapsed.Seconds()
|
|
if tokens > lim.burst {
|
|
tokens = lim.burst
|
|
}
|
|
|
|
// Consume a token.
|
|
tokens--
|
|
|
|
// Update state.
|
|
ok := tokens >= 0
|
|
if ok {
|
|
lim.last = now
|
|
lim.tokens = tokens
|
|
}
|
|
return ok
|
|
}
|