github.com/ulule/limiter/v3@v3.11.3-0.20230613131926-4cb9c1da4633/drivers/store/memory/store.go (about)

     1  package memory
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"time"
     7  
     8  	"github.com/ulule/limiter/v3"
     9  	"github.com/ulule/limiter/v3/drivers/store/common"
    10  )
    11  
    12  // Store is the in-memory store.
    13  type Store struct {
    14  	// Prefix used for the key.
    15  	Prefix string
    16  	// cache used to store values in-memory.
    17  	cache *CacheWrapper
    18  }
    19  
    20  // NewStore creates a new instance of memory store with defaults.
    21  func NewStore() limiter.Store {
    22  	return NewStoreWithOptions(limiter.StoreOptions{
    23  		Prefix:          limiter.DefaultPrefix,
    24  		CleanUpInterval: limiter.DefaultCleanUpInterval,
    25  	})
    26  }
    27  
    28  // NewStoreWithOptions creates a new instance of memory store with options.
    29  func NewStoreWithOptions(options limiter.StoreOptions) limiter.Store {
    30  	return &Store{
    31  		Prefix: options.Prefix,
    32  		cache:  NewCache(options.CleanUpInterval),
    33  	}
    34  }
    35  
    36  // Get returns the limit for given identifier.
    37  func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
    38  	count, expiration := store.cache.Increment(store.getCacheKey(key), 1, rate.Period)
    39  
    40  	lctx := common.GetContextFromState(time.Now(), rate, expiration, count)
    41  	return lctx, nil
    42  }
    43  
    44  // Increment increments the limit by given count & returns the new limit value for given identifier.
    45  func (store *Store) Increment(ctx context.Context, key string, count int64, rate limiter.Rate) (limiter.Context, error) {
    46  	newCount, expiration := store.cache.Increment(store.getCacheKey(key), count, rate.Period)
    47  
    48  	lctx := common.GetContextFromState(time.Now(), rate, expiration, newCount)
    49  	return lctx, nil
    50  }
    51  
    52  // Peek returns the limit for given identifier, without modification on current values.
    53  func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
    54  	count, expiration := store.cache.Get(store.getCacheKey(key), rate.Period)
    55  
    56  	lctx := common.GetContextFromState(time.Now(), rate, expiration, count)
    57  	return lctx, nil
    58  }
    59  
    60  // Reset returns the limit for given identifier.
    61  func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) {
    62  	count, expiration := store.cache.Reset(store.getCacheKey(key), rate.Period)
    63  
    64  	lctx := common.GetContextFromState(time.Now(), rate, expiration, count)
    65  	return lctx, nil
    66  }
    67  
    68  // getCacheKey returns the full path for an identifier.
    69  func (store *Store) getCacheKey(key string) string {
    70  	buffer := strings.Builder{}
    71  	buffer.WriteString(store.Prefix)
    72  	buffer.WriteString(":")
    73  	buffer.WriteString(key)
    74  	return buffer.String()
    75  }