Merge "Add support for experimentally enabling RBE support on specific rules."
diff --git a/android/config.go b/android/config.go
index bab3477..3dae6e2 100644
--- a/android/config.go
+++ b/android/config.go
@@ -816,6 +816,18 @@
 	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()
 }
diff --git a/android/package_ctx.go b/android/package_ctx.go
index 6350667..5a43ea9 100644
--- a/android/package_ctx.go
+++ b/android/package_ctx.go
@@ -232,16 +232,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,
@@ -249,18 +265,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/rule_builder.go b/android/rule_builder.go
index 928ba53..b4f144a 100644
--- a/android/rule_builder.go
+++ b/android/rule_builder.go
@@ -417,10 +417,10 @@
 	}
 
 	var pool blueprint.Pool
-	if ctx.Config().UseGoma() && r.remoteable&SUPPORTS_GOMA != 0 {
+	if ctx.Config().UseGoma() && r.remoteable.Goma {
 		// When USE_GOMA=true is set and the rule is supported by goma, allow jobs to run outside the local pool.
-	} else if ctx.Config().UseRBE() && r.remoteable&SUPPORTS_RBE != 0 {
-		// When USE_GOMA=true is set and the rule is supported by RBE, allow jobs to run outside the local pool.
+	} else if ctx.Config().UseRBE() && r.remoteable.RBE {
+		// When USE_RBE=true is set and the rule is supported by RBE, allow jobs to run outside the local pool.
 	} else if r.highmem {
 		pool = highmemPool
 	} else if ctx.Config().UseRemoteBuild() {
diff --git a/android/variable.go b/android/variable.go
index 7473491..cc42766 100644
--- a/android/variable.go
+++ b/android/variable.go
@@ -204,6 +204,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 024fcbc..1d792ef 100644
--- a/bpf/bpf.go
+++ b/bpf/bpf.go
@@ -33,7 +33,7 @@
 var (
 	pctx = android.NewPackageContext("android/soong/bpf")
 
-	ccRule = pctx.AndroidRemoteStaticRule("ccRule", android.SUPPORTS_GOMA,
+	ccRule = pctx.AndroidRemoteStaticRule("ccRule", android.RemoteRuleSupports{Goma: true},
 		blueprint.RuleParams{
 			Depfile:     "${out}.d",
 			Deps:        blueprint.DepsGCC,
diff --git a/cc/builder.go b/cc/builder.go
index 5f0da5f..b05ca5b 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 26a49ea..f9b5367 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 6da7279..7f446e5 100644
--- a/java/config/config.go
+++ b/java/config/config.go
@@ -148,6 +148,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")
 
 	pctx.HostBinToolVariable("ManifestCheckCmd", "manifest_check")
diff --git a/java/dex.go b/java/dex.go
index ed4dfcf..6afdb6d 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 5319a4f..cb7da20 100644
--- a/java/kotlin.go
+++ b/java/kotlin.go
@@ -26,7 +26,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" "$emptyDir" && ` +
 			`mkdir -p "$classesDir" "$srcJarDir" "$emptyDir" && ` +
@@ -88,7 +88,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 9b19ede..5b9d10a 100644
--- a/ui/build/config.go
+++ b/ui/build/config.go
@@ -787,6 +787,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 c3da38b..ce8f968 100644
--- a/ui/build/dumpvars.go
+++ b/ui/build/dumpvars.go
@@ -191,6 +191,8 @@
 		"CC_WRAPPER",
 		"CXX_WRAPPER",
 		"JAVAC_WRAPPER",
+		"R8_WRAPPER",
+		"D8_WRAPPER",
 
 		// ccache settings
 		"CCACHE_COMPILERCHECK",