llb: avoid concurrent map write on parallel marshal

tonistiigi · tonistiigi · commit 30413b5de707 · 2024-12-06T17:53:59.000-08:00
Calling marshal changes the internal state of the op, for example
addCap() helper adds capability constraints. These can race with
same map being read by another Marshal call. Locking the Marshal
function itself also makes sure that the cache is not recomputed
in this case.

Signed-off-by: Tonis Tiigi &lt;tonistiigi@gmail.com&gt;
diff --git a/client/llb/definition.go b/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
diff --git a/client/llb/diff.go b/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 {
diff --git a/client/llb/exec.go b/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 {
diff --git a/client/llb/fileop.go b/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 {
diff --git a/client/llb/fileop_test.go b/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())
+}
diff --git a/client/llb/llbbuild/llbbuild.go b/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 {
diff --git a/client/llb/marshal.go b/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
diff --git a/client/llb/merge.go b/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 {
diff --git a/client/llb/source.go b/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 {

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@ import (`
`8`	`8`	`)`
`9`	`9`
`10`	`10`	`type DiffOp struct {`
`11`		`- MarshalCache`
	`11`	`+ cache MarshalCache`
`12`	`12`	`lower Output`
`13`	`13`	`upper Output`
`14`	`14`	`output Output`
`@@ -31,7 +31,10 @@ func (m DiffOp) Validate(ctx context.Context, constraints Constraints) error {`
`31`	`31`	`}`
`32`	`32`
`33`	`33`	`func (m DiffOp) Marshal(ctx context.Context, constraints Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`34`		`- if dgst, dt, md, srcs, err := m.Load(constraints); err == nil {`
	`34`	`+ cache := m.cache.Acquire()`
	`35`	`+ defer cache.Release()`
	`36`	`+`
	`37`	`+ if dgst, dt, md, srcs, err := cache.Load(constraints); err == nil {`
`35`	`38`	`return dgst, dt, md, srcs, nil`
`36`	`39`	`}`
`37`	`40`	`if err := m.Validate(ctx, constraints); err != nil {`
`@@ -72,7 +75,7 @@ func (m DiffOp) Marshal(ctx context.Context, constraints Constraints) (digest.`
`72`	`75`	`return "", nil, nil, nil, err`
`73`	`76`	`}`
`74`	`77`
`75`		`- return m.Store(dt, md, m.constraints.SourceLocations, constraints)`
	`78`	`+ return cache.Store(dt, md, m.constraints.SourceLocations, constraints)`
`76`	`79`	`}`
`77`	`80`
`78`	`81`	`func (m *DiffOp) Output() Output {`
Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ type mount struct {`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`type ExecOp struct {`
`54`		`- MarshalCache`
	`54`	`+ cache MarshalCache`
`55`	`55`	`proxyEnv *ProxyEnv`
`56`	`56`	`root Output`
`57`	`57`	`mounts []*mount`
`@@ -63,6 +63,9 @@ type ExecOp struct {`
`63`	`63`	`}`
`64`	`64`
`65`	`65`	`func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Output {`
	`66`	`+ cache := e.cache.Acquire()`
	`67`	`+ defer cache.Release()`
	`68`	`+`
`66`	`69`	`m := &mount{`
`67`	`70`	`target: target,`
`68`	`71`	`source: source,`
`@@ -84,7 +87,7 @@ func (e *ExecOp) AddMount(target string, source Output, opt ...MountOption) Outp`
`84`	`87`	`}`
`85`	`88`	`m.output = o`
`86`	`89`	`}`
`87`		`- e.Store(nil, nil, nil, nil)`
	`90`	`+ cache.Store(nil, nil, nil, nil)`
`88`	`91`	`e.isValidated = false`
`89`	`92`	`return m.output`
`90`	`93`	`}`
`@@ -128,7 +131,10 @@ func (e ExecOp) Validate(ctx context.Context, c Constraints) error {`
`128`	`131`	`}`
`129`	`132`
`130`	`133`	`func (e ExecOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`131`		`- if dgst, dt, md, srcs, err := e.Load(c); err == nil {`
	`134`	`+ cache := e.cache.Acquire()`
	`135`	`+ defer cache.Release()`
	`136`	`+`
	`137`	`+ if dgst, dt, md, srcs, err := cache.Load(c); err == nil {`
`132`	`138`	`return dgst, dt, md, srcs, nil`
`133`	`139`	`}`
`134`	`140`
`@@ -446,7 +452,7 @@ func (e ExecOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []`
`446`	`452`	`if err != nil {`
`447`	`453`	`return "", nil, nil, nil, err`
`448`	`454`	`}`
`449`		`- return e.Store(dt, md, e.constraints.SourceLocations, c)`
	`455`	`+ return cache.Store(dt, md, e.constraints.SourceLocations, c)`
`450`	`456`	`}`
`451`	`457`
`452`	`458`	`func (e *ExecOp) Output() Output {`
Original file line number	Diff line number	Diff line change
`@@ -746,7 +746,10 @@ func (ms marshalState) add(fa FileAction, c Constraints) (fileActionState, e`
`746`	`746`	`}`
`747`	`747`
`748`	`748`	`func (f FileOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`749`		`- if dgst, dt, md, srcs, err := f.Load(c); err == nil {`
	`749`	`+ cache := f.Acquire()`
	`750`	`+ defer cache.Release()`
	`751`	`+`
	`752`	`+ if dgst, dt, md, srcs, err := cache.Load(c); err == nil {`
`750`	`753`	`return dgst, dt, md, srcs, nil`
`751`	`754`	`}`
`752`	`755`
`@@ -816,7 +819,7 @@ func (f FileOp) Marshal(ctx context.Context, c Constraints) (digest.Digest, []`
`816`	`819`	`if err != nil {`
`817`	`820`	`return "", nil, nil, nil, err`
`818`	`821`	`}`
`819`		`- return f.Store(dt, md, f.constraints.SourceLocations, c)`
	`822`	`+ return cache.Store(dt, md, f.constraints.SourceLocations, c)`
`820`	`823`	`}`
`821`	`824`
`822`	`825`	`func normalizePath(parent, p string, keepSlash bool) string {`
Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@ func NewBuildOp(source llb.Output, opt ...BuildOption) llb.Vertex {`
`24`	`24`	`}`
`25`	`25`
`26`	`26`	`type build struct {`
`27`		`- llb.MarshalCache`
	`27`	`+ cache llb.MarshalCache`
`28`	`28`	`source llb.Output`
`29`	`29`	`info *BuildInfo`
`30`	`30`	`constraints llb.Constraints`
`@@ -47,7 +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 dgst, dt, md, srcs, err := b.Load(c); err == nil {`
	`50`	`+ cache := b.cache.Acquire()`
	`51`	`+ defer cache.Release()`
	`52`	`+`
	`53`	`+ if dgst, dt, md, srcs, err := cache.Load(c); err == nil {`
`51`	`54`	`return dgst, dt, md, srcs, nil`
`52`	`55`	`}`
`53`	`56`
`@@ -85,7 +88,7 @@ func (b build) Marshal(ctx context.Context, c llb.Constraints) (digest.Digest,`
`85`	`88`	`if err != nil {`
`86`	`89`	`return "", nil, nil, nil, err`
`87`	`90`	`}`
`88`		`- return b.Store(dt, md, b.constraints.SourceLocations, c)`
	`91`	`+ return cache.Store(dt, md, b.constraints.SourceLocations, c)`
`89`	`92`	`}`
`90`	`93`
`91`	`94`	`func (b *build) Output() llb.Output {`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@ import (`
`9`	`9`	`)`
`10`	`10`
`11`	`11`	`type MergeOp struct {`
`12`		`- MarshalCache`
	`12`	`+ cache MarshalCache`
`13`	`13`	`inputs []Output`
`14`	`14`	`output Output`
`15`	`15`	`constraints Constraints`
`@@ -32,7 +32,10 @@ func (m MergeOp) Validate(ctx context.Context, constraints Constraints) error`
`32`	`32`	`}`
`33`	`33`
`34`	`34`	`func (m MergeOp) Marshal(ctx context.Context, constraints Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`35`		`- if dgst, dt, md, srcs, err := m.Load(constraints); err == nil {`
	`35`	`+ cache := m.cache.Acquire()`
	`36`	`+ defer cache.Release()`
	`37`	`+`
	`38`	`+ if dgst, dt, md, srcs, err := cache.Load(constraints); err == nil {`
`36`	`39`	`return dgst, dt, md, srcs, nil`
`37`	`40`	`}`
`38`	`41`
`@@ -59,7 +62,7 @@ func (m MergeOp) Marshal(ctx context.Context, constraints Constraints) (digest`
`59`	`62`	`return "", nil, nil, nil, err`
`60`	`63`	`}`
`61`	`64`
`62`		`- return m.Store(dt, md, m.constraints.SourceLocations, constraints)`
	`65`	`+ return cache.Store(dt, md, m.constraints.SourceLocations, constraints)`
`63`	`66`	`}`
`64`	`67`
`65`	`68`	`func (m *MergeOp) Output() Output {`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ import (`
`20`	`20`	`)`
`21`	`21`
`22`	`22`	`type SourceOp struct {`
`23`		`- MarshalCache`
	`23`	`+ cache MarshalCache`
`24`	`24`	`id string`
`25`	`25`	`attrs map[string]string`
`26`	`26`	`output Output`
`@@ -49,7 +49,10 @@ func (s SourceOp) Validate(ctx context.Context, c Constraints) error {`
`49`	`49`	`}`
`50`	`50`
`51`	`51`	`func (s SourceOp) Marshal(ctx context.Context, constraints Constraints) (digest.Digest, []byte, pb.OpMetadata, []SourceLocation, error) {`
`52`		`- if dgst, dt, md, srcs, err := s.Load(constraints); err == nil {`
	`52`	`+ cache := s.cache.Acquire()`
	`53`	`+ defer cache.Release()`
	`54`	`+`
	`55`	`+ if dgst, dt, md, srcs, err := cache.Load(constraints); err == nil {`
`53`	`56`	`return dgst, dt, md, srcs, nil`
`54`	`57`	`}`
`55`	`58`
`@@ -82,7 +85,7 @@ func (s SourceOp) Marshal(ctx context.Context, constraints Constraints) (diges`
`82`	`85`	`return "", nil, nil, nil, err`
`83`	`86`	`}`
`84`	`87`
`85`		`- return s.Store(dt, md, s.constraints.SourceLocations, constraints)`
	`88`	`+ return cache.Store(dt, md, s.constraints.SourceLocations, constraints)`
`86`	`89`	`}`
`87`	`90`
`88`	`91`	`func (s *SourceOp) Output() Output {`