From ea69a5901008267e0a57d25c89b0086b0f9a46cd 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>
---
 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 323747937780..4dc968d94d37 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -194,6 +194,7 @@ func TestIntegration(t *testing.T) {
 		testMountStubsDirectory,
 		testMountStubsTimestamp,
 		testSourcePolicy,
+		testLLBMountPerformance,
 	)
 }
 
@@ -8991,3 +8992,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