From d5a0d6dd1db2a8a6bcb52e202ad37a454c7df3ca Mon Sep 17 00:00:00 2001
From: Roshan Piyush <roshan.piyush@traceable.ai>
Date: Mon, 30 Dec 2024 11:40:44 +0530
Subject: [PATCH] chore: Replace sync.Mutex with sync.Map (#1197)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Replace sync.Mutex with sync.Map

* Benchmark AddTransformation

* update benchmark parallelism

---------

Co-authored-by: Juan Pablo Tosso <jptosso@gmail.com>
Co-authored-by: Felipe Zipitría <3012076+fzipi@users.noreply.github.com>
---
 internal/corazawaf/rule.go      | 12 ++++--------
 internal/corazawaf/rule_test.go | 34 +++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/internal/corazawaf/rule.go b/internal/corazawaf/rule.go
index 4e04ff8de..bfc700f7b 100644
--- a/internal/corazawaf/rule.go
+++ b/internal/corazawaf/rule.go
@@ -586,22 +586,18 @@ func (r *Rule) AddVariableNegation(v variables.RuleVariable, key string) error {
 }
 
 var transformationIDToName = []string{""}
-var transformationNameToID = map[string]int{"": 0}
-var transformationIDsLock = sync.Mutex{}
+var transformationNameToID = sync.Map{} // map[string]int
 
 func transformationID(currentID int, transformationName string) int {
-	transformationIDsLock.Lock()
-	defer transformationIDsLock.Unlock()
-
 	currName := transformationIDToName[currentID]
 	nextName := fmt.Sprintf("%s+%s", currName, transformationName)
-	if id, ok := transformationNameToID[nextName]; ok {
-		return id
+	if id, ok := transformationNameToID.Load(nextName); ok {
+		return id.(int)
 	}
 
 	id := len(transformationIDToName)
 	transformationIDToName = append(transformationIDToName, nextName)
-	transformationNameToID[nextName] = id
+	transformationNameToID.Store(nextName, id)
 	return id
 }
 
diff --git a/internal/corazawaf/rule_test.go b/internal/corazawaf/rule_test.go
index d999bb0c8..f96eb4021 100644
--- a/internal/corazawaf/rule_test.go
+++ b/internal/corazawaf/rule_test.go
@@ -384,6 +384,40 @@ func TestAddTransformation(t *testing.T) {
 	}
 }
 
+func BenchmarkAddTransformationUnique(b *testing.B) {
+	transformation := func(input string) (string, bool, error) {
+		return "Test", true, nil
+	}
+	b.ResetTimer()
+	b.RunParallel(func(p *testing.PB) {
+		rule := NewRule()
+		for p.Next() {
+			transformationName := "transformation" + b.Name()
+			err := rule.AddTransformation(transformationName, transformation)
+			if err != nil {
+				b.Fatalf("Failed to add a transformation: %s", err.Error())
+			}
+		}
+	})
+}
+
+func BenchmarkAddTransformationSame(b *testing.B) {
+	transformation := func(input string) (string, bool, error) {
+		return "Test", true, nil
+	}
+	b.ResetTimer()
+	b.RunParallel(func(p *testing.PB) {
+		for p.Next() {
+			rule := NewRule()
+			transformationName := "transformation"
+			err := rule.AddTransformation(transformationName, transformation)
+			if err != nil {
+				b.Fatalf("Failed to add a transformation: %s", err.Error())
+			}
+		}
+	})
+}
+
 func TestAddTransformationEmpty(t *testing.T) {
 	rule := NewRule()
 	transformationName := ""