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

llb: deterministic marshaling for protobuf and store results from multiple constraints
moby · Oct 2, 2024 · 8445ccf · 8445ccf
2 parents a38e4cc + d59218e
commit 8445ccf
Show file tree

Hide file tree

Showing 8 changed files with 63 additions and 49 deletions.
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)
+}