Run non-RBE supported actions in the local pool when USE_RBE is set.

Bug: 143938974
Test: ran CTS build at -j500 successfully.
Change-Id: I55074bd67308cd716972e24fb56a20bc393d5d9d

android/module.go

@@ -1195,9 +1195,9 @@ func (m *moduleContext) Variable(pctx PackageContext, name, value string) {
 func (m *moduleContext) Rule(pctx PackageContext, name string, params blueprint.RuleParams,
 	argNames ...string) blueprint.Rule {
 
-	if m.config.UseGoma() && params.Pool == nil {
+	if (m.config.UseGoma() || m.config.UseRBE()) && params.Pool == nil {
-		// When USE_GOMA=true is set and the rule is not supported by goma, restrict jobs to the
+		// When USE_GOMA=true or USE_RBE=true are set and the rule is not supported by goma/RBE, restrict
-		// local parallelism value
+		// jobs to the local parallelism value
 		params.Pool = localPool
 	}
 

android/package_ctx.go

@@ -115,9 +115,9 @@ func (p PackageContext) RuleFunc(name string,
 		if len(ctx.errors) > 0 {
 			return params, ctx.errors[0]
 		}
-		if ctx.Config().UseGoma() && params.Pool == nil {
+		if (ctx.Config().UseGoma() || ctx.Config().UseRBE()) && params.Pool == nil {
-			// When USE_GOMA=true is set and the rule is not supported by goma, restrict jobs to the
+			// When USE_GOMA=true or USE_RBE=true are set and the rule is not supported by
-			// local parallelism value
+			// goma/RBE, restrict jobs to the local parallelism value
 			params.Pool = localPool
 		}
 		return params, nil
@@ -254,9 +254,35 @@ func (p PackageContext) StaticRule(name string, params blueprint.RuleParams,
 	}, argNames...)
 }
 
-// AndroidGomaStaticRule wraps blueprint.StaticRule but uses goma's parallelism if goma is enabled
+// RemoteRuleSupports selects if a AndroidRemoteStaticRule supports goma, RBE, or both.
-func (p PackageContext) AndroidGomaStaticRule(name string, params blueprint.RuleParams,
+type RemoteRuleSupports int
+
+const (
+	SUPPORTS_NONE = 0
+	SUPPORTS_GOMA = 1 << iota
+	SUPPORTS_RBE  = 1 << iota
+	SUPPORTS_BOTH = SUPPORTS_GOMA | SUPPORTS_RBE
+)
+
+// 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,
 	argNames ...string) blueprint.Rule {
-	// bypass android.PackageContext.StaticRule so that Pool does not get set to local_pool.
+
-	return p.PackageContext.StaticRule(name, params, argNames...)
+	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 {
+			// 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 {
+			// When USE_RBE=true is set and the rule is not supported by RBE, restrict jobs to the
+			// local parallelism value
+			params.Pool = localPool
+		}
+
+		return params, nil
+	}, argNames...)
 }

android/singleton.go

@@ -131,9 +131,9 @@ func (s *singletonContextAdaptor) Variable(pctx PackageContext, name, value stri
 }
 
 func (s *singletonContextAdaptor) Rule(pctx PackageContext, name string, params blueprint.RuleParams, argNames ...string) blueprint.Rule {
-	if s.Config().UseGoma() && params.Pool == nil {
+	if (s.Config().UseGoma() || s.Config().UseRBE()) && params.Pool == nil {
-		// When USE_GOMA=true is set and the rule is not supported by goma, restrict jobs to the
+		// When USE_GOMA=true or USE_RBE=true are set and the rule is not supported by goma/RBE, restrict
-		// local parallelism value
+		// jobs to the local parallelism value
 		params.Pool = localPool
 	}
 	rule := s.SingletonContext.Rule(pctx.PackageContext, name, params, argNames...)

bpf/bpf.go

@@ -33,7 +33,7 @@ func init() {
 var (
 	pctx = android.NewPackageContext("android/soong/bpf")
 
-	ccRule = pctx.AndroidGomaStaticRule("ccRule",
+	ccRule = pctx.AndroidRemoteStaticRule("ccRule", android.SUPPORTS_GOMA,
 		blueprint.RuleParams{
 			Depfile:     "${out}.d",
 			Deps:        blueprint.DepsGCC,

cc/builder.go

@@ -46,7 +46,7 @@ var (
 var (
 	pctx = android.NewPackageContext("android/soong/cc")
 
-	cc = pctx.AndroidGomaStaticRule("cc",
+	cc = pctx.AndroidRemoteStaticRule("cc", android.SUPPORTS_BOTH,
 		blueprint.RuleParams{
 			Depfile:     "${out}.d",
 			Deps:        blueprint.DepsGCC,
@@ -55,7 +55,7 @@ var (
 		},
 		"ccCmd", "cFlags")
 
-	ccNoDeps = pctx.AndroidGomaStaticRule("ccNoDeps",
+	ccNoDeps = pctx.AndroidRemoteStaticRule("ccNoDeps", android.SUPPORTS_GOMA,
 		blueprint.RuleParams{
 			Command:     "$relPwd ${config.CcWrapper}$ccCmd -c $cFlags -o $out $in",
 			CommandDeps: []string{"$ccCmd"},

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.AndroidGomaStaticRule("javac",
+	javac = pctx.AndroidRemoteStaticRule("javac", android.SUPPORTS_GOMA,
 		blueprint.RuleParams{
 			Command: `rm -rf "$outDir" "$annoDir" "$srcJarDir" && mkdir -p "$outDir" "$annoDir" "$srcJarDir" && ` +
 				`${config.ZipSyncCmd} -d $srcJarDir -l $srcJarDir/list -f "*.java" $srcJars && ` +

java/kotlin.go

@@ -26,7 +26,7 @@ import (
 	"github.com/google/blueprint"
 )
 
-var kotlinc = pctx.AndroidGomaStaticRule("kotlinc",
+var kotlinc = pctx.AndroidRemoteStaticRule("kotlinc", android.SUPPORTS_GOMA,
 	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.AndroidGomaStaticRule("kapt",
+var kapt = pctx.AndroidRemoteStaticRule("kapt", android.SUPPORTS_GOMA,
 	blueprint.RuleParams{
 		Command: `rm -rf "$srcJarDir" "$kotlinBuildFile" "$kaptDir" && mkdir -p "$srcJarDir" "$kaptDir" && ` +
 			`${config.ZipSyncCmd} -d $srcJarDir -l $srcJarDir/list -f "*.java" $srcJars && ` +