Refactor RBE support for Javac/R8/D8 to use the remoteexec package.

Bug: b/156613606
Test: build with and without RBE_JAVAC/RBE_R8/RBE_D8 set.
Change-Id: I1607f8cfb4d2c3cbb3672152bbeb561d9968bc60
Merged-In: I1607f8cfb4d2c3cbb3672152bbeb561d9968bc60

8 files changed, 79 insertions, 123 deletions

diff --git a/Android.bp b/Android.bp
index a81676fc..7ee87f78 100644
--- a/Android.bp
+++ b/Android.bp
@@ -349,6 +349,7 @@ bootstrap_go_package {
     deps: [
         "blueprint-proptools",
         "soong-android",
+        "soong-remoteexec",
     ],
     srcs: [
         "java/config/config.go",
diff --git a/android/package_ctx.go b/android/package_ctx.go
index 5a43ea9c..0de356e1 100644
--- a/android/package_ctx.go
+++ b/android/package_ctx.go
@@ -232,30 +232,10 @@ func (p PackageContext) StaticRule(name string, params blueprint.RuleParams,
 	}, argNames...)
 }
 
-// RBEExperimentalFlag indicates which flag should be set for the AndroidRemoteStaticRule
-// to use RBE.
-type RBEExperimentalFlag int
-
-const (
-	// 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
+	Goma bool
+	RBE  bool
 }
 
 // AndroidRemoteStaticRule wraps blueprint.StaticRule but uses goma or RBE's parallelism if goma or RBE are enabled
@@ -277,18 +257,6 @@ func (p PackageContext) AndroidRemoteStaticRule(name string, supports RemoteRule
 			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/cc/config/global.go b/cc/config/global.go
index f4d188e8..aaffcbc1 100644
--- a/cc/config/global.go
+++ b/cc/config/global.go
@@ -257,10 +257,10 @@ func init() {
 		return ""
 	})
 
-	pctx.VariableFunc("RECXXPool", envOverrideFunc("RBE_CXX_POOL", remoteexec.DefaultPool))
-	pctx.VariableFunc("RECXXLinksPool", envOverrideFunc("RBE_CXX_LINKS_POOL", remoteexec.DefaultPool))
-	pctx.VariableFunc("RECXXLinksExecStrategy", envOverrideFunc("RBE_CXX_LINKS_EXEC_STRATEGY", remoteexec.LocalExecStrategy))
-	pctx.VariableFunc("REAbiDumperExecStrategy", envOverrideFunc("RBE_ABI_DUMPER_EXEC_STRATEGY", remoteexec.LocalExecStrategy))
+	pctx.VariableFunc("RECXXPool", remoteexec.EnvOverrideFunc("RBE_CXX_POOL", remoteexec.DefaultPool))
+	pctx.VariableFunc("RECXXLinksPool", remoteexec.EnvOverrideFunc("RBE_CXX_LINKS_POOL", remoteexec.DefaultPool))
+	pctx.VariableFunc("RECXXLinksExecStrategy", remoteexec.EnvOverrideFunc("RBE_CXX_LINKS_EXEC_STRATEGY", remoteexec.LocalExecStrategy))
+	pctx.VariableFunc("REAbiDumperExecStrategy", remoteexec.EnvOverrideFunc("RBE_ABI_DUMPER_EXEC_STRATEGY", remoteexec.LocalExecStrategy))
 }
 
 var HostPrebuiltTag = pctx.VariableConfigMethod("HostPrebuiltTag", android.Config.PrebuiltOS)
