Support jarjar sharding

- If the new option `jarjar_shards` is set, use this as the number
of shards and enable sharding.

- The plan is to use it on "big" jars such as framework-minus-apex, but
for now, we only use it in a "test" target,
`framework-minus-apex_jarjar-sharded`.

- Also add `framework-minus-apex_jarjar-sharded` to the various
hardcoded spacial cases.

Test: m framework-minus-apex framework-minus-apex_jarjar-sharded
  and make sure the javac output jars have the same entries.
  $ diff \
    <(jar tf $ANDROID_BUILD_TOP/out/soong/.intermediates/frameworks/base/framework-minus-apex/android_common/jarjar/framework.jar \
      | sort) \
    <(jar tf $ANDROID_BUILD_TOP/out/soong/.intermediates/frameworks/base/framework-minus-apex_jarjar-sharded/android_common/jarjar/framework-minus-apex_jarjar-sharded.jar \
      | sort)
  (no output)

Bug: 383559945
Flag: EXEMPT New options are not used on production jars
Change-Id: I94c2ad978add25267f22b65c063e9148901b208b

5 files changed, 82 insertions, 16 deletions

diff --git a/android/container.go b/android/container.go
index 27b17ed99..99e6a3075 100644
--- a/android/container.go
+++ b/android/container.go
@@ -197,7 +197,10 @@ var unstableInfoProvider = blueprint.NewProvider[unstableInfo]()
 
 func determineUnstableModule(mctx ModuleContext) bool {
 	module := mctx.Module()
-	unstableModule := module.Name() == "framework-minus-apex"
+
+	// TODO(b/383559945) Remove "framework-minus-apex_jarjar-sharded" once
+	// we remove this module.
+	unstableModule := module.Name() == "framework-minus-apex" || module.Name() == "framework-minus-apex_jarjar-sharded"
 	if installable, ok := module.(InstallableModule); ok {
 		for _, staticDepTag := range installable.StaticDependencyTags() {
 			mctx.VisitDirectDepsWithTag(staticDepTag, func(dep Module) {
diff --git a/java/base.go b/java/base.go
index cd5550a7c..18a875b9c 100644
--- a/java/base.go
+++ b/java/base.go
@@ -109,6 +109,10 @@ type CommonProperties struct {
 	// if not blank, used as prefix to generate repackage rule
 	Jarjar_prefix *string
 
+	// Number of shards for jarjar. It needs to be an integer represented as a string.
+	// TODO(b/383559945) change it to int, once Configurable supports the type.
+	Jarjar_shards proptools.Configurable[string]
+
 	// If not blank, set the java version passed to javac as -source and -target
 	Java_version *string
 
@@ -1273,7 +1277,7 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 		localHeaderJars, combinedHeaderJarFile := j.compileJavaHeader(ctx, uniqueJavaFiles, srcJars, deps, flags, jarName,
 			extraCombinedJars)
 
-		combinedHeaderJarFile, jarjared := j.jarjarIfNecessary(ctx, combinedHeaderJarFile, jarName, "turbine")
+		combinedHeaderJarFile, jarjared := j.jarjarIfNecessary(ctx, combinedHeaderJarFile, jarName, "turbine", false)
 		if jarjared {
 			localHeaderJars = android.Paths{combinedHeaderJarFile}
 			transitiveStaticLibsHeaderJars = nil
@@ -1409,7 +1413,7 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 		shardingHeaderJars = localHeaderJars
 
 		var jarjared bool
-		j.headerJarFile, jarjared = j.jarjarIfNecessary(ctx, combinedHeaderJarFile, jarName, "turbine")
+		j.headerJarFile, jarjared = j.jarjarIfNecessary(ctx, combinedHeaderJarFile, jarName, "turbine", false)
 		if jarjared {
 			// jarjar modifies transitive static dependencies, use the combined header jar and drop the transitive
 			// static libs header jars.
@@ -1643,7 +1647,7 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 	}
 
 	// jarjar implementation jar if necessary
-	jarjarFile, jarjarred := j.jarjarIfNecessary(ctx, outputFile, jarName, "")
+	jarjarFile, jarjarred := j.jarjarIfNecessary(ctx, outputFile, jarName, "", true)
 	if jarjarred {
 		localImplementationJars = android.Paths{jarjarFile}
 		completeStaticLibsImplementationJars = depset.New(depset.PREORDER, localImplementationJars, nil)
@@ -1652,7 +1656,7 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 
 	// jarjar resource jar if necessary
 	if combinedResourceJar != nil {
-		resourceJarJarFile, jarjarred := j.jarjarIfNecessary(ctx, combinedResourceJar, jarName, "resource")
+		resourceJarJarFile, jarjarred := j.jarjarIfNecessary(ctx, combinedResourceJar, jarName, "resource", false)
 		combinedResourceJar = resourceJarJarFile
 		if jarjarred {
 			localResourceJars = android.Paths{resourceJarJarFile}
@@ -2941,12 +2945,23 @@ func (j *Module) repackageFlagsIfNecessary(ctx android.ModuleContext, infile and
 	return repackagedJarjarFile, true
 }
 
-func (j *Module) jarjarIfNecessary(ctx android.ModuleContext, infile android.Path, jarName, info string) (android.Path, bool) {
+func (j *Module) jarjarIfNecessary(ctx android.ModuleContext, infile android.Path, jarName, info string, useShards bool) (android.Path, bool) {
 	if j.expandJarjarRules == nil {
 		return infile, false
 	}
 	jarjarFile := android.PathForModuleOut(ctx, "jarjar", info, jarName)
-	TransformJarJar(ctx, jarjarFile, infile, j.expandJarjarRules)
+
+	totalShards := 1
+	if useShards {
+		totalShardsStr := j.properties.Jarjar_shards.GetOrDefault(ctx, "1")
+		ts, err := strconv.Atoi(totalShardsStr)
+		if err != nil {
+			ctx.PropertyErrorf("jarjar_shards", "jarjar_shards must be an integer represented as a string")
+			return infile, false
+		}
+		totalShards = ts
+	}
+	TransformJarJarWithShards(ctx, jarjarFile, infile, j.expandJarjarRules, totalShards)
 	return jarjarFile, true
 
 }
diff --git a/java/builder.go b/java/builder.go
index 01fbbddf5..f25798d3f 100644
--- a/java/builder.go
+++ b/java/builder.go
@@ -238,12 +238,12 @@ var (
 				// for newly repackaged classes. Dropping @UnsupportedAppUsage on repackaged classes
 				// avoids adding new hiddenapis after jarjar'ing.
 				" -DremoveAndroidCompatAnnotations=true" +
-				" -jar ${config.JarjarCmd} process $rulesFile $in $out && " +
+				" -jar ${config.JarjarCmd} process $rulesFile $in $out $total_shards $shard_index && " +
 				// Turn a missing output file into a ninja error
 				`[ -e ${out} ] || (echo "Missing output file"; exit 1)`,
 			CommandDeps: []string{"${config.JavaCmd}", "${config.JarjarCmd}", "$rulesFile"},
 		},
-		"rulesFile")
+		"rulesFile", "total_shards", "shard_index")
 
 	packageCheck = pctx.AndroidStaticRule("packageCheck",
 		blueprint.RuleParams{
@@ -750,16 +750,58 @@ func convertImplementationJarToHeaderJar(ctx android.ModuleContext, implementati
 
 func TransformJarJar(ctx android.ModuleContext, outputFile android.WritablePath,
 	classesJar android.Path, rulesFile android.Path) {
+	TransformJarJarWithShards(ctx, outputFile, classesJar, rulesFile, 1)
+}
+
+func TransformJarJarWithShards(ctx android.ModuleContext, outputFile android.WritablePath,
+	classesJar android.Path, rulesFile android.Path, totalShards int) {
+
+	// If the total number of shards is 1, just run jarjar as-is, with `total_shards` = 1
+	// and `shard_index` == 0, which effectively disables sharding
+	if totalShards == 1 {
+		ctx.Build(pctx, android.BuildParams{
+			Rule:        jarjar,
+			Description: "jarjar",
+			Output:      outputFile,
+			Input:       classesJar,
+			Implicit:    rulesFile,
+			Args: map[string]string{
+				"rulesFile":    rulesFile.String(),
+				"total_shards": "1",
+				"shard_index":  "0",
+			},
+		})
+		return
+	}
+
+	// Otherwise, run multiple jarjar instances and use merge_zips to combine the output.
+	tempJars := make([]android.Path, 0)
+	totalStr := strconv.Itoa(totalShards)
+	for i := 0; i < totalShards; i++ {
+		iStr := strconv.Itoa(i)
+		tempOut := android.PathForOutput(ctx, outputFile.String()+"-"+iStr+".jar")
+		ctx.Build(pctx, android.BuildParams{
+			Rule:        jarjar,
+			Description: "jarjar (" + iStr + "/" + totalStr + ")",
+			Output:      tempOut,
+			Input:       classesJar,
+			Implicit:    rulesFile,
+			Args: map[string]string{
+				"rulesFile":    rulesFile.String(),
+				"total_shards": totalStr,
+				"shard_index":  iStr,
+			},
+		})
+		tempJars = append(tempJars, tempOut)
+	}
+
 	ctx.Build(pctx, android.BuildParams{
-		Rule:        jarjar,
-		Description: "jarjar",
+		Rule:        combineJar,
+		Description: "merge jarjar shards",
 		Output:      outputFile,
-		Input:       classesJar,
-		Implicit:    rulesFile,
-		Args: map[string]string{
-			"rulesFile": rulesFile.String(),
-		},
+		Inputs:      tempJars,
 	})
+
 }
 
 func CheckJarPackages(ctx android.ModuleContext, outputFile android.WritablePath,
diff --git a/java/config/config.go b/java/config/config.go
index 7c29722a2..19d695291 100644
--- a/java/config/config.go
+++ b/java/config/config.go
@@ -42,6 +42,9 @@ var (
 	InstrumentFrameworkModules = []string{
 		"framework",
 		"framework-minus-apex",
+		// TODO(b/383559945) Remove "framework-minus-apex_jarjar-sharded" once
+		// we remove this module.
+		"framework-minus-apex_jarjar-sharded",
 		"ims-common",
 		"telephony-common",
 		"services",
diff --git a/java/legacy_core_platform_api_usage.go b/java/legacy_core_platform_api_usage.go
index 4be7d0470..6128ce3df 100644
--- a/java/legacy_core_platform_api_usage.go
+++ b/java/legacy_core_platform_api_usage.go
@@ -28,6 +28,9 @@ var legacyCorePlatformApiModules = []string{
 	"FloralClocks",
 	"framework-jobscheduler",
 	"framework-minus-apex",
+	// TODO(b/383559945) Remove "framework-minus-apex_jarjar-sharded" once
+	// we remove this module.
+	"framework-minus-apex_jarjar-sharded",
 	"framework-minus-apex-headers",
 	"framework-minus-apex-intdefs",
 	"FrameworksCoreTests",