llb: add cache for marshalers

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

client/llb/exec.go

@@ -37,9 +37,10 @@ type mount struct {
 }
 
 type ExecOp struct {
-	root   Output
-	mounts []*mount
-	meta   Meta
+	root     Output
+	mounts   []*mount
+	meta     Meta
+	cachedPB []byte
 }
 
 func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Output {
@@ -52,6 +53,7 @@ func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Outp
 	}
 	e.mounts = append(e.mounts, m)
 	m.output = &output{vertex: e, getIndex: e.getMountIndexFn(m)}
+	e.cachedPB = nil
 	return m.output
 }
 
@@ -82,6 +84,9 @@ func (e *ExecOp) Validate() error {
 }
 
 func (e *ExecOp) Marshal() ([]byte, error) {
+	if e.cachedPB != nil {
+		return e.cachedPB, nil
+	}
 	if err := e.Validate(); err != nil {
 		return nil, err
 	}
@@ -139,7 +144,12 @@ func (e *ExecOp) Marshal() ([]byte, error) {
 		peo.Mounts = append(peo.Mounts, pm)
 	}
 
-	return pop.Marshal()
+	dt, err := pop.Marshal()
+	if err != nil {
+		return nil, err
+	}
+	e.cachedPB = dt
+	return dt, nil
 }
 
 func (e *ExecOp) Output() Output {

client/llb/llbbuild/llbbuild.go

@@ -21,8 +21,9 @@ func NewBuildOp(source llb.Output, opt ...BuildOption) llb.Vertex {
 }
 
 type build struct {
-	source llb.Output
-	info   *BuildInfo
+	source   llb.Output
+	info     *BuildInfo
+	cachedPB []byte
 }
 
 func (b *build) ToInput() (*pb.Input, error) {
@@ -43,6 +44,9 @@ func (b *build) Validate() error {
 }
 
 func (b *build) Marshal() ([]byte, error) {
+	if b.cachedPB != nil {
+		return b.cachedPB, nil
+	}
 	pbo := &pb.BuildOp{
 		Builder: pb.LLBBuilder,
 		Inputs: map[string]*pb.BuildInput{
@@ -68,7 +72,12 @@ func (b *build) Marshal() ([]byte, error) {
 
 	pop.Inputs = append(pop.Inputs, inp)
 
-	return pop.Marshal()
+	dt, err := pop.Marshal()
+	if err != nil {
+		return nil, err
+	}
+	b.cachedPB = dt
+	return dt, nil
 }
 
 func (b *build) Output() llb.Output {

client/llb/source.go

@@ -8,9 +8,10 @@ import (
 )
 
 type SourceOp struct {
-	id     string
-	attrs  map[string]string
-	output Output
+	id       string
+	attrs    map[string]string
+	output   Output
+	cachedPB []byte
 }
 
 func NewSource(id string, attrs map[string]string) *SourceOp {
@@ -30,6 +31,9 @@ func (s *SourceOp) Validate() error {
 }
 
 func (s *SourceOp) Marshal() ([]byte, error) {
+	if s.cachedPB != nil {
+		return s.cachedPB, nil
+	}
 	if err := s.Validate(); err != nil {
 		return nil, err
 	}
@@ -39,7 +43,12 @@ func (s *SourceOp) Marshal() ([]byte, error) {
 			Source: &pb.SourceOp{Identifier: s.id, Attrs: s.attrs},
 		},
 	}
-	return proto.Marshal()
+	dt, err := proto.Marshal()
+	if err != nil {
+		return nil, err
+	}
+	s.cachedPB = dt
+	return dt, nil
 }
 
 func (s *SourceOp) Output() Output {

client/llb/state.go

@@ -49,7 +49,7 @@ func (s State) Value(k interface{}) interface{} {
 }
 
 func (s State) Marshal() ([][]byte, error) {
-	list, err := marshal(s.Output().Vertex(), nil, map[digest.Digest]struct{}{})
+	list, err := marshal(s.Output().Vertex(), nil, map[digest.Digest]struct{}{}, map[Vertex]struct{}{})
 	if err != nil {
 		return nil, err
 	}
@@ -66,18 +66,23 @@ func (s State) Marshal() ([][]byte, error) {
 	return list, nil
 }
 
-func marshal(v Vertex, list [][]byte, cache map[digest.Digest]struct{}) (out [][]byte, err error) {
+func marshal(v Vertex, list [][]byte, cache map[digest.Digest]struct{}, vertexCache map[Vertex]struct{}) (out [][]byte, err error) {
 	for _, inp := range v.Inputs() {
 		var err error
-		list, err = marshal(inp.Vertex(), list, cache)
+		list, err = marshal(inp.Vertex(), list, cache, vertexCache)
 		if err != nil {
 			return nil, err
 		}
 	}
+	if _, ok := vertexCache[v]; ok {
+		return list, nil
+	}
+
 	dt, err := v.Marshal()
 	if err != nil {
 		return nil, err
 	}
+	vertexCache[v] = struct{}{}
 	dgst := digest.FromBytes(dt)
 	if _, ok := cache[dgst]; ok {
 		return list, nil