Enhancement: Adding new method Inc under drivers to allow rate limiting based on given variable count

drivers/store/memory/store.go

@@ -45,6 +45,18 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim
 	return lctx, nil
 }
 
+// Inc increments the limit by given count & returns the new limit value for given identifier.
+func (store *Store) Inc(ctx context.Context, key string, count int64, rate limiter.Rate) (limiter.Context, error) {
+	buffer := bytebuffer.New()
+	defer buffer.Close()
+	buffer.Concat(store.Prefix, ":", key)
+
+	newCount, expiration := store.cache.Increment(buffer.String(), count, rate.Period)
+
+	lctx := common.GetContextFromState(time.Now(), rate, expiration, newCount)
+	return lctx, nil
+}
+
 // Peek returns the limit for given identifier, without modification on current values.
 func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
 	buffer := bytebuffer.New()

drivers/store/redis/store.go

@@ -96,22 +96,18 @@ func NewStoreWithOptions(client Client, options limiter.StoreOptions) (limiter.S
 	return store, nil
 }
 
+// Inc increments the limit by given count & gives back the new limit for given identifier
+func (store *Store) Inc(ctx context.Context, key string, count int64, rate limiter.Rate) (limiter.Context, error) {
+	key = fmt.Sprintf("%s:%s", store.Prefix, key)
+	cmd := store.evalSHA(ctx, store.getLuaIncrSHA, []string{key}, count, rate.Period.Milliseconds())
+	return currentContext(cmd, rate)
+}
+
 // Get returns the limit for given identifier.
 func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
 	key = fmt.Sprintf("%s:%s", store.Prefix, key)
 	cmd := store.evalSHA(ctx, store.getLuaIncrSHA, []string{key}, 1, rate.Period.Milliseconds())
-	count, ttl, err := parseCountAndTTL(cmd)
-	if err != nil {
-		return limiter.Context{}, err
-	}
-
-	now := time.Now()
-	expiration := now.Add(rate.Period)
-	if ttl > 0 {
-		expiration = now.Add(time.Duration(ttl) * time.Millisecond)
-	}
-
-	return common.GetContextFromState(now, rate, expiration, count), nil
+	return currentContext(cmd, rate)
 }
 
 // Peek returns the limit for given identifier, without modification on current values.
@@ -253,3 +249,18 @@ func parseCountAndTTL(cmd *libredis.Cmd) (int64, int64, error) {
 
 	return count, ttl, nil
 }
+
+func currentContext(cmd *libredis.Cmd, rate limiter.Rate) (limiter.Context, error) {
+	count, ttl, err := parseCountAndTTL(cmd)
+	if err != nil {
+		return limiter.Context{}, err
+	}
+
+	now := time.Now()
+	expiration := now.Add(rate.Period)
+	if ttl > 0 {
+		expiration = now.Add(time.Duration(ttl) * time.Millisecond)
+	}
+
+	return common.GetContextFromState(now, rate, expiration, count), nil
+}

limiter.go

@@ -58,3 +58,7 @@ func (limiter *Limiter) Peek(ctx context.Context, key string) (Context, error) {
 func (limiter *Limiter) Reset(ctx context.Context, key string) (Context, error) {
 	return limiter.Store.Reset(ctx, key, limiter.Rate)
 }
+
+func (limiter *Limiter) Inc(ctx context.Context, key string, count int64) (Context, error) {
+	return limiter.Store.Inc(ctx, key, count, limiter.Rate)
+}

store.go

@@ -13,6 +13,8 @@ type Store interface {
 	Peek(ctx context.Context, key string, rate Rate) (Context, error)
 	// Reset resets the limit to zero for given identifier.
 	Reset(ctx context.Context, key string, rate Rate) (Context, error)
+	// Inc increments the limit by given count & gives back the new limit for given identifier
+	Inc(ctx context.Context, key string, count int64, rate Rate) (Context, error)
 }
 
 // StoreOptions are options for store.