Enable plz build --shell` to work correctly with remote execution a…

…nd subrepos (thought-machine#3066)

* Enable `plz build `--shell` to work correctly with remote execution and subrepos

Instead of trying to move files around after everything is built, we instead
simply download all the original target's inputs from CAS when running in
remote execution mode.

* appease the linter

* Ensure everything is in CAS before attempting to download the inputs

* simplify

src/build/build_step.go

@@ -123,10 +123,17 @@ func prepareOnly(state *core.BuildState, target *core.BuildTarget) error {
 
 		return fmt.Errorf("can't prepare temporary directory for %s; filegroups don't have temporary directories", target.Label)
 	}
-	// Ensure we have downloaded any previous dependencies if that's relevant.
-	if err := state.DownloadInputsIfNeeded(target, false); err != nil {
-		return err
+
+	if state.RemoteClient != nil {
+		// Targets were built remotely so we can simply download the inputs and place them in the
+		// tmp/ folder and exit.
+		if err := state.DownloadAllInputs(target, target.TmpDir(), false); err != nil {
+			return err
+		}
+
+		return errStop
 	}
+
 	if err := prepareDirectories(state.ProcessExecutor, target); err != nil {
 		return err
 	}

src/core/state.go

@@ -107,6 +107,9 @@ type RemoteClient interface {
 	Run(target *BuildTarget) error
 	// Download downloads the outputs for the given target that has already been built remotely.
 	Download(target *BuildTarget) error
+	// DownloadInputs downloads the whole of inputs folder for the given target that has already
+	// been built remotely, into the target directory
+	DownloadInputs(target *BuildTarget, targetDir string, isTest bool) error
 	// PrintHashes shows the hashes of a target.
 	PrintHashes(target *BuildTarget, isTest bool)
 	// DataRate returns an estimate of the current in/out RPC data rates and totals so far in bytes per second.
@@ -1298,6 +1301,7 @@ 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 {
+	// TODO(jan): Remove this function once `DownloadAllInputs` has been fully implemented across both build & test.
 	if state.RemoteClient != nil {
 		state.LogBuildResult(target, TargetBuilding, "Downloading inputs...")
 		for input := range state.IterInputs(target, runtime) {
@@ -1314,6 +1318,12 @@ func (state *BuildState) DownloadInputsIfNeeded(target *BuildTarget, runtime boo
 	return nil
 }
 
+// DownloadAllInputs downloads all inputs (including sources) for the target. Assumes remote execution.
+func (state *BuildState) DownloadAllInputs(target *BuildTarget, targetDir string, isTest bool) error {
+	state.LogBuildResult(target, TargetBuilding, "Downloading inputs...")
+	return state.RemoteClient.DownloadInputs(target, targetDir, isTest)
+}
+
 // IterInputs returns a channel that iterates all the input files needed for a target.
 func (state *BuildState) IterInputs(target *BuildTarget, test bool) <-chan BuildInput {
 	if !test {

src/remote/remote.go

@@ -515,6 +515,37 @@ func (c *Client) downloadActionOutputs(ctx context.Context, ar *pb.ActionResult,
 	return nil
 }
 
+// DownloadInputs downloads all the inputs of a remotely built target into the target directory.
+func (c *Client) DownloadInputs(target *core.BuildTarget, targetDir string, isTest bool) error {
+	var pbDir *pb.Directory
+	// first ensure all inputs are in the CAS
+	err := c.uploadBlobs(func(ch chan<- *uploadinfo.Entry) error {
+		defer close(ch)
+
+		var err error
+		pbDir, err = c.uploadInputs(ch, target, isTest)
+		return err
+	})
+	if err != nil {
+		return err
+	}
+
+	dirDigest, err := digest.NewFromMessage(pbDir)
+	if err != nil {
+		return fmt.Errorf("could not calculate digest for directory proto: %w", err)
+	}
+
+	if err := os.RemoveAll(targetDir); err != nil {
+		return fmt.Errorf("could not delete target directory %q: %w", targetDir, err)
+	}
+
+	if _, _, err = c.client.DownloadDirectory(context.Background(), dirDigest, targetDir, c.fileMetadataCache); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // moveTmpFilesToOutDir moves files from the target tmp dir to the out dir, handling output directories as well
 func moveTmpFilesToOutDir(target *core.BuildTarget) error {
 	defer os.RemoveAll(target.TmpDir())