diff --git a/java/builder.go b/java/builder.go
index 714d76aa..b8340295 100644
--- a/java/builder.go
+++ b/java/builder.go
@@ -27,6 +27,7 @@ import (
 	"github.com/google/blueprint/proptools"
 
 	"android/soong/android"
+	"android/soong/remoteexec"
 )
 
 var (
@@ -38,12 +39,13 @@ 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.RemoteRuleSupports{Goma: true, RBE: true, RBEFlag: android.RBE_JAVAC},
+	// TODO(b/143658984): goma can't handle the --system argument to javac.
+	javac, javacRE = remoteexec.StaticRules(pctx, "javac",
 		blueprint.RuleParams{
 			Command: `rm -rf "$outDir" "$annoDir" "$srcJarDir" && mkdir -p "$outDir" "$annoDir" "$srcJarDir" && ` +
 				`${config.ZipSyncCmd} -d $srcJarDir -l $srcJarDir/list -f "*.java" $srcJars && ` +
 				`(if [ -s $srcJarDir/list ] || [ -s $out.rsp ] ; then ` +
-				`${config.SoongJavacWrapper} ${config.JavacWrapper}${config.JavacCmd} ` +
+				`${config.SoongJavacWrapper} $reTemplate${config.JavacCmd} ` +
 				`${config.JavacHeapFlags} ${config.JavacVmFlags} ${config.CommonJdkFlags} ` +
 				`$processorpath $processor $javacFlags $bootClasspath $classpath ` +
 				`-source $javaVersion -target $javaVersion ` +
@@ -58,9 +60,12 @@ var (
 			CommandOrderOnly: []string{"${config.SoongJavacWrapper}"},
 			Rspfile:          "$out.rsp",
 			RspfileContent:   "$in",
-		},
-		"javacFlags", "bootClasspath", "classpath", "processorpath", "processor", "srcJars", "srcJarDir",
-		"outDir", "annoDir", "javaVersion")
+		}, &remoteexec.REParams{
+			Labels:       map[string]string{"type": "compile", "lang": "java", "compiler": "javac"},
+			ExecStrategy: "${config.REJavacExecStrategy}",
+			Platform:     map[string]string{remoteexec.PoolKey: "${config.REJavaPool}"},
+		}, []string{"javacFlags", "bootClasspath", "classpath", "processorpath", "processor", "srcJars", "srcJarDir",
+			"outDir", "annoDir", "javaVersion"}, nil)
 
 	_ = pctx.VariableFunc("kytheCorpus",
 		func(ctx android.PackageVarContext) string { return ctx.Config().XrefCorpusName() })
