From 3da3d783d60070f5752413dcd86ee6ff7442898f Mon Sep 17 00:00:00 2001
From: Ryan Lettieri <67934986+RyanLettieri@users.noreply.github.com>
Date: Fri, 12 Jan 2024 12:40:07 -0700
Subject: [PATCH] Adding deleteWithPrefix for in-memory statestore (#3314)

Signed-off-by: Ryan Lettieri <ryanLettieri@microsoft.com>
---
 state/in-memory/in_memory.go | 30 ++++++++++++++++++++++++++++++
 tests/config/state/tests.yml |  2 +-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/state/in-memory/in_memory.go b/state/in-memory/in_memory.go
index f0bc73f3f7..7e8fbe8442 100644
--- a/state/in-memory/in_memory.go
+++ b/state/in-memory/in_memory.go
@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -92,6 +93,7 @@ func (store *inMemoryStore) Features() []state.Feature {
 		state.FeatureETag,
 		state.FeatureTransactional,
 		state.FeatureTTL,
+		state.FeatureDeleteWithPrefix,
 	}
 }
 
@@ -116,6 +118,34 @@ func (store *inMemoryStore) Delete(ctx context.Context, req *state.DeleteRequest
 	return nil
 }
 
+func (store *inMemoryStore) DeleteWithPrefix(ctx context.Context, req state.DeleteWithPrefixRequest) (state.DeleteWithPrefixResponse, error) {
+	// step1: validate parameters
+	err := req.Validate()
+	if err != nil {
+		return state.DeleteWithPrefixResponse{}, err
+	}
+
+	// step2 should be protected by write-lock
+	store.lock.Lock()
+	defer store.lock.Unlock()
+
+	// step2: do really delete
+	// this operation won't fail
+	var count int64
+
+	for key := range store.items {
+		if strings.HasPrefix(key, req.Prefix) {
+			// The string contains the prefix, now we check to make sure there aren't more || after
+			longerPrefix := strings.Contains(key[len(req.Prefix):], "||")
+			if !longerPrefix {
+				delete(store.items, key)
+				count++
+			}
+		}
+	}
+	return state.DeleteWithPrefixResponse{Count: count}, nil
+}
+
 func (store *inMemoryStore) doValidateEtag(key string, etag *string, concurrency string) error {
 	hasEtag := etag != nil && *etag != ""
 
diff --git a/tests/config/state/tests.yml b/tests/config/state/tests.yml
index e281780aa0..77a3234619 100644
--- a/tests/config/state/tests.yml
+++ b/tests/config/state/tests.yml
@@ -90,7 +90,7 @@ components:
   - component: rethinkdb
     operations: []
   - component: in-memory
-    operations: [ "transaction", "etag",  "first-write", "ttl" ]
+    operations: [ "transaction", "etag",  "first-write", "ttl", "delete-with-prefix" ]
   - component: aws.dynamodb.docker
     # In the Docker variant, we do not set ttlAttributeName in the metadata, so TTLs are not enabled
     operations: [ "transaction", "etag", "first-write" ]