Cache helm templates for the same inputs (#1016)

Author: SimKev2

pkg/helm/jsonnet.go

@@ -1,17 +1,24 @@
 package helm
 
 import (
+	"crypto/sha256"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"sync"
 
 	"github.com/google/go-jsonnet"
 	"github.com/google/go-jsonnet/ast"
 	"github.com/grafana/tanka/pkg/kubernetes/manifest"
+	"github.com/rs/zerolog/log"
 )
 
 // DefaultNameFormat to use when no nameFormat is supplied
 const DefaultNameFormat = `{{ print .kind "_" .metadata.name | snakecase }}`
 
+// helmTemplateCache caches the inline environments' rendered helm templates.
+var helmTemplateCache sync.Map
+
 // JsonnetOpts are additional properties the consumer of the native func might
 // pass.
 type JsonnetOpts struct {
@@ -57,6 +64,16 @@ func NativeFunc(h Helm) *jsonnet.NativeFunction {
 				return nil, fmt.Errorf("helmTemplate: Failed to find a chart at '%s': %s. See https://tanka.dev/helm#failed-to-find-chart", chart, err)
 			}
 
+			// check if resources exist in cache
+			helmKey, err := templateKey(name, chartpath, opts.TemplateOpts)
+			if err != nil {
+				return nil, err
+			}
+			if entry, ok := helmTemplateCache.Load(helmKey); ok {
+				log.Debug().Msgf("Using cached template for %s", name)
+				return entry, nil
+			}
+
 			// render resources
 			list, err := h.Template(name, chart, opts.TemplateOpts)
 			if err != nil {
@@ -69,11 +86,25 @@ func NativeFunc(h Helm) *jsonnet.NativeFunction {
 				return nil, err
 			}
 
+			helmTemplateCache.Store(helmKey, out)
 			return out, nil
 		},
 	}
 }
 
