Add functionality to sandbox mixed build actions

The use case for this is for building rules_go's root builder which runs
into issues when built in a directory that contains a symlink to
prebuilts/go

The implementation will involve two changes of working dir
- `sbox` to change the working directory to
__SBOX_SANDBOX_DIR__
- the generated manifest will change the working
directory to mixed build execution root relative to that

Implemenation details
1. Create a unique intermediate path by hashing the outputs of a buildAction.
   "out/bazel/output/execroot/__main__/" was deliberately not chosen as
   the outpuDir for the sandbox because ruleBuilder would wipe it.
   `sbox` will generate the files in __SBOX_SANDBOX_DIR__ and then place
   the files in this intermediate directory.
2. After the files have been generated in (1), copy them to
   out/bazel/output/execroot/__main__/...
3. For bazel depsets that are inputs of an action, copy the direct
   artifacts into the sandbox instead of the phony target
4. Make sandboxing an opt-in. Currently we will only use it for
   `GoToolchainBinaryBuild`

In the current implementation, (3) will increase the size of the ninja
file. With sboxing turned on for only GoToolchainBinaryBuild, this will
increase the size of the ninja file by around 1.3% on aosp's cf

Test: m com.android.neuralnetworks (will build soong_zip from source
using rules_go)
Test: OUT_DIR=out.other m com.android.neuralnetworks
Bug: 289102849

Change-Id: I7addda9af583ba0ff306e50c1dfa16ed16c29799

android/bazel_handler.go

