Install dexpreopt artifacts of apex system server jars in same partition

e.g. If a system server jar providing apex is installed in /system_ext,
the dexpreopt files of that jar will also be installed in /system_ext.
Currently, all of these artifacts are installed in /system

This behavior will be flag guarded by
RELEASE_INSTALL_APEX_SYSTEMSERVER_DEXPREOPT_SAME_PARTITION. This is
necessary because the ART runtime needs to be updated to look in the new
/system_ext location. Since some release configs build with ART
prebuilts, the runtime in those prebuilt apexes will not have the
additional search path introduced in https://r.android.com/3287191

Test: Verified that "Could not check odex file" does not appear for
service-compos (a system_ext apex systemserver jar)

Bug: 369678122
Change-Id: I752bdc7f5f69226b503800ce25726a211302cb07

7 files changed, 108 insertions, 21 deletions

diff --git a/android/config.go b/android/config.go
index 10e43cee9..e51976054 100644
--- a/android/config.go
+++ b/android/config.go
@@ -2105,3 +2105,10 @@ func (c *config) BoardAvbEnable() bool {
 func (c *config) BoardAvbSystemAddHashtreeFooterArgs() []string {
 	return c.productVariables.BoardAvbSystemAddHashtreeFooterArgs
 }
+
+// Returns true if RELEASE_INSTALL_APEX_SYSTEMSERVER_DEXPREOPT_SAME_PARTITION is set to true.
+// If true, dexpreopt files of apex system server jars will be installed in the same partition as the parent apex.
+// If false, all these files will be installed in /system partition.
+func (c Config) InstallApexSystemServerDexpreoptSamePartition() bool {
+	return c.config.productVariables.GetBuildFlagBool("RELEASE_INSTALL_APEX_SYSTEMSERVER_DEXPREOPT_SAME_PARTITION")
+}
diff --git a/android/module.go b/android/module.go
index 20caae2d8..a1a9a4ad2 100644
--- a/android/module.go
+++ b/android/module.go
@@ -81,6 +81,7 @@ type Module interface {
 	InstallInOdm() bool
 	InstallInProduct() bool
 	InstallInVendor() bool
+	InstallInSystemExt() bool
 	InstallForceOS() (*OsType, *ArchType)
 	PartitionTag(DeviceConfig) string
 	HideFromMake()
@@ -1514,6 +1515,10 @@ func (m *ModuleBase) InstallInVendor() bool {
 	return Bool(m.commonProperties.Vendor) || Bool(m.commonProperties.Soc_specific) || Bool(m.commonProperties.Proprietary)
 }
 
+func (m *ModuleBase) InstallInSystemExt() bool {
+	return Bool(m.commonProperties.System_ext_specific)
+}
+
 func (m *ModuleBase) InstallInRoot() bool {
 	return false
 }
diff --git a/dexpreopt/config.go b/dexpreopt/config.go
index fe6317cb2..84d4f10c1 100644
--- a/dexpreopt/config.go
+++ b/dexpreopt/config.go
@@ -191,6 +191,10 @@ type ModuleConfig struct {
 	ForceCreateAppImage bool
 
 	PresignedPrebuilt bool
+
+	// ApexPartition is the partition in which the dexpreopt files of apex system server jars (if any) are installed.
+	// This is a noop unless the module is apex system server jar.
+	ApexPartition string
 }
 
 type globalSoongConfigSingleton struct{}
diff --git a/dexpreopt/dexpreopt.go b/dexpreopt/dexpreopt.go
index 56164834a..7a39fa1d7 100644
--- a/dexpreopt/dexpreopt.go
+++ b/dexpreopt/dexpreopt.go
@@ -219,9 +219,9 @@ func GetSystemServerDexLocation(ctx android.PathContext, global *GlobalConfig, l
 }
 
 // Returns the location to the odex file for the dex file at `path`.
-func ToOdexPath(path string, arch android.ArchType) string {
+func ToOdexPath(path string, arch android.ArchType, partition string) string {
 	if strings.HasPrefix(path, "/apex/") {
-		return filepath.Join("/system/framework/oat", arch.String(),
+		return filepath.Join(partition, "framework/oat", arch.String(),
 			strings.ReplaceAll(path[1:], "/", "@")+"@classes.odex")
 	}
 
@@ -245,7 +245,7 @@ func dexpreoptCommand(ctx android.BuilderContext, globalSoong *GlobalSoongConfig
 
 	odexPath := module.BuildPath.InSameDir(ctx, "oat", arch.String(), pathtools.ReplaceExtension(base, "odex"))
 	odexSymbolsPath := odexPath.ReplaceExtension(ctx, "symbols.odex")
-	odexInstallPath := ToOdexPath(module.DexLocation, arch)
+	odexInstallPath := ToOdexPath(module.DexLocation, arch, module.ApexPartition)
 	if odexOnSystemOther(module, global) {
 		odexInstallPath = filepath.Join(SystemOtherPartition, odexInstallPath)
 	}
diff --git a/dexpreopt/dexpreopt_test.go b/dexpreopt/dexpreopt_test.go
index 6f7d3bb67..7b0f51fbd 100644
--- a/dexpreopt/dexpreopt_test.go
+++ b/dexpreopt/dexpreopt_test.go
@@ -42,12 +42,14 @@ func testModuleConfig(ctx android.PathContext, name, partition string) *ModuleCo
 }
 
 func testApexModuleConfig(ctx android.PathContext, name, apexName string) *ModuleConfig {
-	return createTestModuleConfig(
+	ret := createTestModuleConfig(
 		name,
 		fmt.Sprintf("/apex/%s/javalib/%s.jar", apexName, name),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/dexpreopt/%s.jar", name, name)),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/aligned/%s.jar", name, name)),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/enforce_uses_libraries.status", name)))
+	ret.ApexPartition = "/system"
+	return ret
 }
 
 func testPlatformSystemServerModuleConfig(ctx android.PathContext, name string) *ModuleConfig {
@@ -221,6 +223,49 @@ func TestDexPreoptApexSystemServerJars(t *testing.T) {
 	DexpreoptRunningInSoong = oldDexpreoptRunningInSoong
 }
 
+// Same as `TestDexPreoptApexSystemServerJars`, but the apex jar is in /system_ext
+func TestDexPreoptApexSystemServerJarsSystemExt(t *testing.T) {
+	// modify the global variable for test
+	var oldDexpreoptRunningInSoong = DexpreoptRunningInSoong
+	DexpreoptRunningInSoong = true
+
+	// test begin
+	config := android.TestConfig("out", nil, "", nil)
+	ctx := android.BuilderContextForTesting(config)
+	globalSoong := globalSoongConfigForTests(ctx)
+	global := GlobalConfigForTests(ctx)
+	module := testApexModuleConfig(ctx, "service-A", "com.android.apex1")
+	module.ApexPartition = "/system_ext"
+	productPackages := android.PathForTesting("product_packages.txt")
+
+	global.ApexSystemServerJars = android.CreateTestConfiguredJarList(
+		[]string{"com.android.apex1:service-A"})
+
+	rule, err := GenerateDexpreoptRule(ctx, globalSoong, global, module, productPackages, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wantInstalls := android.RuleBuilderInstalls{
+		{android.PathForOutput(ctx, "service-A/dexpreopt/oat/arm/javalib.odex"), "/system_ext/framework/oat/arm/apex@com.android.apex1@javalib@service-A.jar@classes.odex"},
+		{android.PathForOutput(ctx, "service-A/dexpreopt/oat/arm/javalib.vdex"), "/system_ext/framework/oat/arm/apex@com.android.apex1@javalib@service-A.jar@classes.vdex"},
+	}
+
+	android.AssertStringEquals(t, "installs", wantInstalls.String(), rule.Installs().String())
+
+	android.AssertStringListContains(t, "apex sscp jar copy", rule.Outputs().Strings(), "out/soong/system_server_dexjars/service-A.jar")
+
+	// rule with apex sscp cp as false
+	rule, err = GenerateDexpreoptRule(ctx, globalSoong, global, module, productPackages, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	android.AssertStringListDoesNotContain(t, "apex sscp jar copy", rule.Outputs().Strings(), "out/soong/system_server_dexjars/service-A.jar")
+
+	// cleanup the global variable for test
+	DexpreoptRunningInSoong = oldDexpreoptRunningInSoong
+}
+
 func TestDexPreoptStandaloneSystemServerJars(t *testing.T) {
 	config := android.TestConfig("out", nil, "", nil)
 	ctx := android.BuilderContextForTesting(config)
diff --git a/java/dexpreopt.go b/java/dexpreopt.go
index 63a863497..637da363c 100644
--- a/java/dexpreopt.go
+++ b/java/dexpreopt.go
@@ -494,6 +494,12 @@ func (d *dexpreopter) dexpreopt(ctx android.ModuleContext, libName string, dexJa
 		PresignedPrebuilt: d.isPresignedPrebuilt,
 	}
 
+	if ctx.Config().InstallApexSystemServerDexpreoptSamePartition() {
+		dexpreoptConfig.ApexPartition = android.PathForModuleInstall(ctx).Partition()
+	} else {
+		dexpreoptConfig.ApexPartition = "system"
+	}
+
 	d.configPath = android.PathForModuleOut(ctx, "dexpreopt", dexJarStem, "dexpreopt.config")
 	dexpreopt.WriteModuleConfig(ctx, dexpreoptConfig, d.configPath)
 	ctx.CheckbuildFile(d.configPath)
diff --git a/java/dexpreopt_check.go b/java/dexpreopt_check.go
index 33be60352..c97156541 100644
--- a/java/dexpreopt_check.go
+++ b/java/dexpreopt_check.go
@@ -17,6 +17,8 @@ package java
 import (
 	"strings"
 
+	"github.com/google/blueprint"
+
 	"android/soong/android"
 	"android/soong/dexpreopt"
 
@@ -43,16 +45,12 @@ func RegisterDexpreoptCheckBuildComponents(ctx android.RegistrationContext) {
 type dexpreoptSystemserverCheck struct {
 	android.SingletonModuleBase
 
-	// Mapping from the module name to the install paths to the compilation artifacts.
-	artifactsByModuleName map[string][]string
-
 	// The install paths to the compilation artifacts.
 	artifacts []string
 }
 
 func dexpreoptSystemserverCheckFactory() android.SingletonModule {
 	m := &dexpreoptSystemserverCheck{}
-	m.artifactsByModuleName = make(map[string][]string)
 	android.InitAndroidArchModule(m, android.DeviceSupported, android.MultilibCommon)
 	return m
 }
@@ -62,7 +60,25 @@ func getInstallPath(ctx android.ModuleContext, location string) android.InstallP
 		ctx, "", strings.TrimPrefix(location, "/"))
 }
 
-func (m *dexpreoptSystemserverCheck) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+type systemServerDependencyTag struct {
+	blueprint.BaseDependencyTag
+}
+
+// systemServerJarDepTag willl be used for validation. Skip visiblility.
+func (b systemServerDependencyTag) ExcludeFromVisibilityEnforcement() {
+}
+
+var (
+	// dep tag for platform and apex system server jars
+	systemServerJarDepTag = systemServerDependencyTag{}
+)
+
+var _ android.ExcludeFromVisibilityEnforcementTag = systemServerJarDepTag
+
+// Add a depenendency on the system server jars. The dexpreopt files of those will be emitted to make.
+// The kati packaging system will verify that those files appear in installed files.
+// Adding the dependency allows the singleton module to determine whether an apex system server jar is system_ext specific.
+func (m *dexpreoptSystemserverCheck) DepsMutator(ctx android.BottomUpMutatorContext) {
 	global := dexpreopt.GetGlobalConfig(ctx)
 	targets := ctx.Config().Targets[android.Android]
 
@@ -72,23 +88,27 @@ func (m *dexpreoptSystemserverCheck) GenerateAndroidBuildActions(ctx android.Mod
 		return
 	}
 
-	systemServerJars := global.AllSystemServerJars(ctx)
-	for _, jar := range systemServerJars.CopyOfJars() {
-		dexLocation := dexpreopt.GetSystemServerDexLocation(ctx, global, jar)
-		odexLocation := dexpreopt.ToOdexPath(dexLocation, targets[0].Arch.ArchType)
+	ctx.AddDependency(ctx.Module(), systemServerJarDepTag, global.AllSystemServerJars(ctx).CopyOfJars()...)
+}
+
+func (m *dexpreoptSystemserverCheck) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	global := dexpreopt.GetGlobalConfig(ctx)
+	targets := ctx.Config().Targets[android.Android]
+
+	ctx.VisitDirectDepsWithTag(systemServerJarDepTag, func(systemServerJar android.Module) {
+		partition := "system"
+		if systemServerJar.InstallInSystemExt() && ctx.Config().InstallApexSystemServerDexpreoptSamePartition() {
+			partition = ctx.DeviceConfig().SystemExtPath() // system_ext
+		}
+		dexLocation := dexpreopt.GetSystemServerDexLocation(ctx, global, systemServerJar.Name())
+		odexLocation := dexpreopt.ToOdexPath(dexLocation, targets[0].Arch.ArchType, partition)
 		odexPath := getInstallPath(ctx, odexLocation)
 		vdexPath := getInstallPath(ctx, pathtools.ReplaceExtension(odexLocation, "vdex"))
-		m.artifactsByModuleName[jar] = []string{odexPath.String(), vdexPath.String()}
-	}
+		m.artifacts = append(m.artifacts, odexPath.String(), vdexPath.String())
+	})
 }
 
 func (m *dexpreoptSystemserverCheck) GenerateSingletonBuildActions(ctx android.SingletonContext) {
-	// Only keep modules defined in Soong.
-	ctx.VisitAllModules(func(module android.Module) {
-		if artifacts, ok := m.artifactsByModuleName[module.Name()]; ok {
-			m.artifacts = append(m.artifacts, artifacts...)
-		}
-	})
 }
 
 func (m *dexpreoptSystemserverCheck) MakeVars(ctx android.MakeVarsContext) {