39 lines
799 B
Go
39 lines
799 B
Go
package limiter
|
|
|
|
import (
|
|
"time"
|
|
|
|
"gitea.com/lunny/tango"
|
|
"golang.org/x/net/context"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// Options represents the options to limit requests
|
|
type Options struct {
|
|
limiter *rate.Limiter
|
|
timeout time.Duration
|
|
}
|
|
|
|
// New creates a request limiter
|
|
func New(requestsPerSec, maxRequests int, timeout time.Duration) *Options {
|
|
limiter := rate.NewLimiter(rate.Limit(requestsPerSec), maxRequests)
|
|
return &Options{
|
|
limiter: limiter,
|
|
timeout: timeout,
|
|
}
|
|
}
|
|
|
|
// RequestLimit represents a request limit
|
|
func RequestLimit(opts *Options) tango.HandlerFunc {
|
|
return func(ctx *tango.Context) {
|
|
c, cancel := context.WithTimeout(context.Background(), opts.timeout)
|
|
defer cancel()
|
|
if err := opts.limiter.Wait(c); err != nil {
|
|
ctx.Abort(429)
|
|
return
|
|
}
|
|
|
|
ctx.Next()
|
|
}
|
|
}
|