From ff99df06529e8117fcc8e03716544567c6343c6d Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Tue, 8 Aug 2017 17:02:50 -0700
Subject: [PATCH] llb: add cache for marshalers

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 client/llb/exec.go              | 18 ++++++++++++++----
 client/llb/llbbuild/llbbuild.go | 15 ++++++++++++---
 client/llb/source.go            | 17 +++++++++++++----
 client/llb/state.go             | 11 ++++++++---
 4 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/client/llb/exec.go b/client/llb/exec.go
index f4d7f289..8f81c831 100644
--- a/client/llb/exec.go
+++ b/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 {
diff --git a/client/llb/llbbuild/llbbuild.go b/client/llb/llbbuild/llbbuild.go
index ed1643bf..51e9b4e5 100644
--- a/client/llb/llbbuild/llbbuild.go
+++ b/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 {
diff --git a/client/llb/source.go b/client/llb/source.go
index 97373714..e84909a5 100644
--- a/client/llb/source.go
+++ b/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 {
diff --git a/client/llb/state.go b/client/llb/state.go
index d956b14f..680ceb7d 100644
--- a/client/llb/state.go
+++ b/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