Move downloading of data dependencies

.circleci/rex_test.sh

@@ -33,3 +33,17 @@ plz test -o build.passenv:PATH --profile ci_remote -p -v notice --colour //src/c
 # And run any tests we deem to be pertinent to remote execution
 echo "Testing anything labeled rex..."
 plz test -o build.passenv:PATH --profile ci_remote -p -v notice --colour -i rex
+
+# Check that a target's runtime data gets downloaded along with it.
+# src/plz builds runtime & data deps in parallel with the target itself, so
+# //test/remote_data:slow_data deliberately takes longer to build than the target that
+# depends on it. Both are named explicitly with --rebuild so that slow_data really does
+# execute rather than being served from the action cache, which would let it win the race.
+echo "Testing download of runtime data..."
+rm -rf plz-out/gen/test/remote_data
+plz build -o build.passenv:PATH --profile ci_remote -p -v notice --colour --rebuild \
+    //test/remote_data:needs_data //test/remote_data:slow_data
+if [ ! -f plz-out/gen/test/remote_data/slow_data.txt ]; then
+    echo "Runtime data for //test/remote_data:needs_data wasn't downloaded"
+    exit 1
+fi

src/build/build_step.go

@@ -216,6 +216,9 @@ func buildTarget(state *core.BuildState, target *core.BuildTarget, runRemotely b
 		state.LogBuildResult(target, core.TargetBuilding, "Preparing...")
 
 		// Ensure we have downloaded any previous dependencies if that's relevant.
+		if state.RemoteClient != nil {
+			state.LogBuildResult(target, core.TargetBuilding, "Downloading inputs...")
+		}
 		if err := state.DownloadInputsIfNeeded(target, false); err != nil {
 			return err
 		}

src/cmap/cerrmap.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"iter"
+	"runtime/debug"
 )
 
 // A Limiter is the interface that we use to release/acquire workers while waiting.
