feat: Implement caching for pulling OCI charts (#2171)

Signed-off-by: Ruslan Khizhnyak <mustdiechik@gmail.com>
Co-authored-by: Ruslan Khizhnyak <rkhizhnyak@ptsecurity.com>

pkg/state/state.go

@@ -1191,6 +1191,19 @@ func (st *HelmState) GetRepositoryAndNameFromChartName(chartName string) (*Repos
 	return nil, chartName
 }
 
+var mutexMap sync.Map
+
+// retrieves or creates a sync.Mutex for a given name
+func (st *HelmState) getNamedMutex(name string) *sync.Mutex {
+	mu, ok := mutexMap.Load(name)
+	if ok {
+		return mu.(*sync.Mutex)
+	}
+	newMu := &sync.Mutex{}
+	actualMu, _ := mutexMap.LoadOrStore(name, newMu)
+	return actualMu.(*sync.Mutex)
+}
+
 type PrepareChartKey struct {
 	Namespace, Name, KubeContext string
 }
@@ -1269,7 +1282,7 @@ func (st *HelmState) PrepareCharts(helm helmexec.Interface, dir string, concurre
 				chartFetchedByGoGetter := chartPath != chartName
 
 				if !chartFetchedByGoGetter {
-					ociChartPath, err := st.getOCIChart(release, dir, helm, opts.OutputDirTemplate)
+					ociChartPath, err := st.getOCIChart(release, dir, helm, opts)
 					if err != nil {
 						results <- &chartPrepareResult{err: fmt.Errorf("release %q: %w", release.Name, err)}
 
@@ -4027,7 +4040,10 @@ func (st *HelmState) Reverse() {
 	}
 }
 
-func (st *HelmState) getOCIChart(release *ReleaseSpec, tempDir string, helm helmexec.Interface, outputDirTemplate string) (*string, error) {
+var downloadedOCICharts = make(map[string]bool)
+var downloadedOCIMutex sync.RWMutex
+
+func (st *HelmState) getOCIChart(release *ReleaseSpec, tempDir string, helm helmexec.Interface, opts ChartPrepareOptions) (*string, error) {
 	qualifiedChartName, chartName, chartVersion, err := st.getOCIQualifiedChartName(release, helm)
 	if err != nil {
 		return nil, err
@@ -4037,7 +4053,34 @@ func (st *HelmState) getOCIChart(release *ReleaseSpec, tempDir string, helm helm
 		return nil, nil
 	}
 
-	chartPath, _ := st.getOCIChartPath(tempDir, release, chartName, chartVersion, outputDirTemplate)
+	if opts.OutputDirTemplate == "" {
+		tempDir = remote.CacheDir()
+	}
+
+	chartPath, _ := st.getOCIChartPath(tempDir, release, chartName, chartVersion, opts.OutputDirTemplate)
+
+	mu := st.getNamedMutex(chartPath)
+	mu.Lock()
+	defer mu.Unlock()
+
+	_, err = os.Stat(tempDir)
+	if err != nil {
+		err = os.MkdirAll(tempDir, 0755)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	downloadedOCIMutex.RLock()
+	alreadyDownloadedFlag := downloadedOCICharts[chartPath]
+	downloadedOCIMutex.RUnlock()
+
+	if !opts.SkipDeps && !opts.SkipRefresh && !alreadyDownloadedFlag {
+		err = os.RemoveAll(chartPath)
+		if err != nil {
+			return nil, err
+		}
+	}
 
 	if st.fs.DirectoryExistsAt(chartPath) {
 		st.logger.Debugf("chart already exists at %s", chartPath)
@@ -4064,6 +4107,10 @@ func (st *HelmState) getOCIChart(release *ReleaseSpec, tempDir string, helm helm
 		return nil, err
 	}
 
+	downloadedOCIMutex.Lock()
+	downloadedOCICharts[chartPath] = true
+	downloadedOCIMutex.Unlock()
+
 	chartPath = filepath.Dir(fullChartPath)
 
 	return &chartPath, nil
@@ -4130,15 +4177,15 @@ func (st *HelmState) getOCIChartPath(tempDir string, release *ReleaseSpec, chart
 
 	pathElems := []string{tempDir}
 
-	if release.Namespace != "" {
-		pathElems = append(pathElems, release.Namespace)
-	}
-
-	if release.KubeContext != "" {
-		pathElems = append(pathElems, release.KubeContext)
-	}
-
-	pathElems = append(pathElems, release.Name, chartName, safeVersionPath(chartVersion))
+	replacer := strings.NewReplacer(
+		":", "_",
+		"//", "_",
+		".", "_",
+		"&", "_",
+	)
+	qName := strings.Split(replacer.Replace(release.Chart), "/")
 
+	pathElems = append(pathElems, qName...)
+	pathElems = append(pathElems, safeVersionPath(chartVersion))
 	return filepath.Join(pathElems...), nil
 }

test/e2e/template/helmfile/snapshot_test.go

@@ -30,12 +30,6 @@ var (
 	helmShortVersionRegex = regexp.MustCompile(`v\d+\.\d+\.\d+\+[a-z0-9]+`)
 )
 
-type ociChart struct {
-	name    string
-	version string
-	digest  string
-}
-
 type Config struct {
 	LocalDockerRegistry struct {
 		Enabled  bool   `yaml:"enabled"`
@@ -166,8 +160,6 @@ func testHelmfileTemplateWithBuildCommand(t *testing.T, GoYamlV3 bool) {
 				}
 			})
 
-			// ociCharts holds a list of chart name, version and digest distributed by local oci registry.
-			ociCharts := []ociChart{}
 			// If localDockerRegistry.enabled is set to `true`,
 			// run the docker registry v2 and push the test charts to the registry
 			// so that it can be accessed by helm and helmfile as a oci registry based chart repository.
@@ -201,16 +193,9 @@ func testHelmfileTemplateWithBuildCommand(t *testing.T, GoYamlV3 bool) {
 					if !c.IsDir() {
 						t.Fatalf("%s is not a directory", c)
 					}
-					chartName, chartVersion := execHelmShowChart(t, chartPath)
 					tgzFile := execHelmPackage(t, chartPath)
-					chartDigest, err := execHelmPush(t, tgzFile, fmt.Sprintf("oci://localhost:%d/myrepo", hostPort))
+					_, err := execHelmPush(t, tgzFile, fmt.Sprintf("oci://localhost:%d/myrepo", hostPort))
 					require.NoError(t, err, "Unable to run helm push to local registry: %v", err)
-
-					ociCharts = append(ociCharts, ociChart{
-						name:    chartName,
-						version: chartVersion,
-						digest:  chartDigest,
-					})
 				}
 			}
 
@@ -221,7 +206,7 @@ func testHelmfileTemplateWithBuildCommand(t *testing.T, GoYamlV3 bool) {
 			helmfileCacheHome := filepath.Join(tmpDir, "helmfile_cache")
 			// HELM_CONFIG_HOME contains the registry auth file (registry.json) and the index of all the repos added via helm-repo-add (repositories.yaml).
 			helmConfigHome := filepath.Join(tmpDir, "helm_config")
-			t.Logf("Using HELM_CACHE_HOME=%s, HELMFILE_CACHE_HOME=%s, HELM_CONFIG_HOME=%s", helmCacheHome, helmfileCacheHome, helmConfigHome)
+			t.Logf("Using HELM_CACHE_HOME=%s, HELMFILE_CACHE_HOME=%s, HELM_CONFIG_HOME=%s, WD=%s", helmCacheHome, helmfileCacheHome, helmConfigHome, wd)
 
 			inputFile := filepath.Join(testdataDir, name, "input.yaml.gotmpl")
 			outputFile := ""
@@ -230,6 +215,7 @@ func testHelmfileTemplateWithBuildCommand(t *testing.T, GoYamlV3 bool) {
 			} else {
 				outputFile = filepath.Join(testdataDir, name, "gopkg.in-yaml.v2-output.yaml")
 			}
+			expectedOutputFile := filepath.Join(testdataDir, name, "output.yaml")
 
 			ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
 			defer cancel()
@@ -262,15 +248,8 @@ func testHelmfileTemplateWithBuildCommand(t *testing.T, GoYamlV3 bool) {
 			gotStr = chartGitFullPathRegex.ReplaceAllString(gotStr, `chart=$$GoGetterPath`)
 			// Replace helm version with $HelmVersion
 			gotStr = helmShortVersionRegex.ReplaceAllString(gotStr, `$$HelmVersion`)
-			// Replace all occurrences of HELMFILE_CACHE_HOME with /home/runner/.cache/helmfile
-			// for stable test result
-			gotStr = strings.ReplaceAll(gotStr, helmfileCacheHome, "/home/runner/.cache/helmfile")
 
-			// OCI based helm charts are pulled and exported under temporary directory.
-			// We are not sure the exact name of the temporary directory generated by helmfile,
-			// so redact its base directory name with $TMP.
 			if config.LocalDockerRegistry.Enabled {
-				var releaseName, chartPath string
 				sc := bufio.NewScanner(strings.NewReader(gotStr))
 				for sc.Scan() {
 					if !strings.HasPrefix(sc.Text(), "Templating ") {
@@ -281,28 +260,20 @@ func testHelmfileTemplateWithBuildCommand(t *testing.T, GoYamlV3 bool) {
 					if len(releaseChartParts) != 2 {
 						t.Fatal("Found unexpected log output of templating oci based helm chart, want=\"Templating release=<release_name>, chart=<chart_name>\"")
 					}
-					releaseNamePart, chartPathPart := releaseChartParts[0], releaseChartParts[1]
-					releaseName = strings.TrimPrefix(releaseNamePart, "release=")
-					chartPath = chartPathPart
-				}
-				for _, ociChart := range ociCharts {
-					chartPathWithoutTempDirBase := fmt.Sprintf("/%s/%s/%s/%s", releaseName, ociChart.name, ociChart.version, ociChart.name)
-					var chartPathBase string
-					if strings.HasSuffix(chartPath, chartPathWithoutTempDirBase) {
-						chartPathBase = strings.TrimSuffix(chartPath, chartPathWithoutTempDirBase)
-					}
-					if len(chartPathBase) != 0 {
-						gotStr = strings.ReplaceAll(gotStr, chartPathBase, "chart=$TMP")
-					}
-					gotStr = strings.ReplaceAll(gotStr, fmt.Sprintf("Digest: %s", ociChart.digest), "Digest: $DIGEST")
 				}
 			}
 
+			gotStr = strings.ReplaceAll(gotStr, helmfileCacheHome, "$HELMFILE_CACHE_HOME")
+			gotStr = strings.ReplaceAll(gotStr, wd, "__workingdir__")
+
 			if stat, _ := os.Stat(outputFile); stat != nil {
 				want, err := os.ReadFile(outputFile)
-				wantStr := strings.ReplaceAll(string(want), "__workingdir__", wd)
 				require.NoError(t, err)
-				require.Equal(t, wantStr, gotStr)
+				require.Equal(t, string(want), gotStr)
+			} else if stat, _ := os.Stat(expectedOutputFile); stat != nil {
+				want, err := os.ReadFile(expectedOutputFile)
+				require.NoError(t, err)
+				require.Equal(t, string(want), gotStr)
 			} else {
 				// To update the test golden image(output.yaml), just remove it and rerun this test.
 				// We automatically capture the output to `output.yaml` in the test case directory
@@ -325,23 +296,6 @@ func execDocker(t *testing.T, args ...string) {
 	}
 }
 
-func execHelmShowChart(t *testing.T, localChart string) (string, string) {
-	t.Helper()
-
-	name, version := "", ""
-	out := execHelm(t, "show", "chart", localChart)
-	sc := bufio.NewScanner(strings.NewReader(out))
-	for sc.Scan() {
-		if strings.HasPrefix(sc.Text(), "name:") {
-			name = strings.TrimPrefix(sc.Text(), "name: ")
-		}
-		if strings.HasPrefix(sc.Text(), "version:") {
-			version = strings.TrimPrefix(sc.Text(), "version: ")
-		}
-	}
-	return name, version
-}
-
 func execHelmPackage(t *testing.T, localChart string) string {
 	t.Helper()
 

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull/output.yaml

@@ -1,5 +1,5 @@
 Pulling localhost:5001/myrepo/raw:0.1.0
-Templating release=foo, chart=$TMP/foo/raw/0.1.0/raw
+Templating release=foo, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
 ---
 # Source: raw/templates/resources.yaml
 apiVersion: v1

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_direct/config.yaml

@@ -0,0 +1,6 @@
+localDockerRegistry:
+  enabled: true
+  port: 5001
+chartifyTempDir: temp2
+helmfileArgs:
+- template

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_direct/input.yaml.gotmpl

@@ -0,0 +1,22 @@
+releases:
+- name: foo
+  chart: oci://localhost:5001/myrepo/raw
+  version: 0.1.0
+  values: &oci_chart_pull_direct
+  - templates:
+    - |
+      apiVersion: v1
+      kind: ConfigMap
+      metadata:
+        name: {{`{{ .Release.Name }}`}}
+        namespace: {{`{{ .Release.Namespace }}`}}
+        annotations:
+          chart-version: {{`{{ .Chart.Version }}`}}
+      data:
+        values: {{`{{ .Release.Name }}`}}
+
+- name: bar
+  chart: oci://localhost:5001/myrepo/raw
+  version: 0.1.0
+  namespace: ns2
+  values: *oci_chart_pull_direct

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_direct/output.yaml

@@ -0,0 +1,27 @@
+Pulling localhost:5001/myrepo/raw:0.1.0
+Templating release=foo, chart=$HELMFILE_CACHE_HOME/oci__localhost_5001/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: foo
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: foo
+
+Templating release=bar, chart=$HELMFILE_CACHE_HOME/oci__localhost_5001/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: bar
+  namespace: ns2
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: bar
+

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_once/config.yaml

@@ -0,0 +1,9 @@
+# Templating two versions of the same chart with only one pulling of each version
+localDockerRegistry:
+  enabled: true
+  port: 5001
+chartifyTempDir: temp3
+helmfileArgs:
+# Prevent releases from racing and randomizing the log
+- --concurrency=1
+- template

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_once/input.yaml.gotmpl

@@ -0,0 +1,43 @@
+repositories:
+- name: myrepo
+  url: localhost:5001/myrepo
+  oci: true
+
+releases:
+- name: foo
+  chart: myrepo/raw
+  version: 0.1.0
+  values: &oci_chart_pull_once_values
+  - templates:
+    - |
+      apiVersion: v1
+      kind: ConfigMap
+      metadata:
+        name: {{`{{ .Release.Name }}`}}
+        namespace: {{`{{ .Release.Namespace }}`}}
+        annotations:
+          chart-version: {{`{{ .Chart.Version }}`}}
+      data:
+        values: {{`{{ .Release.Name }}`}}
+
+- name: bar
+  chart: myrepo/raw
+  version: 0.1.0
+  values: *oci_chart_pull_once_values
+
+- name: release-no-default-ns
+  chart: myrepo/raw
+  version: 0.1.0
+  namespace: no-default
+  values: *oci_chart_pull_once_values
+
+- name: second-version-of-chart
+  chart: myrepo/raw
+  version: 0.0.1
+  namespace: foobar
+  values: *oci_chart_pull_once_values
+
+- name: first-release-version-of-chart
+  chart: myrepo/raw
+  version: 0.1.0
+  values: *oci_chart_pull_once_values

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_once/output.yaml

@@ -0,0 +1,67 @@
+Pulling localhost:5001/myrepo/raw:0.1.0
+Pulling localhost:5001/myrepo/raw:0.0.1
+Templating release=foo, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: foo
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: foo
+
+Templating release=bar, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: bar
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: bar
+
+Templating release=release-no-default-ns, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: release-no-default-ns
+  namespace: no-default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: release-no-default-ns
+
+Templating release=second-version-of-chart, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.0.1/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: second-version-of-chart
+  namespace: foobar
+  annotations:
+    chart-version: 0.0.1
+data:
+  values: second-version-of-chart
+
+Templating release=first-release-version-of-chart, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: first-release-version-of-chart
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: first-release-version-of-chart
+

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_once2/config.yaml

@@ -0,0 +1,7 @@
+# Templating few releases with the same chart\version and only one pulling
+localDockerRegistry:
+  enabled: true
+  port: 5001
+chartifyTempDir: temp3
+helmfileArgs:
+- template

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_once2/input.yaml.gotmpl

@@ -0,0 +1,23 @@
+repositories:
+- name: myrepo
+  url: localhost:5001/myrepo
+  oci: true
+
+releases:
+{{- range $i := until 5 }}
+- name: release-{{ $i }}
+  chart: myrepo/raw
+  version: 0.1.0
+  values:
+  - templates:
+    - |
+      apiVersion: v1
+      kind: ConfigMap
+      metadata:
+        name: {{`{{ .Release.Name }}`}}
+        namespace: {{`{{ .Release.Namespace }}`}}
+        annotations:
+          chart-version: {{`{{ .Chart.Version }}`}}
+      data:
+        values: {{`{{ .Release.Name }}`}}
+{{- end }}

test/e2e/template/helmfile/testdata/snapshot/oci_chart_pull_once2/output.yaml

@@ -0,0 +1,66 @@
+Pulling localhost:5001/myrepo/raw:0.1.0
+Templating release=release-0, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: release-0
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: release-0
+
+Templating release=release-1, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: release-1
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: release-1
+
+Templating release=release-2, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: release-2
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: release-2
+
+Templating release=release-3, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: release-3
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: release-3
+
+Templating release=release-4, chart=$HELMFILE_CACHE_HOME/myrepo/raw/0.1.0/raw
+---
+# Source: raw/templates/resources.yaml
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: release-4
+  namespace: default
+  annotations:
+    chart-version: 0.1.0
+data:
+  values: release-4
+