feat: add KRM benchmark example (#982)

* feat: add KRM benchmark example

* move init, init rg to init stage

* move common fns into a package

* use CreateTestVariant for variant generation

Co-authored-by: Morgante Pell <morgantep@google.com>

infra/blueprint-test/pkg/benchmark/krm.go

@@ -0,0 +1,99 @@
+package benchmark
+
+import (
+	"os"
+	"path"
+	"path/filepath"
+	"time"
+
+	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/kpt"
+	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/krmt"
+	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/utils"
+	"github.com/gruntwork-io/terratest/modules/shell"
+	"github.com/mitchellh/go-testing-interface"
+	"github.com/otiai10/copy"
+)
+
+// KubectlWaitForDeletion waits for resources in dir to be deleted.
+// Workaround until https://github.com/GoogleContainerTools/kpt/issues/2374
+func KubectlWaitForDeletion(b testing.TB, dir string, retries int, retryInterval time.Duration) {
+	waitArgs := []string{"get", "-R", "-f", ".", "--ignore-not-found"}
+	kptCmd := shell.Command{
+		Command:    "kubectl",
+		Args:       waitArgs,
+		Logger:     utils.GetLoggerFromT(),
+		WorkingDir: dir,
+	}
+	waitFunction := func() (bool, error) {
+		op, err := shell.RunCommandAndGetStdOutE(b, kptCmd)
+		if err != nil {
+			return false, err
+		}
+		// retry if output is not empty
+		retry := op != ""
+		return retry, nil
+	}
+	utils.Poll(b, waitFunction, retries, retryInterval)
+}
+
+// CreateVariant creates a variant of baseDir blueprint in the buildDir/variantName and upserts any given setters for that variant.
+func CreateVariant(b testing.TB, baseDir string, buildDir string, variantName string, setters map[string]string) {
+	for _, p := range []string{baseDir, buildDir} {
+		_, err := os.Stat(p)
+		if err != nil {
+			b.Fatalf("%s does not exist", p)
+		}
+	}
+	variantPath := path.Join(buildDir, variantName)
+	err := copy.Copy(baseDir, variantPath)
+	if err != nil {
+		b.Fatalf("Error copying resource from %s to %s", baseDir, variantPath)
+	}
+	rs, err := kpt.ReadPkgResources(variantPath)
+	if err != nil {
+		b.Fatalf("unable to read resources in %s :%v", variantPath, err)
+	}
+	kpt.UpsertSetters(rs, setters)
+	err = kpt.WritePkgResources(variantPath, rs)
+	if err != nil {
+		b.Fatalf("unable to write resources in %s :%v", variantPath, err)
+	}
+}
+
+// GetBuildDir creates a directory to store generated variants and cleanup fn.
+func GetBuildDir(b testing.TB) (string, func()) {
+	buildDir := path.Join(".build", b.Name())
+	err := os.MkdirAll(buildDir, 0755)
+	if err != nil {
+		b.Fatalf("unable to create %s :%v", buildDir, err)
+	}
+	abs, err := filepath.Abs(buildDir)
+	if err != nil {
+		b.Fatalf("unable to get absolute path for %s :%v", buildDir, err)
+	}
+	rmBuildDir := func() {
+		os.RemoveAll(buildDir)
+	}
+	return abs, rmBuildDir
+}
+
+// CreateTestVariant creates variants of a baseBlueprint and renders them with variantSetters
+func CreateTestVariant(b testing.TB, baseBlueprintDir string, variantSetters map[string]map[string]string) (*krmt.KRMBlueprintTest, string, func()) {
+	// precreate a custom build directory to generate variants for a given resource blueprint
+	buildDir, cleanup := GetBuildDir(b)
+	// init empty kpt pkg in the build dir
+	kptHelper := kpt.NewCmdConfig(b, kpt.WithDir(buildDir))
+	// generate package variants into the build dir
+	kptHelper.RunCmd("pkg", "init")
+	// generate variants of the base blueprint
+	for name, setters := range variantSetters {
+		CreateVariant(b, baseBlueprintDir, buildDir, name, setters)
+	}
+	// render variants
+	// TODO(bharathkkb): this is currently done in serial by kpt and can be slow for bigger topicCounts
+	// We should look into doing this in parallel possibly bundling variant creation with rendering
+	kptHelper.RunCmd("fn", "render")
+	kptHelper.RunCmd("live", "install-resource-group")
+	kptHelper.RunCmd("live", "init")
+	return krmt.NewKRMBlueprintTest(b, krmt.WithDir(buildDir), krmt.WithBuildDir(buildDir), krmt.WithUpdatePkgs(false)), buildDir, cleanup
+}

infra/blueprint-test/pkg/krmt/krm.go

@@ -54,6 +54,12 @@ func WithDir(dir string) krmtOption {
 	}
 }
 
