perf: optimize vex and progress hot paths

pkg/common/progress.go

@@ -34,6 +34,10 @@ func DisplayProgress(ctx context.Context, eg *errgroup.Group, buildChannel chan
 // ForwardProgressWithPrefix forwards progress from src to dst, prefixing vertex names with the given prefix.
 // This allows multiplexing multiple build progress streams into a single display.
 func ForwardProgressWithPrefix(ctx context.Context, src <-chan *client.SolveStatus, dst chan<- *client.SolveStatus, prefix string) {
+	namePrefix := "[" + prefix + "] "
+	digestPrefix := prefix + ":"
+	remappedDigests := make(map[digest.Digest]digest.Digest)
+
 	for {
 		select {
 		case status, ok := <-src:
@@ -44,35 +48,33 @@ func ForwardProgressWithPrefix(ctx context.Context, src <-chan *client.SolveStat
 				continue
 			}
 
-			// Clone and prefix vertex names
+			// Clone and prefix vertex names.
 			prefixed := &client.SolveStatus{
 				Vertexes: make([]*client.Vertex, len(status.Vertexes)),
-				Statuses: status.Statuses, // Statuses reference vertices by digest, no need to modify
-				Logs:     status.Logs,     // Logs reference vertices by digest, no need to modify
 				Warnings: status.Warnings,
 			}
 
 			for i, v := range status.Vertexes {
-				// Clone the vertex and prefix its name
+				// Clone the vertex and prefix its name.
 				cloned := *v
-				cloned.Name = "[" + prefix + "] " + v.Name
-				// Create a new digest that includes the prefix to avoid collisions
-				cloned.Digest = digest.FromString(prefix + ":" + v.Digest.String())
+				cloned.Name = namePrefix + v.Name
+				// Create a new digest that includes the prefix to avoid collisions.
+				cloned.Digest = prefixedDigest(remappedDigests, digestPrefix, v.Digest)
 				prefixed.Vertexes[i] = &cloned
 			}
 
-			// Also update status and log references to use new digests
+			// Also update status and log references to use new digests.
 			prefixed.Statuses = make([]*client.VertexStatus, len(status.Statuses))
 			for i, s := range status.Statuses {
 				cloned := *s
-				cloned.Vertex = digest.FromString(prefix + ":" + s.Vertex.String())
+				cloned.Vertex = prefixedDigest(remappedDigests, digestPrefix, s.Vertex)
 				prefixed.Statuses[i] = &cloned
 			}
 
 			prefixed.Logs = make([]*client.VertexLog, len(status.Logs))
 			for i, l := range status.Logs {
 				cloned := *l
-				cloned.Vertex = digest.FromString(prefix + ":" + l.Vertex.String())
+				cloned.Vertex = prefixedDigest(remappedDigests, digestPrefix, l.Vertex)
 				prefixed.Logs[i] = &cloned
 			}
 
@@ -87,3 +89,13 @@ func ForwardProgressWithPrefix(ctx context.Context, src <-chan *client.SolveStat
 		}
 	}
 }
+
+func prefixedDigest(remappedDigests map[digest.Digest]digest.Digest, digestPrefix string, original digest.Digest) digest.Digest {
+	if remapped, ok := remappedDigests[original]; ok {
+		return remapped
+	}
+
+	remapped := digest.FromString(digestPrefix + original.String())
+	remappedDigests[original] = remapped
+	return remapped
+}

pkg/common/progress_benchmark_test.go

@@ -0,0 +1,88 @@
+package common
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/moby/buildkit/client"
+	"github.com/opencontainers/go-digest"
+)
+
+func BenchmarkForwardProgressWithPrefixRepeatedDigests(b *testing.B) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	src := make(chan *client.SolveStatus)
+	dst := make(chan *client.SolveStatus)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		ForwardProgressWithPrefix(ctx, src, dst, "linux/amd64")
+	}()
+
+	status := repeatedDigestSolveStatus()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		src <- status
+		prefixed := <-dst
+		if len(prefixed.Statuses) != len(status.Statuses) || len(prefixed.Logs) != len(status.Logs) {
+			b.Fatalf("unexpected forwarded status sizes: statuses=%d logs=%d", len(prefixed.Statuses), len(prefixed.Logs))
+		}
+	}
+	b.StopTimer()
+
+	close(src)
+	<-done
+}
+
+func repeatedDigestSolveStatus() *client.SolveStatus {
+	now := time.Unix(1_700_000_000, 0)
+	digests := []digest.Digest{
+		digest.FromString("vertex-0"),
+		digest.FromString("vertex-1"),
+		digest.FromString("vertex-2"),
+		digest.FromString("vertex-3"),
+	}
+
+	vertexes := make([]*client.Vertex, len(digests))
+	for i, d := range digests {
+		vertexes[i] = &client.Vertex{
+			Digest:  d,
+			Name:    forwardProgressTestVertexName,
+			Started: &now,
+		}
+	}
+
+	statuses := make([]*client.VertexStatus, 64)
+	logs := make([]*client.VertexLog, 64)
+	for i := range statuses {
+		d := digests[i%len(digests)]
+		statuses[i] = &client.VertexStatus{
+			ID:        "status",
+			Vertex:    d,
+			Name:      "status",
+			Total:     100,
+			Current:   int64(i),
+			Timestamp: now,
+			Started:   &now,
+		}
+		logs[i] = &client.VertexLog{
+			Vertex:    d,
+			Stream:    1,
+			Data:      []byte("repeated log data"),
+			Timestamp: now,
+		}
+	}
+
+	return &client.SolveStatus{
+		Vertexes: vertexes,
+		Statuses: statuses,
+		Logs:     logs,
+		Warnings: []*client.VertexWarning{
+			{Vertex: digests[0], Short: []byte("warning")},
+		},
+	}
+}