@@ -59,17 +60,18 @@ func (m *ErrMap[K, V]) Get(key K) (V, error) {
 // GetOrSet returns the value if set, or an error if one has been set.
 // If nothing has been set for the key, it runs the given function to generate the value and then sets it.
 // If the called function panics, the panic will be recovered and treated as though it returned an error.
-func (m *ErrMap[K, V]) GetOrSet(key K, f func() (V, error)) (V, error) {
+func (m *ErrMap[K, V]) GetOrSet(key K, f func() (V, error)) (val V, err error) {
 	v, wait, first := m.m.GetOrWait(key)
 	if v.Err != nil {
 		return v.Val, v.Err
 	} else if first {
 		defer func() {
 			if r := recover(); r != nil {
-				m.m.Set(key, errV[V]{Err: fmt.Errorf("%s", r)})
+				err = panicToErr(r)
+				m.m.Set(key, errV[V]{Err: err})
 			}
 		}()
-		val, err := f()
+		val, err = f()
 		m.m.Set(key, errV[V]{Val: val, Err: err})
 		return val, err
 	} else if wait != nil {
@@ -84,19 +86,29 @@ func (m *ErrMap[K, V]) GetOrSet(key K, f func() (V, error)) (V, error) {
 	return v.Val, v.Err
 }
 
+// panicToErr converts a recovered panic value into an error, retaining its type if it already is an error.
+func panicToErr(r any) error {
+	if e, ok := r.(error); ok {
+		return e
+	}
+	// Maintain the stack; if something panicked that wasn't an error we will probably want it for debugging.
+	return fmt.Errorf("%v\n%s", r, debug.Stack())
+}
+
 // GetOrSetCtx is like GetOrSet but accepts a context that can be cancelled.
 // If the called function panics, the panic will be recovered and treated as though it returned an error.
-func (m *ErrMap[K, V]) GetOrSetCtx(ctx context.Context, key K, f func() (V, error)) (V, error) {
+func (m *ErrMap[K, V]) GetOrSetCtx(ctx context.Context, key K, f func() (V, error)) (val V, err error) {
 	v, wait, first := m.m.GetOrWait(key)
 	if v.Err != nil {
 		return v.Val, v.Err
 	} else if first {
 		defer func() {
 			if r := recover(); r != nil {
-				m.m.Set(key, errV[V]{Err: fmt.Errorf("%s", r)})
+				err = panicToErr(r)
+				m.m.Set(key, errV[V]{Err: err})
 			}
 		}()
-		val, err := f()
+		val, err = f()
 		m.m.Set(key, errV[V]{Val: val, Err: err})
 		return val, err
 	} else if wait != nil {

src/cmap/cerrmap_test.go

@@ -34,6 +34,16 @@ func TestErrWait(t *testing.T) {
 	assert.Error(t, err)
 }
 
+func TestErrMapGetOrSetPanicReturnsErrorToFirstCaller(t *testing.T) {
+	m := NewErrMap[int, int](DefaultShardCount, hashInts, nil)
+	v, err := m.GetOrSet(9, func() (int, error) {
+		panic(fmt.Errorf("it broke"))
+	})
+	assert.Zero(t, v)
+	assert.Error(t, err)
+	assert.Equal(t, "it broke", err.Error())
+}
+
 func TestValues(t *testing.T) {
 	m := NewErrMap[int, int](DefaultShardCount, hashInts, nil)
 	broke := fmt.Errorf("it broke")

src/core/build_target.go

@@ -1422,6 +1422,9 @@ func (target *BuildTarget) AddDebugTool(tool BuildInput) {
 
 // AllTestTools returns all the test tool paths for this rule.
 func (target *BuildTarget) AllTestTools() []BuildInput {
+	if target.Test == nil {
+		return nil
+	}
 	if target.Test.namedTools == nil {
 		return target.Test.tools
 	}

src/core/state.go

@@ -257,7 +257,7 @@ func (state *BuildState) Initialise(subrepo *Subrepo) (err error) {
 		// handled for us already in plz.go
 		if state.CurrentSubrepo != "" {
 			state.RepoConfig = &Configuration{}
-			if err := readSubrepoConfig(state.RepoConfig, subrepo); err != nil {
+			if err = readSubrepoConfig(state.RepoConfig, subrepo); err != nil {
 				return
 			}
 			if err = validateSubrepoNameAndPluginConfig(state.Config, state.RepoConfig, subrepo); err != nil {
@@ -756,9 +756,15 @@ func (state *BuildState) GetPreloadedSubincludes() []BuildLabel {
 	done := map[BuildLabel]struct{}{}
 	includes := make([]BuildLabel, 0, len(state.Config.Parse.PreloadSubincludes)+len(state.RepoConfig.Parse.PreloadSubincludes))
 
-	is := append(state.Config.Parse.PreloadSubincludes, state.RepoConfig.Parse.PreloadSubincludes...)
+	for _, i := range state.Config.Parse.PreloadSubincludes {
+		if _, ok := done[i]; ok {
+			continue
+		}
 
-	for _, i := range is {
+		includes = append(includes, i)
+		done[i] = struct{}{}
+	}
+	for _, i := range state.RepoConfig.Parse.PreloadSubincludes {
 		if _, ok := done[i]; ok {
 			continue
 		}
@@ -772,7 +778,6 @@ func (state *BuildState) GetPreloadedSubincludes() []BuildLabel {
 // DownloadInputsIfNeeded downloads all the inputs (or runtime files) for a target if we are building remotely.
 func (state *BuildState) DownloadInputsIfNeeded(target *BuildTarget, runtime bool) error {
 	if state.RemoteClient != nil {
-		state.LogBuildResult(target, TargetBuilding, "Downloading inputs...")
 		for input := range state.IterInputs(target, runtime) {
 			if l, ok := input.Label(); ok {
 				dep := state.Graph.TargetOrDie(l)

src/plz/plz.go

@@ -179,6 +179,9 @@ func (r *runner) parse(ctx context.Context, label, dependent core.BuildLabel, qu
 		return nil, err
 	}
 	if err := r.ensurePreloads(ctx, state); err != nil {
+		if !quiet {
+			r.state.LogBuildError(label, core.ParseFailed, err, "Failed to resolve preloaded subincludes")
+		}
 		return nil, err
 	}
 	pkg, wait, first, err := r.state.Graph.PackageOrWait(label)
@@ -288,7 +291,7 @@ func (r *runner) registerPreloads(ctx context.Context, state *core.BuildState) e
 		}
 		// Queue them up asynchronously to feed the queues as quickly as possible
 		g.Go(func() error {
-			if _, err := r.Build(gctx, inc, core.OriginalTarget); err != nil {
+			if _, err := r.BuildAndDownload(gctx, inc, core.OriginalTarget); err != nil {
 				return err
 			}
 			return r.parser.PreloadSubinclude(gctx, inc)
@@ -322,6 +325,10 @@ func (r *runner) ensureSubrepo(ctx context.Context, label, dependent core.BuildL
 	if r.state.Graph.Subrepo(label.Subrepo) != nil {
 		return nil // The parse above defined it, we're done.
 	}
+	if arch, ok := couldBeArch(label.Subrepo); ok {
+		r.state.Graph.MaybeAddSubrepo(core.SubrepoForArch(r.state, arch))
+		return nil
+	}
 	if sl.Subrepo != dependent.Subrepo {
 		nested := sl
 		nested.Subrepo = dependent.Subrepo
@@ -334,11 +341,7 @@ func (r *runner) ensureSubrepo(ctx context.Context, label, dependent core.BuildL
 			}
 		}
 	}
-	// Nothing defines it, so the only remaining possibility is an architecture subrepo.
-	if arch, ok := couldBeArch(label.Subrepo); ok {
-		r.state.Graph.MaybeAddSubrepo(core.SubrepoForArch(r.state, arch))
-		return nil
-	} else if err != nil {
+	if err != nil {
 		// This returns the missing build file error that we got from tryParse above.
 		return err
 	}
@@ -477,8 +480,35 @@ func (r *runner) buildDep(ctx context.Context, dep core.BuildLabel, target *core
 	return nil
 }
 
-// buildOne builds a single target (which cannot be a pseudo-label like :all)
+// buildOne builds a single target (which cannot be a pseudo-label like :all) and downloads
+// anything needed to run it, if we're building remotely and it's a target we're downloading.
 func (r *runner) buildOne(ctx context.Context, target *core.BuildTarget) error {
+	if err := r.buildTargetAndDeps(ctx, target); err != nil {
+		return err
+	}
+	return r.downloadRuntimeFiles(target)
+}
+
+// downloadRuntimeFiles downloads the files needed to run a target (i.e. itself and its run-time & data dependencies).
+func (r *runner) downloadRuntimeFiles(target *core.BuildTarget) error {
+	state := r.state.ForTarget(target)
+	if state.RemoteClient == nil || !state.ShouldDownload(target) {
+		return nil
+	}
+	limiter := r.limiter(r.anyRemote)
+	limiter.Acquire()
+	defer limiter.Release()
+
+	log.Debug("Downloading runtime files for %s", target)
+	if err := state.DownloadInputsIfNeeded(target, true); err != nil {
+		state.LogBuildError(target.Label, core.TargetBuildFailed, err, "Failed to download runtime files for %s: %s", target.Label, err)
+		return err
+	}
+	return nil
+}
+
+// buildTargetAndDeps builds a single target and all of its dependencies.
+func (r *runner) buildTargetAndDeps(ctx context.Context, target *core.BuildTarget) error {
 	g, gctx := r.group(ctx)
 	for dep := range target.BuildDependencyLabels() {
 		g.Go(func() error {

src/remote/remote.go

@@ -345,56 +345,9 @@ func (c *Client) Build(target *core.BuildTarget) (*core.BuildMetadata, error) {
 	}
 
 	c.setOutputsFromMetadata(target, metadata)
-
-	if c.state.ShouldDownload(target) {
-		c.state.LogBuildResult(target, core.TargetBuilding, "Downloading")
-		if err := c.Download(target); err != nil {
-			return metadata, err
-		}
-		// TODO(peterebden): Should these not just be part of Download()?
-		if err := c.downloadData(target); err != nil {
-			return metadata, err
-		}
-		if err := c.downloadRuntimeDependencies(target); err != nil {
-			return metadata, err
-		}
-	}
 	return metadata, nil
 }
 
-// downloadData downloads all the runtime data for a target, recursively.
-func (c *Client) downloadData(target *core.BuildTarget) error {
-	var g errgroup.Group
-	for _, datum := range target.AllData() {
-		if l, ok := datum.Label(); ok {
-			t := c.state.Graph.TargetOrDie(l)
-			g.Go(func() error {
-				if err := c.Download(t); err != nil {
-					return err
-				}
-				return c.downloadData(t)
-			})
-		}
-	}
-	return g.Wait()
-}
-
-// downloadRuntimeDependencies downloads all the runtime dependencies for a target.
-func (c *Client) downloadRuntimeDependencies(target *core.BuildTarget) error {
-	var g errgroup.Group
-	for runDep := range target.IterAllRuntimeDependencies(c.state.Graph) {
-		l, _ := runDep.Label()
-		t := c.state.Graph.TargetOrDie(l)
-		g.Go(func() error {
-			if err := c.Download(t); err != nil {
-				return err
-			}
-			return nil
-		})
-	}
-	return g.Wait()
-}
-
 // Run runs a target on the remote executors.
 func (c *Client) Run(target *core.BuildTarget) error {
 	if err := c.CheckInitialised(); err != nil {

src/remote/remote_test.go

@@ -265,6 +265,9 @@ func TestOutDirsSetOutsOnTarget(t *testing.T) {
 	c.state.Graph.AddTarget(outDirTarget)
 	_, err := c.Build(outDirTarget)
 	require.NoError(t, err)
+	// Build does not download anything itself; src/build does that once any post-build
+	// function has had its chance to alter the outputs.
+	require.NoError(t, c.Download(outDirTarget))
 
 	assert.Len(t, outDirTarget.Outputs(c.state.Graph), 2)
 	assert.ElementsMatch(t, []string{"foo.txt", "bar.txt"}, outDirTarget.Outputs(c.state.Graph))

src/test/test_step.go

@@ -584,6 +584,9 @@ func runtimeHash(state *core.BuildState, target *core.BuildTarget, runRemotely b
 	}
 	if target.Local {
 		// If the test is marked to run locally, download the inputs as we need these to calculate the runtime hash.
+		if state.RemoteClient != nil {
+			state.LogTestRunning(target, run, core.TargetTesting, "Downloading inputs...")
+		}
 		if err := state.DownloadInputsIfNeeded(target, true); err != nil {
 			return nil, err
 		}

test/remote_data/BUILD

@@ -0,0 +1,22 @@
+# Regression test for downloading a remotely built target's runtime data.
+#
+# These are only built by .circleci/rex_test.sh, hence "manual"; they're pointless
+# (and slow) without remote execution configured.
+#
+# The data dependency must take longer to build than the target that depends on it.
+# src/plz builds runtime & data deps in parallel with the target itself, so if the
+# data is quicker it wins the race and the ordering never gets exercised.
+genrule(
+    name = "slow_data",
+    outs = ["slow_data.txt"],
+    cmd = "sleep 5 && echo hello > $OUT",
+    labels = ["manual"],
+)
+
+genrule(
+    name = "needs_data",
+    outs = ["needs_data.txt"],
+    cmd = "echo hello > $OUT",
+    data = [":slow_data"],
+    labels = ["manual"],
+)