Reset rate limiter (#2)

* Implementing 'Reset' method for rate limiter

gcra.go

@@ -14,6 +14,11 @@ type gcra struct {
 	rdb   rediser
 }
 
+func (c *gcra) Reset() error {
+	res := c.rdb.Del(c.key)
+	return res.Err()
+}
+
 // Allow is shorthand for AllowN(key, 1).
 func (c *gcra) Allow() (*Result, error) {
 	return c.AllowN(1)

rate.go

@@ -27,6 +27,7 @@ var (
 type (
 	Algorithm interface {
 		Allow() (*Result, error)
+		Reset() error
 		SetKey(string)
 	}
 
@@ -35,6 +36,7 @@ type (
 		EvalSha(sha1 string, keys []string, args ...interface{}) *redis.Cmd
 		ScriptExists(hashes ...string) *redis.BoolSliceCmd
 		ScriptLoad(script string) *redis.StringCmd
+		Del(key ...string) *redis.IntCmd
 	}
 
 	Limit struct {
@@ -90,14 +92,9 @@ func NewLimiter(rdb rediser) *Limiter {
 
 func (l *Limiter) Allow(key string, limit *Limit) (*Result, error) {
 	var algo Algorithm
-
-	switch limit.Algorithm {
-	case SlidingWindowAlgorithm:
-		algo = &slidingWindow{limit: limit, rdb: l.rdb}
-	case GCRAAlgorithm:
-		algo = &gcra{limit: limit, rdb: l.rdb}
-	default:
-		return nil, errors.New("algorithm is not supported")
+	algo, err := l.getAlgorithm(limit)
+	if err != nil {
+		return nil, err
 	}
 
 	name, _ := GetAlgorithmName(limit.Algorithm)
@@ -107,6 +104,30 @@ func (l *Limiter) Allow(key string, limit *Limit) (*Result, error) {
 	return algo.Allow()
 }
 
+func (l *Limiter) Reset(key string, limit *Limit) error {
+	var algo Algorithm
+	algo, err := l.getAlgorithm(limit)
+	if err != nil {
+		return err
+	}
+
+	name, _ := GetAlgorithmName(limit.Algorithm)
+
+	algo.SetKey(l.Prefix + ":" + name + ":" + key)
+	return algo.Reset()
+}
+
+func (l *Limiter) getAlgorithm(limit *Limit) (Algorithm, error) {
+	switch limit.Algorithm {
+	case SlidingWindowAlgorithm:
+		return &slidingWindow{limit: limit, rdb: l.rdb}, nil
+	case GCRAAlgorithm:
+		return &gcra{limit: limit, rdb: l.rdb}, nil
+	default:
+		return nil, errors.New("algorithm is not supported")
+	}
+}
+
 func GetAlgorithmName(a uint) (string, bool) {
 	if name, ok := algorithmNames[a]; ok {
 		return name, ok

rate_test.go

@@ -55,6 +55,54 @@ func TestLimiter_Allow(t *testing.T) {
 	})
 }
 
+func TestLimiter_Reset(t *testing.T) {
+	l := rateLimiter()
+
+	limit := &Limit{
+		Rate:      1,
+		Period:    time.Minute,
+		Burst:     1,
+	}
+
+	t.Run("reset-sliding-window", func(t *testing.T) {
+		limit.Algorithm = SlidingWindowAlgorithm
+
+		res, err := l.Allow("sliding-reset_me", limit)
+		assert.Nil(t, err)
+		assert.True(t, res.Allowed)
+
+		res, err = l.Allow("sliding-reset_me", limit)
+		assert.Nil(t, err)
+		assert.False(t, res.Allowed)
+
+		err = l.Reset("sliding-reset_me", limit)
+		assert.Nil(t, err)
+
+		res, err = l.Allow("sliding-reset_me", limit)
+		assert.Nil(t, err)
+		assert.True(t, res.Allowed)
+	})
+
+	t.Run("reset-gcra-window", func(t *testing.T) {
+		limit.Algorithm = GCRAAlgorithm
+
+		res, err := l.Allow("gcra-reset_me", limit)
+		assert.Nil(t, err)
+		assert.True(t, res.Allowed)
+
+		res, err = l.Allow("gcra-reset_me", limit)
+		assert.Nil(t, err)
+		assert.False(t, res.Allowed)
+
+		err = l.Reset("gcra-reset_me", limit)
+		assert.Nil(t, err)
+
+		res, err = l.Allow("gcra-reset_me", limit)
+		assert.Nil(t, err)
+		assert.True(t, res.Allowed)
+	})
+}
+
 func BenchmarkAllow_Simple(b *testing.B) {
 	l := rateLimiter()
 	limit := &Limit{

sliding_window.go

@@ -12,6 +12,11 @@ type slidingWindow struct {
 	rdb   rediser
 }
 
+func (c *slidingWindow) Reset() error {
+	res := c.rdb.Del(c.key)
+	return res.Err()
+}
+
 func (c *slidingWindow) SetKey(key string) {
 	c.key = key
 }