pkg/common/progress_test.go

@@ -13,6 +13,8 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+const forwardProgressTestVertexName = "vertex"
+
 func TestDisplayProgress(t *testing.T) {
 	ctx := context.Background()
 	eg, egCtx := errgroup.WithContext(ctx)
@@ -172,3 +174,119 @@ func TestDisplayProgressWithDebugMode(t *testing.T) {
 	err := eg.Wait()
 	assert.NoError(t, err)
 }
+
+func TestForwardProgressWithPrefixRemapsRepeatedDigestsWithoutMutatingInput(t *testing.T) {
+	ctx := context.Background()
+	src := make(chan *client.SolveStatus)
+	dst := make(chan *client.SolveStatus)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		ForwardProgressWithPrefix(ctx, src, dst, "linux/arm64")
+	}()
+
+	now := time.Unix(1_700_000_001, 0)
+	vertexA := digest.FromString("vertex-a")
+	vertexB := digest.FromString("vertex-b")
+	status := &client.SolveStatus{
+		Vertexes: []*client.Vertex{
+			{Digest: vertexA, Name: "install packages", Started: &now},
+			{Digest: vertexB, Name: "copy files", Completed: &now},
+		},
+		Statuses: []*client.VertexStatus{
+			{ID: "status-a-1", Vertex: vertexA, Name: "download", Timestamp: now},
+			{ID: "status-a-2", Vertex: vertexA, Name: "extract", Timestamp: now},
+			{ID: "status-b", Vertex: vertexB, Name: "copy", Timestamp: now},
+		},
+		Logs: []*client.VertexLog{
+			{Vertex: vertexA, Stream: 1, Data: []byte("log a"), Timestamp: now},
+			{Vertex: vertexA, Stream: 2, Data: []byte("log a again"), Timestamp: now},
+			{Vertex: vertexB, Stream: 1, Data: []byte("log b"), Timestamp: now},
+		},
+		Warnings: []*client.VertexWarning{
+			{Vertex: vertexA, Short: []byte("warning")},
+		},
+	}
+
+	src <- status
+	prefixed := <-dst
+	close(src)
+	<-done
+
+	remappedA := digest.FromString("linux/arm64:" + vertexA.String())
+	remappedB := digest.FromString("linux/arm64:" + vertexB.String())
+
+	assert.Equal(t, remappedA, prefixed.Vertexes[0].Digest)
+	assert.Equal(t, "[linux/arm64] install packages", prefixed.Vertexes[0].Name)
+	assert.Equal(t, remappedB, prefixed.Vertexes[1].Digest)
+	assert.Equal(t, "[linux/arm64] copy files", prefixed.Vertexes[1].Name)
+
+	assert.Equal(t, remappedA, prefixed.Statuses[0].Vertex)
+	assert.Equal(t, remappedA, prefixed.Statuses[1].Vertex)
+	assert.Equal(t, remappedB, prefixed.Statuses[2].Vertex)
+	assert.Equal(t, remappedA, prefixed.Logs[0].Vertex)
+	assert.Equal(t, remappedA, prefixed.Logs[1].Vertex)
+	assert.Equal(t, remappedB, prefixed.Logs[2].Vertex)
+
+	// Warnings are forwarded unchanged, matching the previous behavior.
+	assert.Equal(t, vertexA, prefixed.Warnings[0].Vertex)
+
+	// The input status and nested vertex/status/log fields are not mutated.
+	assert.Equal(t, vertexA, status.Vertexes[0].Digest)
+	assert.Equal(t, "install packages", status.Vertexes[0].Name)
+	assert.Equal(t, vertexB, status.Vertexes[1].Digest)
+	assert.Equal(t, "copy files", status.Vertexes[1].Name)
+	assert.Equal(t, vertexA, status.Statuses[0].Vertex)
+	assert.Equal(t, vertexA, status.Statuses[1].Vertex)
+	assert.Equal(t, vertexB, status.Statuses[2].Vertex)
+	assert.Equal(t, vertexA, status.Logs[0].Vertex)
+	assert.Equal(t, vertexA, status.Logs[1].Vertex)
+	assert.Equal(t, vertexB, status.Logs[2].Vertex)
+}
+
+func TestForwardProgressWithPrefixSkipsNilStatus(t *testing.T) {
+	ctx := context.Background()
+	src := make(chan *client.SolveStatus)
+	dst := make(chan *client.SolveStatus)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		ForwardProgressWithPrefix(ctx, src, dst, "target")
+	}()
+
+	vertex := digest.FromString("vertex")
+	status := &client.SolveStatus{
+		Vertexes: []*client.Vertex{{Digest: vertex, Name: forwardProgressTestVertexName}},
+	}
+
+	src <- nil
+	src <- status
+	prefixed := <-dst
+	close(src)
+	<-done
+
+	assert.Len(t, prefixed.Vertexes, 1)
+	assert.Equal(t, digest.FromString("target:"+vertex.String()), prefixed.Vertexes[0].Digest)
+}
+
+func TestForwardProgressWithPrefixReturnsOnCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	src := make(chan *client.SolveStatus)
+	dst := make(chan *client.SolveStatus)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		ForwardProgressWithPrefix(ctx, src, dst, "target")
+	}()
+
+	src <- &client.SolveStatus{
+		Vertexes: []*client.Vertex{{Digest: digest.FromString(forwardProgressTestVertexName), Name: forwardProgressTestVertexName}},
+	}
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("ForwardProgressWithPrefix did not return after context cancellation")
+	}
+}

