llb: deterministic marshaling for protobuf and store results from mul…

…tiple constraints This fixes a problem with the new protobuf marshaling with the standard library. LLB digests are now forced into deterministic marshaling to ensure they produce the same digest when marshaled multiple times. In addition, the marshal cache has also been fixed to work in multi-threaded frontends with multiple different constraints. Previously, if an LLB vertex was used in multiple goroutines and marshaled concurrently, the cache would be broken. This could cause certain problems when a specific node was used multiple times in the same LLB tree. Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com>
moby · Sep 30, 2024 · 5776658 · 5776658
1 parent 185751b
commit 5776658
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)
+}