Trigger on-device AOT compilation for system_server on mainline update.

Before this change, AOT compilation for system_server are triggered only
on ART mainline update, and only the system_server components in /system are compiled.

After this change, AOT compilation for system_server are triggered on
any mainline update that touches $BOOTCLASSPATH or
$SYSTEMSERVERCLASSPATH, and the system_server components in mainline
modules are also compiled.

Test: manual - 1. Modify a mainline module (such as com.android.wifi).
    2. Install the modified module on a device and reboot the device.
    3. See if system_server components are recompiled (while boot
    classpath components are not).
Test: atest odsign_e2e_tests
Test: art/tools/run-gtests.sh -j8 apex/com.android.art/bin/art/x86_64/art_libartbase_tests
Test: atest art_odrefresh_tests
Bug: 189467174
Change-Id: Icaba5471e9f62fd035f0a24d662ecfaf2e064cab
diff --git a/libartbase/base/file_utils.cc b/libartbase/base/file_utils.cc
index cb0023e..5650457 100644
--- a/libartbase/base/file_utils.cc
+++ b/libartbase/base/file_utils.cc
@@ -430,13 +430,14 @@
 
 static std::string GetApexDataDalvikCacheFilename(std::string_view dex_location,
                                                   InstructionSet isa,
-                                                  bool encode_location,
+                                                  bool is_boot_classpath_location,
                                                   std::string_view file_extension) {
-  if (LocationIsOnApex(dex_location)) {
+  if (LocationIsOnApex(dex_location) && is_boot_classpath_location) {
+    // We don't compile boot images for updatable APEXes.
     return {};
   }
   std::string apex_data_dalvik_cache = GetApexDataDalvikCacheDirectory(isa);
-  if (encode_location) {
+  if (!is_boot_classpath_location) {
     // Arguments: "/system/framework/xyz.jar", "arm", true, "odex"
     // Result:
     // "/data/misc/apexdata/com.android.art/dalvik-cache/arm/system@framework@xyz.jar@classes.odex"
@@ -455,24 +456,25 @@
 }
 
 std::string GetApexDataOatFilename(std::string_view location, InstructionSet isa) {
-  return GetApexDataDalvikCacheFilename(location, isa, /*encode_location=*/false, "oat");
+  return GetApexDataDalvikCacheFilename(location, isa, /*is_boot_classpath_location=*/true, "oat");
 }
 
 std::string GetApexDataOdexFilename(std::string_view location, InstructionSet isa) {
-  return GetApexDataDalvikCacheFilename(location, isa, /*encode_location=*/true, "odex");
+  return GetApexDataDalvikCacheFilename(
+      location, isa, /*is_boot_classpath_location=*/false, "odex");
 }
 
 std::string GetApexDataBootImage(std::string_view dex_location) {
   return GetApexDataDalvikCacheFilename(dex_location,
                                         InstructionSet::kNone,
-                                        /*encode_location=*/false,
+                                        /*is_boot_classpath_location=*/true,
                                         kArtImageExtension);
 }
 
 std::string GetApexDataImage(std::string_view dex_location) {
   return GetApexDataDalvikCacheFilename(dex_location,
                                         InstructionSet::kNone,
-                                        /*encode_location=*/true,
+                                        /*is_boot_classpath_location=*/false,
                                         kArtImageExtension);
 }
 
@@ -480,7 +482,7 @@
                                            InstructionSet isa,
                                            std::string_view file_extension) {
   return GetApexDataDalvikCacheFilename(
-      dex_location, isa, /*encode_location=*/true, file_extension);
+      dex_location, isa, /*is_boot_classpath_location=*/false, file_extension);
 }
 
 std::string GetVdexFilename(const std::string& oat_location) {
diff --git a/libartbase/base/file_utils.h b/libartbase/base/file_utils.h
index c1b0095..337640e 100644
--- a/libartbase/base/file_utils.h
+++ b/libartbase/base/file_utils.h
@@ -93,20 +93,20 @@
 // than in an APEX. Returns the oat filename if `location` is valid, empty string otherwise.
 std::string GetApexDataOatFilename(std::string_view location, InstructionSet isa);
 
-// Gets the odex location in the ART APEX data directory for a DEX file installed anywhere other
-// than in an APEX. Returns the odex filename if `location` is valid, empty string otherwise.
+// Gets the odex location in the ART APEX data directory for a DEX file. Returns the odex filename
+// if `location` is valid, empty string otherwise.
 std::string GetApexDataOdexFilename(std::string_view location, InstructionSet isa);
 
 // Gets the boot image in the ART APEX data directory for a DEX file installed anywhere other
 // than in an APEX. Returns the image location if `dex_location` is valid, empty string otherwise.
 std::string GetApexDataBootImage(std::string_view dex_location);
 
-// Gets the image in the ART APEX data directory for a DEX file installed installed anywhere other
-// than in an APEX. Returns the image location if `dex_location` is valid, empty string otherwise.
+// Gets the image in the ART APEX data directory for a DEX file. Returns the image location if
+// `dex_location` is valid, empty string otherwise.
 std::string GetApexDataImage(std::string_view dex_location);
 
 // Gets the name of a file in the ART APEX directory dalvik-cache. This method assumes the
-// `dex_location` is for an application and that the `dex_location` is not within an APEX.
+// `dex_location` is for an application.
 // Returns the location of the file in the dalvik-cache
 std::string GetApexDataDalvikCacheFilename(std::string_view dex_location,
                                            InstructionSet isa,
diff --git a/libartbase/base/file_utils_test.cc b/libartbase/base/file_utils_test.cc
index 4dce3dc..e9093b2 100644
--- a/libartbase/base/file_utils_test.cc
+++ b/libartbase/base/file_utils_test.cc
@@ -20,12 +20,15 @@
 #include <stdlib.h>
 
 #include <optional>
+#include <vector>
 
 #include "base/stl_util.h"
 #include "common_art_test.h"
 
 namespace art {
 
+static constexpr const char kAndroidWifiApexDefaultPath[] = "/apex/com.android.wifi";
+
 class FileUtilsTest : public CommonArtTest {};
 
 TEST_F(FileUtilsTest, GetDalvikCacheFilename) {
@@ -169,11 +172,11 @@
             GetApexDataOatFilename("/product/javalib/beep.jar", InstructionSet::kArm));
 
   const std::string art_apex_jar = std::string {kAndroidArtApexDefaultPath} + "/javalib/some.jar";
-  EXPECT_EQ(std::string{}, GetApexDataOatFilename(art_apex_jar.c_str(), InstructionSet::kArm));
+  EXPECT_EQ(std::string {}, GetApexDataOatFilename(art_apex_jar.c_str(), InstructionSet::kArm));
 
   const std::string i18n_jar =
       std::string {kAndroidI18nApexDefaultPath} + "/javalib/core-icu4j.jar";
-  EXPECT_EQ(std::string{}, GetApexDataOatFilename(i18n_jar, InstructionSet::kArm));
+  EXPECT_EQ(std::string {}, GetApexDataOatFilename(i18n_jar, InstructionSet::kArm));
 
   const std::string system_jar_apexdata_oat = GetArtApexData() + "/dalvik-cache/x86/boot-lace.oat";
   EXPECT_EQ(system_jar_apexdata_oat,
@@ -188,11 +191,15 @@
             GetApexDataOdexFilename("/data/some/code.dex", InstructionSet::kArm));
 
   const std::string art_apex_jar = std::string {kAndroidArtApexDefaultPath} + "/javalib/some.jar";
-  EXPECT_EQ(std::string{}, GetApexDataOdexFilename(art_apex_jar.c_str(), InstructionSet::kArm));
+  EXPECT_EQ(
+      GetArtApexData() + "/dalvik-cache/arm/apex@com.android.art@javalib@some.jar@classes.odex",
+      GetApexDataOdexFilename(art_apex_jar.c_str(), InstructionSet::kArm));
 
   const std::string i18n_jar =
       std::string {kAndroidI18nApexDefaultPath} + "/javalib/core-icu4j.jar";
-  EXPECT_EQ(std::string{}, GetApexDataOdexFilename(i18n_jar.c_str(), InstructionSet::kArm));
+  EXPECT_EQ(GetArtApexData() +
+                "/dalvik-cache/arm/apex@com.android.i18n@javalib@core-icu4j.jar@classes.odex",
+            GetApexDataOdexFilename(i18n_jar.c_str(), InstructionSet::kArm));
 
   const std::string system_jar_apexdata_odex =
       GetArtApexData() + "/dalvik-cache/x86/system@framework@cookie.jar@classes.odex";
@@ -204,7 +211,7 @@
   ScopedUnsetEnvironmentVariable android_root("ANDROID_ROOT");
   ScopedUnsetEnvironmentVariable art_apex_data("ART_APEX_DATA");
 
-  EXPECT_EQ(std::string{},
+  EXPECT_EQ(std::string {},
             GetApexDataBootImage(std::string {kAndroidI18nApexDefaultPath} + "/javalib/bar.jar"));
 
   // Check image location has the prefix "boot-" in front of the basename of dex location and
@@ -224,8 +231,9 @@
   ScopedUnsetEnvironmentVariable android_root("ANDROID_ROOT");
   ScopedUnsetEnvironmentVariable art_apex_data("ART_APEX_DATA");
 
-  EXPECT_EQ(std::string{},
-            GetApexDataImage(std::string {kAndroidI18nApexDefaultPath} + "/lib/javalib/bar.jar"));
+  EXPECT_EQ(
+      GetArtApexData() + "/dalvik-cache/apex@com.android.wifi@lib@javalib@bar.jar@classes.art",
+      GetApexDataImage(std::string {kAndroidWifiApexDefaultPath} + "/lib/javalib/bar.jar"));
 
   // Check image has basename of dex location with the .art suffix.
   const char* jar = "/system/framework/mcguffin/test.jar";
@@ -241,9 +249,9 @@
 }
 
 TEST_F(FileUtilsTest, GetApexDataDalvikCacheFilename) {
-  // Check /apex inputs return empty string
-  const std::string apex_jar = std::string {kAndroidI18nApexDefaultPath} + "/lib/javalib/bar.jar";
-  EXPECT_EQ(std::string{},
+  const std::string apex_jar = std::string {kAndroidWifiApexDefaultPath} + "/lib/javalib/bar.jar";
+  EXPECT_EQ(GetArtApexData() +
+                "/dalvik-cache/x86_64/apex@com.android.wifi@lib@javalib@bar.jar@classes.art",
             GetApexDataDalvikCacheFilename(apex_jar, InstructionSet::kX86_64, "art"));
 
   // Check dalvik-cache filename follows convention.
diff --git a/odrefresh/CacheInfo.xsd b/odrefresh/CacheInfo.xsd
index b58453a..485c6b6 100644
--- a/odrefresh/CacheInfo.xsd
+++ b/odrefresh/CacheInfo.xsd
@@ -26,6 +26,7 @@
     <xs:complexType>
     <xs:sequence>
       <xs:element name="artModuleInfo" minOccurs="1" maxOccurs="1" type="t:artModuleInfo" />
+      <xs:element name="bootClasspath" type="t:bootClasspath" />
       <xs:element name="dex2oatBootClasspath" type="t:dex2oatBootClasspath" />
       <xs:element name="systemServerClasspath" type="t:systemServerClasspath" />
     </xs:sequence>
@@ -42,6 +43,13 @@
     <xs:attribute name="lastUpdateMillis" type="xs:long" use="required" />
   </xs:complexType>
 
+  <!-- Components of the `BOOTCLASSPATH`. -->
+  <xs:complexType name="bootClasspath">
+    <xs:sequence>
+      <xs:element name="component" type="t:component" />
+    </xs:sequence>
+  </xs:complexType>
+
   <!-- Components of the `DEX2OATBOOTCLASSPATH`. -->
   <xs:complexType name="dex2oatBootClasspath">
     <xs:sequence>
diff --git a/odrefresh/odr_config.h b/odrefresh/odr_config.h
index 89ca8c7..23d5cb6 100644
--- a/odrefresh/odr_config.h
+++ b/odrefresh/odr_config.h
@@ -55,6 +55,7 @@
   std::string updatable_bcp_packages_file_;
   ZygoteKind zygote_kind_;
   std::string compilation_os_address_;
+  std::string boot_classpath_;
 
  public:
   explicit OdrConfig(const char* program_name)
@@ -139,6 +140,10 @@
   void SetUpdatableBcpPackagesFile(const std::string& file) { updatable_bcp_packages_file_ = file; }
   void SetZygoteKind(ZygoteKind zygote_kind) { zygote_kind_ = zygote_kind; }
 
+  const std::string& GetBootClasspath() const { return boot_classpath_; }
+
+  void SetBootClasspath(const std::string& classpath) { boot_classpath_ = classpath; }
+
  private:
   // Returns a pair for the possible instruction sets for the configured instruction set
   // architecture. The first item is the 32-bit architecture and the second item is the 64-bit
diff --git a/odrefresh/odrefresh.cc b/odrefresh/odrefresh.cc
index 6f8fbb8..fbf30c4 100644
--- a/odrefresh/odrefresh.cc
+++ b/odrefresh/odrefresh.cc
@@ -251,6 +251,10 @@
   // List of system_server components that should be compiled.
   std::vector<std::string> systemserver_compilable_jars_;
 
+  // List of all boot classpath components. Used as the dependencies for compiling the
+  // system_server.
+  std::vector<std::string> boot_classpath_jars_;
+
   const time_t start_time_;
 
  public:
@@ -267,13 +271,8 @@
       }
     }
 
-    for (const std::string& jar : android::base::Split(config_.GetSystemServerClasspath(), ":")) {
-      // Only consider DEX files on the SYSTEMSERVERCLASSPATH for compilation that do not reside
-      // in APEX modules. Otherwise, we'll recompile on boot any time one of these APEXes updates.
-      if (!LocationIsOnApex(jar)) {
-        systemserver_compilable_jars_.emplace_back(jar);
-      }
-    }
+    systemserver_compilable_jars_ = android::base::Split(config_.GetSystemServerClasspath(), ":");
+    boot_classpath_jars_ = android::base::Split(config_.GetBootClasspath(), ":");
   }
 
   time_t GetExecutionTimeUsed() const { return time(nullptr) - start_time_; }
@@ -331,9 +330,16 @@
     const std::vector<art_apex::ArtModuleInfo> art_module_infos { art_module_info.value() };
 
     std::optional<std::vector<art_apex::Component>> bcp_components =
-        GenerateBootExtensionComponents();
+        GenerateBootClasspathComponents();
     if (!bcp_components.has_value()) {
-      LOG(ERROR) << "No boot classpath extension components.";
+      LOG(ERROR) << "No boot classpath components.";
+      return;
+    }
+
+    std::optional<std::vector<art_apex::Component>> bcp_compilable_components =
+        GenerateBootExtensionCompilableComponents();
+    if (!bcp_compilable_components.has_value()) {
+      LOG(ERROR) << "No boot classpath extension compilable components.";
       return;
     }
 
@@ -346,7 +352,8 @@
 
     std::ofstream out(cache_info_filename_.c_str());
     art_apex::CacheInfo info{art_module_infos,
-                             {{art_apex::Dex2oatBootClasspath{bcp_components.value()}}},
+                             {{art_apex::BootClasspath{bcp_components.value()}}},
+                             {{art_apex::Dex2oatBootClasspath{bcp_compilable_components.value()}}},
                              {{art_apex::SystemServerClasspath{system_server_components.value()}}}};
 
     art_apex::write(out, info);
@@ -439,7 +446,11 @@
     return components;
   }
 
-  std::vector<art_apex::Component> GenerateBootExtensionComponents() const {
+  std::vector<art_apex::Component> GenerateBootClasspathComponents() const {
+    return GenerateComponents(boot_classpath_jars_);
+  }
+
+  std::vector<art_apex::Component> GenerateBootExtensionCompilableComponents() const {
     return GenerateComponents(boot_extension_compilable_jars_);
   }
 
@@ -579,9 +590,9 @@
     //
     // The boot class components may change unexpectedly, for example an OTA could update
     // framework.jar.
-    const std::vector<art_apex::Component> expected_bcp_components =
-        GenerateBootExtensionComponents();
-    if (expected_bcp_components.size() != 0 &&
+    const std::vector<art_apex::Component> expected_bcp_compilable_components =
+        GenerateBootExtensionCompilableComponents();
+    if (expected_bcp_compilable_components.size() != 0 &&
         (!cache_info->hasDex2oatBootClasspath() ||
          !cache_info->getFirstDex2oatBootClasspath()->hasComponent())) {
       LOG(INFO) << "Missing Dex2oatBootClasspath components.";
@@ -590,9 +601,10 @@
     }
 
     std::string error_msg;
-    const std::vector<art_apex::Component>& bcp_components =
+    const std::vector<art_apex::Component>& bcp_compilable_components =
         cache_info->getFirstDex2oatBootClasspath()->getComponent();
-    if (!CheckComponents(expected_bcp_components, bcp_components, &error_msg)) {
+    if (!CheckComponents(
+            expected_bcp_compilable_components, bcp_compilable_components, &error_msg)) {
       LOG(INFO) << "Dex2OatClasspath components mismatch: " << error_msg;
       metrics.SetTrigger(OdrMetrics::Trigger::kDexFilesChanged);
       return cleanup_return(ExitCode::kCompilationRequired);
@@ -629,6 +641,25 @@
       return cleanup_system_server_return(ExitCode::kCompilationRequired);
     }
 
+    const std::vector<art_apex::Component> expected_bcp_components =
+        GenerateBootClasspathComponents();
+    if (expected_bcp_components.size() != 0 &&
+        (!cache_info->hasBootClasspath() || !cache_info->getFirstBootClasspath()->hasComponent())) {
+      LOG(INFO) << "Missing BootClasspath components.";
+      metrics.SetTrigger(OdrMetrics::Trigger::kDexFilesChanged);
+      return cleanup_system_server_return(ExitCode::kCompilationRequired);
+    }
+
+    const std::vector<art_apex::Component>& bcp_components =
+        cache_info->getFirstBootClasspath()->getComponent();
+    if (!CheckComponents(expected_bcp_components, bcp_components, &error_msg)) {
+      LOG(INFO) << "BootClasspath components mismatch: " << error_msg;
+      metrics.SetTrigger(OdrMetrics::Trigger::kDexFilesChanged);
+      // Boot classpath components can be dependencies of system_server components, so system_server
+      // components need to be recompiled if boot classpath components are changed.
+      return cleanup_system_server_return(ExitCode::kCompilationRequired);
+    }
+
     // Cache info looks good, check all compilation artifacts exist.
     auto cleanup_boot_extensions_return = [this](ExitCode exit_code, InstructionSet isa) {
       return RemoveBootExtensionArtifactsFromData(isa) ? exit_code : ExitCode::kCleanupFailed;
@@ -785,7 +816,7 @@
       args.emplace_back(
           Concatenate({"--isa=", GetInstructionSetString(config_.GetSystemServerIsa())}));
       args.emplace_back("--runtime-arg");
-      args.emplace_back(Concatenate({"-Xbootclasspath:", config_.GetDex2oatBootClasspath()}));
+      args.emplace_back(Concatenate({"-Xbootclasspath:", config_.GetBootClasspath()}));
       args.emplace_back(Concatenate(
           {"--class-loader-context=PCL[", android::base::Join(classloader_context, ':'), "]"}));
 
@@ -1252,8 +1283,8 @@
       }
 
       args.emplace_back("--runtime-arg");
-      args.emplace_back(Concatenate({"-Xbootclasspath:", config_.GetDex2oatBootClasspath()}));
-      auto bcp_jars = android::base::Split(config_.GetDex2oatBootClasspath(), ":");
+      args.emplace_back(Concatenate({"-Xbootclasspath:", config_.GetBootClasspath()}));
+      auto bcp_jars = android::base::Split(config_.GetBootClasspath(), ":");
       if (!AddBootClasspathFds(args, readonly_files_raii, bcp_jars)) {
         return false;
       }
@@ -1508,6 +1539,7 @@
   static int InitializeTargetConfig(int argc, const char** argv, OdrConfig* config) {
     config->SetApexInfoListFile("/apex/apex-info-list.xml");
     config->SetArtBinDir(GetArtBinDir());
+    config->SetBootClasspath(GetEnvironmentVariableOrDie("BOOTCLASSPATH"));
     config->SetDex2oatBootclasspath(GetEnvironmentVariableOrDie("DEX2OATBOOTCLASSPATH"));
     config->SetSystemServerClasspath(GetEnvironmentVariableOrDie("SYSTEMSERVERCLASSPATH"));
     config->SetIsa(kRuntimeISA);
diff --git a/odrefresh/schema/current.txt b/odrefresh/schema/current.txt
index e6933f6..9a38400 100644
--- a/odrefresh/schema/current.txt
+++ b/odrefresh/schema/current.txt
@@ -11,12 +11,20 @@
     method public void setVersionName(String);
   }
 
+  public class BootClasspath {
+    ctor public BootClasspath();
+    method public com.android.art.Component getComponent();
+    method public void setComponent(com.android.art.Component);
+  }
+
   public class CacheInfo {
     ctor public CacheInfo();
     method public com.android.art.ArtModuleInfo getArtModuleInfo();
+    method public com.android.art.BootClasspath getBootClasspath();
     method public com.android.art.Dex2oatBootClasspath getDex2oatBootClasspath();
     method public com.android.art.SystemServerClasspath getSystemServerClasspath();
     method public void setArtModuleInfo(com.android.art.ArtModuleInfo);
+    method public void setBootClasspath(com.android.art.BootClasspath);
     method public void setDex2oatBootClasspath(com.android.art.Dex2oatBootClasspath);
     method public void setSystemServerClasspath(com.android.art.SystemServerClasspath);
   }
diff --git a/test/odsign/test-src/com/android/tests/odsign/OnDeviceSigningHostTest.java b/test/odsign/test-src/com/android/tests/odsign/OnDeviceSigningHostTest.java
index 5245698..362e33d 100644
--- a/test/odsign/test-src/com/android/tests/odsign/OnDeviceSigningHostTest.java
+++ b/test/odsign/test-src/com/android/tests/odsign/OnDeviceSigningHostTest.java
@@ -39,7 +39,10 @@
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.Optional;
 import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.regex.Matcher;
 import java.util.stream.Collectors;
 
 @RunWith(DeviceJUnit4ClassRunner.class)
@@ -54,6 +57,8 @@
     private static final String ODREFRESH_COMPILATION_LOG =
             "/data/misc/odrefresh/compilation-log.txt";
 
+    private static final String CACHE_INFO_FILE = ART_APEX_DALVIK_CACHE_DIRNAME + "/cache-info.xml";
+
     private final String[] APP_ARTIFACT_EXTENSIONS = new String[] {".art", ".odex", ".vdex"};
 
     private final String[] BCP_ARTIFACT_EXTENSIONS = new String[] {".art", ".oat", ".vdex"};
@@ -65,6 +70,8 @@
 
     private static final Duration BOOT_COMPLETE_TIMEOUT = Duration.ofMinutes(2);
 
+    private final String[] ZYGOTE_NAMES = new String[] {"zygote", "zygote64"};
+
     @Before
     public void setUp() throws Exception {
         assumeTrue("Updating APEX is not supported", mInstallUtils.isApexUpdateSupported());
@@ -113,10 +120,36 @@
         return mappedFiles;
     }
 
+    /**
+     * Returns the mapped artifacts of the Zygote process, or {@code Optional.empty()} if the
+     * process does not exist.
+     */
+    private Optional<Set<String>> getZygoteLoadedArtifacts(String zygoteName) throws Exception {
+        final CommandResult pgrepResult = getDevice().executeShellV2Command("pgrep " + zygoteName);
+        if (pgrepResult.getExitCode() != 0) {
+            return Optional.empty();
+        }
+        final String zygotePid = pgrepResult.getStdout();
+
+        final String bootExtensionName = "boot-framework";
+        return Optional.of(getMappedArtifacts(zygotePid, bootExtensionName));
+    }
+
+    private Set<String> getSystemServerLoadedArtifacts() throws Exception {
+        String systemServerPid = getDevice().executeShellCommand("pgrep system_server");
+        assertTrue(systemServerPid != null);
+
+        // system_server artifacts are in the APEX data dalvik cache and names all contain
+        // the word "@classes". Look for mapped files that match this pattern in the proc map for
+        // system_server.
+        final String grepPattern = ART_APEX_DALVIK_CACHE_DIRNAME + ".*@classes";
+        return getMappedArtifacts(systemServerPid, grepPattern);
+    }
+
     private String[] getSystemServerClasspath() throws Exception {
         String systemServerClasspath =
                 getDevice().executeShellCommand("echo $SYSTEMSERVERCLASSPATH");
-        return systemServerClasspath.split(":");
+        return systemServerClasspath.trim().split(":");
     }
 
     private String getSystemServerIsa(String mappedArtifact) {
@@ -132,56 +165,34 @@
         String[] classpathElements = getSystemServerClasspath();
         assertTrue("SYSTEMSERVERCLASSPATH is empty", classpathElements.length > 0);
 
-        String systemServerPid = getDevice().executeShellCommand("pgrep system_server");
-        assertTrue(systemServerPid != null);
-
-        // system_server artifacts are in the APEX data dalvik cache and names all contain
-        // the word "@classes". Look for mapped files that match this pattern in the proc map for
-        // system_server.
-        final String grepPattern = ART_APEX_DALVIK_CACHE_DIRNAME + ".*@classes";
-        final Set<String> mappedArtifacts = getMappedArtifacts(systemServerPid, grepPattern);
+        final Set<String> mappedArtifacts = getSystemServerLoadedArtifacts();
         assertTrue(
                 "No mapped artifacts under " + ART_APEX_DALVIK_CACHE_DIRNAME,
                 mappedArtifacts.size() > 0);
         final String isa = getSystemServerIsa(mappedArtifacts.iterator().next());
         final String isaCacheDirectory = String.format("%s/%s", ART_APEX_DALVIK_CACHE_DIRNAME, isa);
 
-        // Check the non-APEX components in the system_server classpath have mapped artifacts.
+        // Check components in the system_server classpath have mapped artifacts.
         for (String element : classpathElements) {
-            // Skip system_server classpath elements from APEXes as these are not currently
-            // compiled.
-            if (element.startsWith("/apex")) {
-                continue;
-            }
-            String escapedPath = element.substring(1).replace('/', '@');
-            for (String extension : APP_ARTIFACT_EXTENSIONS) {
-                final String fullArtifactPath =
-                        String.format("%s/%s@classes%s", isaCacheDirectory, escapedPath, extension);
-                assertTrue(
-                        "Missing " + fullArtifactPath, mappedArtifacts.contains(fullArtifactPath));
-            }
+          String escapedPath = element.substring(1).replace('/', '@');
+          for (String extension : APP_ARTIFACT_EXTENSIONS) {
+            final String fullArtifactPath =
+                String.format("%s/%s@classes%s", isaCacheDirectory, escapedPath, extension);
+            assertTrue("Missing " + fullArtifactPath, mappedArtifacts.contains(fullArtifactPath));
+          }
         }
 
         for (String mappedArtifact : mappedArtifacts) {
-            // Check no APEX JAR artifacts are mapped for system_server since if there
-            // are, then the policy around not compiling APEX jars for system_server has
-            // changed and this test needs updating here and in the system_server classpath
-            // check above.
-            assertTrue(
-                    "Unexpected mapped artifact: " + mappedArtifact,
-                    mappedArtifact.contains("/apex"));
-
-            // Check the mapped artifact has a .art, .odex or .vdex extension.
-            final boolean knownArtifactKind =
-                    Arrays.stream(APP_ARTIFACT_EXTENSIONS)
-                            .anyMatch(e -> mappedArtifact.endsWith(e));
-            assertTrue("Unknown artifact kind: " + mappedArtifact, knownArtifactKind);
+          // Check the mapped artifact has a .art, .odex or .vdex extension.
+          final boolean knownArtifactKind =
+              Arrays.stream(APP_ARTIFACT_EXTENSIONS).anyMatch(e -> mappedArtifact.endsWith(e));
+          assertTrue("Unknown artifact kind: " + mappedArtifact, knownArtifactKind);
         }
     }
 
-    private void verifyZygoteLoadedArtifacts(String zygoteName, String zygotePid) throws Exception {
+    private void verifyZygoteLoadedArtifacts(String zygoteName, Set<String> mappedArtifacts)
+            throws Exception {
         final String bootExtensionName = "boot-framework";
-        final Set<String> mappedArtifacts = getMappedArtifacts(zygotePid, bootExtensionName);
 
         assertTrue("Expect 3 boot-framework artifacts", mappedArtifacts.size() == 3);
 
@@ -198,14 +209,12 @@
         // instances 32-bit and 64-bit unspecialized app_process processes.
         // (frameworks/base/cmds/app_process).
         int zygoteCount = 0;
-        for (String zygoteName : new String[] {"zygote", "zygote64"}) {
-            final CommandResult pgrepResult =
-                    getDevice().executeShellV2Command("pgrep " + zygoteName);
-            if (pgrepResult.getExitCode() != 0) {
+        for (String zygoteName : ZYGOTE_NAMES) {
+            final Optional<Set<String>> mappedArtifacts = getZygoteLoadedArtifacts(zygoteName);
+            if (mappedArtifacts.isEmpty()) {
                 continue;
             }
-            final String zygotePid = pgrepResult.getStdout();
-            verifyZygoteLoadedArtifacts(zygoteName, zygotePid);
+            verifyZygoteLoadedArtifacts(zygoteName, mappedArtifacts.get());
             zygoteCount += 1;
         }
         assertTrue("No zygote processes found", zygoteCount > 0);
@@ -284,6 +293,59 @@
         verifyGeneratedArtifactsLoaded();
     }
 
+    /**
+     * A workaround to simulate that an APEX has been upgraded. We could install a real APEX, but
+     * that would introduce an extra dependency to this test, which we want to avoid.
+     */
+    void mutateCacheInfo() throws Exception {
+        String cacheInfo = getDevice().pullFileContents(CACHE_INFO_FILE);
+        StringBuffer output = new StringBuffer();
+        // com.android.wifi is a module in $BOOTCLASSPATH, but not in $DEX2OATBOOTCLASSPATH.
+        Pattern p = Pattern.compile("(.*/apex/com\\.android\\.wifi.*checksums=\\\").*?(\\\".*)");
+        for (String line : cacheInfo.split("\n")) {
+            Matcher m = p.matcher(line);
+            if (m.matches()) {
+                m.appendReplacement(output, "$1aaaaaaaa$2");
+                output.append("\n");
+            } else {
+                output.append(line + "\n");
+            }
+        }
+        getDevice().pushString(output.toString(), CACHE_INFO_FILE);
+    }
+
+    long getModifiedTimeSec(String filename) throws Exception {
+        String timeStr = getDevice()
+                .executeShellCommand(String.format("stat -c '%%Y' '%s'", filename))
+                .trim();
+        return Long.parseLong(timeStr);
+    }
+
+    @Test
+    public void verifyApexUpgradeTriggersCompilation() throws Exception {
+        Set<String> zygoteArtifacts = new HashSet<>();
+        for (String zygoteName : ZYGOTE_NAMES) {
+            zygoteArtifacts.addAll(getZygoteLoadedArtifacts(zygoteName).orElse(new HashSet<>()));
+        }
+        Set<String> systemServerArtifacts = getSystemServerLoadedArtifacts();
+        long timeSec = Math.floorDiv(getDevice().getDeviceDate(), 1000);
+
+        mutateCacheInfo();
+        removeCompilationLogToAvoidBackoff();
+        CommandResult result =
+                getDevice().executeShellV2Command("odrefresh --compile");
+
+        for (String artifact : zygoteArtifacts) {
+            assertTrue("Boot classpath artifact " + artifact + " is re-compiled",
+                    getModifiedTimeSec(artifact) < timeSec);
+        }
+
+        for (String artifact : systemServerArtifacts) {
+            assertTrue("System server artifact " + artifact + " is not re-compiled",
+                    getModifiedTimeSec(artifact) >= timeSec);
+        }
+    }
+
     private boolean haveCompilationLog() throws Exception {
         CommandResult result =
                 getDevice().executeShellV2Command("stat " + ODREFRESH_COMPILATION_LOG);