From b189652eaec8038100b747445799ecf38e92c73e Mon Sep 17 00:00:00 2001
From: Jonathan Poole <jpoole@thoughtmachine.net>
Date: Thu, 28 Nov 2024 16:19:54 +0000
Subject: [PATCH] Fix litter deadlock with queuing data for tests

---
 src/build/build_step.go  |  3 ---
 src/core/build_target.go |  4 ++--
 src/core/state.go        |  4 ++--
 src/plz/plz.go           | 44 ++++++++++++++++++++++++++++++----------
 4 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/src/build/build_step.go b/src/build/build_step.go
index 8a99d0eb9..f01e1a5e9 100644
--- a/src/build/build_step.go
+++ b/src/build/build_step.go
@@ -80,9 +80,6 @@ func Build(state *core.BuildState, target *core.BuildTarget, remote bool) {
 	}
 	// Mark the target as having finished building.
 	target.FinishBuild()
-	if target.IsTest() && state.NeedTests && state.IsOriginalTarget(target) {
-		state.QueueTestTarget(target)
-	}
 }
 
 func validateBuildTargetBeforeBuild(state *core.BuildState, target *core.BuildTarget) error {
diff --git a/src/core/build_target.go b/src/core/build_target.go
index 927878e43..9978e2f85 100644
--- a/src/core/build_target.go
+++ b/src/core/build_target.go
@@ -712,8 +712,8 @@ func (target *BuildTarget) FinishBuild() {
 }
 
 // WaitForBuild blocks until this target has finished building.
-func (target *BuildTarget) WaitForBuild() {
-	<-target.finishedBuilding
+func (target *BuildTarget) WaitForBuild(dependant BuildLabel) {
+	waitOnChan(target.finishedBuilding, "Still waiting on (target %v).WaitForBuild(dependant %v)", target.Label, dependant)
 }
 
 // DeclaredOutputs returns the outputs from this target's original declaration.
diff --git a/src/core/state.go b/src/core/state.go
index f53861cf4..f98b22881 100644
--- a/src/core/state.go
+++ b/src/core/state.go
@@ -904,7 +904,7 @@ func (state *BuildState) WaitForBuiltTarget(l, dependent BuildLabel, mode ParseM
 		return make(chan struct{})
 	}); !inserted {
 		// Something's already registered for this, get on the train
-		waitOnChan(ch, "Still waiting on WaitForBuiltTarget(%v, %v, %v)", l, dependent, mode)
+		waitOnChan(ch, "Still waiting on WaitForBuiltTarget(label %v, dependant %v, ParseMode(%v))", l, dependent, mode)
 		return state.Graph.Target(l)
 	}
 	if err := state.queueTarget(l, dependent, mode.IsForSubinclude(), mode); err != nil {
@@ -1187,7 +1187,7 @@ func (state *BuildState) queueTargetAsync(target *BuildTarget, forceBuild, build
 		// Wait for these targets to actually build.
 		if building {
 			for _, t := range target.Dependencies() {
-				t.WaitForBuild()
+				t.WaitForBuild(target.Label)
 				if t.State() >= DependencyFailed { // Either the target failed or its dependencies failed
 					// Give up and set the original target as dependency failed
 					target.SetState(DependencyFailed)
diff --git a/src/plz/plz.go b/src/plz/plz.go
index 0eba00951..c164ce9c2 100644
--- a/src/plz/plz.go
+++ b/src/plz/plz.go
@@ -45,6 +45,32 @@ func Run(targets, preTargets []core.BuildLabel, state *core.BuildState, config *
 	remoteLimiter := make(limiter, config.NumRemoteExecutors())
 	anyRemote := config.NumRemoteExecutors() > 0
 
+	completeAction := func(remote bool, task core.Task) {
+		if remote {
+			remoteLimiter.Release()
+		} else {
+			localLimiter.Release()
+		}
+
+		if task.Type != core.BuildTask {
+			state.TaskDone()
+			return
+		}
+
+		if state.NeedTests && task.Target.IsTest() && state.IsOriginalTarget(task.Target) {
+			state.QueueTestTarget(task.Target)
+		}
+		state.TaskDone()
+	}
+
+	startAction := func(remote bool) {
+		if remote {
+			remoteLimiter.Acquire()
+		} else {
+			localLimiter.Acquire()
+		}
+	}
+
 	// Start up all the build workers
 	var wg sync.WaitGroup
 	wg.Add(2)
@@ -64,21 +90,17 @@ func Run(targets, preTargets []core.BuildLabel, state *core.BuildState, config *
 			wg.Add(1)
 			go func(task core.Task) {
 				defer wg.Done()
-				remote := anyRemote && !task.Target.Local
-				if remote {
-					remoteLimiter.Acquire()
-					defer remoteLimiter.Release()
-				} else {
-					localLimiter.Acquire()
-					defer localLimiter.Release()
-				}
+
+				isRemote := anyRemote && !task.Target.Local
+				startAction(isRemote)
+				defer completeAction(isRemote, task)
+
 				switch task.Type {
 				case core.TestTask:
-					test.Test(state, task.Target, remote, int(task.Run))
+					test.Test(state, task.Target, isRemote, int(task.Run))
 				case core.BuildTask:
-					build.Build(state, task.Target, remote)
+					build.Build(state, task.Target, isRemote)
 				}
-				state.TaskDone()
 			}(task)
 		}
 		wg.Done()