llb: Cache DefinitionOp inputs to support memoization during Marshal

llb.State.Marshal uses a vertexCache map (keyed by the llb.Vertex interface
value) to memoize which vertexes it has already visited while marshalling.

However, before this change, llb.DefinitionOp was constructing new pointers for
each value in the return slice of Inputs(), which meant that each input had a
different key in the vertexCache (due to being different pointer values). This
meant no memoization actually occured. I noticed this while using
llb.DefinitionOp with a fairly large graph when my program crashed after using
>16GB of RSS during a call to llb.State.Marshal.

The fix here avoids changing vertexCache as that would impact many other
vertex implementations. Instead, just DefinitionOp.Inputs() is updated to cache
the pointers it creates and pass that cache to its descendents. This fix
resulted in the program that was previously crashing out-of-memory to run
without any perceivable increase in RSS.

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

client/llb/definition.go

@@ -16,14 +16,15 @@ import (
 // LLB state can be reconstructed from the definition.
 type DefinitionOp struct {
 	MarshalCache
-	mu        sync.Mutex
-	ops       map[digest.Digest]*pb.Op
-	defs      map[digest.Digest][]byte
-	metas     map[digest.Digest]pb.OpMetadata
-	sources   map[digest.Digest][]*SourceLocation
-	platforms map[digest.Digest]*specs.Platform
-	dgst      digest.Digest
-	index     pb.OutputIndex
+	mu         sync.Mutex
+	ops        map[digest.Digest]*pb.Op
+	defs       map[digest.Digest][]byte
+	metas      map[digest.Digest]pb.OpMetadata
+	sources    map[digest.Digest][]*SourceLocation
+	platforms  map[digest.Digest]*specs.Platform
+	dgst       digest.Digest
+	index      pb.OutputIndex
+	inputCache map[digest.Digest][]*DefinitionOp
 }
 
 // NewDefinitionOp returns a new operation from a marshalled definition.
@@ -89,13 +90,14 @@ func NewDefinitionOp(def *pb.Definition) (*DefinitionOp, error) {
 	}
 
 	return &DefinitionOp{
-		ops:       ops,
-		defs:      defs,
-		metas:     def.Metadata,
-		sources:   srcs,
-		platforms: platforms,
-		dgst:      dgst,
-		index:     index,
+		ops:        ops,
+		defs:       defs,
+		metas:      def.Metadata,
+		sources:    srcs,
+		platforms:  platforms,
+		dgst:       dgst,
+		index:      index,
+		inputCache: make(map[digest.Digest][]*DefinitionOp),
 	}, nil
 }
 
@@ -188,14 +190,34 @@ func (d *DefinitionOp) Inputs() []Output {
 	d.mu.Unlock()
 
 	for _, input := range op.Inputs {
-		vtx := &DefinitionOp{
-			ops:       d.ops,
-			defs:      d.defs,
-			metas:     d.metas,
-			platforms: d.platforms,
-			dgst:      input.Digest,
-			index:     input.Index,
+		var vtx *DefinitionOp
+		d.mu.Lock()
+		if existingIndexes, ok := d.inputCache[input.Digest]; ok {
+			if int(input.Index) < len(existingIndexes) && existingIndexes[input.Index] != nil {
+				vtx = existingIndexes[input.Index]
+			}
 		}
+		if vtx == nil {
+			vtx = &DefinitionOp{
+				ops:        d.ops,
+				defs:       d.defs,
+				metas:      d.metas,
+				platforms:  d.platforms,
+				dgst:       input.Digest,
+				index:      input.Index,
+				inputCache: d.inputCache,
+			}
+			existingIndexes := d.inputCache[input.Digest]
+			indexDiff := int(input.Index) - len(existingIndexes)
+			if indexDiff >= 0 {
+				// make room in the slice for the new index being set
+				existingIndexes = append(existingIndexes, make([]*DefinitionOp, indexDiff+1)...)
+			}
+			existingIndexes[input.Index] = vtx
+			d.inputCache[input.Digest] = existingIndexes
+		}
+		d.mu.Unlock()
+
 		inputs = append(inputs, &output{vertex: vtx, platform: platform, getIndex: func() (pb.OutputIndex, error) {
 			return pb.OutputIndex(vtx.index), nil
 		}})

client/llb/definition_test.go

@@ -6,6 +6,8 @@ import (
 	"testing"
 
 	"github.com/containerd/containerd/platforms"
+	"github.com/moby/buildkit/solver/pb"
+	digest "github.com/opencontainers/go-digest"
 	"github.com/stretchr/testify/require"
 )
 
@@ -69,3 +71,50 @@ func TestDefinitionEquivalence(t *testing.T) {
 		})
 	}
 }
+
+func TestDefinitionInputCache(t *testing.T) {
+	src := HTTP("url")
+
+	stA := Scratch().Run(
+		Shlex("A"),
+		AddMount("/mnt", src),
+	)
+
+	stB := Scratch().Run(
+		Shlex("B"),
+		AddMount("/mnt", src),
+	)
+
+	st := Scratch().Run(
+		Shlex("args"),
+		AddMount("/a", stA.Root()),
+		AddMount("/a2", stA.GetMount("/mnt")),
+		AddMount("/b", stB.Root()),
+		AddMount("/b2", stB.GetMount("/mnt")),
+	).Root()
+
+	ctx := context.TODO()
+
+	def, err := st.Marshal(context.TODO())
+	require.NoError(t, err)
+
+	op, err := NewDefinitionOp(def.ToPB())
+	require.NoError(t, err)
+
+	err = op.Validate(ctx)
+	require.NoError(t, err)
+
+	st2 := NewState(op.Output())
+	marshalDef := &Definition{
+		Metadata: make(map[digest.Digest]pb.OpMetadata, 0),
+	}
+	constraints := &Constraints{}
+	smc := newSourceMapCollector()
+
+	// verify the expected number of vertexes gets marshalled
+	vertexCache := make(map[Vertex]struct{})
+	_, err = marshal(ctx, st2.Output().Vertex(ctx), marshalDef, smc, map[digest.Digest]struct{}{}, vertexCache, constraints)
+	require.NoError(t, err)
+	// 1 exec + 2x2 mounts from stA and stB + 1 src = 6 vertexes
+	require.Equal(t, 6, len(vertexCache))
+}