Merge pull request moby#5374 from jsternberg/marshal-constraints-cache

tonistiigi · web-flow · commit 8445ccf1cb0b · 2024-10-02T13:48:25.000-07:00
llb: deterministic marshaling for protobuf and store results from multiple constraints
diff --git a/client/llb/diff.go b/client/llb/diff.go
@@ -5,7 +5,6 @@ import (
 
 	"github.com/moby/buildkit/solver/pb"
 	digest "github.com/opencontainers/go-digest"
-	protobuf "google.golang.org/protobuf/proto"
 )
 
 type DiffOp struct {
@@ -32,8 +31,8 @@ func (m *DiffOp) Validate(ctx context.Context, constraints *Constraints) error {
 }
 
 func (m *DiffOp) Marshal(ctx context.Context, constraints *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
-	if m.Cached(constraints) {
-		return m.Load()
+	if dgst, dt, md, srcs, err := m.Load(constraints); err == nil {
+		return dgst, dt, md, srcs, nil
 	}
 	if err := m.Validate(ctx, constraints); err != nil {
 		return "", nil, nil, nil, err
@@ -68,13 +67,12 @@ func (m *DiffOp) Marshal(ctx context.Context, constraints *Constraints) (digest.
 
 	proto.Op = &pb.Op_Diff{Diff: op}
 
-	dt, err := protobuf.Marshal(proto)
+	dt, err := deterministicMarshal(proto)
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
 
-	m.Store(dt, md, m.constraints.SourceLocations, constraints)
-	return m.Load()
+	return m.Store(dt, md, m.constraints.SourceLocations, constraints)
 }
 
 func (m *DiffOp) Output() Output {
diff --git a/client/llb/exec.go b/client/llb/exec.go
@@ -129,9 +129,10 @@ func (e *ExecOp) Validate(ctx context.Context, c *Constraints) error {
 }
 
 func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
-	if e.Cached(c) {
-		return e.Load()
+	if dgst, dt, md, srcs, err := e.Load(c); err == nil {
+		return dgst, dt, md, srcs, nil
 	}
+
 	if err := e.Validate(ctx, c); err != nil {
 		return "", nil, nil, nil, err
 	}
@@ -442,12 +443,11 @@ func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		peo.Mounts = append(peo.Mounts, pm)
 	}
 
-	dt, err := proto.Marshal(pop)
+	dt, err := deterministicMarshal(pop)
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
-	e.Store(dt, md, e.constraints.SourceLocations, c)
-	return e.Load()
+	return e.Store(dt, md, e.constraints.SourceLocations, c)
 }
 
 func (e *ExecOp) Output() Output {
diff --git a/client/llb/fileop.go b/client/llb/fileop.go
@@ -730,9 +730,10 @@ func (ms *marshalState) add(fa *FileAction, c *Constraints) (*fileActionState, e
 }
 
 func (f *FileOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
-	if f.Cached(c) {
-		return f.Load()
+	if dgst, dt, md, srcs, err := f.Load(c); err == nil {
+		return dgst, dt, md, srcs, nil
 	}
+
 	if err := f.Validate(ctx, c); err != nil {
 		return "", nil, nil, nil, err
 	}
@@ -795,12 +796,11 @@ func (f *FileOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		})
 	}
 
-	dt, err := proto.Marshal(pop)
+	dt, err := deterministicMarshal(pop)
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
-	f.Store(dt, md, f.constraints.SourceLocations, c)
-	return f.Load()
+	return f.Store(dt, md, f.constraints.SourceLocations, c)
 }
 
 func normalizePath(parent, p string, keepSlash bool) string {
diff --git a/client/llb/llbbuild/llbbuild.go b/client/llb/llbbuild/llbbuild.go
@@ -47,9 +47,10 @@ func (b *build) Validate(context.Context, *llb.Constraints) error {
 }
 
 func (b *build) Marshal(ctx context.Context, c *llb.Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*llb.SourceLocation, error) {
-	if b.Cached(c) {
-		return b.Load()
+	if dgst, dt, md, srcs, err := b.Load(c); err == nil {
+		return dgst, dt, md, srcs, nil
 	}
+
 	pbo := &pb.BuildOp{
 		Builder: int64(pb.LLBBuilder),
 		Inputs: map[string]*pb.BuildInput{
@@ -84,8 +85,7 @@ func (b *build) Marshal(ctx context.Context, c *llb.Constraints) (digest.Digest,
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
-	b.Store(dt, md, b.constraints.SourceLocations, c)
-	return b.Load()
+	return b.Store(dt, md, b.constraints.SourceLocations, c)
 }
 
 func (b *build) Output() llb.Output {
diff --git a/client/llb/marshal.go b/client/llb/marshal.go
@@ -2,7 +2,9 @@ package llb
 
 import (
 	"io"
+	"sync"
 
+	cerrdefs "github.com/containerd/errdefs"
 	"github.com/containerd/platforms"
 	"github.com/moby/buildkit/solver/pb"
 	digest "github.com/opencontainers/go-digest"
@@ -115,25 +117,37 @@ func MarshalConstraints(base, override *Constraints) (*pb.Op, *pb.OpMetadata) {
 }
 
 type MarshalCache struct {
-	digest      digest.Digest
-	dt          []byte
-	md          *pb.OpMetadata
-	srcs        []*SourceLocation
-	constraints *Constraints
+	cache sync.Map
 }
 
-func (mc *MarshalCache) Cached(c *Constraints) bool {
-	return mc.dt != nil && mc.constraints == c
+func (mc *MarshalCache) Load(c *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
+	v, ok := mc.cache.Load(c)
+	if !ok {
+		return "", nil, nil, nil, cerrdefs.ErrNotFound
+	}
+
+	res := v.(*marshalCacheResult)
+	return res.digest, res.dt, res.md, res.srcs, nil
+}
+
+func (mc *MarshalCache) Store(dt []byte, md *pb.OpMetadata, srcs []*SourceLocation, c *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
+	res := &marshalCacheResult{
+		digest: digest.FromBytes(dt),
+		dt:     dt,
+		md:     md,
+		srcs:   srcs,
+	}
+	mc.cache.Store(c, res)
+	return res.digest, res.dt, res.md, res.srcs, nil
 }
 
-func (mc *MarshalCache) Load() (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
-	return mc.digest, mc.dt, mc.md, mc.srcs, nil
+type marshalCacheResult struct {
+	digest digest.Digest
+	dt     []byte
+	md     *pb.OpMetadata
+	srcs   []*SourceLocation
 }
 
-func (mc *MarshalCache) Store(dt []byte, md *pb.OpMetadata, srcs []*SourceLocation, c *Constraints) {
-	mc.digest = digest.FromBytes(dt)
-	mc.dt = dt
-	mc.md = md
-	mc.constraints = c
-	mc.srcs = srcs
+func deterministicMarshal[Message proto.Message](m Message) ([]byte, error) {
+	return proto.MarshalOptions{Deterministic: true}.Marshal(m)
 }
diff --git a/client/llb/merge.go b/client/llb/merge.go
@@ -6,7 +6,6 @@ import (
 	"github.com/moby/buildkit/solver/pb"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
-	"google.golang.org/protobuf/proto"
 )
 
 type MergeOp struct {
@@ -33,9 +32,10 @@ func (m *MergeOp) Validate(ctx context.Context, constraints *Constraints) error
 }
 
 func (m *MergeOp) Marshal(ctx context.Context, constraints *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
-	if m.Cached(constraints) {
-		return m.Load()
+	if dgst, dt, md, srcs, err := m.Load(constraints); err == nil {
+		return dgst, dt, md, srcs, nil
 	}
+
 	if err := m.Validate(ctx, constraints); err != nil {
 		return "", nil, nil, nil, err
 	}
@@ -54,13 +54,12 @@ func (m *MergeOp) Marshal(ctx context.Context, constraints *Constraints) (digest
 	}
 	pop.Op = &pb.Op_Merge{Merge: op}
 
-	dt, err := proto.Marshal(pop)
+	dt, err := deterministicMarshal(pop)
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
 
-	m.Store(dt, md, m.constraints.SourceLocations, constraints)
-	return m.Load()
+	return m.Store(dt, md, m.constraints.SourceLocations, constraints)
 }
 
 func (m *MergeOp) Output() Output {
diff --git a/client/llb/source.go b/client/llb/source.go
@@ -17,7 +17,6 @@ import (
 	"github.com/moby/buildkit/util/sshutil"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
-	protobuf "google.golang.org/protobuf/proto"
 )
 
 type SourceOp struct {
@@ -50,9 +49,10 @@ func (s *SourceOp) Validate(ctx context.Context, c *Constraints) error {
 }
 
 func (s *SourceOp) Marshal(ctx context.Context, constraints *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
-	if s.Cached(constraints) {
-		return s.Load()
+	if dgst, dt, md, srcs, err := s.Load(constraints); err == nil {
+		return dgst, dt, md, srcs, nil
 	}
+
 	if err := s.Validate(ctx, constraints); err != nil {
 		return "", nil, nil, nil, err
 	}
@@ -77,13 +77,12 @@ func (s *SourceOp) Marshal(ctx context.Context, constraints *Constraints) (diges
 		proto.Platform = nil
 	}
 
-	dt, err := protobuf.Marshal(proto)
+	dt, err := deterministicMarshal(proto)
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
 
-	s.Store(dt, md, s.constraints.SourceLocations, constraints)
-	return s.Load()
+	return s.Store(dt, md, s.constraints.SourceLocations, constraints)
 }
 
 func (s *SourceOp) Output() Output {
diff --git a/solver/llbsolver/vertex.go b/solver/llbsolver/vertex.go
@@ -228,7 +228,7 @@ func recomputeDigests(ctx context.Context, all map[digest.Digest]*pb.Op, visited
 		return dgst, nil
 	}
 
-	dt, err := proto.Marshal(op)
+	dt, err := deterministicMarshal(op)
 	if err != nil {
 		return "", err
 	}
@@ -263,7 +263,7 @@ func loadLLB(ctx context.Context, def *pb.Definition, polEngine SourcePolicyEval
 				return solver.Edge{}, errors.Wrap(err, "error evaluating the source policy")
 			}
 			if mutated {
-				dtMutated, err := proto.Marshal(&op)
+				dtMutated, err := deterministicMarshal(&op)
 				if err != nil {
 					return solver.Edge{}, err
 				}
@@ -397,3 +397,7 @@ func fileOpName(actions []*pb.FileAction) string {
 
 	return strings.Join(names, ", ")
 }
+
+func deterministicMarshal[Message proto.Message](m Message) ([]byte, error) {
+	return proto.MarshalOptions{Deterministic: true}.Marshal(m)
+}

Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,6 @@ import (`
`5`	`5`
`6`	`6`	`"github.com/moby/buildkit/solver/pb"`
`7`	`7`	`digest "github.com/opencontainers/go-digest"`
`8`		`- protobuf "google.golang.org/protobuf/proto"`
`9`	`8`	`)`
`10`	`9`
`11`	`10`	`type DiffOp struct {`
`@@ -32,8 +31,8 @@ func (m DiffOp) Validate(ctx context.Context, constraints Constraints) error {`
`32`	`31`	`}`
`33`	`32`
`34`	`33`	`func (m DiffOp) Marshal(ctx context.Context, constraints Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`35`		`- if m.Cached(constraints) {`
`36`		`- return m.Load()`
	`34`	`+ if dgst, dt, md, srcs, err := m.Load(constraints); err == nil {`
	`35`	`+ return dgst, dt, md, srcs, nil`
`37`	`36`	`}`
`38`	`37`	`if err := m.Validate(ctx, constraints); err != nil {`
`39`	`38`	`return "", nil, nil, nil, err`
`@@ -68,13 +67,12 @@ func (m DiffOp) Marshal(ctx context.Context, constraints Constraints) (digest.`
`68`	`67`
`69`	`68`	`proto.Op = &pb.Op_Diff{Diff: op}`
`70`	`69`
`71`		`- dt, err := protobuf.Marshal(proto)`
	`70`	`+ dt, err := deterministicMarshal(proto)`
`72`	`71`	`if err != nil {`
`73`	`72`	`return "", nil, nil, nil, err`
`74`	`73`	`}`
`75`	`74`
`76`		`- m.Store(dt, md, m.constraints.SourceLocations, constraints)`
`77`		`- return m.Load()`
	`75`	`+ return m.Store(dt, md, m.constraints.SourceLocations, constraints)`
`78`	`76`	`}`
`79`	`77`
`80`	`78`	`func (m *DiffOp) Output() Output {`
Original file line number	Diff line number	Diff line change
`@@ -129,9 +129,10 @@ func (e ExecOp) Validate(ctx context.Context, c Constraints) error {`
`129`	`129`	`}`
`130`	`130`
`131`	`131`	`func (e ExecOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`132`		`- if e.Cached(c) {`
`133`		`- return e.Load()`
	`132`	`+ if dgst, dt, md, srcs, err := e.Load(c); err == nil {`
	`133`	`+ return dgst, dt, md, srcs, nil`
`134`	`134`	`}`
	`135`	`+`
`135`	`136`	`if err := e.Validate(ctx, c); err != nil {`
`136`	`137`	`return "", nil, nil, nil, err`
`137`	`138`	`}`
`@@ -442,12 +443,11 @@ func (e ExecOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []`
`442`	`443`	`peo.Mounts = append(peo.Mounts, pm)`
`443`	`444`	`}`
`444`	`445`
`445`		`- dt, err := proto.Marshal(pop)`
	`446`	`+ dt, err := deterministicMarshal(pop)`
`446`	`447`	`if err != nil {`
`447`	`448`	`return "", nil, nil, nil, err`
`448`	`449`	`}`
`449`		`- e.Store(dt, md, e.constraints.SourceLocations, c)`
`450`		`- return e.Load()`
	`450`	`+ return e.Store(dt, md, e.constraints.SourceLocations, c)`
`451`	`451`	`}`
`452`	`452`
`453`	`453`	`func (e *ExecOp) Output() Output {`
Original file line number	Diff line number	Diff line change
`@@ -730,9 +730,10 @@ func (ms marshalState) add(fa FileAction, c Constraints) (fileActionState, e`
`730`	`730`	`}`
`731`	`731`
`732`	`732`	`func (f FileOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`733`		`- if f.Cached(c) {`
`734`		`- return f.Load()`
	`733`	`+ if dgst, dt, md, srcs, err := f.Load(c); err == nil {`
	`734`	`+ return dgst, dt, md, srcs, nil`
`735`	`735`	`}`
	`736`	`+`
`736`	`737`	`if err := f.Validate(ctx, c); err != nil {`
`737`	`738`	`return "", nil, nil, nil, err`
`738`	`739`	`}`
`@@ -795,12 +796,11 @@ func (f FileOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []`
`795`	`796`	`})`
`796`	`797`	`}`
`797`	`798`
`798`		`- dt, err := proto.Marshal(pop)`
	`799`	`+ dt, err := deterministicMarshal(pop)`
`799`	`800`	`if err != nil {`
`800`	`801`	`return "", nil, nil, nil, err`
`801`	`802`	`}`
`802`		`- f.Store(dt, md, f.constraints.SourceLocations, c)`
`803`		`- return f.Load()`
	`803`	`+ return f.Store(dt, md, f.constraints.SourceLocations, c)`
`804`	`804`	`}`
`805`	`805`
`806`	`806`	`func normalizePath(parent, p string, keepSlash bool) string {`
Original file line number	Diff line number	Diff line change
`@@ -47,9 +47,10 @@ func (b build) Validate(context.Context, llb.Constraints) error {`
`47`	`47`	`}`
`48`	`48`
`49`	`49`	`func (b build) Marshal(ctx context.Context, c llb.Constraints) (digest.Digest, []byte, pb.OpMetadata, []llb.SourceLocation, error) {`
`50`		`- if b.Cached(c) {`
`51`		`- return b.Load()`
	`50`	`+ if dgst, dt, md, srcs, err := b.Load(c); err == nil {`
	`51`	`+ return dgst, dt, md, srcs, nil`
`52`	`52`	`}`
	`53`	`+`
`53`	`54`	`pbo := &pb.BuildOp{`
`54`	`55`	`Builder: int64(pb.LLBBuilder),`
`55`	`56`	`Inputs: map[string]*pb.BuildInput{`
`@@ -84,8 +85,7 @@ func (b build) Marshal(ctx context.Context, c llb.Constraints) (digest.Digest,`
`84`	`85`	`if err != nil {`
`85`	`86`	`return "", nil, nil, nil, err`
`86`	`87`	`}`
`87`		`- b.Store(dt, md, b.constraints.SourceLocations, c)`
`88`		`- return b.Load()`
	`88`	`+ return b.Store(dt, md, b.constraints.SourceLocations, c)`
`89`	`89`	`}`
`90`	`90`
`91`	`91`	`func (b *build) Output() llb.Output {`
Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,6 @@ import (`
`6`	`6`	`"github.com/moby/buildkit/solver/pb"`
`7`	`7`	`digest "github.com/opencontainers/go-digest"`
`8`	`8`	`"github.com/pkg/errors"`
`9`		`- "google.golang.org/protobuf/proto"`
`10`	`9`	`)`
`11`	`10`
`12`	`11`	`type MergeOp struct {`
`@@ -33,9 +32,10 @@ func (m MergeOp) Validate(ctx context.Context, constraints Constraints) error`
`33`	`32`	`}`
`34`	`33`
`35`	`34`	`func (m MergeOp) Marshal(ctx context.Context, constraints Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`36`		`- if m.Cached(constraints) {`
`37`		`- return m.Load()`
	`35`	`+ if dgst, dt, md, srcs, err := m.Load(constraints); err == nil {`
	`36`	`+ return dgst, dt, md, srcs, nil`
`38`	`37`	`}`
	`38`	`+`
`39`	`39`	`if err := m.Validate(ctx, constraints); err != nil {`
`40`	`40`	`return "", nil, nil, nil, err`
`41`	`41`	`}`
`@@ -54,13 +54,12 @@ func (m MergeOp) Marshal(ctx context.Context, constraints Constraints) (digest`
`54`	`54`	`}`
`55`	`55`	`pop.Op = &pb.Op_Merge{Merge: op}`
`56`	`56`
`57`		`- dt, err := proto.Marshal(pop)`
	`57`	`+ dt, err := deterministicMarshal(pop)`
`58`	`58`	`if err != nil {`
`59`	`59`	`return "", nil, nil, nil, err`
`60`	`60`	`}`
`61`	`61`
`62`		`- m.Store(dt, md, m.constraints.SourceLocations, constraints)`
`63`		`- return m.Load()`
	`62`	`+ return m.Store(dt, md, m.constraints.SourceLocations, constraints)`
`64`	`63`	`}`
`65`	`64`
`66`	`65`	`func (m *MergeOp) Output() Output {`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,6 @@ import (`
`17`	`17`	`"github.com/moby/buildkit/util/sshutil"`
`18`	`18`	`digest "github.com/opencontainers/go-digest"`
`19`	`19`	`"github.com/pkg/errors"`
`20`		`- protobuf "google.golang.org/protobuf/proto"`
`21`	`20`	`)`
`22`	`21`
`23`	`22`	`type SourceOp struct {`
`@@ -50,9 +49,10 @@ func (s SourceOp) Validate(ctx context.Context, c Constraints) error {`
`50`	`49`	`}`
`51`	`50`
`52`	`51`	`func (s SourceOp) Marshal(ctx context.Context, constraints Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`53`		`- if s.Cached(constraints) {`
`54`		`- return s.Load()`
	`52`	`+ if dgst, dt, md, srcs, err := s.Load(constraints); err == nil {`
	`53`	`+ return dgst, dt, md, srcs, nil`
`55`	`54`	`}`
	`55`	`+`
`56`	`56`	`if err := s.Validate(ctx, constraints); err != nil {`
`57`	`57`	`return "", nil, nil, nil, err`
`58`	`58`	`}`
`@@ -77,13 +77,12 @@ func (s SourceOp) Marshal(ctx context.Context, constraints Constraints) (diges`
`77`	`77`	`proto.Platform = nil`
`78`	`78`	`}`
`79`	`79`
`80`		`- dt, err := protobuf.Marshal(proto)`
	`80`	`+ dt, err := deterministicMarshal(proto)`
`81`	`81`	`if err != nil {`
`82`	`82`	`return "", nil, nil, nil, err`
`83`	`83`	`}`
`84`	`84`
`85`		`- s.Store(dt, md, s.constraints.SourceLocations, constraints)`
`86`		`- return s.Load()`
	`85`	`+ return s.Store(dt, md, s.constraints.SourceLocations, constraints)`
`87`	`86`	`}`
`88`	`87`
`89`	`88`	`func (s *SourceOp) Output() Output {`
Original file line number	Diff line number	Diff line change
`@@ -228,7 +228,7 @@ func recomputeDigests(ctx context.Context, all map[digest.Digest]*pb.Op, visited`
`228`	`228`	`return dgst, nil`
`229`	`229`	`}`
`230`	`230`
`231`		`- dt, err := proto.Marshal(op)`
	`231`	`+ dt, err := deterministicMarshal(op)`
`232`	`232`	`if err != nil {`
`233`	`233`	`return "", err`
`234`	`234`	`}`
`@@ -263,7 +263,7 @@ func loadLLB(ctx context.Context, def *pb.Definition, polEngine SourcePolicyEval`
`263`	`263`	`return solver.Edge{}, errors.Wrap(err, "error evaluating the source policy")`
`264`	`264`	`}`
`265`	`265`	`if mutated {`
`266`		`- dtMutated, err := proto.Marshal(&op)`
	`266`	`+ dtMutated, err := deterministicMarshal(&op)`
`267`	`267`	`if err != nil {`
`268`	`268`	`return solver.Edge{}, err`
`269`	`269`	`}`
`@@ -397,3 +397,7 @@ func fileOpName(actions []*pb.FileAction) string {`
`397`	`397`
`398`	`398`	`return strings.Join(names, ", ")`
`399`	`399`	`}`
	`400`	`+`
	`401`	`+func deterministicMarshal[Message proto.Message](m Message) ([]byte, error) {`
	`402`	`+ return proto.MarshalOptions{Deterministic: true}.Marshal(m)`
	`403`	`+}`