+func WithBuildDir(buildDir string) krmtOption {
+	return func(f *KRMBlueprintTest) {
+		f.buildDir = buildDir
+	}
+}
+
 func WithUpdatePkgs(update bool) krmtOption {
 	return func(f *KRMBlueprintTest) {
 		f.updatePkgs = update
@@ -225,12 +231,12 @@ func (b *KRMBlueprintTest) DefaultInit(assert *assert.Assertions) {
 	}
 	b.updateSetters()
 	kpt.NewCmdConfig(b.t, kpt.WithDir(b.buildDir)).RunCmd("fn", "render")
+	b.kpt.RunCmd("live", "install-resource-group")
+	b.kpt.RunCmd("live", "init")
 }
 
 // DefaultApply installs resource-group, initializes inventory, applies pkg and polls resource statuses until current.
 func (b *KRMBlueprintTest) DefaultApply(assert *assert.Assertions) {
-	b.kpt.RunCmd("live", "install-resource-group")
-	b.kpt.RunCmd("live", "init")
 	b.kpt.RunCmd("live", "apply")
 	b.kpt.RunCmd("live", "status", "--output", "json", "--poll-until", "current", "--timeout", b.timeout)
 }

infra/blueprint-test/test/benchmark_fixtures/simple_pubsub_krm/Kptfile

@@ -0,0 +1,8 @@
+apiVersion: kpt.dev/v1
+kind: Kptfile
+metadata:
+  name: pubsub-topic
+pipeline:
+  mutators:
+    - image: gcr.io/kpt-fn/apply-setters:v0.1
+      configPath: setters.yaml

infra/blueprint-test/test/benchmark_fixtures/simple_pubsub_krm/pubsub.yaml

@@ -0,0 +1,21 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: pubsub.cnrm.cloud.google.com/v1beta1
+kind: PubSubTopic
+metadata: # kpt-merge: config-control/test-topic
+  namespace: config-control # kpt-set: ${namespace}
+  annotations:
+    cnrm.cloud.google.com/project-id: project-id # kpt-set: ${project-id}
+  name: test-topic # kpt-set: ${topic-name}

infra/blueprint-test/test/benchmark_fixtures/simple_pubsub_krm/setters.yaml

@@ -0,0 +1,22 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: v1
+kind: ConfigMap
+metadata: # kpt-merge: /setters
+  name: setters
+data:
+  namespace: config-control
+  topic-name: test-topic
+  project-id: project-id

infra/blueprint-test/test/krm_simple_pubsub_bench_test.go

@@ -0,0 +1,52 @@
+package test
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/benchmark"
+	"github.com/GoogleCloudPlatform/cloud-foundation-toolkit/infra/blueprint-test/pkg/utils"
+)
+
+// generateNTopics generates a slice of topicCount topic names
+func generateNTopics(topicCount int) []string {
+	var m []string
+	for i := 0; i < topicCount; i++ {
+		m = append(m, fmt.Sprintf("topic-%d", i))
+	}
+	return m
+}
+
+func BenchmarkKRMPubSubBlueprint(b *testing.B) {
+	projectID := utils.ValFromEnv(b, "PROJECT_ID")
+	// base blueprint dir
+	blueprintDir := "benchmark_fixtures/simple_pubsub_krm"
+	// topicCounts := []int{10, 50, 100, 500, 1000}
+	topicCounts := []int{10}
+	for _, topicCount := range topicCounts {
+		b.Run(fmt.Sprintf("PubSub Bench mark with %d topics", topicCount), func(b *testing.B) {
+			// generate setters for each variants
+			topicNames := generateNTopics(topicCount)
+			variantSetters := make(map[string]map[string]string)
+			for _, name := range topicNames {
+				variantSetters[name] = map[string]string{"topic-name": name, "project-id": projectID}
+			}
+			pubsubTest, buildDir, cleanup := benchmark.CreateTestVariant(b, blueprintDir, variantSetters)
+			defer cleanup()
+			b.ResetTimer()
+			// start benchmark
+			for n := 0; n < b.N; n++ {
+				pubsubTest.Apply(nil)
+				b.StopTimer()
+				// stop benchmark
+				pubsubTest.Teardown(nil)
+				// confirm resources are deleted
+				benchmark.KubectlWaitForDeletion(b, buildDir, 50, 5*time.Second)
+				// restart timer
+				b.StartTimer()
+			}
+
+		})
+	}
+}