diff --git a/redis/redis.go b/redis/redis.go index 86958ed2e67e6e3d5b41a0e9b71add25ae2a88d2..e4f6ccc5e65a069fa8703fb681d99c86a43ea279 100644 --- a/redis/redis.go +++ b/redis/redis.go @@ -27,6 +27,13 @@ type ClientWithHelpers struct { Available bool } +type SetLockKeyOptions struct { + Value *string + MaxRetries *int + RetryInterval *time.Duration + FallbackResult *bool +} + func GetRedisClient(isDebug bool) *ClientWithHelpers { if redisClient != nil && redisClient.IsConnected() { return redisClient @@ -232,3 +239,54 @@ func (r ClientWithHelpers) RateLimit(key string, limitFn func(int) redis_rate.Li return res.Allowed == 1, nil } } + +// SetLockKey attempts to set a lock on the specified key. +func (r ClientWithHelpers) SetLockKey(key string, expiration time.Duration, lockOptions ...SetLockKeyOptions) (success bool, err error) { + fallbackResult := false + if len(lockOptions) > 0 && lockOptions[0].FallbackResult != nil { + fallbackResult = *lockOptions[0].FallbackResult + } + + if !r.IsConnected() { + return fallbackResult, errors.Error("Redis is not connected") + } + + value := "1" + retries := 0 + retryInterval := 100 * time.Millisecond + if len(lockOptions) > 0 { + // Only use the values that were set + if lockOptions[0].Value != nil { + value = *lockOptions[0].Value + } + if lockOptions[0].MaxRetries != nil { + retries = *lockOptions[0].MaxRetries + } + if lockOptions[0].RetryInterval != nil { + retryInterval = *lockOptions[0].RetryInterval + } + } + + for retries >= 0 { + success, err = r.Client.SetNX(ctx, key, value, expiration).Result() + if err != nil { + logs.ErrorWithMsg(fmt.Sprintf("Error setting lock key %s", key), err) + + // Prevent further calls in this execution from trying to connect and also timeout + if strings.HasSuffix(err.Error(), "i/o timeout") { + r.Available = false + } + + return fallbackResult, err + } + + if success || retries == 0 { + break + } + + retries-- + time.Sleep(retryInterval) + } + + return success, nil +}