[DOM-49106] Emit image size bytes as part of build (#131)

[DOM-49106] Emit image size bytes as part of build (#131)

deployments/crds/hephaestus.dominodatalab.com_imagebuilds.yaml

@@ -150,6 +150,10 @@ spec:
                 description: BuilderAddr is the routable address to the buildkit pod
                   used during the image build process.
                 type: string
+              compressedImageSizeBytes:
+                description: CompressedImageSizeBytes is the total size of all the
+                  compressed layers in the image.
+                type: string
               conditions:
                 items:
                   description: "Condition contains details for one aspect of the current

go.mod

@@ -24,6 +24,7 @@ require (
 	github.com/dominodatalab/controller-util v0.1.0
 	github.com/go-logr/logr v1.3.0
 	github.com/go-logr/zapr v1.3.0
+	github.com/google/go-containerregistry v0.14.0
 	github.com/h2non/filetype v1.1.3
 	github.com/hashicorp/go-retryablehttp v0.7.5
 	github.com/moby/buildkit v0.12.4
@@ -117,6 +118,7 @@ require (
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/klauspost/compress v1.17.2 // indirect
 	github.com/kylelemons/godebug v1.1.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 // indirect

go.sum

@@ -256,6 +256,8 @@ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-containerregistry v0.14.0 h1:z58vMqHxuwvAsVwvKEkmVBz2TlgBgH5k6koEXBtlYkw=
+github.com/google/go-containerregistry v0.14.0/go.mod h1:aiJ2fp/SXvkWgmYHioXnbMdlgB8eXiiYOY55gfN91Wk=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/gofuzz v1.1.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=

pkg/api/hephaestus/v1/imagebuild_types.go

@@ -49,6 +49,8 @@ type ImageBuildStatus struct {
 	BuildTime string `json:"buildTime,omitempty"`
 	// BuilderAddr is the routable address to the buildkit pod used during the image build process.
 	BuilderAddr string `json:"builderAddr,omitempty"`
+	// CompressedImageSizeBytes is the total size of all the compressed layers in the image.
+	CompressedImageSizeBytes string `json:"compressedImageSizeBytes,omitempty"`
 
 	Conditions  []metav1.Condition     `json:"conditions,omitempty"`
 	Transitions []ImageBuildTransition `json:"transitions,omitempty"`

pkg/buildkit/buildkit.go

@@ -13,6 +13,9 @@ import (
 	"github.com/containerd/console"
 	"github.com/docker/cli/cli/config"
 	"github.com/go-logr/logr"
+	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/name"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
 	bkclient "github.com/moby/buildkit/client"
 	"github.com/moby/buildkit/cmd/buildctl/build"
 	"github.com/moby/buildkit/session"
@@ -144,11 +147,11 @@ func validateCompression(compression string, name string) map[string]string {
 	return attrs
 }
 
-func (c *Client) Build(ctx context.Context, opts BuildOptions) error {
+func (c *Client) Build(ctx context.Context, opts BuildOptions) (int64, error) {
 	// setup build directory
 	buildDir, err := os.MkdirTemp("", "hephaestus-build-")
 	if err != nil {
-		return fmt.Errorf("failed to create build dir: %w", err)
+		return 0, fmt.Errorf("failed to create build dir: %w", err)
 	}
 
 	defer func(path string) {
@@ -172,7 +175,7 @@ func (c *Client) Build(ctx context.Context, opts BuildOptions) error {
 		c.log.Info("Fetching remote context", "url", opts.Context)
 		extract, err := archive.FetchAndExtract(ctx, c.log, opts.Context, buildDir, opts.FetchAndExtractTimeout)
 		if err != nil {
-			return fmt.Errorf("cannot fetch remote context: %w", err)
+			return 0, fmt.Errorf("cannot fetch remote context: %w", err)
 		}
 
 		contentsDir = extract.ContentsDir
@@ -182,13 +185,13 @@ func (c *Client) Build(ctx context.Context, opts BuildOptions) error {
 	// verify manifest is present
 	dockerfile := filepath.Join(contentsDir, "Dockerfile")
 	if _, err := os.Stat(dockerfile); errors.Is(err, os.ErrNotExist) {
-		return fmt.Errorf("build requires a Dockerfile inside context dir: %w", err)
+		return 0, fmt.Errorf("build requires a Dockerfile inside context dir: %w", err)
 	}
 
 	if l := c.log.V(1); l.Enabled() {
 		bs, err := os.ReadFile(dockerfile)
 		if err != nil {
-			return fmt.Errorf("cannot read Dockerfile: %w", err)
+			return 0, fmt.Errorf("cannot read Dockerfile: %w", err)
 		}
 		l.Info("Dockerfile contents:\n" + string(bs))
 	}
@@ -199,7 +202,7 @@ func (c *Client) Build(ctx context.Context, opts BuildOptions) error {
 	for name, path := range opts.Secrets {
 		contents, err := os.ReadFile(path)
 		if err != nil {
-			return err
+			return 0, err
 		}
 
 		secrets[name] = contents
@@ -256,7 +259,7 @@ func (c *Client) Build(ctx context.Context, opts BuildOptions) error {
 
 		attrs, err := build.ParseOpt(args)
 		if err != nil {
-			return fmt.Errorf("cannot parse build args: %w", err)
+			return 0, fmt.Errorf("cannot parse build args: %w", err)
 		}
 
 		for k, v := range attrs {
@@ -333,10 +336,58 @@ func (c *Client) solveWith(ctx context.Context, modify func(buildDir string, sol
 		return err
 	}
 
-	return c.runSolve(ctx, solveOpt)
+	_, err = c.runSolve(ctx, solveOpt)
+	return err
+}
+
+func (c *Client) ResolveAuth(registryHostname string) (authn.Authenticator, error) {
+	cf, err := config.Load(c.dockerConfigDir)
+	if err != nil {
+		return nil, err
+	}
+	cfg, err := cf.GetAuthConfig(registryHostname)
+	if err != nil {
+		return nil, err
+	}
+
+	return authn.FromConfig(authn.AuthConfig{
+		Username:      cfg.Username,
+		Password:      cfg.Password,
+		Auth:          cfg.Auth,
+		IdentityToken: cfg.IdentityToken,
+		RegistryToken: cfg.RegistryToken,
+	}), nil
 }
 
-func (c *Client) runSolve(ctx context.Context, so bkclient.SolveOpt) error {
+func (c *Client) retrieveImageSize(ctx context.Context, imageName string, sizePtr *int64) error {
+	ref, err := name.ParseReference(imageName)
+	if err != nil {
+		return err
+	}
+	registryName := ref.Context().RegistryStr()
+	auth, err := c.ResolveAuth(registryName)
+	if err != nil {
+		return err
+	}
+	img, err := remote.Image(ref, remote.WithContext(ctx), remote.WithAuth(auth))
+	if err != nil {
+		return err
+	}
+	layers, err := img.Layers()
+	if err != nil {
+		return err
+	}
+	for _, layer := range layers {
+		compressedSize, err := layer.Size()
+		if err != nil {
+			return err
+		}
+		*sizePtr += compressedSize
+	}
+	return nil
+}
+
+func (c *Client) runSolve(ctx context.Context, so bkclient.SolveOpt) (int64, error) {
 	lw := &LogWriter{Logger: c.log}
 	ch := make(chan *bkclient.SolveStatus)
 	eg, ctx := errgroup.WithContext(ctx)
@@ -354,19 +405,29 @@ func (c *Client) runSolve(ctx context.Context, so bkclient.SolveOpt) error {
 		return err
 	})
 
+	var size int64
+
 	eg.Go(func() error {
-		if _, err := c.bk.Solve(ctx, nil, so, ch); err != nil {
+		res, err := c.bk.Solve(ctx, nil, so, ch)
+		if err != nil {
 			return err
 		}
 
 		c.log.Info("Solve complete")
+		imageName := res.ExporterResponse["image.name"]
+		err = c.retrieveImageSize(ctx, imageName, &size)
+		if err != nil {
+			c.log.Error(err, "Cannot retrieve image size from registry", "imageName", imageName)
+		}
+
 		return nil
 	})
 
 	if err := eg.Wait(); err != nil {
 		c.log.Info(fmt.Sprintf("Build failed: %s", err.Error()))
-		return fmt.Errorf("buildkit solve issue: %w", err)
+		return 0, fmt.Errorf("buildkit solve issue: %w", err)
 	}
 
-	return nil
+	c.log.Info(fmt.Sprintf("Final image size: %d", size))
+	return size, nil
 }

pkg/controller/imagebuild/component/builddispatcher.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"strconv"
 	"sync"
 	"time"
 
@@ -230,7 +231,8 @@ func (c *BuildDispatcherComponent) Reconcile(coreCtx *core.Context) (ctrl.Result
 
 	// best effort phase change regardless if the original context is "done"
 	coreCtx.Context = context.Background()
-	if err = bk.Build(buildCtx, buildOpts); err != nil {
+	imageSize, err := bk.Build(buildCtx, buildOpts)
+	if err != nil {
 		// if the underlying buildkit pod is terminated via resource delete, then buildCtx will be closed and there will
 		// be an error on it. otherwise, some external event (e.g. pod terminated) cancelled the build, so we should
 		// mark the build as failed.
@@ -253,6 +255,7 @@ func (c *BuildDispatcherComponent) Reconcile(coreCtx *core.Context) (ctrl.Result
 	obj.Status.BuildTime = time.Since(start).Truncate(time.Millisecond).String()
 	buildSeg.End()
 
+	obj.Status.CompressedImageSizeBytes = strconv.FormatInt(imageSize, 10)
 	c.phase.SetSucceeded(coreCtx, obj)
 	return ctrl.Result{}, nil
 }

pkg/controller/imagebuildmessage/component/amqpmessenger.go

@@ -28,7 +28,10 @@ import (
 	"github.com/dominodatalab/hephaestus/pkg/config"
 )
 
-const publishContentType = "application/json"
+const (
+	publishContentType                 = "application/json"
+	compressedImageSizeBytesAnnotation = "imagebuilder.dominodatalab.com/compressed-image-size-bytes"
+)
 
 type AMQPMessengerComponent struct {
 	cfg      config.Messaging
@@ -57,6 +60,7 @@ func (c *AMQPMessengerComponent) Initialize(_ *core.Context, bldr *ctrl.Builder)
 	return nil
 }
 
+//nolint:maintidx
 func (c *AMQPMessengerComponent) Reconcile(ctx *core.Context) (ctrl.Result, error) {
 	log := ctx.Log
 	obj := ctx.Object
@@ -201,6 +205,7 @@ func (c *AMQPMessengerComponent) Reconcile(ctx *core.Context) (ctrl.Result, erro
 				images = append(images, reference.TagNameOnly(named).String())
 			}
 			message.ImageURLs = images
+			message.Annotations[compressedImageSizeBytesAnnotation] = ib.Status.CompressedImageSizeBytes
 		case hephv1.PhaseFailed:
 			if ib.Status.Conditions == nil {
 				return ctrl.Result{Requeue: true}, nil