diff --git a/drivers/store/memory/store.go b/drivers/store/memory/store.go index 21b12f6..3199831 100644 --- a/drivers/store/memory/store.go +++ b/drivers/store/memory/store.go @@ -2,11 +2,11 @@ package memory import ( "context" + "strings" "time" "github.com/ulule/limiter/v3" "github.com/ulule/limiter/v3/drivers/store/common" - "github.com/ulule/limiter/v3/internal/bytebuffer" ) // Store is the in-memory store. @@ -35,11 +35,7 @@ func NewStoreWithOptions(options limiter.StoreOptions) limiter.Store { // Get returns the limit for given identifier. func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) { - buffer := bytebuffer.New() - defer buffer.Close() - buffer.Concat(store.Prefix, ":", key) - - count, expiration := store.cache.Increment(buffer.String(), 1, rate.Period) + count, expiration := store.cache.Increment(store.getCacheKey(key), 1, rate.Period) lctx := common.GetContextFromState(time.Now(), rate, expiration, count) return lctx, nil @@ -47,11 +43,7 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim // Increment increments the limit by given count & returns the new limit value for given identifier. func (store *Store) Increment(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) + newCount, expiration := store.cache.Increment(store.getCacheKey(key), count, rate.Period) lctx := common.GetContextFromState(time.Now(), rate, expiration, newCount) return lctx, nil @@ -59,11 +51,7 @@ func (store *Store) Increment(ctx context.Context, key string, count int64, rate // 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() - defer buffer.Close() - buffer.Concat(store.Prefix, ":", key) - - count, expiration := store.cache.Get(buffer.String(), rate.Period) + count, expiration := store.cache.Get(store.getCacheKey(key), rate.Period) lctx := common.GetContextFromState(time.Now(), rate, expiration, count) return lctx, nil @@ -71,12 +59,17 @@ func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (li // Reset returns the limit for given identifier. func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) { - buffer := bytebuffer.New() - defer buffer.Close() - buffer.Concat(store.Prefix, ":", key) - - count, expiration := store.cache.Reset(buffer.String(), rate.Period) + count, expiration := store.cache.Reset(store.getCacheKey(key), rate.Period) lctx := common.GetContextFromState(time.Now(), rate, expiration, count) return lctx, nil } + +// getCacheKey returns the full path for an identifier. +func (store *Store) getCacheKey(key string) string { + buffer := strings.Builder{} + buffer.WriteString(store.Prefix) + buffer.WriteString(":") + buffer.WriteString(key) + return buffer.String() +} diff --git a/drivers/store/redis/store.go b/drivers/store/redis/store.go index 13a4aa3..195b0d7 100644 --- a/drivers/store/redis/store.go +++ b/drivers/store/redis/store.go @@ -2,7 +2,6 @@ package redis import ( "context" - "fmt" "strings" "sync" "sync/atomic" @@ -98,22 +97,19 @@ func NewStoreWithOptions(client Client, options limiter.StoreOptions) (limiter.S // Increment increments the limit by given count & gives back the new limit for given identifier func (store *Store) Increment(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()) + cmd := store.evalSHA(ctx, store.getLuaIncrSHA, []string{store.getCacheKey(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()) + cmd := store.evalSHA(ctx, store.getLuaIncrSHA, []string{store.getCacheKey(key)}, 1, rate.Period.Milliseconds()) return currentContext(cmd, rate) } // 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) { - key = fmt.Sprintf("%s:%s", store.Prefix, key) - cmd := store.evalSHA(ctx, store.getLuaPeekSHA, []string{key}) + cmd := store.evalSHA(ctx, store.getLuaPeekSHA, []string{store.getCacheKey(key)}) count, ttl, err := parseCountAndTTL(cmd) if err != nil { return limiter.Context{}, err @@ -130,8 +126,7 @@ func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (li // Reset returns the limit for given identifier which is set to zero. func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) { - key = fmt.Sprintf("%s:%s", store.Prefix, key) - _, err := store.client.Del(ctx, key).Result() + _, err := store.client.Del(ctx, store.getCacheKey(key)).Result() if err != nil { return limiter.Context{}, err } @@ -143,6 +138,15 @@ func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (l return common.GetContextFromState(now, rate, expiration, count), nil } +// getCacheKey returns the full path for an identifier. +func (store *Store) getCacheKey(key string) string { + buffer := strings.Builder{} + buffer.WriteString(store.Prefix) + buffer.WriteString(":") + buffer.WriteString(key) + return buffer.String() +} + // preloadLuaScripts preloads the "incr" and "peek" lua scripts. func (store *Store) preloadLuaScripts(ctx context.Context) error { // Verify if we need to load lua scripts.