[TT-16259] [Security] Unable to delete Session that uses a quota and is under heavy load [Plan A]

temporal/Taskfile.yml

@@ -38,4 +38,4 @@ tasks:
     status:
       - type mockery
     cmds:
-      - go install github.com/vektra/mockery/v2@v2.38.0
+      - go install github.com/vektra/mockery/v2@v2.53.6

temporal/internal/driver/redisv9/keyvalue.go

@@ -472,6 +472,14 @@ func (r *RedisV9) SetIfNotExist(ctx context.Context, key, value string, expirati
 	return res.Val(), nil
 }
 
+func (r *RedisV9) SetIfExist(ctx context.Context, key, value string, expiration time.Duration) (bool, error) {
+	if key == "" {
+		return false, temperr.KeyEmpty
+	}
+
+	return r.client.SetXX(ctx, key, value, expiration).Result()
+}
+
 func fetchKeysWithCursor(ctx context.Context,
 	client redis.UniversalClient,
 	pattern string,

temporal/keyvalue/keyvalue_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/TykTechnologies/storage/temporal/model"
 	"github.com/TykTechnologies/storage/temporal/temperr"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestKeyValue_Set(t *testing.T) {
@@ -1351,3 +1352,51 @@ func TestKeyValue_SetIfNotExist(t *testing.T) {
 		}
 	}
 }
+
+func TestKeyValue_SetIfExist(t *testing.T) {
+	connectors := testutil.TestConnectors(t)
+	defer testutil.CloseConnectors(t, connectors)
+
+	for _, connector := range connectors {
+		t.Run(connector.Type(), func(t *testing.T) {
+			ctx := context.Background()
+			kv, err := NewKeyValue(connector)
+			require.Nil(t, err)
+
+			flusher, err := flusher.NewFlusher(connector)
+			assert.Nil(t, err)
+			defer assert.Nil(t, flusher.FlushAll(ctx))
+
+			t.Run("does not set key if key does not exist", func(t *testing.T) {
+				const key = "non_existent_key"
+				ok, err := kv.SetIfExist(t.Context(), key, "value", 0)
+				assert.False(t, ok)
+				assert.NoError(t, err)
+
+				_, err = kv.Get(t.Context(), key)
+				assert.ErrorIs(t, err, temperr.KeyNotFound)
+			})
+
+			t.Run("overrides key of already exists", func(t *testing.T) {
+				const key = "existent_key"
+
+				err := kv.Set(t.Context(), key, "value1", 0)
+				assert.NoError(t, err)
+
+				ok, err := kv.SetIfExist(t.Context(), key, "value2", 0)
+				assert.True(t, ok)
+				assert.NoError(t, err)
+
+				v, err := kv.Get(t.Context(), key)
+				assert.Equal(t, "value2", v)
+				assert.NoError(t, err)
+			})
+
+			t.Run("returns err if empty key was given", func(t *testing.T) {
+				ok, err := kv.SetIfExist(t.Context(), "", "1", 0)
+				assert.False(t, ok)
+				assert.ErrorIs(t, err, temperr.KeyEmpty)
+			})
+		})
+	}
+}

temporal/model/types.go

@@ -58,6 +58,9 @@ type KeyValue interface {
 	// SetIfNotExist sets the string value of a key if the key does not exist.
 	// Returns true if the key was set, false otherwise.
 	SetIfNotExist(ctx context.Context, key, value string, expiration time.Duration) (bool, error)
+	// SetIfExist sets the string value of a key if the key exists.
+	// Returns true if the key was set, false otherwise.
+	SetIfExist(ctx context.Context, key, value string, expiration time.Duration) (bool, error)
 	// Delete removes the specified keys
 	Delete(ctx context.Context, key string) error
 	// Increment atomically increments the integer value of a key by one