@@ -16,6 +16,8 @@ package android
 
 import (
 	"bytes"
+	"crypto/sha1"
+	"encoding/hex"
 	"fmt"
 	"os"
 	"path"
@@ -1222,7 +1224,11 @@ func (c *bazelSingleton) GenerateBuildActions(ctx SingletonContext) {
 		ctx.AddNinjaFileDeps(file)
 	}
 
+	depsetHashToDepset := map[string]bazel.AqueryDepset{}
+
 	for _, depset := range ctx.Config().BazelContext.AqueryDepsets() {
+		depsetHashToDepset[depset.ContentHash] = depset
+
 		var outputs []Path
 		var orderOnlies []Path
 		for _, depsetDepHash := range depset.TransitiveDepSetHashes {
@@ -1257,7 +1263,30 @@ func (c *bazelSingleton) GenerateBuildActions(ctx SingletonContext) {
 		}
 		if len(buildStatement.Command) > 0 {
 			rule := NewRuleBuilder(pctx, ctx)
-			createCommand(rule.Command(), buildStatement, executionRoot, bazelOutDir, ctx)
+			intermediateDir, intermediateDirHash := intermediatePathForSboxMixedBuildAction(ctx, buildStatement)
+			if buildStatement.ShouldRunInSbox {
+				// Create a rule to build the output inside a sandbox
+				// This will create two changes of working directory
+				// 1. From ANDROID_BUILD_TOP to sbox top
+				// 2. From sbox top to a a synthetic mixed build execution root relative to it
+				// Finally, the outputs will be copied to intermediateDir
+				rule.Sbox(intermediateDir,
+					PathForOutput(ctx, "mixed_build_sbox_intermediates", intermediateDirHash+".textproto")).
+					SandboxInputs().
+					// Since we will cd to mixed build execution root, set sbox's out subdir to empty
+					// Without this, we will try to copy from $SBOX_SANDBOX_DIR/out/out/bazel/output/execroot/__main__/...
+					SetSboxOutDirDirAsEmpty()
+
+				// Create another set of rules to copy files from the intermediate dir to mixed build execution root
+				for _, outputPath := range buildStatement.OutputPaths {
+					ctx.Build(pctx, BuildParams{
+						Rule:   CpIfChanged,
+						Input:  intermediateDir.Join(ctx, executionRoot, outputPath),
+						Output: PathForBazelOut(ctx, outputPath),
+					})
+				}
+			}
+			createCommand(rule.Command(), buildStatement, executionRoot, bazelOutDir, ctx, depsetHashToDepset)
 			desc := fmt.Sprintf("%s: %s", buildStatement.Mnemonic, buildStatement.OutputPaths)
 			rule.Build(fmt.Sprintf("bazel %d", index), desc)
 			continue
@@ -1304,10 +1333,25 @@ func (c *bazelSingleton) GenerateBuildActions(ctx SingletonContext) {
 	}
 }
 
+// Returns a out dir path for a sandboxed mixed build action
+func intermediatePathForSboxMixedBuildAction(ctx PathContext, statement *bazel.BuildStatement) (OutputPath, string) {
+	// An artifact can be generated by a single buildstatement.
+	// Use the hash of the first artifact to create a unique path
+	uniqueDir := sha1.New()
+	uniqueDir.Write([]byte(statement.OutputPaths[0]))
+	uniqueDirHashString := hex.EncodeToString(uniqueDir.Sum(nil))
+	return PathForOutput(ctx, "mixed_build_sbox_intermediates", uniqueDirHashString), uniqueDirHashString
+}
+
 // Register bazel-owned build statements (obtained from the aquery invocation).
-func createCommand(cmd *RuleBuilderCommand, buildStatement *bazel.BuildStatement, executionRoot string, bazelOutDir string, ctx BuilderContext) {
+func createCommand(cmd *RuleBuilderCommand, buildStatement *bazel.BuildStatement, executionRoot string, bazelOutDir string, ctx BuilderContext, depsetHashToDepset map[string]bazel.AqueryDepset) {
 	// executionRoot is the action cwd.
-	cmd.Text(fmt.Sprintf("cd '%s' &&", executionRoot))
+	if buildStatement.ShouldRunInSbox {
+		// mkdir -p ensures that the directory exists when run via sbox
+		cmd.Text(fmt.Sprintf("mkdir -p '%s' && cd '%s' &&", executionRoot, executionRoot))
+	} else {
+		cmd.Text(fmt.Sprintf("cd '%s' &&", executionRoot))
+	}
 
 	// Remove old outputs, as some actions might not rerun if the outputs are detected.
 	if len(buildStatement.OutputPaths) > 0 {
@@ -1334,14 +1378,30 @@ func createCommand(cmd *RuleBuilderCommand, buildStatement *bazel.BuildStatement
 	}
 
 	for _, outputPath := range buildStatement.OutputPaths {
-		cmd.ImplicitOutput(PathForBazelOut(ctx, outputPath))
+		if buildStatement.ShouldRunInSbox {
+			// The full path has three components that get joined together
+			// 1. intermediate output dir that `sbox` will place the artifacts at
+			// 2. mixed build execution root
+			// 3. artifact path returned by aquery
+			intermediateDir, _ := intermediatePathForSboxMixedBuildAction(ctx, buildStatement)
+			cmd.ImplicitOutput(intermediateDir.Join(ctx, executionRoot, outputPath))
+		} else {
+			cmd.ImplicitOutput(PathForBazelOut(ctx, outputPath))
+		}
 	}
 	for _, inputPath := range buildStatement.InputPaths {
 		cmd.Implicit(PathForBazelOut(ctx, inputPath))
 	}
 	for _, inputDepsetHash := range buildStatement.InputDepsetHashes {
-		otherDepsetName := bazelDepsetName(inputDepsetHash)
-		cmd.Implicit(PathForPhony(ctx, otherDepsetName))
+		if buildStatement.ShouldRunInSbox {
+			// Bazel depsets are phony targets that are used to group files.
+			// We need to copy the grouped files into the sandbox
+			ds, _ := depsetHashToDepset[inputDepsetHash]
+			cmd.Implicits(PathsForBazelOut(ctx, ds.DirectArtifacts))
+		} else {
+			otherDepsetName := bazelDepsetName(inputDepsetHash)
+			cmd.Implicit(PathForPhony(ctx, otherDepsetName))
+		}
 	}
 
 	if depfile := buildStatement.Depfile; depfile != nil {

android/bazel_handler_test.go

@@ -181,13 +181,62 @@ func TestInvokeBazelPopulatesBuildStatements(t *testing.T) {
 
 		cmd := RuleBuilderCommand{}
 		ctx := builderContextForTests{PathContextForTesting(TestConfig("out", nil, "", nil))}
-		createCommand(&cmd, got[0], "test/exec_root", "test/bazel_out", ctx)
+		createCommand(&cmd, got[0], "test/exec_root", "test/bazel_out", ctx, map[string]bazel.AqueryDepset{})
 		if actual, expected := cmd.buf.String(), testCase.command; expected != actual {
 			t.Errorf("expected: [%s], actual: [%s]", expected, actual)
 		}
 	}
 }
 
+func TestMixedBuildSandboxedAction(t *testing.T) {
+	input := `{
+ "artifacts": [
+   { "id": 1, "path_fragment_id": 1 },
+   { "id": 2, "path_fragment_id": 2 }],
+ "actions": [{
+   "target_Id": 1,
+   "action_Key": "x",
+   "mnemonic": "x",
+   "arguments": ["touch", "foo"],
+   "input_dep_set_ids": [1],
+   "output_Ids": [1],
+   "primary_output_id": 1
+ }],
+ "dep_set_of_files": [
+   { "id": 1, "direct_artifact_ids": [1, 2] }],
+ "path_fragments": [
+   { "id": 1, "label": "one" },
+   { "id": 2, "label": "two" }]
+}`
+	data, err := JsonToActionGraphContainer(input)
+	if err != nil {
+		t.Error(err)
+	}
+	bazelContext, _ := testBazelContext(t, map[bazelCommand]string{aqueryCmd: string(data)})
+
+	err = bazelContext.InvokeBazel(testConfig, &testInvokeBazelContext{})
+	if err != nil {
+		t.Fatalf("TestMixedBuildSandboxedAction did not expect error invoking Bazel, but got %s", err)
+	}
+
+	statement := bazelContext.BuildStatementsToRegister()[0]
+	statement.ShouldRunInSbox = true
+
+	cmd := RuleBuilderCommand{}
+	ctx := builderContextForTests{PathContextForTesting(TestConfig("out", nil, "", nil))}
+	createCommand(&cmd, statement, "test/exec_root", "test/bazel_out", ctx, map[string]bazel.AqueryDepset{})
+	// Assert that the output is generated in an intermediate directory
+	// fe05bcdcdc4928012781a5f1a2a77cbb5398e106 is the sha1 checksum of "one"
+	if actual, expected := cmd.outputs[0].String(), "out/soong/mixed_build_sbox_intermediates/fe05bcdcdc4928012781a5f1a2a77cbb5398e106/test/exec_root/one"; expected != actual {
+		t.Errorf("expected: [%s], actual: [%s]", expected, actual)
+	}
+
+	// Assert the actual command remains unchanged inside the sandbox
+	if actual, expected := cmd.buf.String(), "mkdir -p 'test/exec_root' && cd 'test/exec_root' && rm -rf 'one' && touch foo"; expected != actual {
+		t.Errorf("expected: [%s], actual: [%s]", expected, actual)
+	}
+}
+
 func TestCoverageFlagsAfterInvokeBazel(t *testing.T) {
 	testConfig.productVariables.ClangCoverage = boolPtr(true)
 

android/rule_builder.go

@@ -53,6 +53,7 @@ type RuleBuilder struct {
 	remoteable       RemoteRuleSupports
 	rbeParams        *remoteexec.REParams
 	outDir           WritablePath
+	sboxOutSubDir    string
 	sboxTools        bool
 	sboxInputs       bool
 	sboxManifestPath WritablePath
@@ -65,9 +66,18 @@ func NewRuleBuilder(pctx PackageContext, ctx BuilderContext) *RuleBuilder {
 		pctx:           pctx,
 		ctx:            ctx,
 		temporariesSet: make(map[WritablePath]bool),
+		sboxOutSubDir:  sboxOutSubDir,
 	}
 }
 
+// SetSboxOutDirDirAsEmpty sets the out subdirectory to an empty string
+// This is useful for sandboxing actions that change the execution root to a path in out/ (e.g mixed builds)
+// For such actions, SetSboxOutDirDirAsEmpty ensures that the path does not become $SBOX_SANDBOX_DIR/out/out/bazel/output/execroot/__main__/...
+func (rb *RuleBuilder) SetSboxOutDirDirAsEmpty() *RuleBuilder {
+	rb.sboxOutSubDir = ""
+	return rb
+}
+
 // RuleBuilderInstall is a tuple of install from and to locations.
 type RuleBuilderInstall struct {
 	From Path
@@ -585,7 +595,7 @@ func (r *RuleBuilder) Build(name string, desc string) {
 		for _, output := range outputs {
 			rel := Rel(r.ctx, r.outDir.String(), output.String())
 			command.CopyAfter = append(command.CopyAfter, &sbox_proto.Copy{
-				From: proto.String(filepath.Join(sboxOutSubDir, rel)),
+				From: proto.String(filepath.Join(r.sboxOutSubDir, rel)),
 				To:   proto.String(output.String()),
 			})
 		}

bazel/aquery.go

@@ -116,6 +116,9 @@ type BuildStatement struct {
 	InputDepsetHashes []string
 	InputPaths        []string
 	FileContents      string
+	// If ShouldRunInSbox is true, Soong will use sbox to created an isolated environment
+	// and run the mixed build action there
+	ShouldRunInSbox bool
 }
 
 // A helper type for aquery processing which facilitates retrieval of path IDs from their
@@ -496,6 +499,12 @@ func (a *aqueryArtifactHandler) normalActionBuildStatement(actionEntry *analysis
 		Env:               actionEntry.EnvironmentVariables,
 		Mnemonic:          actionEntry.Mnemonic,
 	}
+	if buildStatement.Mnemonic == "GoToolchainBinaryBuild" {
+		// Unlike b's execution root, mixed build execution root contains a symlink to prebuilts/go
+		// This causes issues for `GOCACHE=$(mktemp -d) go build ...`
+		// To prevent this, sandbox this action in mixed builds as well
+		buildStatement.ShouldRunInSbox = true
+	}
 	return buildStatement, nil
 }