Naming, layout and comments: descriptive full-word names, one item per line in wrapped calls and literals, and doc comments on exported names
Short and single-letter names in application code become full words (tests keep theirs). Wrapped calls, signatures and composite literals put one item per line, and one-line function bodies are expanded. Exported names and groups of constants get doc comments, and comments refer to the renamed identifiers.
This commit is contained in:
1 parent
043e307663
commit
465358609a
78 files changed
+2374
-1402
No files matched your search
@@ -24,6 +24,7 @@ type Limiter interface {
|
||||
// LimiterFunc adapts a function to the Limiter interface.
|
||||
type LimiterFunc func(key string) (bool, time.Duration)
|
||||
|
||||
// Allow implements Limiter by calling f.
|
||||
func (f LimiterFunc) Allow(key string) (bool, time.Duration) { return f(key) }
|
||||
|
||||
// KeyFunc derives the rate limit key from a request.
|
||||
@@ -34,22 +35,24 @@ func KeyByIP(r *http.Request) string { return ClientIP(r) }
|
||||
|
||||
// RateLimit turns a Limiter into middleware. Rejected requests get 429 with
|
||||
// a Retry-After header. A nil limiter yields a pass-through middleware.
|
||||
func RateLimit(l Limiter, key KeyFunc) func(http.Handler) http.Handler {
|
||||
if l == nil {
|
||||
return func(next http.Handler) http.Handler { return next }
|
||||
func RateLimit(limiter Limiter, key KeyFunc) func(http.Handler) http.Handler {
|
||||
if limiter == nil {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return next
|
||||
}
|
||||
}
|
||||
if key == nil {
|
||||
key = KeyByIP
|
||||
}
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
ok, retry := l.Allow(key(r))
|
||||
ok, retry := limiter.Allow(key(r))
|
||||
if !ok {
|
||||
secs := int(math.Ceil(retry.Seconds()))
|
||||
if secs < 1 {
|
||||
secs = 1
|
||||
seconds := int(math.Ceil(retry.Seconds()))
|
||||
if seconds < 1 {
|
||||
seconds = 1
|
||||
}
|
||||
w.Header().Set("Retry-After", strconv.Itoa(secs))
|
||||
w.Header().Set("Retry-After", strconv.Itoa(seconds))
|
||||
httpx.WriteProblem(w, r, http.StatusTooManyRequests, "rate limit exceeded")
|
||||
return
|
||||
}
|
||||
@@ -74,22 +77,32 @@ type limiterShard struct {
|
||||
}
|
||||
|
||||
type limiterEntry struct {
|
||||
lim *rate.Limiter
|
||||
last time.Time
|
||||
limiter *rate.Limiter
|
||||
last time.Time
|
||||
}
|
||||
|
||||
// NewKeyedLimiter creates a limiter allowing rps sustained requests per
|
||||
// second with the given burst. A janitor goroutine evicts entries idle for
|
||||
// longer than ttl until ctx is cancelled.
|
||||
func NewKeyedLimiter(ctx context.Context, rps float64, burst int, ttl time.Duration) *KeyedLimiter {
|
||||
k := &KeyedLimiter{limit: rate.Limit(rps), burst: burst, ttl: ttl, seed: maphash.MakeSeed()}
|
||||
for i := range k.shards {
|
||||
k.shards[i].entries = make(map[string]*limiterEntry)
|
||||
// NewKeyedLimiter creates a limiter allowing requestsPerSecond sustained
|
||||
// requests per second with the given burst. A janitor goroutine evicts
|
||||
// entries idle for longer than ttl until ctx is cancelled.
|
||||
func NewKeyedLimiter(
|
||||
ctx context.Context,
|
||||
requestsPerSecond float64,
|
||||
burst int,
|
||||
ttl time.Duration,
|
||||
) *KeyedLimiter {
|
||||
limiter := &KeyedLimiter{
|
||||
limit: rate.Limit(requestsPerSecond),
|
||||
burst: burst,
|
||||
ttl: ttl,
|
||||
seed: maphash.MakeSeed(),
|
||||
}
|
||||
for i := range limiter.shards {
|
||||
limiter.shards[i].entries = make(map[string]*limiterEntry)
|
||||
}
|
||||
if ttl > 0 {
|
||||
go k.janitor(ctx)
|
||||
go limiter.janitor(ctx)
|
||||
}
|
||||
return k
|
||||
return limiter
|
||||
}
|
||||
|
||||
func (k *KeyedLimiter) shard(key string) *limiterShard {
|
||||
@@ -99,45 +112,47 @@ func (k *KeyedLimiter) shard(key string) *limiterShard {
|
||||
// Allow implements Limiter.
|
||||
func (k *KeyedLimiter) Allow(key string) (bool, time.Duration) {
|
||||
now := time.Now()
|
||||
s := k.shard(key)
|
||||
s.mu.Lock()
|
||||
e := s.entries[key]
|
||||
if e == nil {
|
||||
e = &limiterEntry{lim: rate.NewLimiter(k.limit, k.burst)}
|
||||
s.entries[key] = e
|
||||
shard := k.shard(key)
|
||||
shard.mu.Lock()
|
||||
entry := shard.entries[key]
|
||||
if entry == nil {
|
||||
entry = &limiterEntry{
|
||||
limiter: rate.NewLimiter(k.limit, k.burst),
|
||||
}
|
||||
shard.entries[key] = entry
|
||||
}
|
||||
e.last = now
|
||||
lim := e.lim
|
||||
s.mu.Unlock()
|
||||
entry.last = now
|
||||
limiter := entry.limiter
|
||||
shard.mu.Unlock()
|
||||
|
||||
res := lim.ReserveN(now, 1)
|
||||
res := limiter.ReserveN(now, 1)
|
||||
if !res.OK() {
|
||||
return false, time.Second
|
||||
}
|
||||
if d := res.DelayFrom(now); d > 0 {
|
||||
if delay := res.DelayFrom(now); delay > 0 {
|
||||
res.CancelAt(now)
|
||||
return false, d
|
||||
return false, delay
|
||||
}
|
||||
return true, 0
|
||||
}
|
||||
|
||||
func (k *KeyedLimiter) janitor(ctx context.Context) {
|
||||
t := time.NewTicker(k.ttl / 2)
|
||||
defer t.Stop()
|
||||
ticker := time.NewTicker(k.ttl / 2)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case now := <-t.C:
|
||||
case now := <-ticker.C:
|
||||
for i := range k.shards {
|
||||
s := &k.shards[i]
|
||||
s.mu.Lock()
|
||||
for key, e := range s.entries {
|
||||
if now.Sub(e.last) > k.ttl {
|
||||
delete(s.entries, key)
|
||||
shard := &k.shards[i]
|
||||
shard.mu.Lock()
|
||||
for key, entry := range shard.entries {
|
||||
if now.Sub(entry.last) > k.ttl {
|
||||
delete(shard.entries, key)
|
||||
}
|
||||
}
|
||||
s.mu.Unlock()
|
||||
shard.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user