github.com/cozy/cozy-stack@v0.0.0-20240603063001-31110fa4cae1/pkg/limits/impl_redis.go (about)

     1  package limits
     2  
     3  import (
     4  	"context"
     5  	"strconv"
     6  	"time"
     7  
     8  	"github.com/redis/go-redis/v9"
     9  )
    10  
    11  // Redis implementation of [Counter].
    12  //
    13  // This implementation is safe to use in multi-instances installation.
    14  type Redis struct {
    15  	Client redis.UniversalClient
    16  	ctx    context.Context
    17  }
    18  
    19  // NewRedis returns a counter that can be mutualized between several
    20  // cozy-stack processes by using redis.
    21  func NewRedis(client redis.UniversalClient) Counter {
    22  	return &Redis{client, context.Background()}
    23  }
    24  
    25  // incrWithTTL is a lua script for redis to increment a counter and sets a TTL
    26  // if it doesn't have one.
    27  const incrWithTTL = `
    28  local n = redis.call("INCR", KEYS[1])
    29  if redis.call("TTL", KEYS[1]) == -1 then
    30    redis.call("EXPIRE", KEYS[1], KEYS[2])
    31  end
    32  return n
    33  `
    34  
    35  func (r *Redis) Increment(key string, timeLimit time.Duration) (int64, error) {
    36  	ttl := strconv.FormatInt(int64(timeLimit/time.Second), 10)
    37  	count, err := r.Client.Eval(r.ctx, incrWithTTL, []string{key, ttl}).Result()
    38  	if err != nil {
    39  		return 0, err
    40  	}
    41  	return count.(int64), nil
    42  }
    43  
    44  func (r *Redis) Reset(key string) error {
    45  	_, err := r.Client.Del(r.ctx, key).Result()
    46  	return err
    47  }