From 8ea054a81e3249d33ef29672c8a67025d90b4551 Mon Sep 17 00:00:00 2001
From: Ramy Medhat <abdelaal@google.com>
Date: Mon, 27 Jan 2020 14:19:44 -0500
Subject: [PATCH] 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.
Change-Id: Ic82f3627944f6a5ee7b9f3228170c2709b1bfcb8
---
 android/config.go       | 12 +++++++++++
 android/package_ctx.go  | 44 +++++++++++++++++++++++++++++++++--------
 android/rule_builder.go |  6 +++---
 android/variable.go     |  3 +++
 bpf/bpf.go              |  2 +-
 cc/builder.go           |  4 ++--
 java/builder.go         |  2 +-
 java/config/config.go   | 14 +++++++++++++
 java/dex.go             |  8 ++++----
 java/kotlin.go          |  4 ++--
 ui/build/config.go      | 42 +++++++++++++++++++++++++++++++++++++++
 ui/build/dumpvars.go    |  2 ++
 12 files changed, 122 insertions(+), 21 deletions(-)

diff --git a/android/config.go b/android/config.go
index f907de66a..ffd774417 100644
--- a/android/config.go
+++ b/android/config.go
@@ -827,6 +827,18 @@ func (c *config) UseRBE() bool {
 	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 635066767..5a43ea9c7 100644
--- a/android/package_ctx.go
+++ b/android/package_ctx.go
@@ -232,16 +232,32 @@ func (p PackageContext) StaticRule(name string, params blueprint.RuleParams,
 	}, 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 @@ func (p PackageContext) AndroidRemoteStaticRule(name string, supports RemoteRule
 
 	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 928ba532d..b4f144a0d 100644
--- a/android/rule_builder.go
+++ b/android/rule_builder.go
@@ -417,10 +417,10 @@ func (r *RuleBuilder) Build(pctx PackageContext, ctx BuilderContext, name string
 	}
 
 	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 74734913b..cc42766f5 100644
--- a/android/variable.go
+++ b/android/variable.go
@@ -204,6 +204,9 @@ type productVariables struct {
 	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 024fcbccc..1d792ef1f 100644
--- a/bpf/bpf.go
+++ b/bpf/bpf.go
@@ -33,7 +33,7 @@ func init() {
 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 5f0da5ff7..b05ca5bdd 100644
--- a/cc/builder.go
+++ b/cc/builder.go
@@ -46,7 +46,7 @@ var (
 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 @@ var (
 		},
 		"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 26a49ea53..f9b53674d 100644
--- a/java/builder.go
+++ b/java/builder.go
@@ -38,7 +38,7 @@ var (
 	// 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 6da7279c3..7f446e5db 100644
--- a/java/config/config.go
+++ b/java/config/config.go
@@ -148,6 +148,20 @@ func init() {
 		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 cd6d90de6..cd97563cd 100644
--- a/java/dex.go
+++ b/java/dex.go
@@ -22,10 +22,10 @@ import (
 	"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 @@ var d8 = pctx.AndroidStaticRule("d8",
 	},
 	"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 5319a4ff9..cb7da20f8 100644
--- a/java/kotlin.go
+++ b/java/kotlin.go
@@ -26,7 +26,7 @@ import (
 	"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 @@ func kotlinCompile(ctx android.ModuleContext, outputFile android.WritablePath,
 	})
 }
 
-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 9b19ede3a..5b9d10aa9 100644
--- a/ui/build/config.go
+++ b/ui/build/config.go
@@ -787,6 +787,48 @@ func (c *configImpl) UseRBE() bool {
 	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 c3da38bcd..ce8f968d7 100644
--- a/ui/build/dumpvars.go
+++ b/ui/build/dumpvars.go
@@ -191,6 +191,8 @@ func runMakeProductConfig(ctx Context, config Config) {
 		"CC_WRAPPER",
 		"CXX_WRAPPER",
 		"JAVAC_WRAPPER",
+		"R8_WRAPPER",
+		"D8_WRAPPER",
 
 		// ccache settings
 		"CCACHE_COMPILERCHECK",