@@ -199,6 +204,7 @@ var (
 func init() {
 	pctx.Import("android/soong/android")
 	pctx.Import("android/soong/java/config")
+	pctx.Import("android/soong/remoteexec")
 }
 
 type javaBuilderFlags struct {
@@ -409,8 +415,12 @@ func transformJavaToClasses(ctx android.ModuleContext, outputFile android.Writab
 		outDir = filepath.Join(shardDir, outDir)
 		annoDir = filepath.Join(shardDir, annoDir)
 	}
+	rule := javac
+	if ctx.Config().IsEnvTrue("RBE_JAVAC") {
+		rule = javacRE
+	}
 	ctx.Build(pctx, android.BuildParams{
-		Rule:        javac,
+		Rule:        rule,
 		Description: desc,
 		Output:      outputFile,
 		Inputs:      srcFiles,
diff --git a/java/config/config.go b/java/config/config.go
index 9ac5a50c..b25d14ae 100644
--- a/java/config/config.go
+++ b/java/config/config.go
@@ -22,6 +22,7 @@ import (
 	_ "github.com/google/blueprint/bootstrap"
 
 	"android/soong/android"
+	"android/soong/remoteexec"
 )
 
 var (
@@ -139,30 +140,16 @@ func init() {
 	pctx.HostJavaToolVariable("MetalavaJar", "metalava.jar")
 	pctx.HostJavaToolVariable("DokkaJar", "dokka.jar")
 	pctx.HostJavaToolVariable("JetifierJar", "jetifier.jar")
+	pctx.HostJavaToolVariable("R8Jar", "r8-compat-proguard.jar")
+	pctx.HostJavaToolVariable("D8Jar", "d8.jar")
 
 	pctx.HostBinToolVariable("SoongJavacWrapper", "soong_javac_wrapper")
 	pctx.HostBinToolVariable("DexpreoptGen", "dexpreopt_gen")
 
-	pctx.VariableFunc("JavacWrapper", func(ctx android.PackageVarContext) string {
-		if override := ctx.Config().Getenv("JAVAC_WRAPPER"); override != "" {
-			return override + " "
-		}
-		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.VariableFunc("REJavaPool", remoteexec.EnvOverrideFunc("RBE_JAVA_POOL", "java16"))
+	pctx.VariableFunc("REJavacExecStrategy", remoteexec.EnvOverrideFunc("RBE_JAVAC_EXEC_STRATEGY", remoteexec.LocalExecStrategy))
+	pctx.VariableFunc("RED8ExecStrategy", remoteexec.EnvOverrideFunc("RBE_D8_EXEC_STRATEGY", remoteexec.LocalExecStrategy))
+	pctx.VariableFunc("RER8ExecStrategy", remoteexec.EnvOverrideFunc("RBE_R8_EXEC_STRATEGY", remoteexec.LocalExecStrategy))
 
 	pctx.HostJavaToolVariable("JacocoCLIJar", "jacoco-cli.jar")
 
diff --git a/java/dex.go b/java/dex.go
index 6afdb6dc..27ec6ee4 100644
--- a/java/dex.go
+++ b/java/dex.go
@@ -20,12 +20,13 @@ import (
 	"github.com/google/blueprint"
 
 	"android/soong/android"
+	"android/soong/remoteexec"
 )
 
-var d8 = pctx.AndroidRemoteStaticRule("d8", android.RemoteRuleSupports{RBE: true, RBEFlag: android.RBE_D8},
+var d8, d8RE = remoteexec.StaticRules(pctx, "d8",
 	blueprint.RuleParams{
 		Command: `rm -rf "$outDir" && mkdir -p "$outDir" && ` +
-			`${config.D8Wrapper}${config.D8Cmd} ${config.DexFlags} --output $outDir $d8Flags $in && ` +
+			`$reTemplate${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{
@@ -33,14 +34,19 @@ var d8 = pctx.AndroidRemoteStaticRule("d8", android.RemoteRuleSupports{RBE: true
 			"${config.SoongZipCmd}",
 			"${config.MergeZipsCmd}",
 		},
-	},
-	"outDir", "d8Flags", "zipFlags")
-
-var r8 = pctx.AndroidRemoteStaticRule("r8", android.RemoteRuleSupports{RBE: true, RBEFlag: android.RBE_R8},
+	}, &remoteexec.REParams{
+		Labels:          map[string]string{"type": "compile", "compiler": "d8"},
+		Inputs:          []string{"${config.D8Jar}"},
+		ExecStrategy:    "${config.RED8ExecStrategy}",
+		ToolchainInputs: []string{"${config.JavaCmd}"},
+		Platform:        map[string]string{remoteexec.PoolKey: "${config.REJavaPool}"},
+	}, []string{"outDir", "d8Flags", "zipFlags"}, nil)
+
+var r8, r8RE = remoteexec.StaticRules(pctx, "r8",
 	blueprint.RuleParams{
 		Command: `rm -rf "$outDir" && mkdir -p "$outDir" && ` +
 			`rm -f "$outDict" && ` +
-			`${config.R8Wrapper}${config.R8Cmd} ${config.DexFlags} -injars $in --output $outDir ` +
+			`$reTemplate${config.R8Cmd} ${config.DexFlags} -injars $in --output $outDir ` +
 			`--force-proguard-compatibility ` +
 			`--no-data-resources ` +
 			`-printmapping $outDict ` +
@@ -53,8 +59,13 @@ var r8 = pctx.AndroidRemoteStaticRule("r8", android.RemoteRuleSupports{RBE: true
 			"${config.SoongZipCmd}",
 			"${config.MergeZipsCmd}",
 		},
-	},
-	"outDir", "outDict", "r8Flags", "zipFlags")
+	}, &remoteexec.REParams{
+		Labels:          map[string]string{"type": "compile", "compiler": "r8"},
+		Inputs:          []string{"$implicits", "${config.R8Jar}"},
+		ExecStrategy:    "${config.RER8ExecStrategy}",
+		ToolchainInputs: []string{"${config.JavaCmd}"},
+		Platform:        map[string]string{remoteexec.PoolKey: "${config.REJavaPool}"},
+	}, []string{"outDir", "outDict", "r8Flags", "zipFlags"}, []string{"implicits"})
 
 func (j *Module) dexCommonFlags(ctx android.ModuleContext) []string {
 	flags := j.deviceProperties.Dxflags
@@ -185,24 +196,34 @@ func (j *Module) compileDex(ctx android.ModuleContext, flags javaBuilderFlags,
 		proguardDictionary := android.PathForModuleOut(ctx, "proguard_dictionary")
 		j.proguardDictionary = proguardDictionary
 		r8Flags, r8Deps := j.r8Flags(ctx, flags)
+		rule := r8
+		args := map[string]string{
+			"r8Flags":  strings.Join(r8Flags, " "),
+			"zipFlags": zipFlags,
+			"outDict":  j.proguardDictionary.String(),
+			"outDir":   outDir.String(),
+		}
+		if ctx.Config().IsEnvTrue("RBE_R8") {
+			rule = r8RE
+			args["implicits"] = strings.Join(r8Deps.Strings(), ",")
+		}
 		ctx.Build(pctx, android.BuildParams{
-			Rule:           r8,
+			Rule:           rule,
 			Description:    "r8",
 			Output:         javalibJar,
 			ImplicitOutput: proguardDictionary,
 			Input:          classesJar,
 			Implicits:      r8Deps,
-			Args: map[string]string{
-				"r8Flags":  strings.Join(r8Flags, " "),
-				"zipFlags": zipFlags,
-				"outDict":  j.proguardDictionary.String(),
-				"outDir":   outDir.String(),
-			},
+			Args:           args,
 		})
 	} else {
 		d8Flags, d8Deps := j.d8Flags(ctx, flags)
+		rule := d8
+		if ctx.Config().IsEnvTrue("RBE_D8") {
+			rule = d8RE
+		}
 		ctx.Build(pctx, android.BuildParams{
-			Rule:        d8,
+			Rule:        rule,
 			Description: "d8",
 			Output:      javalibJar,
 			Input:       classesJar,
diff --git a/remoteexec/remoteexec.go b/remoteexec/remoteexec.go
index f5119226..860db556 100644
--- a/remoteexec/remoteexec.go
+++ b/remoteexec/remoteexec.go
@@ -154,3 +154,14 @@ func StaticRules(ctx android.PackageContext, name string, ruleParams blueprint.R
 	return ctx.AndroidStaticRule(name, ruleParams, commonArgs...),
 		ctx.AndroidRemoteStaticRule(name+"RE", android.RemoteRuleSupports{RBE: true}, ruleParamsRE, append(commonArgs, reArgs...)...)
 }
+
+// EnvOverrideFunc retrieves a variable func that evaluates to the value of the given environment
+// variable if set, otherwise the given default.
+func EnvOverrideFunc(envVar, defaultVal string) func(ctx android.PackageVarContext) string {
+	return func(ctx android.PackageVarContext) string {
+		if override := ctx.Config().Getenv(envVar); override != "" {
+			return override
+		}
+		return defaultVal
+	}
+}
diff --git a/ui/build/config.go b/ui/build/config.go
index 5b9d10aa..9b19ede3 100644
--- a/ui/build/config.go
+++ b/ui/build/config.go
@@ -787,48 +787,6 @@ 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