From c575077cd68784b5aa6af7698ea57e37a8d633b0 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Sat, 21 Jul 2018 21:00:37 -0700
Subject: [PATCH] solver: fix edge merge skipping issue

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 solver/edge.go           |  3 ++-
 solver/scheduler_test.go | 55 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/solver/edge.go b/solver/edge.go
index 1547cf4bd..0de01f881 100644
--- a/solver/edge.go
+++ b/solver/edge.go
@@ -413,7 +413,7 @@ func (e *edge) processUpdate(upt pipe.Receiver) (depChanged bool) {
 	}
 
 	// response for requests to dependencies
-	if dep, ok := e.depRequests[upt]; ok { // TODO: ignore canceled
+	if dep, ok := e.depRequests[upt]; ok {
 		if err := upt.Status().Err; !upt.Status().Canceled && upt.Status().Completed && err != nil {
 			if e.err == nil {
 				e.err = err
@@ -427,6 +427,7 @@ func (e *edge) processUpdate(upt pipe.Receiver) (depChanged bool) {
 			newKeys := state.keys[len(dep.keys):]
 			if e.cacheMap != nil {
 				e.probeCache(dep, withSelector(newKeys, e.cacheMap.Deps[dep.index].Selector))
+				dep.edgeState.keys = state.keys
 				if e.allDepsHaveKeys() {
 					e.keysDidChange = true
 				}
diff --git a/solver/scheduler_test.go b/solver/scheduler_test.go
index fe3367df1..18911700a 100644
--- a/solver/scheduler_test.go
+++ b/solver/scheduler_test.go
@@ -2799,6 +2799,61 @@ func TestCacheExportingMergedKey(t *testing.T) {
 	require.Equal(t, len(expTarget.records), 5)
 }
 
+// moby/buildkit#434
+func TestMergedEdgesLookup(t *testing.T) {
+	t.Parallel()
+
+	rand.Seed(time.Now().UnixNano())
+
+	// this test requires multiple runs to trigger the race
+	for i := 0; i < 20; i++ {
+		func() {
+			ctx := context.TODO()
+
+			cacheManager := newTrackingCacheManager(NewInMemoryCacheManager())
+
+			l := NewSolver(SolverOpt{
+				ResolveOpFunc: testOpResolver,
+				DefaultCache:  cacheManager,
+			})
+			defer l.Close()
+
+			j0, err := l.NewJob("j0")
+			require.NoError(t, err)
+
+			defer func() {
+				if j0 != nil {
+					j0.Discard()
+				}
+			}()
+
+			g := Edge{
+				Vertex: vtxSum(3, vtxOpt{inputs: []Edge{
+					{Vertex: vtxSum(0, vtxOpt{inputs: []Edge{
+						{Vertex: vtxSum(2, vtxOpt{inputs: []Edge{
+							{Vertex: vtxConst(2, vtxOpt{})},
+						}})},
+						{Vertex: vtxConst(0, vtxOpt{})},
+					}})},
+					{Vertex: vtxSum(2, vtxOpt{inputs: []Edge{
+						{Vertex: vtxConst(2, vtxOpt{})},
+					}})},
+				}}),
+			}
+			g.Vertex.(*vertexSum).setupCallCounters()
+
+			res, err := j0.Build(ctx, g)
+			require.NoError(t, err)
+			require.Equal(t, unwrapInt(res), 11)
+			require.Equal(t, int64(7), *g.Vertex.(*vertexSum).cacheCallCount)
+			require.Equal(t, int64(0), cacheManager.loadCounter)
+
+			require.NoError(t, j0.Discard())
+			j0 = nil
+		}()
+	}
+}
+
 func generateSubGraph(nodes int) (Edge, int) {
 	if nodes == 1 {
 		value := rand.Int() % 500