[DO NOT MERGE] Add support for experimentally enabling RBE support on specific rules.

This CL adds RBE support to javac, r8, and d8 rules which is only
enabled if respective environment variables are set.

Test: an aosp_crosshatch build with and without the new variables.

Bug: b/166182389
Change-Id: Ic82f3627944f6a5ee7b9f3228170c2709b1bfcb8
Merged-In: Ic82f3627944f6a5ee7b9f3228170c2709b1bfcb8
diff --git a/android/config.go b/android/config.go
index 7899119..2e0c247 100644
--- a/android/config.go
+++ b/android/config.go
@@ -718,6 +718,22 @@
 	return Bool(c.productVariables.UseRBE)
 }
 
+func (c *config) UseRBEJAVAC() bool {
+	return Bool(c.productVariables.UseRBEJAVAC)
+}
+
+func (c *config) UseRBER8() bool {
+	return Bool(c.productVariables.UseRBER8)
+}
+
+func (c *config) UseRBED8() bool {
+	return Bool(c.productVariables.UseRBED8)
+}
+
+func (c *config) UseRemoteBuild() bool {
+	return c.UseGoma() || c.UseRBE()
+}
+
 func (c *config) RunErrorProne() bool {
 	return c.IsEnvTrue("RUN_ERROR_PRONE")
 }
diff --git a/android/package_ctx.go b/android/package_ctx.go
index cf8face..ff10fe7 100644
--- a/android/package_ctx.go
+++ b/android/package_ctx.go
@@ -254,16 +254,32 @@
 	}, argNames...)
 }
 
-// RemoteRuleSupports selects if a AndroidRemoteStaticRule supports goma, RBE, or both.
-type RemoteRuleSupports int
+// RBEExperimentalFlag indicates which flag should be set for the AndroidRemoteStaticRule
+// to use RBE.
+type RBEExperimentalFlag int
 
 const (
-	SUPPORTS_NONE = 0
-	SUPPORTS_GOMA = 1 << iota
-	SUPPORTS_RBE  = 1 << iota
-	SUPPORTS_BOTH = SUPPORTS_GOMA | SUPPORTS_RBE
+	// RBE_NOT_EXPERIMENTAL indicates the rule should use RBE in every build that has
+	// UseRBE set.
+	RBE_NOT_EXPERIMENTAL RBEExperimentalFlag = iota
+	// RBE_JAVAC indicates the rule should use RBE only if the RBE_JAVAC variable is
+	// set in an RBE enabled build.
+	RBE_JAVAC
+	// RBE_R8 indicates the rule should use RBE only if the RBE_R8 variable is set in
+	// an RBE enabled build.
+	RBE_R8
+	// RBE_D8 indicates the rule should use RBE only if the RBE_D8 variable is set in
+	// an RBE enabled build.
+	RBE_D8
 )
 
