Merge pull request #5577 from tonistiigi/llb-constraints-race-fix

llb: avoid concurrent map write on parallel marshal

client/llb/definition.go

@@ -16,7 +16,6 @@ import (
 // For example, after marshalling a LLB state and sending over the wire, the
 // LLB state can be reconstructed from the definition.
 type DefinitionOp struct {
-	MarshalCache
 	mu         sync.Mutex
 	ops        map[digest.Digest]*pb.Op
 	defs       map[digest.Digest][]byte

client/llb/diff.go

@@ -8,7 +8,7 @@ import (
 )
 
 type DiffOp struct {
-	MarshalCache
+	cache       MarshalCache
 	lower       Output
 	upper       Output
 	output      Output
@@ -31,7 +31,10 @@ 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 dgst, dt, md, srcs, err := m.Load(constraints); err == nil {
+	cache := m.cache.Acquire()
+	defer cache.Release()
+
+	if dgst, dt, md, srcs, err := cache.Load(constraints); err == nil {
 		return dgst, dt, md, srcs, nil
 	}
 	if err := m.Validate(ctx, constraints); err != nil {
@@ -72,7 +75,7 @@ func (m *DiffOp) Marshal(ctx context.Context, constraints *Constraints) (digest.
 		return "", nil, nil, nil, err
 	}
 
-	return m.Store(dt, md, m.constraints.SourceLocations, constraints)
+	return cache.Store(dt, md, m.constraints.SourceLocations, constraints)
 }
 
 func (m *DiffOp) Output() Output {

client/llb/exec.go

@@ -51,7 +51,7 @@ type mount struct {
 }
 
 type ExecOp struct {
-	MarshalCache
+	cache       MarshalCache
 	proxyEnv    *ProxyEnv
 	root        Output
 	mounts      []*mount
@@ -63,6 +63,9 @@ type ExecOp struct {
 }
 
 func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Output {
+	cache := e.cache.Acquire()
+	defer cache.Release()
+
 	m := &mount{
 		target: target,
 		source: source,
@@ -84,7 +87,7 @@ func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Outp
 		}
 		m.output = o
 	}
-	e.Store(nil, nil, nil, nil)
+	cache.Store(nil, nil, nil, nil)
 	e.isValidated = false
 	return m.output
 }
@@ -128,7 +131,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 dgst, dt, md, srcs, err := e.Load(c); err == nil {
+	cache := e.cache.Acquire()
+	defer cache.Release()
+
+	if dgst, dt, md, srcs, err := cache.Load(c); err == nil {
 		return dgst, dt, md, srcs, nil
 	}
 
@@ -446,7 +452,7 @@ func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
-	return e.Store(dt, md, e.constraints.SourceLocations, c)
+	return cache.Store(dt, md, e.constraints.SourceLocations, c)
 }
 
 func (e *ExecOp) Output() Output {

client/llb/fileop.go

@@ -746,7 +746,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 dgst, dt, md, srcs, err := f.Load(c); err == nil {
+	cache := f.Acquire()
+	defer cache.Release()
+
+	if dgst, dt, md, srcs, err := cache.Load(c); err == nil {
 		return dgst, dt, md, srcs, nil
 	}
 
@@ -816,7 +819,7 @@ func (f *FileOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
-	return f.Store(dt, md, f.constraints.SourceLocations, c)
+	return cache.Store(dt, md, f.constraints.SourceLocations, c)
 }
 
 func normalizePath(parent, p string, keepSlash bool) string {

client/llb/fileop_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/moby/buildkit/solver/pb"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
 )
 
 func TestFileMkdir(t *testing.T) {
@@ -737,3 +738,15 @@ func TestFileOpMarshalConsistency(t *testing.T) {
 		prevDef = def.Def
 	}
 }
+
+func TestParallelMarshal(t *testing.T) {
+	st := Scratch().File(Mkfile("/tmp", 0644, []byte("tmp 1")))
+	eg, ctx := errgroup.WithContext(context.Background())
+	for i := 0; i < 100; i++ {
+		eg.Go(func() error {
+			_, err := st.Marshal(ctx)
+			return err
+		})
+	}
+	require.NoError(t, eg.Wait())
+}

client/llb/llbbuild/llbbuild.go

@@ -24,7 +24,7 @@ func NewBuildOp(source llb.Output, opt ...BuildOption) llb.Vertex {
 }
 
 type build struct {
-	llb.MarshalCache
+	cache       llb.MarshalCache
 	source      llb.Output
 	info        *BuildInfo
 	constraints llb.Constraints
@@ -47,7 +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 dgst, dt, md, srcs, err := b.Load(c); err == nil {
+	cache := b.cache.Acquire()
+	defer cache.Release()
+
+	if dgst, dt, md, srcs, err := cache.Load(c); err == nil {
 		return dgst, dt, md, srcs, nil
 	}
 
@@ -85,7 +88,7 @@ func (b *build) Marshal(ctx context.Context, c *llb.Constraints) (digest.Digest,
 	if err != nil {
 		return "", nil, nil, nil, err
 	}
-	return b.Store(dt, md, b.constraints.SourceLocations, c)
+	return cache.Store(dt, md, b.constraints.SourceLocations, c)
 }
 
 func (b *build) Output() llb.Output {

client/llb/marshal.go

@@ -117,30 +117,45 @@ func MarshalConstraints(base, override *Constraints) (*pb.Op, *pb.OpMetadata) {
 }
 
 type MarshalCache struct {
-	cache sync.Map
+	mu    sync.Mutex
+	cache map[*Constraints]*marshalCacheResult
 }
 
-func (mc *MarshalCache) Load(c *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
-	v, ok := mc.cache.Load(c)
+type MarshalCacheInstance struct {
+	*MarshalCache
+}
+
+func (mc *MarshalCache) Acquire() *MarshalCacheInstance {
+	mc.mu.Lock()
+	return &MarshalCacheInstance{mc}
+}
+
+func (mc *MarshalCacheInstance) Load(c *Constraints) (digest.Digest, []byte, *pb.OpMetadata, []*SourceLocation, error) {
+	res, ok := mc.cache[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) {
+func (mc *MarshalCacheInstance) 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)
+	if mc.cache == nil {
+		mc.cache = make(map[*Constraints]*marshalCacheResult)
+	}
+	mc.cache[c] = res
 	return res.digest, res.dt, res.md, res.srcs, nil
 }
 
+func (mc *MarshalCacheInstance) Release() {
+	mc.mu.Unlock()
+}
+
 type marshalCacheResult struct {
 	digest digest.Digest
 	dt     []byte

client/llb/merge.go

@@ -9,7 +9,7 @@ import (
 )
 
 type MergeOp struct {
-	MarshalCache
+	cache       MarshalCache
 	inputs      []Output
 	output      Output
 	constraints Constraints
@@ -32,7 +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 dgst, dt, md, srcs, err := m.Load(constraints); err == nil {
+	cache := m.cache.Acquire()
+	defer cache.Release()
+
+	if dgst, dt, md, srcs, err := cache.Load(constraints); err == nil {
 		return dgst, dt, md, srcs, nil
 	}
 
@@ -59,7 +62,7 @@ func (m *MergeOp) Marshal(ctx context.Context, constraints *Constraints) (digest
 		return "", nil, nil, nil, err
 	}
 
-	return m.Store(dt, md, m.constraints.SourceLocations, constraints)
+	return cache.Store(dt, md, m.constraints.SourceLocations, constraints)
 }
 
 func (m *MergeOp) Output() Output {

client/llb/source.go

@@ -20,7 +20,7 @@ import (
 )
 
 type SourceOp struct {
-	MarshalCache
+	cache       MarshalCache
 	id          string
 	attrs       map[string]string
 	output      Output
@@ -49,7 +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 dgst, dt, md, srcs, err := s.Load(constraints); err == nil {
+	cache := s.cache.Acquire()
+	defer cache.Release()
+
+	if dgst, dt, md, srcs, err := cache.Load(constraints); err == nil {
 		return dgst, dt, md, srcs, nil
 	}
 
@@ -82,7 +85,7 @@ func (s *SourceOp) Marshal(ctx context.Context, constraints *Constraints) (diges
 		return "", nil, nil, nil, err
 	}
 
-	return s.Store(dt, md, s.constraints.SourceLocations, constraints)
+	return cache.Store(dt, md, s.constraints.SourceLocations, constraints)
 }
 
 func (s *SourceOp) Output() Output {