From 3b418db323f6cf90e5366fa382d313da28219bc7 Mon Sep 17 00:00:00 2001 From: Erik Sipsma Date: Mon, 10 Aug 2020 13:51:58 -0700 Subject: [PATCH] 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 --- client/llb/definition.go | 66 +++++++++++++++++++++++------------ client/llb/definition_test.go | 49 ++++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 22 deletions(-) diff --git a/client/llb/definition.go b/client/llb/definition.go index fe9f7c17f..99af7c687 100644 --- a/client/llb/definition.go +++ b/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 }}) diff --git a/client/llb/definition_test.go b/client/llb/definition_test.go index 4f69b0a69..449855b37 100644 --- a/client/llb/definition_test.go +++ b/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)) +}