[Ravenwood] Poorman's sharding

Framework-minus-apex is rather slow to process.

Ideally HostStubGen should be internally parallelize the processing,
but that means now we'll have to carefully synchronize concurrent access
to the common data structure.

So instead, let's just run multiple HSG instances in parallel.

Now the build trace looks like this: http://screen/9B5NUtPjaaQ8jfA

Bug: 313930116
Test: $ANDROID_BUILD_TOP/frameworks/base/ravenwood/scripts/run-ravenwood-tests.sh
Flag: EXEMPT host side test change only
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:a0ad8aa151aa4015ef1b63b4f4e0e1abb38f1891)
Merged-In: I50fd932a8e61196eda66c074eb2f8089e455b515
Change-Id: I50fd932a8e61196eda66c074eb2f8089e455b515
diff --git a/Ravenwood.bp b/Ravenwood.bp
index cebeb18..7faa33f 100644
--- a/Ravenwood.bp
+++ b/Ravenwood.bp
@@ -26,24 +26,13 @@
 }
 
 // Generate the stub/impl from framework-all, with hidden APIs.
-java_genrule {
-    name: "framework-minus-apex.ravenwood-base",
+// This step takes several tens of seconds, so we manually shard it to multiple modules.
+// All the copies have to be kept in sync.
+// TODO: Do the sharding better.
+
+genrule_defaults {
+    name: "framework-minus-apex.ravenwood-base_defaults",
     tools: ["hoststubgen"],
-    cmd: "$(location hoststubgen) " +
-        "@$(location :ravenwood-standard-options) " +
-
-        "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
-        "--stats-file $(location hoststubgen_framework-minus-apex_stats.csv) " +
-        "--supported-api-list-file $(location hoststubgen_framework-minus-apex_apis.csv) " +
-
-        "--out-impl-jar $(location ravenwood.jar) " +
-
-        "--gen-keep-all-file $(location hoststubgen_framework-minus-apex_keep_all.txt) " +
-        "--gen-input-dump-file $(location hoststubgen_framework-minus-apex_dump.txt) " +
-
-        "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
-        "--policy-override-file $(location :ravenwood-framework-policies) " +
-        "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) ",
     srcs: [
         ":framework-minus-apex-for-hoststubgen",
         ":ravenwood-framework-policies",
@@ -61,24 +50,221 @@
         "hoststubgen_framework-minus-apex_stats.csv",
         "hoststubgen_framework-minus-apex_apis.csv",
     ],
-    defaults: ["ravenwood-internal-only-visibility-genrule"],
+    visibility: ["//visibility:private"],
 }
 
-// Extract the impl jar from "framework-minus-apex.ravenwood-base" for subsequent build rules.
-// Note this emits a "device side" output, so that ravenwood tests can (implicitly)
-// depend on it.
+java_genrule {
+    name: "framework-minus-apex.ravenwood-base_X0",
+    defaults: ["framework-minus-apex.ravenwood-base_defaults"],
+    cmd: "$(location hoststubgen) " +
+        "--num-shards 6 --shard-index 0 " + // Only this line differs
+
+        "@$(location :ravenwood-standard-options) " +
+
+        "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
+        "--stats-file $(location hoststubgen_framework-minus-apex_stats.csv) " +
+        "--supported-api-list-file $(location hoststubgen_framework-minus-apex_apis.csv) " +
+
+        "--out-impl-jar $(location ravenwood.jar) " +
+
+        "--gen-keep-all-file $(location hoststubgen_framework-minus-apex_keep_all.txt) " +
+        "--gen-input-dump-file $(location hoststubgen_framework-minus-apex_dump.txt) " +
+
+        "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
+        "--policy-override-file $(location :ravenwood-framework-policies) " +
+        "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) ",
+}
+
+java_genrule {
+    name: "framework-minus-apex.ravenwood-base_X1",
+    defaults: ["framework-minus-apex.ravenwood-base_defaults"],
+    cmd: "$(location hoststubgen) " +
+        "--num-shards 6 --shard-index 1 " + // Only this line differs
+
+        "@$(location :ravenwood-standard-options) " +
+
+        "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
+        "--stats-file $(location hoststubgen_framework-minus-apex_stats.csv) " +
+        "--supported-api-list-file $(location hoststubgen_framework-minus-apex_apis.csv) " +
+
+        "--out-impl-jar $(location ravenwood.jar) " +
+
+        "--gen-keep-all-file $(location hoststubgen_framework-minus-apex_keep_all.txt) " +
+        "--gen-input-dump-file $(location hoststubgen_framework-minus-apex_dump.txt) " +
+
+        "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
+        "--policy-override-file $(location :ravenwood-framework-policies) " +
+        "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) ",
+}
+
+java_genrule {
+    name: "framework-minus-apex.ravenwood-base_X2",
+    defaults: ["framework-minus-apex.ravenwood-base_defaults"],
+    cmd: "$(location hoststubgen) " +
+        "--num-shards 6 --shard-index 2 " + // Only this line differs
+
+        "@$(location :ravenwood-standard-options) " +
+
+        "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
+        "--stats-file $(location hoststubgen_framework-minus-apex_stats.csv) " +
+        "--supported-api-list-file $(location hoststubgen_framework-minus-apex_apis.csv) " +
+
+        "--out-impl-jar $(location ravenwood.jar) " +
+
+        "--gen-keep-all-file $(location hoststubgen_framework-minus-apex_keep_all.txt) " +
+        "--gen-input-dump-file $(location hoststubgen_framework-minus-apex_dump.txt) " +
+
+        "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
+        "--policy-override-file $(location :ravenwood-framework-policies) " +
+        "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) ",
+}
+
+java_genrule {
+    name: "framework-minus-apex.ravenwood-base_X3",
+    defaults: ["framework-minus-apex.ravenwood-base_defaults"],
+    cmd: "$(location hoststubgen) " +
+        "--num-shards 6 --shard-index 3 " + // Only this line differs
+
+        "@$(location :ravenwood-standard-options) " +
+
+        "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
+        "--stats-file $(location hoststubgen_framework-minus-apex_stats.csv) " +
+        "--supported-api-list-file $(location hoststubgen_framework-minus-apex_apis.csv) " +
+
+        "--out-impl-jar $(location ravenwood.jar) " +
+
+        "--gen-keep-all-file $(location hoststubgen_framework-minus-apex_keep_all.txt) " +
+        "--gen-input-dump-file $(location hoststubgen_framework-minus-apex_dump.txt) " +
+
+        "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
+        "--policy-override-file $(location :ravenwood-framework-policies) " +
+        "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) ",
+}
+
+java_genrule {
+    name: "framework-minus-apex.ravenwood-base_X4",
+    defaults: ["framework-minus-apex.ravenwood-base_defaults"],
+    cmd: "$(location hoststubgen) " +
+        "--num-shards 6 --shard-index 4 " + // Only this line differs
+
+        "@$(location :ravenwood-standard-options) " +
+
+        "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
+        "--stats-file $(location hoststubgen_framework-minus-apex_stats.csv) " +
+        "--supported-api-list-file $(location hoststubgen_framework-minus-apex_apis.csv) " +
+
+        "--out-impl-jar $(location ravenwood.jar) " +
+
+        "--gen-keep-all-file $(location hoststubgen_framework-minus-apex_keep_all.txt) " +
+        "--gen-input-dump-file $(location hoststubgen_framework-minus-apex_dump.txt) " +
+
+        "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
+        "--policy-override-file $(location :ravenwood-framework-policies) " +
+        "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) ",
+}
+
+java_genrule {
+    name: "framework-minus-apex.ravenwood-base_X5",
+    defaults: ["framework-minus-apex.ravenwood-base_defaults"],
+    cmd: "$(location hoststubgen) " +
+        "--num-shards 6 --shard-index 5 " + // Only this line differs
+
+        "@$(location :ravenwood-standard-options) " +
+
+        "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
+        "--stats-file $(location hoststubgen_framework-minus-apex_stats.csv) " +
+        "--supported-api-list-file $(location hoststubgen_framework-minus-apex_apis.csv) " +
+
+        "--out-impl-jar $(location ravenwood.jar) " +
+
+        "--gen-keep-all-file $(location hoststubgen_framework-minus-apex_keep_all.txt) " +
+        "--gen-input-dump-file $(location hoststubgen_framework-minus-apex_dump.txt) " +
+
+        "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
+        "--policy-override-file $(location :ravenwood-framework-policies) " +
+        "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) ",
+}
+
+// Marge all the sharded jars
 java_genrule {
     name: "framework-minus-apex.ravenwood",
-    defaults: ["ravenwood-internal-only-visibility-genrule"],
-    cmd: "cp $(in) $(out)",
+    defaults: ["ravenwood-internal-only-visibility-java"],
+    cmd: "$(location merge_zips) $(out) $(in)",
+    tools: ["merge_zips"],
     srcs: [
-        ":framework-minus-apex.ravenwood-base{ravenwood.jar}",
+        ":framework-minus-apex.ravenwood-base_X0{ravenwood.jar}",
+        ":framework-minus-apex.ravenwood-base_X1{ravenwood.jar}",
+        ":framework-minus-apex.ravenwood-base_X2{ravenwood.jar}",
+        ":framework-minus-apex.ravenwood-base_X3{ravenwood.jar}",
+        ":framework-minus-apex.ravenwood-base_X4{ravenwood.jar}",
+        ":framework-minus-apex.ravenwood-base_X5{ravenwood.jar}",
     ],
     out: [
         "framework-minus-apex.ravenwood.jar",
     ],
 }
 
+// Merge the sharded text files
+genrule {
+    name: "hoststubgen_framework-minus-apex_stats.csv",
+    defaults: ["ravenwood-internal-only-visibility-genrule"],
+    cmd: "cat $(in) > $(out)",
+    srcs: [
+        ":framework-minus-apex.ravenwood-base_X0{hoststubgen_framework-minus-apex_stats.csv}",
+        ":framework-minus-apex.ravenwood-base_X1{hoststubgen_framework-minus-apex_stats.csv}",
+        ":framework-minus-apex.ravenwood-base_X2{hoststubgen_framework-minus-apex_stats.csv}",
+        ":framework-minus-apex.ravenwood-base_X3{hoststubgen_framework-minus-apex_stats.csv}",
+        ":framework-minus-apex.ravenwood-base_X4{hoststubgen_framework-minus-apex_stats.csv}",
+        ":framework-minus-apex.ravenwood-base_X5{hoststubgen_framework-minus-apex_stats.csv}",
+    ],
+    out: ["hoststubgen_framework-minus-apex_stats.csv"],
+}
+
+genrule {
+    name: "hoststubgen_framework-minus-apex_apis.csv",
+    defaults: ["ravenwood-internal-only-visibility-genrule"],
+    cmd: "cat $(in) > $(out)",
+    srcs: [
+        ":framework-minus-apex.ravenwood-base_X0{hoststubgen_framework-minus-apex_apis.csv}",
+        ":framework-minus-apex.ravenwood-base_X1{hoststubgen_framework-minus-apex_apis.csv}",
+        ":framework-minus-apex.ravenwood-base_X2{hoststubgen_framework-minus-apex_apis.csv}",
+        ":framework-minus-apex.ravenwood-base_X3{hoststubgen_framework-minus-apex_apis.csv}",
+        ":framework-minus-apex.ravenwood-base_X4{hoststubgen_framework-minus-apex_apis.csv}",
+        ":framework-minus-apex.ravenwood-base_X5{hoststubgen_framework-minus-apex_apis.csv}",
+    ],
+    out: ["hoststubgen_framework-minus-apex_apis.csv"],
+}
+
+genrule {
+    name: "hoststubgen_framework-minus-apex_keep_all.txt",
+    defaults: ["ravenwood-internal-only-visibility-genrule"],
+    cmd: "cat $(in) > $(out)",
+    srcs: [
+        ":framework-minus-apex.ravenwood-base_X0{hoststubgen_framework-minus-apex_keep_all.txt}",
+        ":framework-minus-apex.ravenwood-base_X1{hoststubgen_framework-minus-apex_keep_all.txt}",
+        ":framework-minus-apex.ravenwood-base_X2{hoststubgen_framework-minus-apex_keep_all.txt}",
+        ":framework-minus-apex.ravenwood-base_X3{hoststubgen_framework-minus-apex_keep_all.txt}",
+        ":framework-minus-apex.ravenwood-base_X4{hoststubgen_framework-minus-apex_keep_all.txt}",
+        ":framework-minus-apex.ravenwood-base_X5{hoststubgen_framework-minus-apex_keep_all.txt}",
+    ],
+    out: ["hoststubgen_framework-minus-apex_keep_all.txt"],
+}
+
+genrule {
+    name: "hoststubgen_framework-minus-apex_dump.txt",
+    defaults: ["ravenwood-internal-only-visibility-genrule"],
+    cmd: "cat $(in) > $(out)",
+    srcs: [
+        ":framework-minus-apex.ravenwood-base_X0{hoststubgen_framework-minus-apex_dump.txt}",
+        ":framework-minus-apex.ravenwood-base_X1{hoststubgen_framework-minus-apex_dump.txt}",
+        ":framework-minus-apex.ravenwood-base_X2{hoststubgen_framework-minus-apex_dump.txt}",
+        ":framework-minus-apex.ravenwood-base_X3{hoststubgen_framework-minus-apex_dump.txt}",
+        ":framework-minus-apex.ravenwood-base_X4{hoststubgen_framework-minus-apex_dump.txt}",
+        ":framework-minus-apex.ravenwood-base_X5{hoststubgen_framework-minus-apex_dump.txt}",
+    ],
+    out: ["hoststubgen_framework-minus-apex_dump.txt"],
+}
+
 java_library {
     name: "services.core-for-hoststubgen",
     installable: false, // host only jar.
diff --git a/ravenwood/Android.bp b/ravenwood/Android.bp
index 2de3c5e..4faf03c 100644
--- a/ravenwood/Android.bp
+++ b/ravenwood/Android.bp
@@ -280,10 +280,10 @@
     src: "scripts/ravenwood-stats-checker.sh",
     test_suites: ["general-tests"],
     data: [
-        ":framework-minus-apex.ravenwood-base{hoststubgen_framework-minus-apex_stats.csv}",
-        ":framework-minus-apex.ravenwood-base{hoststubgen_framework-minus-apex_apis.csv}",
-        ":framework-minus-apex.ravenwood-base{hoststubgen_framework-minus-apex_keep_all.txt}",
-        ":framework-minus-apex.ravenwood-base{hoststubgen_framework-minus-apex_dump.txt}",
+        ":hoststubgen_framework-minus-apex_stats.csv",
+        ":hoststubgen_framework-minus-apex_apis.csv",
+        ":hoststubgen_framework-minus-apex_keep_all.txt",
+        ":hoststubgen_framework-minus-apex_dump.txt",
         ":services.core.ravenwood-base{hoststubgen_services.core_stats.csv}",
         ":services.core.ravenwood-base{hoststubgen_services.core_apis.csv}",
         ":services.core.ravenwood-base{hoststubgen_services.core_keep_all.txt}",
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
index 73cfd3f..5dde265 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
@@ -85,6 +85,8 @@
                 errors,
                 stats,
                 policyFileRemapper,
+                options.numShards.get,
+                options.shard.get,
         )
 
         // Dump statistics, if specified.
@@ -213,6 +215,8 @@
             errors: HostStubGenErrors,
             stats: HostStubGenStats,
             remapper: Remapper?,
+            numShards: Int,
+            shard: Int,
             ) {
         log.i("Converting %s into [stub: %s, impl: %s] ...", inJar, outStubJar, outImplJar)
         log.i("ASM CheckClassAdapter is %s", if (enableChecker) "enabled" else "disabled")
@@ -221,17 +225,32 @@
 
         val packageRedirector = PackageRedirectRemapper(options.packageRedirects)
 
+        var itemIndex = 0
+        var numItemsProcessed = 0
+        var numItems = -1 // == Unknown
+
         log.withIndent {
             // Open the input jar file and process each entry.
             ZipFile(inJar).use { inZip ->
+
+                numItems = inZip.size()
+                val shardStart = numItems * shard / numShards
+                val shardNextStart = numItems * (shard + 1) / numShards
+
                 maybeWithZipOutputStream(outStubJar) { stubOutStream ->
                     maybeWithZipOutputStream(outImplJar) { implOutStream ->
                         val inEntries = inZip.entries()
                         while (inEntries.hasMoreElements()) {
                             val entry = inEntries.nextElement()
+                            val inShard = (shardStart <= itemIndex) && (itemIndex < shardNextStart)
+                            itemIndex++
+                            if (!inShard) {
+                                continue
+                            }
                             convertSingleEntry(inZip, entry, stubOutStream, implOutStream,
                                     filter, packageRedirector, remapper,
                                     enableChecker, classes, errors, stats)
+                            numItemsProcessed++
                         }
                         log.i("Converted all entries.")
                     }
@@ -241,7 +260,8 @@
             }
         }
         val end = System.currentTimeMillis()
-        log.i("Done transforming the jar in %.1f second(s).", (end - start) / 1000.0)
+        log.i("Done transforming the jar in %.1f second(s). %d / %d item(s) processed.",
+            (end - start) / 1000.0, numItemsProcessed, numItems)
     }
 
     private fun <T> maybeWithZipOutputStream(filename: String?, block: (ZipOutputStream?) -> T): T {
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
index e192516..2f833a8 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
@@ -112,6 +112,9 @@
         var statsFile: SetOnce<String?> = SetOnce(null),
 
         var apiListFile: SetOnce<String?> = SetOnce(null),
+
+        var numShards: SetOnce<Int> = SetOnce(1),
+        var shard: SetOnce<Int> = SetOnce(0),
 ) {
     companion object {
 
@@ -162,6 +165,13 @@
                 fun SetOnce<String?>.setNextStringArg(): String = nextArg().also { this.set(it) }
                 fun MutableSet<String>.addUniqueAnnotationArg(): String =
                         nextArg().also { this += ensureUniqueAnnotation(it) }
+                fun SetOnce<Int>.setNextIntArg(): String = nextArg().also {
+                    try {
+                        this.set(it.toInt())
+                    } catch (e: NumberFormatException) {
+                        throw ArgumentsException("Invalid integer for $arg: $it")
+                    }
+                }
 
                 try {
                     when (arg) {
@@ -259,6 +269,9 @@
                         "--stats-file" -> ret.statsFile.setNextStringArg()
                         "--supported-api-list-file" -> ret.apiListFile.setNextStringArg()
 
+                        "--num-shards" -> ret.numShards.setNextIntArg()
+                        "--shard-index" -> ret.shard.setNextIntArg()
+
                         else -> throw ArgumentsException("Unknown option: $arg")
                     }
                 } catch (e: SetOnce.SetMoreThanOnceException) {
@@ -396,6 +409,8 @@
               enableNonStubMethodCallDetection=$enableNonStubMethodCallDetection,
               statsFile=$statsFile,
               apiListFile=$apiListFile,
+              numShards=$numShards,
+              shard=$shard,
             }
             """.trimIndent()
     }