Repository: chainguard-dev/apko
Pull request #1735 (2025-06-30 21:44:04)

Lazily compress layers, maintain diffID -> digest cache
Carrying on from #1667 

This adds the compression cache, so that if `Digest()` is called and there's already a cached copy of the diffID -> digest available, we'll use it, instead of computing it again. Same with `Size()`.

If the pre-computed compressed digest is already present in the registry for example, the resulting `HEAD` request's `200` will tell us we don't need to compress anything at all, and we can skip a lot of work.

If the digest isn't in the registry already, we'll compress and compute it like we did before, but that should happen a lot less often with this cache.


Patch:
diff --git a/pkg/build/build.go b/pkg/build/build.go
index 24a0554..0006902 100644
--- a/pkg/build/build.go
+++ b/pkg/build/build.go
@@ -15,8 +15,9 @@
 package build
 
 import (
-	"compress/gzip"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -25,6 +26,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -44,6 +46,10 @@ import (
 	"chainguard.dev/apko/pkg/s6"
 )
 
+// compressionCache stores descriptor information for already-compressed layers,
+// keyed by diffID. This avoids recompressing identical layers.
+var compressionCache sync.Map // map[string]*v1.Descriptor
+
 // Context contains all of the information necessary to build an
 // OCI image. Includes the configuration for the build,
 // the path to the config file, the executor for root jails and
@@ -349,9 +355,67 @@ func (f *notAFile) Close() error {
 // layer implements v1.Layer from go-containerregistry to avoid re-computing
 // digests and diffids.
 type layer struct {
-	filename string
-	diffid   *v1.Hash
-	desc     *v1.Descriptor
+	uncompressed string
+	compressed   string
+	diffid       *v1.Hash
+	desc         *v1.Descriptor
+}
+
+func (l *layer) compress() error {
+	if l.compressed != "" {
+		return nil
+	}
+
+	in, err := l.Uncompressed()
+	if err != nil {
+		return err
+	}
+	defer in.Close()
+
+	out, err := os.Create(l.uncompressed + ".gz")
+	if err != nil {
+		return err
+	}
+
+	buf := pooledBufioWriter(out)
+	defer bufioPool.Put(buf)
+
+	digest := sha256.New()
+	gzw := pooledGzipWriter(io.MultiWriter(digest, buf))
+	defer pgzipPool.Put(gzw)
+
+	if _, err := io.Copy(gzw, in); err != nil {
+		return err
+	}
+
+	if err := gzw.Close(); err != nil {
+		return fmt.Errorf("closing gzip writer: %w", err)
+	}
+
+	if err := buf.Flush(); err != nil {
+		return fmt.Errorf("flushing %s: %w", out.Name(), err)
+	}
+
+	stat, err := out.Stat()
+	if err != nil {
+		return fmt.Errorf("statting %s: %w", out.Name(), err)
+	}
+
+	h := v1.Hash{
+		Algorithm: "sha256",
+		Hex:       hex.EncodeToString(digest.Sum(make([]byte, 0, digest.Size()))),
+	}
+
+	l.desc.Digest = h
+	l.desc.Size = stat.Size()
+
+	// Store in cache for future use
+	descCopy := *l.desc
+	compressionCache.Store(l.diffid.String(), &descCopy)
+
+	l.compressed = l.uncompressed + ".gz"
+
+	return out.Close()
 }
 
 func (l *layer) DiffID() (v1.Hash, error) {
@@ -359,11 +423,25 @@ func (l *layer) DiffID() (v1.Hash, error) {
 }
 
 func (l *layer) Digest() (v1.Hash, error) {
+	// Check if we've already compressed a layer with this diffID
+	if cached, ok := compressionCache.Load(l.diffid.String()); ok {
+		cachedDesc := cached.(*v1.Descriptor)
+		l.desc.Digest = cachedDesc.Digest
+		l.desc.Size = cachedDesc.Size
+		return l.desc.Digest, nil
+	}
+
+	if err := l.compress(); err != nil {
+		return v1.Hash{}, err
+	}
 	return l.desc.Digest, nil
 }
 
 func (l *layer) Compressed() (io.ReadCloser, error) {
-	f, err := os.Open(l.filename)
+	if err := l.compress(); err != nil {
+		return nil, err
+	}
+	f, err := os.Open(l.compressed)
 	if err != nil {
 		return nil, err
 	}
@@ -374,20 +452,21 @@ func (l *layer) Compressed() (io.ReadCloser, error) {
 }
 
 func (l *layer) Uncompressed() (io.ReadCloser, error) {
-	rc, err := l.Compressed()
-	if err != nil {
-		return nil, err
-	}
-
-	// In practice, this won't be called, but this should work anyway.
-	zr, err := gzip.NewReader(rc)
-	if err != nil {
-		return nil, err
-	}
-	return zr, nil
+	return os.Open(l.uncompressed)
 }
 
 func (l *layer) Size() (int64, error) {
+	// Check if we've already compressed a layer with this diffID
+	if cached, ok := compressionCache.Load(l.diffid.String()); ok {
+		cachedDesc := cached.(*v1.Descriptor)
+		l.desc.Digest = cachedDesc.Digest
+		l.desc.Size = cachedDesc.Size
+		return l.desc.Size, nil
+	}
+
+	if err := l.compress(); err != nil {
+		return 0, err
+	}
 	return l.desc.Size, nil
 }
 
diff --git a/pkg/build/build_implementation.go b/pkg/build/build_implementation.go
index 01dc44d..f5d53aa 100644
--- a/pkg/build/build_implementation.go
+++ b/pkg/build/build_implementation.go
@@ -92,15 +92,11 @@ type layerWriter struct {
 // 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()
+	diffid := sha256.New()
 
 	buf := pooledBufioWriter(out)
 
-	gzw := pooledGzipWriter(io.MultiWriter(digest, buf))
-
-	diffid := sha256.New()
-
-	w := tar.NewWriter(io.MultiWriter(diffid, gzw))
+	w := tar.NewWriter(io.MultiWriter(diffid, buf))
 
 	// Just capturing everything in a closure here is more straightforward
 	// to read (as a translation from what used to implement this) than
@@ -108,36 +104,19 @@ func newLayerWriter(out *os.File) *layerWriter {
 	return &layerWriter{
 		w: w,
 		finalize: func() (*layer, error) {
-			defer pgzipPool.Put(gzw)
 			defer bufioPool.Put(buf)
 
 			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)
-			}
-
 			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(),
+				uncompressed: out.Name(),
 				desc: &v1.Descriptor{
-					Digest:    h,
-					Size:      stat.Size(),
 					MediaType: v1types.OCILayer,
 				},
 				diffid: &v1.Hash{


Test patch:
diff --git a/pkg/build/layer_cache_test.go b/pkg/build/layer_cache_test.go
new file mode 100644
index 0000000..ffb25d5
--- /dev/null
+++ b/pkg/build/layer_cache_test.go
@@ -0,0 +1,232 @@
+// Copyright 2025 Chainguard, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package build
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"io"
+	"os"
+	"path/filepath"
+	"strconv"
+	"testing"
+
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	v1types "github.com/google/go-containerregistry/pkg/v1/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLayerCompressionCache(t *testing.T) {
+	// Create a temporary directory for test files
+	tmpDir := t.TempDir()
+
+	// Create test content
+	testContent := []byte("test layer content for compression cache")
+
+	// Calculate the diffID (hash of uncompressed content)
+	h := sha256.Sum256(testContent)
+	diffID := v1.Hash{
+		Algorithm: "sha256",
+		Hex:       hex.EncodeToString(h[:]),
+	}
+
+	// Create first layer with test content
+	file1 := filepath.Join(tmpDir, "layer1.tar")
+	err := os.WriteFile(file1, testContent, 0644)
+	require.NoError(t, err)
+
+	layer1 := &layer{
+		uncompressed: file1,
+		diffid:       &diffID,
+		desc: &v1.Descriptor{
+			MediaType: v1types.OCILayer,
+		},
+	}
+
+	// Get digest for first layer - this should compress the layer
+	digest1, err := layer1.Digest()
+	require.NoError(t, err)
+	require.NotEmpty(t, digest1.String())
+
+	// Verify the compressed file was created
+	require.FileExists(t, file1+".gz")
+
+	// Get size for first layer - should use cached values
+	size1, err := layer1.Size()
+	require.NoError(t, err)
+	require.Greater(t, size1, int64(0))
+
+	// Create second layer with identical content (same diffID)
+	file2 := filepath.Join(tmpDir, "layer2.tar")
+	err = os.WriteFile(file2, testContent, 0644)
+	require.NoError(t, err)
+
+	layer2 := &layer{
+		uncompressed: file2,
+		diffid:       &diffID, // Same diffID as layer1
+		desc: &v1.Descriptor{
+			MediaType: v1types.OCILayer,
+		},
+	}
+
+	// Get digest for second layer - should use cached descriptor
+	digest2, err := layer2.Digest()
+	require.NoError(t, err)
+	require.Equal(t, digest1, digest2, "Cached digest should match")
+
+	// Get size for second layer - should use cached descriptor
+	size2, err := layer2.Size()
+	require.NoError(t, err)
+	require.Equal(t, size1, size2, "Cached size should match")
+
+	// Verify that layer2's descriptor was populated from cache
+	require.Equal(t, layer1.desc.Digest, layer2.desc.Digest)
+	require.Equal(t, layer1.desc.Size, layer2.desc.Size)
+
+	// The compressed file for layer2 should NOT exist yet since we only used cached values
+	require.NoFileExists(t, file2+".gz")
+
+	// Now actually compress layer2 by calling Compressed()
+	rc, err := layer2.Compressed()
+	require.NoError(t, err)
+	rc.Close()
+
+	// Now the compressed file should exist
+	require.FileExists(t, file2+".gz")
+
+	// Create third layer with different content
+	differentContent := []byte("different content that will have different diffID")
+	h3 := sha256.Sum256(differentContent)
+	diffID3 := v1.Hash{
+		Algorithm: "sha256",
+		Hex:       hex.EncodeToString(h3[:]),
+	}
+
+	file3 := filepath.Join(tmpDir, "layer3.tar")
+	err = os.WriteFile(file3, differentContent, 0644)
+	require.NoError(t, err)
+
+	layer3 := &layer{
+		uncompressed: file3,
+		diffid:       &diffID3,
+		desc: &v1.Descriptor{
+			MediaType: v1types.OCILayer,
+		},
+	}
+
+	// Get digest for third layer - should NOT use cache (different diffID)
+	digest3, err := layer3.Digest()
+	require.NoError(t, err)
+	require.NotEqual(t, digest1, digest3, "Different content should have different digest")
+
+	// Verify the compressed file was created for layer3
+	require.FileExists(t, file3+".gz")
+}
+
+func TestLayerCompressionCacheConsistency(t *testing.T) {
+	// This test verifies that the cache returns consistent results
+	tmpDir := t.TempDir()
+
+	// Create test content
+	testContent := []byte("consistency test content")
+	h := sha256.Sum256(testContent)
+	diffID := v1.Hash{
+		Algorithm: "sha256",
+		Hex:       hex.EncodeToString(h[:]),
+	}
+
+	// Create and compress first layer
+	file1 := filepath.Join(tmpDir, "consistent1.tar")
+	err := os.WriteFile(file1, testContent, 0644)
+	require.NoError(t, err)
+
+	layer1 := &layer{
+		uncompressed: file1,
+		diffid:       &diffID,
+		desc: &v1.Descriptor{
+			MediaType: v1types.OCILayer,
+		},
+	}
+
+	// Compress the layer
+	err = layer1.compress()
+	require.NoError(t, err)
+
+	originalDigest := layer1.desc.Digest
+	originalSize := layer1.desc.Size
+
+	// Create multiple layers with same diffID and verify they all get same cached values
+	for i := 2; i <= 5; i++ {
+		file := filepath.Join(tmpDir, "consistent"+strconv.Itoa(i)+".tar")
+		err := os.WriteFile(file, testContent, 0644)
+		require.NoError(t, err)
+
+		layer := &layer{
+			uncompressed: file,
+			diffid:       &diffID,
+			desc: &v1.Descriptor{
+				MediaType: v1types.OCILayer,
+			},
+		}
+
+		// Check digest uses cache
+		digest, err := layer.Digest()
+		require.NoError(t, err)
+		require.Equal(t, originalDigest, digest, "Cached digest should be consistent across all layers")
+
+		// Check size uses cache
+		size, err := layer.Size()
+		require.NoError(t, err)
+		require.Equal(t, originalSize, size, "Cached size should be consistent across all layers")
+	}
+}
+
+func TestLayerUncompressedAccess(t *testing.T) {
+	// Test that we can still access uncompressed content
+	tmpDir := t.TempDir()
+
+	testContent := []byte("uncompressed access test")
+	h := sha256.Sum256(testContent)
+	diffID := v1.Hash{
+		Algorithm: "sha256",
+		Hex:       hex.EncodeToString(h[:]),
+	}
+
+	file := filepath.Join(tmpDir, "uncompressed.tar")
+	err := os.WriteFile(file, testContent, 0644)
+	require.NoError(t, err)
+
+	layer := &layer{
+		uncompressed: file,
+		diffid:       &diffID,
+		desc: &v1.Descriptor{
+			MediaType: v1types.OCILayer,
+		},
+	}
+
+	// Access uncompressed content
+	rc, err := layer.Uncompressed()
+	require.NoError(t, err)
+	defer rc.Close()
+
+	content, err := io.ReadAll(rc)
+	require.NoError(t, err)
+	require.Equal(t, testContent, content)
+
+	// Verify DiffID
+	gotDiffID, err := layer.DiffID()
+	require.NoError(t, err)
+	require.Equal(t, diffID, gotDiffID)
+}