pkg/generate/generate.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"sort"
 	"strings"
 	"time"
 
@@ -420,6 +421,7 @@ LABEL sh.copa.image.patched="%s"
 	// Extract and rewrite the patch layer tar with proper paths
 	tr := tar.NewReader(bytes.NewReader(patchLayer))
 	hardlinks := make(map[string]*tar.Header)
+	hardlinkTargets := make(map[string]struct{})
 
 	for {
 		hdr, err := tr.Next()
@@ -434,6 +436,7 @@ LABEL sh.copa.image.patched="%s"
 		if hdr.Typeflag == tar.TypeLink {
 			// Store hardlink info for later
 			hardlinks[hdr.Name] = hdr
+			hardlinkTargets[hdr.Linkname] = struct{}{}
 			continue
 		}
 
@@ -461,7 +464,7 @@ LABEL sh.copa.image.patched="%s"
 	if len(hardlinks) > 0 {
 		// Re-read the tar to find hardlink targets
 		tr = tar.NewReader(bytes.NewReader(patchLayer))
-		fileContents := make(map[string][]byte)
+		fileContents := make(map[string][]byte, len(hardlinkTargets))
 
 		for {
 			hdr, err := tr.Next()
@@ -472,8 +475,10 @@ LABEL sh.copa.image.patched="%s"
 				return errors.Wrap(err, "failed to read patch layer tar for hardlinks")
 			}
 
-			// Store file contents for hardlink targets
-			if hdr.Typeflag == tar.TypeReg && hdr.Size > 0 {
+			// Store only file contents that are actual hardlink targets. The first
+			// pass has already copied every regular file into the output context;
+			// this second pass only needs bytes for materializing hardlinks as files.
+			if _, needed := hardlinkTargets[hdr.Name]; needed && hdr.Typeflag == tar.TypeReg && hdr.Size > 0 {
 				// Limit file size to 1GB to prevent decompression bombs
 				if hdr.Size > maxFileSize {
 					return errors.Errorf("file %s exceeds maximum allowed size of 1GB", hdr.Name)
@@ -486,8 +491,15 @@ LABEL sh.copa.image.patched="%s"
 			}
 		}
 
-		// Write hardlinks as regular files
-		for name, hdr := range hardlinks {
+		// Write hardlinks as regular files. Sort names so identical inputs produce
+		// deterministic build contexts instead of inheriting map iteration order.
+		hardlinkNames := make([]string, 0, len(hardlinks))
+		for name := range hardlinks {
+			hardlinkNames = append(hardlinkNames, name)
+		}
+		sort.Strings(hardlinkNames)
+		for _, name := range hardlinkNames {
+			hdr := hardlinks[name]
 			if content, ok := fileContents[hdr.Linkname]; ok {
 				// Convert to regular file
 				newHdr := &tar.Header{