+// templateKey returns the key identifier used in the template cache for the given helm chart.
+func templateKey(chartName string, chartPath string, opts TemplateOpts) (string, error) {
+	hasher := sha256.New()
+	hasher.Write([]byte(chartName))
+	hasher.Write([]byte(chartPath))
+	valuesBytes, err := json.Marshal(opts)
+	if err != nil {
+		return "", err
+	}
+	hasher.Write(valuesBytes)
+	return base64.URLEncoding.EncodeToString(hasher.Sum(nil)), nil
+}
+
 func parseOpts(data interface{}) (*JsonnetOpts, error) {
 	c, err := json.Marshal(data)
 	if err != nil {

pkg/helm/jsonnet_test.go

@@ -10,6 +10,7 @@ import (
 )
 
 const calledFrom = "/my/path/here"
+const kubeVersion = "1.18.0"
 
 type MockHelm struct {
 	mock.Mock
@@ -88,8 +89,6 @@ func callNativeFunction(t *testing.T, expectedHelmTemplateOptions TemplateOpts,
 // TestDefaultCommandineFlagsIncludeCrds tests that the includeCrds flag is set
 // to true by default
 func TestDefaultCommandLineFlagsIncludeCrds(t *testing.T) {
-	kubeVersion := "1.18.0"
-
 	// we will check that the template function is called with these options,
 	// i.e. that includeCrds got set to true. This is not us passing an input,
 	// we are asserting here that the template function is called with these
@@ -116,8 +115,6 @@ func TestDefaultCommandLineFlagsIncludeCrds(t *testing.T) {
 // TestIncludeCrdsFalse tests that the includeCrds flag is can be set to false,
 // and this makes it to the helm.Template() method call
 func TestIncludeCrdsFalse(t *testing.T) {
-	kubeVersion := "1.18.0"
-
 	// we will check that the template function is called with these options,
 	// i.e. that includeCrds got set to false. This is not us passing an input,
 	// we are asserting here that the template function is called with these
@@ -140,3 +137,121 @@ func TestIncludeCrdsFalse(t *testing.T) {
 	// `helm template` don't contain the --include-crds flag
 	require.NotContains(t, args, "--include-crds")
 }
+
+// TestTemplateCachingWorks tests that calling template with the same chart and values twice will
+// use the cached version the second time.
+func TestTemplateCachingWorks(t *testing.T) {
+	values := map[string]interface{}{"testkey": "testvalue"}
+	expectedHelmTemplateOptions := TemplateOpts{
+		Values:      values,
+		KubeVersion: kubeVersion,
+		IncludeCRDs: true,
+	}
+	inputOptionsFromJsonnet := make(map[string]interface{})
+	inputOptionsFromJsonnet["values"] = values
+	inputOptionsFromJsonnet["calledFrom"] = calledFrom
+	inputOptionsFromJsonnet["kubeVersion"] = kubeVersion
+
+	helmMock := &MockHelm{}
+	// ChartExists called on both function calls, only template is cached
+	helmMock.On(
+		"ChartExists",
+		"exampleChartPath",
+		mock.AnythingOfType("*helm.JsonnetOpts")).
+		Return("/full/chart/path", nil).
+		Twice()
+	// this verifies that the helmMock.Template() method is called with the
+	// correct arguments and only a single time.
+	helmMock.On("Template", "exampleChartName", "/full/chart/path", expectedHelmTemplateOptions).
+		Return(manifest.List{}, nil).
+		Once()
+
+	nf := NativeFunc(helmMock)
+	require.NotNil(t, nf)
+
+	// the mandatory parameters to helm.template() in Jsonnet
+	params := []string{
+		"exampleChartName",
+		"exampleChartPath",
+	}
+
+	// mandatory parameters + the k-v pairs from the Jsonnet input
+	paramsInterface := make([]interface{}, 3)
+	paramsInterface[0] = params[0]
+	paramsInterface[1] = params[1]
+	paramsInterface[2] = inputOptionsFromJsonnet
+
+	_, err := nf.Func(paramsInterface)
+	firstCommandArgs := helmMock.TestData().Get("templateCommandArgs").StringSlice()
+	require.NoError(t, err)
+	_, err = nf.Func(paramsInterface)
+	secondCommandArgs := helmMock.TestData().Get("templateCommandArgs").StringSlice()
+	require.NoError(t, err)
+
+	helmMock.AssertExpectations(t)
+
+	// Verify command line args are same between the two calls
+	require.Equal(t, firstCommandArgs, secondCommandArgs)
+}
+
+type templateData struct {
+	chartName string
+	chartPath string
+	opts      TemplateOpts
+}
+
+var templateTestCases = []struct {
+	name       string
+	data       templateData
+	errMessage string
+}{
+	{
+		name: "emptyData",
+		data: templateData{
+			chartName: "testChart",
+			chartPath: "./chart/path",
+			opts:      TemplateOpts{},
+		},
+	},
+	{
+		name: "fullData",
+		data: templateData{
+			chartName: "bigChart",
+			chartPath: "./chart/bigPath",
+			opts: TemplateOpts{
+				map[string]interface{}{
+					"installCRDs": true,
+					"multitenancy": map[string]interface{}{
+						"enabled":               false,
+						"defaultServiceAccount": "default",
+						"privileged":            false,
+					},
+					"clusterDomain": "cluster.local",
+					"cli": map[string]interface{}{
+						"image":        "test-image.io",
+						"nodeSelector": map[string]interface{}{},
+						"tolerations":  []interface{}{},
+					},
+					"baz": []int32{12, 13},
+				},
+				[]string{"asdf", "qwer", "zxcv"},
+				true,
+				false,
+				"version",
+				"namespace",
+				false,
+			},
+		},
+	},
+}
+
+func BenchmarkTemplateKey(b *testing.B) {
+	for _, c := range templateTestCases {
+		b.Run(c.name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				// nolint:errcheck
+				templateKey(c.data.chartName, c.data.chartPath, c.data.opts)
+			}
+		})
+	}
+}