From 90ff220dded60ef7aafde29f2e5274cb62afe71d Mon Sep 17 00:00:00 2001
From: Erik Sipsma <erik@sipsma.dev>
Date: Wed, 22 Mar 2023 12:32:38 -0700
Subject: [PATCH] llbsolver: Fix performance of recomputeDigests

Before this, in the case where nothing was mutated the visited memo
would never be updated, thus causing exponential complexity.

Now the memo is updated even when nothing is mutated, just setting old
and new to be the same digest.

Signed-off-by: Erik Sipsma <erik@sipsma.dev>
(cherry picked from commit ea69a5901008267e0a57d25c89b0086b0f9a46cd)
---
 client/client_test.go      | 29 +++++++++++++++++++++++++++++
 solver/llbsolver/vertex.go |  3 ++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/client/client_test.go b/client/client_test.go
index 6ca36d2a71a6..348d81070f96 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -195,6 +195,7 @@ func TestIntegration(t *testing.T) {
 		testMountStubsDirectory,
 		testMountStubsTimestamp,
 		testSourcePolicy,
+		testLLBMountPerformance,
 	)
 }
 
@@ -8945,3 +8946,31 @@ func testSourcePolicy(t *testing.T, sb integration.Sandbox) {
 		require.ErrorContains(t, err, sourcepolicy.ErrSourceDenied.Error())
 	})
 }
+
+func testLLBMountPerformance(t *testing.T, sb integration.Sandbox) {
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	mntInput := llb.Image("busybox:latest")
+	st := llb.Image("busybox:latest")
+	var mnts []llb.State
+	for i := 0; i < 20; i++ {
+		execSt := st.Run(
+			llb.Args([]string{"true"}),
+		)
+		mnts = append(mnts, mntInput)
+		for j := range mnts {
+			mnts[j] = execSt.AddMount(fmt.Sprintf("/tmp/bin%d", j), mnts[j], llb.SourcePath("/bin"))
+		}
+		st = execSt.Root()
+	}
+
+	def, err := st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	timeoutCtx, cancel := context.WithTimeout(sb.Context(), time.Minute)
+	defer cancel()
+	_, err = c.Solve(timeoutCtx, def, SolveOpt{}, nil)
+	require.NoError(t, err)
+}
diff --git a/solver/llbsolver/vertex.go b/solver/llbsolver/vertex.go
index 6901332d2b65..41a31bb9bbba 100644
--- a/solver/llbsolver/vertex.go
+++ b/solver/llbsolver/vertex.go
@@ -210,6 +210,7 @@ func recomputeDigests(ctx context.Context, all map[digest.Digest]*pb.Op, visited
 	}
 
 	if !mutated {
+		visited[dgst] = dgst
 		return dgst, nil
 	}
 
@@ -274,7 +275,7 @@ func loadLLB(ctx context.Context, def *pb.Definition, polEngine SourcePolicyEval
 
 	for {
 		newDgst, ok := mutatedDigests[lastDgst]
-		if !ok {
+		if !ok || newDgst == lastDgst {
 			break
 		}
 		lastDgst = newDgst