From 3b706fde7f3ce35a3a8d06de86aa0f9718406f5b Mon Sep 17 00:00:00 2001
From: Colin Cross <ccross@android.com>
Date: Thu, 5 Sep 2019 16:44:18 -0700
Subject: [PATCH] Shard kythe invocations when javac is sharded

The kythe pipeline has trouble injesting all 6000 files of
framework.jar at once.  Shared invocations of the kythe extractor
when javac is sharded to produce multiple smaller kzip files.

Bug: 140426870
Test: no change to build.ninja
Test: m out/soong/.intermediates/frameworks/base/framework/android_common/framework0.kzip
Test: TestSharding in java_test.go
Change-Id: I867db4ef5cb1e7f3ce8359a46aac2c00ed8a8912
---
 java/builder.go | 10 +++++++---
 java/java.go    | 39 +++++++++++++++++++++++++++------------
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/java/builder.go b/java/builder.go
index f174cf0c6..9e068fa64 100644
--- a/java/builder.go
+++ b/java/builder.go
@@ -229,10 +229,9 @@ func RunErrorProne(ctx android.ModuleContext, outputFile android.WritablePath,
 
 // Emits the rule to generate Xref input file (.kzip file) for the given set of source files and source jars
 // to compile with given set of builder flags, etc.
-func emitXrefRule(ctx android.ModuleContext, xrefFile android.WritablePath,
+func emitXrefRule(ctx android.ModuleContext, xrefFile android.WritablePath, idx int,
 	srcFiles, srcJars android.Paths,
-	flags javaBuilderFlags, deps android.Paths,
-	intermediatesDir string) {
+	flags javaBuilderFlags, deps android.Paths) {
 
 	deps = append(deps, srcJars...)
 
@@ -260,6 +259,11 @@ func emitXrefRule(ctx android.ModuleContext, xrefFile android.WritablePath,
 		processor = "-processor " + flags.processor
 	}
 
+	intermediatesDir := "xref"
+	if idx >= 0 {
+		intermediatesDir += strconv.Itoa(idx)
+	}
+
 	ctx.Build(pctx,
 		android.BuildParams{
 			Rule:        kytheExtract,
diff --git a/java/java.go b/java/java.go
index 87ee2f29c..b05d7bbd5 100644
--- a/java/java.go
+++ b/java/java.go
@@ -25,6 +25,7 @@ import (
 	"strings"
 
 	"github.com/google/blueprint"
+	"github.com/google/blueprint/pathtools"
 	"github.com/google/blueprint/proptools"
 
 	"android/soong/android"
@@ -1151,27 +1152,20 @@ func (j *Module) compile(ctx android.ModuleContext, aaptSrcJar android.Path) {
 			if len(uniqueSrcFiles) > 0 {
 				shardSrcs = shardPaths(uniqueSrcFiles, shardSize)
 				for idx, shardSrc := range shardSrcs {
-					classes := android.PathForModuleOut(ctx, "javac", jarName+strconv.Itoa(idx))
-					TransformJavaToClasses(ctx, classes, idx, shardSrc, nil, flags, extraJarDeps)
+					classes := j.compileJavaClasses(ctx, jarName, idx, shardSrc,
+						nil, flags, extraJarDeps)
 					jars = append(jars, classes)
 				}
 			}
 			if len(srcJars) > 0 {
-				classes := android.PathForModuleOut(ctx, "javac", jarName+strconv.Itoa(len(shardSrcs)))
-				TransformJavaToClasses(ctx, classes, len(shardSrcs), nil, srcJars, flags, extraJarDeps)
+				classes := j.compileJavaClasses(ctx, jarName, len(shardSrcs),
+					nil, srcJars, flags, extraJarDeps)
 				jars = append(jars, classes)
 			}
 		} else {
-			classes := android.PathForModuleOut(ctx, "javac", jarName)
-			TransformJavaToClasses(ctx, classes, -1, uniqueSrcFiles, srcJars, flags, extraJarDeps)
+			classes := j.compileJavaClasses(ctx, jarName, -1, uniqueSrcFiles, srcJars, flags, extraJarDeps)
 			jars = append(jars, classes)
 		}
-		if ctx.Config().EmitXrefRules() {
-			extractionFile := android.PathForModuleOut(ctx, ctx.ModuleName()+".kzip")
-			emitXrefRule(ctx, extractionFile, uniqueSrcFiles, srcJars, flags, extraJarDeps, "xref")
-			j.kytheFiles = append(j.kytheFiles, extractionFile)
-
-		}
 		if ctx.Failed() {
 			return
 		}
@@ -1392,6 +1386,27 @@ func (j *Module) compile(ctx android.ModuleContext, aaptSrcJar android.Path) {
 	j.outputFile = outputFile.WithoutRel()
 }
 
+func (j *Module) compileJavaClasses(ctx android.ModuleContext, jarName string, idx int,
+	srcFiles, srcJars android.Paths, flags javaBuilderFlags, extraJarDeps android.Paths) android.WritablePath {
+
+	kzipName := pathtools.ReplaceExtension(jarName, "kzip")
+	if idx >= 0 {
+		kzipName = strings.TrimSuffix(jarName, filepath.Ext(jarName)) + strconv.Itoa(idx) + ".kzip"
+		jarName += strconv.Itoa(idx)
+	}
+
+	classes := android.PathForModuleOut(ctx, "javac", jarName)
+	TransformJavaToClasses(ctx, classes, idx, srcFiles, srcJars, flags, extraJarDeps)
+
+	if ctx.Config().EmitXrefRules() {
+		extractionFile := android.PathForModuleOut(ctx, kzipName)
+		emitXrefRule(ctx, extractionFile, idx, srcFiles, srcJars, flags, extraJarDeps)
+		j.kytheFiles = append(j.kytheFiles, extractionFile)
+	}
+
+	return classes
+}
+
 // Check for invalid kotlinc flags. Only use this for flags explicitly passed by the user,
 // since some of these flags may be used internally.
 func CheckKotlincFlags(ctx android.ModuleContext, flags []string) {