Merge pull request #237 from stefanprodan/mod-cache

Implement caching for modules pulled from registries

cmd/timoni/apply.go

@@ -167,6 +167,7 @@ func runApplyCmd(cmd *cobra.Command, args []string) error {
 		applyArgs.module,
 		version,
 		tmpDir,
+		rootArgs.cacheDir,
 		applyArgs.creds.String(),
 	)
 	mod, err := fetcher.Fetch()

cmd/timoni/build.go

@@ -118,6 +118,7 @@ func runBuildCmd(cmd *cobra.Command, args []string) error {
 		buildArgs.module,
 		version,
 		tmpDir,
+		rootArgs.cacheDir,
 		buildArgs.creds.String(),
 	)
 	mod, err := fetcher.Fetch()

cmd/timoni/bundle_apply.go

@@ -234,6 +234,7 @@ func fetchBundleInstanceModule(ctx context.Context, instance *engine.BundleInsta
 		instance.Module.Repository,
 		moduleVersion,
 		modDir,
+		rootArgs.cacheDir,
 		bundleApplyArgs.creds.String(),
 	)
 	mod, err := fetcher.Fetch()

cmd/timoni/main.go

@@ -19,6 +19,7 @@ package main
 import (
 	"context"
 	"os"
+	"path"
 	"time"
 
 	"github.com/fatih/color"
@@ -57,6 +58,7 @@ type rootFlags struct {
 	timeout    time.Duration
 	prettyLog  bool
 	coloredLog bool
+	cacheDir   string
 }
 
 var (
@@ -76,6 +78,8 @@ func init() {
 		"Adds timestamps to the logs.")
 	rootCmd.PersistentFlags().BoolVar(&rootArgs.coloredLog, "log-color", rootArgs.coloredLog,
 		"Adds colorized output to the logs. (defaults to false when no tty)")
+	rootCmd.PersistentFlags().StringVar(&rootArgs.cacheDir, "cache-dir", "",
+		"Artifacts cache dir, can be disable with 'TIMONI_CACHING=false' env var. (defaults to \"$HOME/.timoni/cache\")")
 
 	addKubeConfigFlags(rootCmd)
 
@@ -85,6 +89,7 @@ func init() {
 }
 
 func main() {
+	setCacheDir()
 	if err := rootCmd.Execute(); err != nil {
 		// Ensure a logger is initialized even if the rootCmd
 		// failed before running its hooks.
@@ -99,6 +104,26 @@ func main() {
 	}
 }
 
+func setCacheDir() {
+	caching := os.Getenv("TIMONI_CACHING")
+	if caching == "false" || caching == "0" {
+		rootArgs.cacheDir = ""
+		return
+	}
+	if rootArgs.cacheDir == "" {
+		home, err := os.UserHomeDir()
+		if err != nil {
+			return
+		}
+		rootArgs.cacheDir = path.Join(home, ".timoni/cache")
+	}
+
+	if err := os.MkdirAll(rootArgs.cacheDir, os.ModePerm); err != nil {
+		// disable caching if target dir is not writable
+		rootArgs.cacheDir = ""
+	}
+}
+
 // addKubeConfigFlags maps the kubectl config flags to the given persistent flags.
 // The default namespace is set to the value found in current kubeconfig context.
 func addKubeConfigFlags(cmd *cobra.Command) {

cmd/timoni/main_test.go

@@ -75,6 +75,7 @@ func TestMain(m *testing.M) {
 	}
 
 	kubeconfigArgs.KubeConfig = &tmpFilename
+	rootArgs.cacheDir = tmpDir
 
 	code := m.Run()
 	testEnv.Stop()

cmd/timoni/mod_vet.go

@@ -92,6 +92,7 @@ func runVetModCmd(cmd *cobra.Command, args []string) error {
 		vetModArgs.path,
 		apiv1.LatestVersion,
 		tmpDir,
+		rootArgs.cacheDir,
 		"",
 	)
 	mod, err := fetcher.Fetch()

internal/engine/fetcher.go

@@ -30,21 +30,23 @@ import (
 
 // Fetcher downloads a module and extracts it locally.
 type Fetcher struct {
-	ctx     context.Context
-	src     string
-	dst     string
-	version string
-	creds   string
+	ctx      context.Context
+	src      string
+	dst      string
+	cacheDir string
+	version  string
+	creds    string
 }
 
 // NewFetcher creates a Fetcher for the given module.
-func NewFetcher(ctx context.Context, src, version, dst, creds string) *Fetcher {
+func NewFetcher(ctx context.Context, src, version, dst, cacheDir, creds string) *Fetcher {
 	return &Fetcher{
-		ctx:     ctx,
-		src:     src,
-		dst:     dst,
-		version: version,
-		creds:   creds,
+		ctx:      ctx,
+		src:      src,
+		dst:      dst,
+		version:  version,
+		cacheDir: cacheDir,
+		creds:    creds,
 	}
 }
 
@@ -103,5 +105,5 @@ func (f *Fetcher) fetchRemoteModule(dstDir string) (*apiv1.ModuleReference, erro
 	}
 
 	opts := oci.Options(f.ctx, f.creds)
-	return oci.PullModule(ociURL, dstDir, opts)
+	return oci.PullModule(ociURL, dstDir, f.cacheDir, opts)
 }

internal/oci/module_test.go

@@ -19,6 +19,7 @@ package oci
 import (
 	"context"
 	"fmt"
+	"os"
 	"path/filepath"
 	"testing"
 
@@ -90,7 +91,8 @@ func TestModuleOperations(t *testing.T) {
 	}
 
 	dstPath := filepath.Join(tmpDir, "artifact")
-	modRef, err := PullModule(digestURL, dstPath, opts)
+	cacheDir := t.TempDir()
+	modRef, err := PullModule(digestURL, dstPath, cacheDir, opts)
 	g.Expect(err).ToNot(HaveOccurred())
 	g.Expect(modRef.Version).To(BeEquivalentTo(imgVersion))
 	g.Expect(filepath.Join(dstPath, "timoni.ignore")).ToNot(BeAnExistingFile())
@@ -106,4 +108,7 @@ func TestModuleOperations(t *testing.T) {
 	} {
 		g.Expect(filepath.Join(dstPath, entry)).To(Or(BeAnExistingFile(), BeADirectory()))
 	}
+	cachedLayers, err := os.ReadDir(cacheDir)
+	g.Expect(err).ToNot(HaveOccurred())
+	g.Expect(len(cachedLayers)).To(BeEquivalentTo(2))
 }

internal/oci/pull_module.go

@@ -19,6 +19,9 @@ package oci
 import (
 	"bytes"
 	"fmt"
+	"io"
+	"os"
+	"path"
 
 	"github.com/fluxcd/pkg/tar"
 	"github.com/google/go-containerregistry/pkg/crane"
@@ -31,9 +34,10 @@ import (
 // - determines the artifact digest corresponding to the module version
 // - fetches the manifest of the remote artifact
 // - verifies that artifact config matches Timoni's media type
-// - download all the compressed layer matching Timoni's media type
+// - downloads all the compressed layer matching Timoni's media type (if not cached)
+// - stores the compressed layers in the local cache (if caching is enabled)
 // - extracts the module contents to the destination directory
-func PullModule(ociURL, dstPath string, opts []crane.Option) (*apiv1.ModuleReference, error) {
+func PullModule(ociURL, dstPath, cacheDir string, opts []crane.Option) (*apiv1.ModuleReference, error) {
 	ref, err := parseArtifactRef(ociURL)
 	if err != nil {
 		return nil, err
@@ -43,7 +47,7 @@ func PullModule(ociURL, dstPath string, opts []crane.Option) (*apiv1.ModuleRefer
 
 	digest, err := crane.Digest(ref.String(), opts...)
 	if err != nil {
-		return nil, fmt.Errorf("resolving the digest for '%s' failed: %w", ociURL, err)
+		return nil, fmt.Errorf("resolving digest of '%s' failed: %w", ociURL, err)
 	}
 
 	manifestJSON, err := crane.Manifest(ref.String(), opts...)
@@ -76,25 +80,72 @@ func PullModule(ociURL, dstPath string, opts []crane.Option) (*apiv1.ModuleRefer
 		Digest:     digest,
 	}
 
+	// If caching is disable, download the compressed layers to an ephemeral tmp dir.
+	if cacheDir == "" {
+		tmpDir, err := os.MkdirTemp("", apiv1.FieldManager)
+		if err != nil {
+			return nil, err
+		}
+		defer os.RemoveAll(tmpDir)
+		cacheDir = tmpDir
+	}
+
 	var foundLayer bool
 	for _, layer := range manifest.Layers {
 		if layer.MediaType == apiv1.ContentMediaType {
 			foundLayer = true
 			layerDigest := layer.Digest.String()
 			blobURL := fmt.Sprintf("%s@%s", repoURL, layerDigest)
-			layer, err := crane.PullLayer(blobURL, opts...)
-			if err != nil {
-				return nil, fmt.Errorf("pulling layer %s failed: %w", layerDigest, err)
+
+			isCached := false
+			cachedLayer := path.Join(cacheDir, fmt.Sprintf("%s.tgz", layer.Digest.Hex))
+			if _, err := os.Stat(cachedLayer); err == nil {
+				isCached = true
 			}
 
-			blob, err := layer.Compressed()
+			// Pull the compressed layer from the registry and persist the gzip tarball
+			// in the cache at '<cache-dir>/<layer-digest-hex>.tgz'.
+			if !isCached {
+				layer, err := crane.PullLayer(blobURL, opts...)
+				if err != nil {
+					return nil, fmt.Errorf("pulling layer %s failed: %w", layerDigest, err)
+				}
+
+				remote, err := layer.Compressed()
+				if err != nil {
+					return nil, fmt.Errorf("pulling layer %s failed: %w", layerDigest, err)
+				}
+
+				local, err := os.Create(cachedLayer)
+				if err != nil {
+					return nil, fmt.Errorf("writing layer to storage failed: %w", err)
+				}
+
+				if _, err := io.Copy(local, remote); err != nil {
+					return nil, fmt.Errorf("writing layer to storage failed: %w", err)
+				}
+
+				if err := local.Close(); err != nil {
+					return nil, fmt.Errorf("writing layer to storage failed: %w", err)
+				}
+			}
+
+			reader, err := os.Open(cachedLayer)
 			if err != nil {
-				return nil, fmt.Errorf("extracting layer %s failed: %w", layerDigest, err)
+				return nil, fmt.Errorf("reading layer from storage failed: %w", err)
 			}
 
-			if err = tar.Untar(blob, dstPath, tar.WithMaxUntarSize(-1)); err != nil {
+			// Extract the contents from the gzip tarball stored in cache.
+			// If extraction fails, the gzip tarball is removed from cache.
+			if err = tar.Untar(reader, dstPath, tar.WithMaxUntarSize(-1)); err != nil {
+				_ = reader.Close()
+				_ = os.Remove(cachedLayer)
 				return nil, fmt.Errorf("extracting layer %s failed: %w", layerDigest, err)
 			}
+
+			if err := reader.Close(); err != nil {
+				return nil, fmt.Errorf("reading layer from storage failed: %w", err)
+			}
 		}
 	}