From 8b23be8ef06ebdb44c8119b1679a2f730f18aacb Mon Sep 17 00:00:00 2001
From: ddl-retornam <56278673+ddl-retornam@users.noreply.github.com>
Date: Mon, 26 Jun 2023 13:31:19 -0700
Subject: [PATCH] [PLAT-6701] Enable zstd hephaestus

---
 pkg/buildkit/buildkit.go   | 30 ++++++++++++++++++++++++------
 pkg/cmd/controller/root.go |  4 ++++
 pkg/config/config.go       |  2 ++
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/pkg/buildkit/buildkit.go b/pkg/buildkit/buildkit.go
index 94aad1c4..a6d28a0f 100644
--- a/pkg/buildkit/buildkit.go
+++ b/pkg/buildkit/buildkit.go
@@ -23,6 +23,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/wait"
 
 	"github.com/dominodatalab/hephaestus/pkg/buildkit/archive"
+	hephconfig "github.com/dominodatalab/hephaestus/pkg/config"
 )
 
 var clientCheckBackoff = wait.Backoff{ // retries after 500ms 1s 2s 4s 8s 16s 32s 64s with jitter
@@ -118,6 +119,26 @@ type Client struct {
 	dockerAuthConfig string
 }
 
+func validateCompression(compression string, name string) map[string]string {
+	attrs := make(map[string]string)
+	attrs["name"] = name
+	switch compression {
+	case "estargz":
+		attrs["push"] = "true"
+		attrs["compression"] = "estargz"
+		attrs["force-compression"] = "true"
+		attrs["oci-mediatypes"] = "true"
+	case "zstd":
+		attrs["compression"] = "zstd"
+		attrs["force-compression"] = "true"
+		attrs["push"] = "true"
+	// default is gzip
+	default:
+		attrs["push"] = "true"
+	}
+	return attrs
+}
+
 func (c *Client) Build(ctx context.Context, opts BuildOptions) error {
 	// setup build directory
 	buildDir, err := os.MkdirTemp("", "hephaestus-build-")
@@ -230,14 +251,11 @@ func (c *Client) Build(ctx context.Context, opts BuildOptions) error {
 			solveOpt.FrontendAttrs[k] = v
 		}
 	}
-
 	for _, name := range opts.Images {
+		bkclientattrs := validateCompression(hephconfig.CompressionMethod, name)
 		solveOpt.Exports = append(solveOpt.Exports, bkclient.ExportEntry{
-			Type: bkclient.ExporterImage,
-			Attrs: map[string]string{
-				"name": name,
-				"push": "true",
-			},
+			Type:  bkclient.ExporterImage,
+			Attrs: bkclientattrs,
 		})
 	}
 
diff --git a/pkg/cmd/controller/root.go b/pkg/cmd/controller/root.go
index ef522db1..f838b446 100644
--- a/pkg/cmd/controller/root.go
+++ b/pkg/cmd/controller/root.go
@@ -2,6 +2,7 @@ package controller
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/spf13/cobra"
 
@@ -18,6 +19,7 @@ func NewCommand() *cobra.Command {
 		Short: "OCI image build controller using buildkit",
 	}
 	cmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "hephaestus.yaml", "configuration file")
+	cmd.PersistentFlags().StringVarP(&config.CompressionMethod, "compression", "d", "gzip", "Compression method options: zstd,estargz")
 	cmd.AddCommand(
 		newStartCommand(),
 		newRunGCCommand(),
@@ -33,6 +35,8 @@ func newStartCommand() *cobra.Command {
 		Use:   "start",
 		Short: "Start controller",
 		RunE: func(cmd *cobra.Command, _ []string) error {
+			config.CompressionMethod, _ = cmd.Flags().GetString("compression")
+			fmt.Printf("BuildKit compression method: %s enabled\n", config.CompressionMethod)
 			cfgFile, err := cmd.Flags().GetString("config")
 			if err != nil {
 				return err
diff --git a/pkg/config/config.go b/pkg/config/config.go
index d76133d0..a3eb7f23 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -12,6 +12,8 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+var CompressionMethod string
+
 type Controller struct {
 	Logging   Logging   `json:"logging" yaml:"logging"`
 	Manager   Manager   `json:"manager" yaml:"manager"`