Repository: chainguard-dev/apko
Pull request #1616 (2025-04-10 23:01:39)

Use a lazy layerWriter for writing tar files
Instead of writing out the entire layer in one shot, layerWriter exposes a tar.Writer so that it's possible to construct a layerWriter, write to it a bunch of times, then finalize the layerWriter to produce a layer.

This PR should be a no-op, but we'll take advantage of this for multiple layers when we walk the filesystem and have multiple possible layerWriters to choose from.


Patch:
diff --git a/docs/build-process.md b/docs/build-process.md
index 045a6d5..ce6b0ce 100644
--- a/docs/build-process.md
+++ b/docs/build-process.md
@@ -32,7 +32,6 @@ As described above, after everything is setup, the actual build occurs inside th
 The build is in [`build.Context.BuildLayer()`](../pkg/build/build.go#L80-109), which consists of:
 
 1. `Context.BuildImage()`: building the image
-1. `Context.BuildTarball()`: build the tarball for the layer
 1. `Context.GenerateSBOM()` optionally generate the SBoM
 
 The actual building of the image via `BuildImage()` just wraps [`buildImage()`](../pkg/build/build_implementation.go#L195-247).
diff --git a/pkg/build/build.go b/pkg/build/build.go
index a7972a4..590bb61 100644
--- a/pkg/build/build.go
+++ b/pkg/build/build.go
@@ -17,11 +17,11 @@ package build
 import (
 	"compress/gzip"
 	"context"
-	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
 	"os"
+	"path/filepath"
 	"runtime"
 	"strconv"
 	"strings"
@@ -161,31 +161,34 @@ func (bc *Context) ImageLayoutToLayer(ctx context.Context) (string, v1.Layer, er
 		return "", nil, err
 	}
 
-	layerTarGZ, diffid, digest, size, err := bc.BuildTarball(ctx)
-	// build layer tarball
+	var (
+		outfile *os.File
+		err     error
+	)
+
+	if bc.o.TarballPath != "" {
+		outfile, err = os.Create(bc.o.TarballPath)
+	} else {
+		outfile, err = os.Create(filepath.Join(bc.o.TempDir(), bc.o.TarballFileName()))
+	}
 	if err != nil {
-		return "", nil, err
+		return "", nil, fmt.Errorf("creating tarball file: %w", err)
 	}
+	bc.o.TarballPath = outfile.Name()
+	defer outfile.Close()
+
+	lw := newLayerWriter(outfile)
 
-	h := v1.Hash{
-		Algorithm: "sha256",
-		Hex:       hex.EncodeToString(digest.Sum(make([]byte, 0, digest.Size()))),
+	if err := writeTar(ctx, lw.w, bc.fs); err != nil {
+		return "", nil, fmt.Errorf("generating tarball: %w", err)
 	}
 
-	l := &layer{
-		filename: layerTarGZ,
-		desc: &v1.Descriptor{
-			Digest:    h,
-			Size:      size,
-			MediaType: v1types.OCILayer,
-		},
-		diffid: &v1.Hash{
-			Algorithm: "sha256",
-			Hex:       hex.EncodeToString(diffid.Sum(make([]byte, 0, diffid.Size()))),
-		},
+	l, err := lw.finalize()
+	if err != nil {
+		return "", nil, fmt.Errorf("finalizing layer: %w", err)
 	}
 
-	return layerTarGZ, l, nil
+	return outfile.Name(), l, nil
 }
 
 func (bc *Context) checkPaths(ctx context.Context) error {
diff --git a/pkg/build/build_implementation.go b/pkg/build/build_implementation.go
index 80c85fc..7088fd8 100644
--- a/pkg/build/build_implementation.go
+++ b/pkg/build/build_implementation.go
@@ -15,21 +15,23 @@
 package build
 
 import (
+	"archive/tar"
 	"bufio"
 	"context"
 	"crypto/sha256"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
-	"hash"
 	"io"
 	"os"
 	"path/filepath"
 	"runtime"
 	"sync"
 
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	v1types "github.com/google/go-containerregistry/pkg/v1/types"
 	gzip "github.com/klauspost/pgzip"
 	"github.com/sigstore/cosign/v2/pkg/oci"
-	"go.opentelemetry.io/otel"
 
 	"github.com/chainguard-dev/clog"
 
@@ -77,56 +79,75 @@ func pooledBufioWriter(w io.Writer) *bufio.Writer {
 	return bw
 }
 
-// BuildTarball takes the fully populated working directory and saves it to
-// an OCI image layer tar.gz file.
-func (bc *Context) BuildTarball(ctx context.Context) (string, hash.Hash, hash.Hash, int64, error) {
-	log := clog.FromContext(ctx)
-
-	ctx, span := otel.Tracer("apko").Start(ctx, "BuildTarball")
-	defer span.End()
-
-	var outfile *os.File
-	var err error
-
-	if bc.o.TarballPath != "" {
-		outfile, err = os.Create(bc.o.TarballPath)
-	} else {
-		outfile, err = os.Create(filepath.Join(bc.o.TempDir(), bc.o.TarballFileName()))
-	}
-	if err != nil {
-		return "", nil, nil, 0, fmt.Errorf("opening the build context tarball path failed: %w", err)
-	}
-	bc.o.TarballPath = outfile.Name()
-	defer outfile.Close()
+// layerWriter allows lazily writing files to a tarball instead
+// of doing everything in one pass, this is necessary for multi-layer
+// images where we are writing to multiple layers at the same time.
+type layerWriter struct {
+	w        *tar.Writer
+	finalize func() (*layer, error)
+}
 
+// newLayerWriter wraps a file with a gzipping tar writer that computes
+// everything we need to know to implement a v1.Layer, which it will
+// produce when finalize() is called.
+func newLayerWriter(out *os.File) *layerWriter {
 	digest := sha256.New()
 
-	buf := pooledBufioWriter(outfile)
-	defer bufioPool.Put(buf)
+	buf := pooledBufioWriter(out)
 
 	gzw := pooledGzipWriter(io.MultiWriter(digest, buf))
-	defer pgzipPool.Put(gzw)
 
 	diffid := sha256.New()
 
-	if err := writeTar(ctx, io.MultiWriter(diffid, gzw), bc.fs); err != nil {
-		return "", nil, nil, 0, fmt.Errorf("failed to generate tarball for image: %w", err)
-	}
-	if err := gzw.Close(); err != nil {
-		return "", nil, nil, 0, fmt.Errorf("closing gzip writer: %w", err)
-	}
+	w := tar.NewWriter(io.MultiWriter(diffid, gzw))
 
-	if err := buf.Flush(); err != nil {
-		return "", nil, nil, 0, fmt.Errorf("flushing %s: %w", outfile.Name(), err)
-	}
+	// Just capturing everything in a closure here is more straightforward
+	// to read (as a translation from what used to implement this) than
+	// adding a bunch of fields to layerWriter.
+	return &layerWriter{
+		w: w,
+		finalize: func() (*layer, error) {
+			defer pgzipPool.Put(gzw)
+			defer bufioPool.Put(buf)
 
-	stat, err := outfile.Stat()
-	if err != nil {
-		return "", nil, nil, 0, fmt.Errorf("stat(%q): %w", outfile.Name(), err)
-	}
+			if err := w.Close(); err != nil {
+				return nil, fmt.Errorf("closing tar writer: %w", err)
+			}
+
+			if err := gzw.Close(); err != nil {
+				return nil, fmt.Errorf("closing gzip writer: %w", err)
+			}
 
-	log.Infof("built image layer tarball as %s", outfile.Name())
-	return outfile.Name(), diffid, digest, stat.Size(), nil
+			if err := buf.Flush(); err != nil {
+				return nil, fmt.Errorf("flushing %s: %w", out.Name(), err)
+			}
+
+			stat, err := out.Stat()
+			if err != nil {
+				return nil, fmt.Errorf("statting %s: %w", out.Name(), err)
+			}
+
+			h := v1.Hash{
+				Algorithm: "sha256",
+				Hex:       hex.EncodeToString(digest.Sum(make([]byte, 0, digest.Size()))),
+			}
+
+			l := &layer{
+				filename: out.Name(),
+				desc: &v1.Descriptor{
+					Digest:    h,
+					Size:      stat.Size(),
+					MediaType: v1types.OCILayer,
+				},
+				diffid: &v1.Hash{
+					Algorithm: "sha256",
+					Hex:       hex.EncodeToString(diffid.Sum(make([]byte, 0, diffid.Size()))),
+				},
+			}
+
+			return l, nil
+		},
+	}
 }
 
 func (bc *Context) buildImage(ctx context.Context) error {
diff --git a/pkg/build/tarball.go b/pkg/build/tarball.go
index d1745cf..50d3a06 100644
--- a/pkg/build/tarball.go
+++ b/pkg/build/tarball.go
@@ -34,11 +34,10 @@ const xattrTarPAXRecordsPrefix = "SCHILY.xattr."
 
 // writeTar writes a tarball to the provided io.Writer from the provided fs.FS.
 // The etc/passwd and etc/group file provide username and group name mappings for the tar.
-func writeTar(ctx context.Context, dst io.Writer, fsys apkfs.FullFS) error { //nolint:gocyclo
+func writeTar(ctx context.Context, tw *tar.Writer, fsys apkfs.FullFS) error { //nolint:gocyclo
 	ctx, span := otel.Tracer("go-apk").Start(ctx, "writeTar")
 	defer span.End()
 
-	tw := tar.NewWriter(dst)
 	buf := make([]byte, 1<<20)
 
 	for f, err := range walkFS(ctx, fsys) {


Test patch:
diff --git a/pkg/build/tarball_test.go b/pkg/build/tarball_test.go
index ac97901..cad9958 100644
--- a/pkg/build/tarball_test.go
+++ b/pkg/build/tarball_test.go
@@ -28,8 +28,11 @@ func TestWriteTar(t *testing.T) {
 	require.NoError(t, err, "error setting xattr on %s", dir)
 	err = m.SetXattr(file, "user.file", []byte("bar"))
 	require.NoError(t, err, "error setting xattr on %s", file)
-	err = writeTar(context.Background(), &buf, m)
+	tw := tar.NewWriter(&buf)
+	err = writeTar(context.Background(), tw, m)
 	require.NoError(t, err, "error writing tar")
+	err = tw.Close()
+	require.NoError(t, err, "error closing tar writer")
 
 	// now should be able to read the tar and check the xattrs
 	tr := tar.NewReader(&buf)