Avoid panic on parallel walking on DefinitionOp

Signed-off-by: Kohei Tokunaga <[email protected]>

Author: ktock

client/llb/definition.go

@@ -24,7 +24,7 @@ type DefinitionOp struct {
 	platforms  map[digest.Digest]*ocispecs.Platform
 	dgst       digest.Digest
 	index      pb.OutputIndex
-	inputCache map[digest.Digest][]*DefinitionOp
+	inputCache *sync.Map // shared and written among DefinitionOps so avoid race on this map using sync.Map
 }
 
 // NewDefinitionOp returns a new operation from a marshalled definition.
@@ -101,7 +101,7 @@ func NewDefinitionOp(def *pb.Definition) (*DefinitionOp, error) {
 		platforms:  platforms,
 		dgst:       dgst,
 		index:      index,
-		inputCache: make(map[digest.Digest][]*DefinitionOp),
+		inputCache: new(sync.Map),
 	}, nil
 }
 
@@ -180,6 +180,18 @@ func (d *DefinitionOp) Output() Output {
 	}}
 }
 
+func (d *DefinitionOp) loadInputCache(dgst digest.Digest) ([]*DefinitionOp, bool) {
+	a, ok := d.inputCache.Load(dgst.String())
+	if ok {
+		return a.([]*DefinitionOp), true
+	}
+	return nil, false
+}
+
+func (d *DefinitionOp) storeInputCache(dgst digest.Digest, c []*DefinitionOp) {
+	d.inputCache.Store(dgst.String(), c)
+}
+
 func (d *DefinitionOp) Inputs() []Output {
 	if d.dgst == "" {
 		return nil
@@ -195,7 +207,7 @@ func (d *DefinitionOp) Inputs() []Output {
 	for _, input := range op.Inputs {
 		var vtx *DefinitionOp
 		d.mu.Lock()
-		if existingIndexes, ok := d.inputCache[input.Digest]; ok {
+		if existingIndexes, ok := d.loadInputCache(input.Digest); ok {
 			if int(input.Index) < len(existingIndexes) && existingIndexes[input.Index] != nil {
 				vtx = existingIndexes[input.Index]
 			}
@@ -211,14 +223,14 @@ func (d *DefinitionOp) Inputs() []Output {
 				inputCache: d.inputCache,
 				sources:    d.sources,
 			}
-			existingIndexes := d.inputCache[input.Digest]
+			existingIndexes, _ := d.loadInputCache(input.Digest)
 			indexDiff := int(input.Index) - len(existingIndexes)
 			if indexDiff >= 0 {
 				// make room in the slice for the new index being set
 				existingIndexes = append(existingIndexes, make([]*DefinitionOp, indexDiff+1)...)
 			}
 			existingIndexes[input.Index] = vtx
-			d.inputCache[input.Digest] = existingIndexes
+			d.storeInputCache(input.Digest, existingIndexes)
 		}
 		d.mu.Unlock()
 

client/llb/definition_test.go

@@ -3,12 +3,14 @@ package llb
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/containerd/containerd/platforms"
 	"github.com/moby/buildkit/solver/pb"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
 )
 
 func TestDefinitionEquivalence(t *testing.T) {
@@ -117,10 +119,36 @@ func TestDefinitionInputCache(t *testing.T) {
 	require.NoError(t, err)
 	// 1 exec + 2x2 mounts from stA and stB + 1 src = 6 vertexes
 	require.Equal(t, 6, len(vertexCache))
+
+	// make sure that walking vertices in parallel doesn't cause panic
+	var all []RunOption
+	for i := 0; i < 100; i++ {
+		var sts []RunOption
+		for j := 0; j < 100; j++ {
+			sts = append(sts, AddMount("/mnt", Scratch().Run(Shlex(fmt.Sprintf("%d-%d", i, j))).Root()))
+		}
+		all = append(all, AddMount("/mnt", Scratch().Run(append([]RunOption{Shlex("args")}, sts...)...).Root()))
+	}
+	def, err = Scratch().Run(append([]RunOption{Shlex("args")}, all...)...).Root().Marshal(context.TODO())
+	require.NoError(t, err)
+	op, err = NewDefinitionOp(def.ToPB())
+	require.NoError(t, err)
+	require.NoError(t, testParallelWalk(context.Background(), op.Output()))
 }
 
 func TestDefinitionNil(t *testing.T) {
 	// should be an error, not a panic
 	_, err := NewDefinitionOp(nil)
 	require.Error(t, err)
 }
+
+func testParallelWalk(ctx context.Context, out Output) error {
+	eg, egCtx := errgroup.WithContext(ctx)
+	for _, o := range out.Vertex(ctx, nil).Inputs() {
+		o := o
+		eg.Go(func() error {
+			return testParallelWalk(egCtx, o)
+		})
+	}
+	return eg.Wait()
+}