+// RemoteRuleSupports configures rules with whether they have Goma and/or RBE support.
+type RemoteRuleSupports struct {
+	Goma    bool
+	RBE     bool
+	RBEFlag RBEExperimentalFlag
+}
+
 // AndroidRemoteStaticRule wraps blueprint.StaticRule but uses goma or RBE's parallelism if goma or RBE are enabled
 // and the appropriate SUPPORTS_* flag is set.
 func (p PackageContext) AndroidRemoteStaticRule(name string, supports RemoteRuleSupports, params blueprint.RuleParams,
@@ -271,18 +287,30 @@
 
 	return p.PackageContext.RuleFunc(name, func(config interface{}) (blueprint.RuleParams, error) {
 		ctx := &configErrorWrapper{p, config.(Config), nil}
-		if ctx.Config().UseGoma() && supports&SUPPORTS_GOMA == 0 {
+		if ctx.Config().UseGoma() && !supports.Goma {
 			// When USE_GOMA=true is set and the rule is not supported by goma, restrict jobs to the
 			// local parallelism value
 			params.Pool = localPool
 		}
 
-		if ctx.Config().UseRBE() && supports&SUPPORTS_RBE == 0 {
+		if ctx.Config().UseRBE() && !supports.RBE {
 			// When USE_RBE=true is set and the rule is not supported by RBE, restrict jobs to the
 			// local parallelism value
 			params.Pool = localPool
 		}
 
+		if ctx.Config().UseRBE() && supports.RBE {
+			if supports.RBEFlag == RBE_JAVAC && !ctx.Config().UseRBEJAVAC() {
+				params.Pool = localPool
+			}
+			if supports.RBEFlag == RBE_R8 && !ctx.Config().UseRBER8() {
+				params.Pool = localPool
+			}
+			if supports.RBEFlag == RBE_D8 && !ctx.Config().UseRBED8() {
+				params.Pool = localPool
+			}
+		}
+
 		return params, nil
 	}, argNames...)
 }
diff --git a/android/variable.go b/android/variable.go
index 28a8d0c..a88fc9d 100644
--- a/android/variable.go
+++ b/android/variable.go
@@ -193,6 +193,9 @@
 	Binder32bit                      *bool `json:",omitempty"`
 	UseGoma                          *bool `json:",omitempty"`
 	UseRBE                           *bool `json:",omitempty"`
+	UseRBEJAVAC                      *bool `json:",omitempty"`
+	UseRBER8                         *bool `json:",omitempty"`
+	UseRBED8                         *bool `json:",omitempty"`
 	Debuggable                       *bool `json:",omitempty"`
 	Eng                              *bool `json:",omitempty"`
 	Treble_linker_namespaces         *bool `json:",omitempty"`
diff --git a/bpf/bpf.go b/bpf/bpf.go
index 85eb7da..0047636 100644
--- a/bpf/bpf.go
+++ b/bpf/bpf.go
@@ -33,7 +33,7 @@
 var (
 	pctx = android.NewPackageContext("android/soong/bpf")
 
-	cc = pctx.AndroidRemoteStaticRule("cc", android.SUPPORTS_GOMA,
+	cc = pctx.AndroidRemoteStaticRule("cc", android.RemoteRuleSupports{Goma: true},
 		blueprint.RuleParams{
 			Depfile:     "${out}.d",
 			Deps:        blueprint.DepsGCC,
diff --git a/cc/builder.go b/cc/builder.go
index 773e7e8..972f967 100644
--- a/cc/builder.go
+++ b/cc/builder.go
@@ -46,7 +46,7 @@
 var (
 	pctx = android.NewPackageContext("android/soong/cc")
 
-	cc = pctx.AndroidRemoteStaticRule("cc", android.SUPPORTS_BOTH,
+	cc = pctx.AndroidRemoteStaticRule("cc", android.RemoteRuleSupports{Goma: true, RBE: true},
 		blueprint.RuleParams{
 			Depfile:     "${out}.d",
 			Deps:        blueprint.DepsGCC,
@@ -55,7 +55,7 @@
 		},
 		"ccCmd", "cFlags")
 
-	ccNoDeps = pctx.AndroidRemoteStaticRule("ccNoDeps", android.SUPPORTS_GOMA,
+	ccNoDeps = pctx.AndroidRemoteStaticRule("ccNoDeps", android.RemoteRuleSupports{Goma: true},
 		blueprint.RuleParams{
 			Command:     "$relPwd ${config.CcWrapper}$ccCmd -c $cFlags -o $out $in",
 			CommandDeps: []string{"$ccCmd"},
diff --git a/java/builder.go b/java/builder.go
index a8e54cc..e5ac7b0 100644
--- a/java/builder.go
+++ b/java/builder.go
@@ -38,7 +38,7 @@
 	// this, all java rules write into separate directories and then are combined into a .jar file
 	// (if the rule produces .class files) or a .srcjar file (if the rule produces .java files).
 	// .srcjar files are unzipped into a temporary directory when compiled with javac.
-	javac = pctx.AndroidRemoteStaticRule("javac", android.SUPPORTS_GOMA,
+	javac = pctx.AndroidRemoteStaticRule("javac", android.RemoteRuleSupports{Goma: true, RBE: true, RBEFlag: android.RBE_JAVAC},
 		blueprint.RuleParams{
 			Command: `rm -rf "$outDir" "$annoDir" "$srcJarDir" && mkdir -p "$outDir" "$annoDir" "$srcJarDir" && ` +
 				`${config.ZipSyncCmd} -d $srcJarDir -l $srcJarDir/list -f "*.java" $srcJars && ` +
diff --git a/java/config/config.go b/java/config/config.go
index f6688a2..f76a393 100644
--- a/java/config/config.go
+++ b/java/config/config.go
@@ -129,6 +129,20 @@
 		return ""
 	})
 
+	pctx.VariableFunc("R8Wrapper", func(ctx android.PackageVarContext) string {
+		if override := ctx.Config().Getenv("R8_WRAPPER"); override != "" {
+			return override + " "
+		}
+		return ""
+	})
+
+	pctx.VariableFunc("D8Wrapper", func(ctx android.PackageVarContext) string {
+		if override := ctx.Config().Getenv("D8_WRAPPER"); override != "" {
+			return override + " "
+		}
+		return ""
+	})
+
 	pctx.HostJavaToolVariable("JacocoCLIJar", "jacoco-cli.jar")
 
 	hostBinToolVariableWithPrebuilt := func(name, prebuiltDir, tool string) {
diff --git a/java/dex.go b/java/dex.go
index c8a4fa8..86a28fc 100644
--- a/java/dex.go
+++ b/java/dex.go
@@ -22,10 +22,10 @@
 	"android/soong/android"
 )
 
-var d8 = pctx.AndroidStaticRule("d8",
+var d8 = pctx.AndroidRemoteStaticRule("d8", android.RemoteRuleSupports{RBE: true, RBEFlag: android.RBE_D8},
 	blueprint.RuleParams{
 		Command: `rm -rf "$outDir" && mkdir -p "$outDir" && ` +
-			`${config.D8Cmd} ${config.DexFlags} --output $outDir $d8Flags $in && ` +
+			`${config.D8Wrapper}${config.D8Cmd} ${config.DexFlags} --output $outDir $d8Flags $in && ` +
 			`${config.SoongZipCmd} $zipFlags -o $outDir/classes.dex.jar -C $outDir -f "$outDir/classes*.dex" && ` +
 			`${config.MergeZipsCmd} -D -stripFile "**/*.class" $out $outDir/classes.dex.jar $in`,
 		CommandDeps: []string{
@@ -36,11 +36,11 @@
 	},
 	"outDir", "d8Flags", "zipFlags")
 
-var r8 = pctx.AndroidStaticRule("r8",
+var r8 = pctx.AndroidRemoteStaticRule("r8", android.RemoteRuleSupports{RBE: true, RBEFlag: android.RBE_R8},
 	blueprint.RuleParams{
 		Command: `rm -rf "$outDir" && mkdir -p "$outDir" && ` +
 			`rm -f "$outDict" && ` +
-			`${config.R8Cmd} ${config.DexFlags} -injars $in --output $outDir ` +
+			`${config.R8Wrapper}${config.R8Cmd} ${config.DexFlags} -injars $in --output $outDir ` +
 			`--force-proguard-compatibility ` +
 			`--no-data-resources ` +
 			`-printmapping $outDict ` +
diff --git a/java/kotlin.go b/java/kotlin.go
index 7256030..58dc64c 100644
--- a/java/kotlin.go
+++ b/java/kotlin.go
@@ -25,7 +25,7 @@
 	"github.com/google/blueprint"
 )
 
-var kotlinc = pctx.AndroidRemoteStaticRule("kotlinc", android.SUPPORTS_GOMA,
+var kotlinc = pctx.AndroidRemoteStaticRule("kotlinc", android.RemoteRuleSupports{Goma: true},
 	blueprint.RuleParams{
 		Command: `rm -rf "$classesDir" "$srcJarDir" "$kotlinBuildFile" && mkdir -p "$classesDir" "$srcJarDir" && ` +
 			`${config.ZipSyncCmd} -d $srcJarDir -l $srcJarDir/list -f "*.java" $srcJars && ` +
@@ -74,7 +74,7 @@
 	})
 }
 
-var kapt = pctx.AndroidRemoteStaticRule("kapt", android.SUPPORTS_GOMA,
+var kapt = pctx.AndroidRemoteStaticRule("kapt", android.RemoteRuleSupports{Goma: true},
 	blueprint.RuleParams{
 		Command: `rm -rf "$srcJarDir" "$kotlinBuildFile" "$kaptDir" && mkdir -p "$srcJarDir" "$kaptDir" && ` +
 			`${config.ZipSyncCmd} -d $srcJarDir -l $srcJarDir/list -f "*.java" $srcJars && ` +
diff --git a/ui/build/config.go b/ui/build/config.go
index b105dc0..daefb2f 100644
--- a/ui/build/config.go
+++ b/ui/build/config.go
@@ -503,6 +503,48 @@
 	return false
 }
 
+func (c *configImpl) UseRBEJAVAC() bool {
+	if !c.UseRBE() {
+		return false
+	}
+
+	if v, ok := c.environ.Get("RBE_JAVAC"); ok {
+		v = strings.TrimSpace(v)
+		if v != "" && v != "false" {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *configImpl) UseRBER8() bool {
+	if !c.UseRBE() {
+		return false
+	}
+
+	if v, ok := c.environ.Get("RBE_R8"); ok {
+		v = strings.TrimSpace(v)
+		if v != "" && v != "false" {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *configImpl) UseRBED8() bool {
+	if !c.UseRBE() {
+		return false
+	}
+
+	if v, ok := c.environ.Get("RBE_D8"); ok {
+		v = strings.TrimSpace(v)
+		if v != "" && v != "false" {
+			return true
+		}
+	}
+	return false
+}
+
 func (c *configImpl) StartRBE() bool {
 	if !c.UseRBE() {
 		return false
diff --git a/ui/build/dumpvars.go b/ui/build/dumpvars.go
index 3e387c1..2a6a9ca 100644
--- a/ui/build/dumpvars.go
+++ b/ui/build/dumpvars.go
@@ -178,6 +178,8 @@
 		"CC_WRAPPER",
 		"CXX_WRAPPER",
 		"JAVAC_WRAPPER",
+		"R8_WRAPPER",
+		"D8_WRAPPER",
 
 		// ccache settings
 		"CCACHE_COMPILERCHECK",