godoc the (new) rate.Limit middleware

Former-commit-id: fd5b4504ff873b55870b6966851157b8c641e587

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
+		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).
 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)
+}