mirror of
https://github.com/kataras/iris.git
synced 2025-03-15 04:06:25 +01:00
godoc the (new) rate.Limit middleware
Former-commit-id: fd5b4504ff873b55870b6966851157b8c641e587
This commit is contained in:
parent
dbd6fcd2d7
commit
71e9a84442
|
@ -1,4 +1,5 @@
|
||||||
// TODO: godoc and add tests.
|
// Package rate implements rate limiter for Iris client requests.
|
||||||
|
// Example can be found at: _examples/miscellaneous/ratelimit/main.go.
|
||||||
package rate
|
package rate
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
@ -15,24 +16,41 @@ func init() {
|
||||||
context.SetHandlerName("iris/middleware/rate.(*Limiter).serveHTTP-fm", "iris.ratelimit")
|
context.SetHandlerName("iris/middleware/rate.(*Limiter).serveHTTP-fm", "iris.ratelimit")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Option delcares a function which can be passed on `Limit` package-level
|
||||||
|
// to modify its internal fields. Available Options are:
|
||||||
|
// * ExceedHandler
|
||||||
|
// * ClientData
|
||||||
|
// * PurgeEvery
|
||||||
type Option func(*Limiter)
|
type Option func(*Limiter)
|
||||||
|
|
||||||
|
// ExceedHandler is an `Option` that can be passed at the `Limit` package-level function.
|
||||||
|
// It accepts a handler that will be executed every time a client tries to reach a page/resource
|
||||||
|
// which is not accessible for that moment.
|
||||||
func ExceedHandler(handler context.Handler) Option {
|
func ExceedHandler(handler context.Handler) Option {
|
||||||
return func(l *Limiter) {
|
return func(l *Limiter) {
|
||||||
l.exceedHandler = handler
|
l.exceedHandler = handler
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ClientData is an `Option` that can be passed at the `Limit` package-level function.
|
||||||
|
// It accepts a function which provides the Iris Context and should return custom data
|
||||||
|
// that will be stored to the Client and be retrieved as `Get(ctx).Client.Data` later on.
|
||||||
func ClientData(clientDataFunc func(ctx context.Context) interface{}) Option {
|
func ClientData(clientDataFunc func(ctx context.Context) interface{}) Option {
|
||||||
return func(l *Limiter) {
|
return func(l *Limiter) {
|
||||||
l.clientDataFunc = clientDataFunc
|
l.clientDataFunc = clientDataFunc
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// PurgeEvery is an `Option` that can be passed at the `Limit` package-level function.
|
||||||
|
// This function will check for old entries and remove them.
|
||||||
|
//
|
||||||
|
// E.g. Limit(..., PurgeEvery(time.Minute, 5*time.Minute)) to
|
||||||
|
// check every 1 minute if a client's last visit was 5 minutes ago ("old" entry)
|
||||||
|
// and remove it from the memory.
|
||||||
func PurgeEvery(every time.Duration, maxLifetime time.Duration) Option {
|
func PurgeEvery(every time.Duration, maxLifetime time.Duration) Option {
|
||||||
condition := func(c *Client) bool {
|
condition := func(c *Client) bool {
|
||||||
// for a custom purger the end-developer may use the c.Data filled from a `ClientData` option.
|
// for a custom purger the end-developer may use the c.Data filled from a `ClientData` option.
|
||||||
return time.Since(c.LastSeen) > maxLifetime
|
return time.Since(c.LastSeen()) > maxLifetime
|
||||||
}
|
}
|
||||||
|
|
||||||
return func(l *Limiter) {
|
return func(l *Limiter) {
|
||||||
|
@ -47,6 +65,10 @@ func PurgeEvery(every time.Duration, maxLifetime time.Duration) Option {
|
||||||
}
|
}
|
||||||
|
|
||||||
type (
|
type (
|
||||||
|
// Limiter is featured with the necessary functions to limit requests per second.
|
||||||
|
// It has a single exported method `Purge` which helps to manually remove
|
||||||
|
// old clients from the memory. Limiter is not exposed by a function,
|
||||||
|
// callers should use it inside an `Option` for the `Limit` package-level function.
|
||||||
Limiter struct {
|
Limiter struct {
|
||||||
clientDataFunc func(ctx context.Context) interface{} // fill the Client's Data field.
|
clientDataFunc func(ctx context.Context) interface{} // fill the Client's Data field.
|
||||||
exceedHandler context.Handler // when too many requests.
|
exceedHandler context.Handler // when too many requests.
|
||||||
|
@ -57,17 +79,27 @@ type (
|
||||||
mu sync.RWMutex // mutex for clients.
|
mu sync.RWMutex // mutex for clients.
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Client holds some request information and the rate limiter itself.
|
||||||
|
// It can be retrieved by the `Get` package-level function.
|
||||||
|
// It can be used to manually add RateLimit response headers.
|
||||||
Client struct {
|
Client struct {
|
||||||
limiter *rate.Limiter
|
Limiter *rate.Limiter
|
||||||
LastSeen time.Time
|
IP string
|
||||||
IP string
|
Data interface{}
|
||||||
Data interface{}
|
|
||||||
|
lastSeen time.Time
|
||||||
|
mu sync.RWMutex // mutex for lastSeen.
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
// Inf is the infinite rate limit; it allows all events (even if burst is zero).
|
// Inf is the infinite rate limit; it allows all events (even if burst is zero).
|
||||||
const Inf = math.MaxFloat64
|
const Inf = math.MaxFloat64
|
||||||
|
|
||||||
|
// Limit returns a new rate limiter handler that allows requests up to rate "limit" and permits
|
||||||
|
// bursts of at most "burst" tokens.
|
||||||
|
// E.g. Limit(1, 5) to allow 1 request per second, with a maximum burst size of 5.
|
||||||
|
//
|
||||||
|
// See `ExceedHandler`, `ClientData` and `PurgeEvery` for the available "options".
|
||||||
func Limit(limit float64, burst int, options ...Option) context.Handler {
|
func Limit(limit float64, burst int, options ...Option) context.Handler {
|
||||||
l := &Limiter{
|
l := &Limiter{
|
||||||
clients: make(map[string]*Client),
|
clients: make(map[string]*Client),
|
||||||
|
@ -85,6 +117,7 @@ func Limit(limit float64, burst int, options ...Option) context.Handler {
|
||||||
return l.serveHTTP
|
return l.serveHTTP
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Purge removes client entries from the memory based on the given "condition".
|
||||||
func (l *Limiter) Purge(condition func(*Client) bool) {
|
func (l *Limiter) Purge(condition func(*Client) bool) {
|
||||||
l.mu.Lock()
|
l.mu.Lock()
|
||||||
for ip, client := range l.clients {
|
for ip, client := range l.clients {
|
||||||
|
@ -103,7 +136,7 @@ func (l *Limiter) serveHTTP(ctx context.Context) {
|
||||||
|
|
||||||
if !ok {
|
if !ok {
|
||||||
client = &Client{
|
client = &Client{
|
||||||
limiter: rate.NewLimiter(l.limit, l.burstSize),
|
Limiter: rate.NewLimiter(l.limit, l.burstSize),
|
||||||
IP: ip,
|
IP: ip,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -118,10 +151,15 @@ func (l *Limiter) serveHTTP(ctx context.Context) {
|
||||||
l.mu.Unlock()
|
l.mu.Unlock()
|
||||||
}
|
}
|
||||||
|
|
||||||
client.LastSeen = time.Now()
|
client.mu.Lock()
|
||||||
|
client.lastSeen = time.Now()
|
||||||
|
client.mu.Unlock()
|
||||||
|
|
||||||
ctx.Values().Set(clientContextKey, client)
|
ctx.Values().Set(clientContextKey, client)
|
||||||
|
|
||||||
if client.limiter.Allow() {
|
// reserve := client.Limiter.Reserve()
|
||||||
|
// if reserve.OK() {
|
||||||
|
if client.Limiter.Allow() {
|
||||||
ctx.Next()
|
ctx.Next()
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
@ -133,6 +171,12 @@ func (l *Limiter) serveHTTP(ctx context.Context) {
|
||||||
|
|
||||||
const clientContextKey = "iris.ratelimit.client"
|
const clientContextKey = "iris.ratelimit.client"
|
||||||
|
|
||||||
|
// Get returns the current rate limited `Client`.
|
||||||
|
// Use it when you want to log or add response headers based on the current request limitation.
|
||||||
|
//
|
||||||
|
// You can read more about X-RateLimit response headers at:
|
||||||
|
// https://tools.ietf.org/id/draft-polli-ratelimit-headers-00.html.
|
||||||
|
// A good example of that is the GitHub API itself: https://developer.github.com/v3/#rate-limiting
|
||||||
func Get(ctx context.Context) *Client {
|
func Get(ctx context.Context) *Client {
|
||||||
if v := ctx.Values().Get(clientContextKey); v != nil {
|
if v := ctx.Values().Get(clientContextKey); v != nil {
|
||||||
if c, ok := v.(*Client); ok {
|
if c, ok := v.(*Client); ok {
|
||||||
|
@ -142,3 +186,29 @@ func Get(ctx context.Context) *Client {
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// LastSeen reports the last Client's visit.
|
||||||
|
func (c *Client) LastSeen() (t time.Time) {
|
||||||
|
c.mu.RLock()
|
||||||
|
t = c.lastSeen
|
||||||
|
c.mu.RUnlock()
|
||||||
|
return t
|
||||||
|
}
|
||||||
|
|
||||||
|
// TokensFromDuration is a unit conversion function from a time duration to the number of tokens
|
||||||
|
// which could be accumulated during that duration at a rate of limit tokens per second.
|
||||||
|
func (c *Client) TokensFromDuration(d time.Duration) float64 {
|
||||||
|
// rate.go#tokensFromDuration
|
||||||
|
limit := float64(c.Limiter.Limit())
|
||||||
|
sec := float64(d/time.Second) * limit
|
||||||
|
nsec := float64(d%time.Second) * limit
|
||||||
|
return sec + nsec/1e9
|
||||||
|
}
|
||||||
|
|
||||||
|
// DurationFromTokens is a unit conversion function from the number of tokens to the duration
|
||||||
|
// of time it takes to accumulate them at a rate of limit tokens per second.
|
||||||
|
func (c *Client) DurationFromTokens(tokens float64) time.Duration {
|
||||||
|
// rate.go#durationFromTokens
|
||||||
|
seconds := tokens / float64(c.Limiter.Limit())
|
||||||
|
return time.Nanosecond * time.Duration(1e9*seconds)
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue
Block a user