From 71e9a84442c31c7e8d56e0ec3af8d9174866a6c3 Mon Sep 17 00:00:00 2001 From: "Gerasimos (Makis) Maropoulos" Date: Sat, 2 May 2020 22:21:42 +0300 Subject: [PATCH] godoc the (new) rate.Limit middleware Former-commit-id: fd5b4504ff873b55870b6966851157b8c641e587 --- middleware/rate/rate.go | 88 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 79 insertions(+), 9 deletions(-) diff --git a/middleware/rate/rate.go b/middleware/rate/rate.go index ce29ba9f..04036053 100644 --- a/middleware/rate/rate.go +++ b/middleware/rate/rate.go @@ -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 import ( @@ -15,24 +16,41 @@ func init() { 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) +// 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 { return func(l *Limiter) { 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 { return func(l *Limiter) { 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 { condition := func(c *Client) bool { // 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) { @@ -47,6 +65,10 @@ func PurgeEvery(every time.Duration, maxLifetime time.Duration) Option { } 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 { clientDataFunc func(ctx context.Context) interface{} // fill the Client's Data field. exceedHandler context.Handler // when too many requests. @@ -57,17 +79,27 @@ type ( 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 { - limiter *rate.Limiter - LastSeen time.Time - IP string - Data interface{} + Limiter *rate.Limiter + IP string + 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). 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 { l := &Limiter{ clients: make(map[string]*Client), @@ -85,6 +117,7 @@ func Limit(limit float64, burst int, options ...Option) context.Handler { return l.serveHTTP } +// Purge removes client entries from the memory based on the given "condition". func (l *Limiter) Purge(condition func(*Client) bool) { l.mu.Lock() for ip, client := range l.clients { @@ -103,7 +136,7 @@ func (l *Limiter) serveHTTP(ctx context.Context) { if !ok { client = &Client{ - limiter: rate.NewLimiter(l.limit, l.burstSize), + Limiter: rate.NewLimiter(l.limit, l.burstSize), IP: ip, } @@ -118,10 +151,15 @@ func (l *Limiter) serveHTTP(ctx context.Context) { l.mu.Unlock() } - client.LastSeen = time.Now() + client.mu.Lock() + client.lastSeen = time.Now() + client.mu.Unlock() + ctx.Values().Set(clientContextKey, client) - if client.limiter.Allow() { + // reserve := client.Limiter.Reserve() + // if reserve.OK() { + if client.Limiter.Allow() { ctx.Next() return } @@ -133,6 +171,12 @@ func (l *Limiter) serveHTTP(ctx context.Context) { 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 { if v := ctx.Values().Get(clientContextKey); v != nil { if c, ok := v.(*Client); ok { @@ -142,3 +186,29 @@ func Get(ctx context.Context) *Client { 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) +}