72 lines
2 KiB
Go
72 lines
2 KiB
Go
package ratelimiter
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/mailgun/ttlmap"
|
|
"github.com/rs/zerolog"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
type inMemoryRateLimiter struct {
|
|
rate rate.Limit // reqs/s
|
|
burst int64
|
|
// maxDelay is the maximum duration we're willing to wait for a bucket reservation to become effective, in nanoseconds.
|
|
// For now it is somewhat arbitrarily set to 1/(2*rate).
|
|
maxDelay time.Duration
|
|
// Each rate limiter for a given source is stored in the buckets ttlmap.
|
|
// To keep this ttlmap constrained in size,
|
|
// each ratelimiter is "garbage collected" when it is considered expired.
|
|
// It is considered expired after it hasn't been used for ttl seconds.
|
|
ttl int
|
|
buckets *ttlmap.TtlMap // actual buckets, keyed by source.
|
|
|
|
logger *zerolog.Logger
|
|
}
|
|
|
|
func newInMemoryRateLimiter(rate rate.Limit, burst int64, maxDelay time.Duration, ttl int, logger *zerolog.Logger) (*inMemoryRateLimiter, error) {
|
|
buckets, err := ttlmap.NewConcurrent(maxSources)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("creating ttlmap: %w", err)
|
|
}
|
|
|
|
return &inMemoryRateLimiter{
|
|
rate: rate,
|
|
burst: burst,
|
|
maxDelay: maxDelay,
|
|
ttl: ttl,
|
|
logger: logger,
|
|
buckets: buckets,
|
|
}, nil
|
|
}
|
|
|
|
func (i *inMemoryRateLimiter) Allow(_ context.Context, source string) (*time.Duration, error) {
|
|
// Get bucket which contains limiter information.
|
|
var bucket *rate.Limiter
|
|
if rlSource, exists := i.buckets.Get(source); exists {
|
|
bucket = rlSource.(*rate.Limiter)
|
|
} else {
|
|
bucket = rate.NewLimiter(i.rate, int(i.burst))
|
|
}
|
|
|
|
// We Set even in the case where the source already exists,
|
|
// because we want to update the expiryTime everytime we get the source,
|
|
// as the expiryTime is supposed to reflect the activity (or lack thereof) on that source.
|
|
if err := i.buckets.Set(source, bucket, i.ttl); err != nil {
|
|
return nil, fmt.Errorf("setting buckets: %w", err)
|
|
}
|
|
|
|
res := bucket.Reserve()
|
|
if !res.OK() {
|
|
return nil, nil
|
|
}
|
|
|
|
delay := res.Delay()
|
|
if delay > i.maxDelay {
|
|
res.Cancel()
|
|
}
|
|
|
|
return &delay, nil
|
|
}
|