From 100d3cb6b6903be50f7a3e5dba193515aa9530fa Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Wed, 27 Sep 2023 17:10:19 -0700
Subject: [PATCH] solver: fix issue with double merged edges

"Merged edges" is an optimization in the solve graph where two different active LLB edges can be combined into one after cache key computation finds that they generated equivalent cache chain.

In that case, one edge is released and is set to point to another. The reference count for the second one is increased.

An issue was discovered where an edge that was already pointing to another edge became a target to the third one. The current implementation did not handle the case where an edge that already had a different target itself became a target edge as well. This resulted in an inconsistent graph state where edges that were thought to be released could get scheduled again.

Instead of setting the same edge value to two different maps, the new logic is to chain the edges into a linked list that should handle multiple levels of targets. This slightly increases the number of lookups, but "merged edges" case is very rare anyway, and a couple of extra pointer dereferences do not affect build speed.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 solver/edge.go | 13 ++++++++++---
 solver/jobs.go | 21 +++++++++++++++------
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/solver/edge.go b/solver/edge.go
index 3e4ec18242fd..7595f317ee03 100644
--- a/solver/edge.go
+++ b/solver/edge.go
@@ -64,6 +64,7 @@ type edge struct {
 	hasActiveOutgoing         bool
 
 	releaserCount int
+	owner         *edge
 	keysDidChange bool
 	index         *edgeIndex
 
@@ -116,10 +117,12 @@ type edgeRequest struct {
 	currentKeys  int
 }
 
-// incrementReferenceCount increases the number of times release needs to be
+// takeOwnership increases the number of times release needs to be
 // called to release the edge. Called on merging edges.
-func (e *edge) incrementReferenceCount() {
-	e.releaserCount++
+func (e *edge) takeOwnership(old *edge) {
+	e.releaserCount += old.releaserCount + 1
+	old.owner = e
+	old.releaseResult()
 }
 
 // release releases the edge resources
@@ -128,6 +131,10 @@ func (e *edge) release() {
 		e.releaserCount--
 		return
 	}
+	e.releaseResult()
+}
+
+func (e *edge) releaseResult() {
 	e.index.Release(e)
 	if e.result != nil {
 		go e.result.Release(context.TODO())
diff --git a/solver/jobs.go b/solver/jobs.go
index 5ef8205953f7..67d53e2acccb 100644
--- a/solver/jobs.go
+++ b/solver/jobs.go
@@ -141,6 +141,9 @@ func (s *state) getEdge(index Index) *edge {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if e, ok := s.edges[index]; ok {
+		for e.owner != nil {
+			e = e.owner
+		}
 		return e
 	}
 
@@ -153,19 +156,22 @@ func (s *state) getEdge(index Index) *edge {
 	return e
 }
 
-func (s *state) setEdge(index Index, newEdge *edge) {
+func (s *state) setEdge(index Index, targetEdge *edge) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	e, ok := s.edges[index]
 	if ok {
-		if e == newEdge {
+		for e.owner != nil {
+			e = e.owner
+		}
+		if e == targetEdge {
 			return
 		}
-		e.release()
+	} else {
+		e = newEdge(Edge{Index: index, Vertex: s.vtx}, s.op, s.index)
+		s.edges[index] = e
 	}
-
-	newEdge.incrementReferenceCount()
-	s.edges[index] = newEdge
+	targetEdge.takeOwnership(e)
 }
 
 func (s *state) combinedCacheManager() CacheManager {
@@ -186,6 +192,9 @@ func (s *state) combinedCacheManager() CacheManager {
 
 func (s *state) Release() {
 	for _, e := range s.edges {
+		for e.owner != nil {
+			e = e.owner
+		}
 		e.release()
 	}
 	if s.op != nil {