From 938edcd0c5e459e0b56cd3597080eb950a6c9ea9 Mon Sep 17 00:00:00 2001
From: Will Beason <willbeason@google.com>
Date: Fri, 16 Jul 2021 22:56:28 -0500
Subject: [PATCH] [mutation] Cache tester in AssignMetadata like we do in
 Assign mutators (#1442)

---
 pkg/mutation/mutators/assignmeta_mutator.go   | 18 +++--
 .../assignmeta_mutator_benchmark_test.go      | 65 +++++++++++++++++++
 2 files changed, 76 insertions(+), 7 deletions(-)
 create mode 100644 pkg/mutation/mutators/assignmeta_mutator_benchmark_test.go

diff --git a/pkg/mutation/mutators/assignmeta_mutator.go b/pkg/mutation/mutators/assignmeta_mutator.go
index ddb3b9ddd4a..acc2d773f7e 100644
--- a/pkg/mutation/mutators/assignmeta_mutator.go
+++ b/pkg/mutation/mutators/assignmeta_mutator.go
@@ -42,6 +42,7 @@ var (
 // AssignMeta instance.
 type AssignMetadataMutator struct {
 	id             types.ID
+	tester         *tester.Tester
 	assignMetadata *mutationsv1alpha1.AssignMetadata
 	path           parser.Path
 }
@@ -59,13 +60,7 @@ func (m *AssignMetadataMutator) Matches(obj client.Object, ns *corev1.Namespace)
 }
 
 func (m *AssignMetadataMutator) Mutate(obj *unstructured.Unstructured) (bool, error) {
-	t, err := tester.New(m.Path(), []tester.Test{
-		{SubPath: m.Path(), Condition: tester.MustNotExist},
-	})
-	if err != nil {
-		return false, err
-	}
-	return core.Mutate(m, t, nil, obj)
+	return core.Mutate(m, m.tester, nil, obj)
 }
 
 func (m *AssignMetadataMutator) ID() types.ID {
@@ -96,6 +91,7 @@ func (m *AssignMetadataMutator) HasDiff(mutator types.Mutator) bool {
 func (m *AssignMetadataMutator) DeepCopy() types.Mutator {
 	res := &AssignMetadataMutator{
 		id:             m.id,
+		tester:         m.tester.DeepCopy(),
 		assignMetadata: m.assignMetadata.DeepCopy(),
 		path:           m.path.DeepCopy(),
 	}
@@ -142,8 +138,16 @@ func MutatorForAssignMetadata(assignMeta *mutationsv1alpha1.AssignMetadata) (*As
 		return nil, errors.New("spec.parameters.assign.value field must be a string for AssignMetadata " + assignMeta.GetName())
 	}
 
+	t, err := tester.New(path, []tester.Test{
+		{SubPath: path, Condition: tester.MustNotExist},
+	})
+	if err != nil {
+		return nil, err
+	}
+
 	return &AssignMetadataMutator{
 		id:             types.MakeID(assignMeta),
+		tester:         t,
 		assignMetadata: assignMeta.DeepCopy(),
 		path:           path,
 	}, nil
diff --git a/pkg/mutation/mutators/assignmeta_mutator_benchmark_test.go b/pkg/mutation/mutators/assignmeta_mutator_benchmark_test.go
new file mode 100644
index 00000000000..72a10543cbc
--- /dev/null
+++ b/pkg/mutation/mutators/assignmeta_mutator_benchmark_test.go
@@ -0,0 +1,65 @@
+package mutators
+
+import (
+	"testing"
+
+	"github.com/open-policy-agent/gatekeeper/apis/mutations/v1alpha1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+)
+
+func assignMetadata(value interface{}, location string) *v1alpha1.AssignMetadata {
+	result := &v1alpha1.AssignMetadata{
+		Spec: v1alpha1.AssignMetadataSpec{
+			Location: location,
+			Parameters: v1alpha1.MetadataParameters{
+				Assign: makeValue(value),
+			},
+		},
+	}
+
+	return result
+}
+
+func BenchmarkAssignMetadataMutator_Always(b *testing.B) {
+	mutator, err := MutatorForAssignMetadata(assignMetadata("bar", "metadata.labels.foo"))
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// Create a fresh object each time as AssignMetadata.Mutate does nothing if
+		// a label/annotation already exists. Thus, this test is for when we do
+		// actually make the mutation.
+
+		// The performance cost of instantiating the Unstructured is negligible
+		// compared to the time to perform Mutate.
+		obj := &unstructured.Unstructured{
+			Object: make(map[string]interface{}),
+		}
+		_, _ = mutator.Mutate(obj)
+	}
+}
+
+func BenchmarkAssignMetadataMutator_Never(b *testing.B) {
+	mutator, err := MutatorForAssignMetadata(assignMetadata("bar", "metadata.labels.foo"))
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	obj := &unstructured.Unstructured{
+		Object: make(map[string]interface{}),
+	}
+	_, err = mutator.Mutate(obj)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		// Use the same object each time as AssignMetadata.Mutate does nothing if
+		// a label/annotation already exists. Thus, this test is for the case where
+		// no mutation is necessary.
+		_, _ = mutator.Mutate(obj)
+	}
+}