Merge pull request go-redis#8 from heynemann/v4

Added support for verifying the allow limit without incrementing it.

Author: vmihailenco

README.md

@@ -35,6 +35,17 @@ func handler(w http.ResponseWriter, req *http.Request, rateLimiter *rate.Limiter
 	fmt.Fprint(w, "Rate limit remaining: ", strconv.FormatInt(limit-rate, 10))
 }
 
+func statusHandler(w http.ResponseWriter, req *http.Request, rateLimiter *rate.Limiter) {
+	userID := "user-12345"
+	limit := int64(5)
+
+	//With increment 0, we just retrieve the current limit
+	rate, reset, allowed := rateLimiter.AllowN(userID, limit, time.Minute, 0)
+	fmt.Fprintf(w, "Current rate: %d", rate)
+	fmt.Fprintf(w, "Reset: %d", reset)
+	fmt.Fprintf(w, "Allowed: %v", allowed)
+}
+
 func main() {
 	ring := redis.NewRing(&redis.RingOptions{
 		Addrs: map[string]string{
@@ -48,6 +59,10 @@ func main() {
 		handler(w, req, limiter)
 	})
 
+	http.HandleFunc("/status", func(w http.ResponseWriter, req *http.Request) {
+		statusHandler(w, req, limiter)
+	})
+
 	http.HandleFunc("/favicon.ico", http.NotFound)
 	log.Println("listening on localhost:8888...")
 	log.Println(http.ListenAndServe("localhost:8888", nil))

rate.go

@@ -1,12 +1,12 @@
-package rate // import "gopkg.in/go-redis/rate.v4"
+package rate
 
 import (
 	"fmt"
 	"time"
 
-	timerate "golang.org/x/time/rate"
+	redis "gopkg.in/redis.v4"
 
-	"gopkg.in/redis.v4"
+	timerate "golang.org/x/time/rate"
 )
 
 const redisPrefix = "rate"
@@ -15,38 +15,46 @@ type rediser interface {
 	Pipelined(func(pipe *redis.Pipeline) error) ([]redis.Cmder, error)
 }
 
+//Limiter type
 type Limiter struct {
 	fallbackLimiter *timerate.Limiter
 	redis           rediser
 }
 
-// A Limiter controls how frequently events are allowed to happen. It uses
-// the redis to store data and fallbacks to the fallbackLimiter
-// when Redis Server is not available.
+// NewLimiter creates a limiter that controls how frequently events
+// are allowed to happen. It uses redis to store data and fallbacks
+// to the fallbackLimiter when Redis Server is not available.
 func NewLimiter(redis rediser, fallbackLimiter *timerate.Limiter) *Limiter {
 	return &Limiter{
 		fallbackLimiter: fallbackLimiter,
 		redis:           redis,
 	}
 }
 
-// Allow reports whether an event with given name may happen at time now.
-// It allows up to maxn events within duration dur.
-func (l *Limiter) Allow(name string, maxn int64, dur time.Duration) (count, reset int64, allow bool) {
+// AllowN reports whether an event with given name may happen at time now.
+// It allows up to maxn events within duration dur, with each interaction incrementing
+// the limit by n.
+func (l *Limiter) AllowN(name string, maxn int64, dur time.Duration, n int64) (count, reset int64, allow bool) {
 	udur := int64(dur / time.Second)
 	slot := time.Now().Unix() / udur
 	reset = (slot + 1) * udur
 	allow = l.fallbackLimiter.Allow()
 
 	name = fmt.Sprintf("%s:%s-%d", redisPrefix, name, slot)
-	count, err := l.incr(name, dur)
+	count, err := l.incr(name, dur, n)
 	if err == nil {
 		allow = count <= maxn
 	}
 
 	return count, reset, allow
 }
 
+// Allow reports whether an event with given name may happen at time now.
+// It allows up to maxn events within duration dur.
+func (l *Limiter) Allow(name string, maxn int64, dur time.Duration) (count, reset int64, allow bool) {
+	return l.AllowN(name, maxn, dur, 1)
+}
+
 // AllowMinute is shorthand for Allow(name, maxn, time.Minute).
 func (l *Limiter) AllowMinute(name string, maxn int64) (int64, int64, bool) {
 	return l.Allow(name, maxn, time.Minute)
@@ -79,7 +87,7 @@ func (l *Limiter) AllowRate(name string, rateLimit timerate.Limit) (delay time.D
 	allow = l.fallbackLimiter.Allow()
 
 	name = fmt.Sprintf("%s:%s-%d-%d", redisPrefix, name, dur, slot)
-	count, err := l.incr(name, dur)
+	count, err := l.incr(name, dur, 1)
 	if err == nil {
 		allow = count <= limit
 	}
@@ -91,10 +99,10 @@ func (l *Limiter) AllowRate(name string, rateLimit timerate.Limit) (delay time.D
 	return delay, allow
 }
 
-func (l *Limiter) incr(name string, dur time.Duration) (int64, error) {
+func (l *Limiter) incr(name string, dur time.Duration, n int64) (int64, error) {
 	var incr *redis.IntCmd
 	_, err := l.redis.Pipelined(func(pipe *redis.Pipeline) error {
-		incr = pipe.Incr(name)
+		incr = pipe.IncrBy(name, n)
 		pipe.Expire(name, dur)
 		return nil
 	})

rate_test.go

@@ -1,24 +1,23 @@
-package rate_test
+package rate
 
 import (
 	"testing"
 	"time"
 
-	timerate "golang.org/x/time/rate"
-	"gopkg.in/redis.v4"
+	redis "gopkg.in/redis.v4"
 
-	"gopkg.in/go-redis/rate.v4"
+	timerate "golang.org/x/time/rate"
 )
 
-func rateLimiter() *rate.Limiter {
+func rateLimiter() *Limiter {
 	ring := redis.NewRing(&redis.RingOptions{
 		Addrs: map[string]string{"server0": ":6379"},
 	})
 	if err := ring.FlushDb().Err(); err != nil {
 		panic(err)
 	}
 	fallbackLimiter := timerate.NewLimiter(timerate.Every(time.Millisecond), 100)
-	return rate.NewLimiter(ring, fallbackLimiter)
+	return NewLimiter(ring, fallbackLimiter)
 }
 
 func TestAllow(t *testing.T) {
@@ -90,7 +89,7 @@ func TestAllowRateSecond(t *testing.T) {
 func TestRedisIsDown(t *testing.T) {
 	ring := redis.NewRing(&redis.RingOptions{})
 	limiter := timerate.NewLimiter(timerate.Every(time.Second), 1)
-	l := rate.NewLimiter(ring, limiter)
+	l := NewLimiter(ring, limiter)
 
 	rate, _, allow := l.AllowMinute("test_id", 1)
 	if !allow {
@@ -109,6 +108,36 @@ func TestRedisIsDown(t *testing.T) {
 	}
 }
 
+func TestAllowN(t *testing.T) {
+	l := rateLimiter()
+
+	rate, reset, allow := l.AllowN("test_allow_n", 1, time.Minute, 1)
+	if !allow {
+		t.Fatalf("rate limited with rate %d", rate)
+	}
+	if rate != 1 {
+		t.Fatalf("got %d, wanted 1", rate)
+	}
+	dur := time.Duration(reset-time.Now().Unix()) * time.Second
+	if dur > time.Minute {
+		t.Fatalf("got %s, wanted <= %s", dur, time.Minute)
+	}
+
+	l.AllowN("test_allow_n", 1, time.Minute, 2)
+
+	rate, reset, allow = l.AllowN("test_allow_n", 1, time.Minute, 0)
+	if allow {
+		t.Fatalf("should rate limit with rate %d", rate)
+	}
+	if rate != 3 {
+		t.Fatalf("got %d, wanted 3", rate)
+	}
+	dur = time.Duration(reset-time.Now().Unix()) * time.Second
+	if dur > time.Minute {
+		t.Fatalf("got %s, wanted <= %s", dur, time.Minute)
+	}
+}
+
 func durEqual(got, wanted time.Duration) bool {
 	return got > 0 && got < wanted
 }