Do not call getStagedApexInfos repeatedly

IApexService.getStagedApexInfos(session) is a costly operation. The
current implementation calls the method more than necessary.

- StagingManager calls it to get the list of apex names for each
  observer.
- Each observer calls it to get the info for each staged apex name.

In this change,, StagingManager passes the result to all observers and
each observer just use the result.

Bug: 370712193
Test: StagingManagerTest
Test: StagedInstallInternalTest
Change-Id: Ica92dbad7e72ee3c853a0fc34069c89f47f072ee

7 files changed, 67 insertions, 124 deletions

diff --git a/services/core/java/com/android/server/BinaryTransparencyService.java b/services/core/java/com/android/server/BinaryTransparencyService.java
index 05d07ae761c1..485bf31dfb64 100644
--- a/services/core/java/com/android/server/BinaryTransparencyService.java
+++ b/services/core/java/com/android/server/BinaryTransparencyService.java
@@ -1523,7 +1523,7 @@ public class BinaryTransparencyService extends SystemService {
                 @Override
                 public void onApexStaged(ApexStagedEvent event) throws RemoteException {
                     Slog.d(TAG, "A new APEX has been staged for update. There are currently "
-                            + event.stagedApexModuleNames.length + " APEX(s) staged for update. "
+                            + event.stagedApexInfos.length + " APEX(s) staged for update. "
                             + "Scheduling measurement...");
                     UpdateMeasurementsJobService.scheduleBinaryMeasurements(mContext,
                             BinaryTransparencyService.this);
diff --git a/services/core/java/com/android/server/pm/DexOptHelper.java b/services/core/java/com/android/server/pm/DexOptHelper.java
index e34bdc60cfdb..07f069f98edc 100644
--- a/services/core/java/com/android/server/pm/DexOptHelper.java
+++ b/services/core/java/com/android/server/pm/DexOptHelper.java
@@ -87,6 +87,7 @@ import java.io.File;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
@@ -813,7 +814,8 @@ public final class DexOptHelper {
 
         @Override
         public void onApexStaged(@NonNull ApexStagedEvent event) {
-            mArtManager.onApexStaged(event.stagedApexModuleNames);
+            mArtManager.onApexStaged(Arrays.stream(event.stagedApexInfos)
+                    .map(info -> info.moduleName).toArray(String[]::new));
         }
     }
 }
diff --git a/services/core/java/com/android/server/pm/PackageManagerNative.java b/services/core/java/com/android/server/pm/PackageManagerNative.java
index 66ecd6e67e56..7d8573e35522 100644
--- a/services/core/java/com/android/server/pm/PackageManagerNative.java
+++ b/services/core/java/com/android/server/pm/PackageManagerNative.java
@@ -20,7 +20,6 @@ import static android.content.pm.PackageManager.CERT_INPUT_SHA256;
 
 import static com.android.server.pm.PackageManagerService.TAG;
 
-import android.annotation.Nullable;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.IPackageManagerNative;
 import android.content.pm.IStagedApexObserver;
@@ -184,14 +183,8 @@ final class PackageManagerNative extends IPackageManagerNative.Stub {
     }
 
     @Override
-    public String[] getStagedApexModuleNames() {
-        return mPm.mInstallerService.getStagingManager()
-                .getStagedApexModuleNames().toArray(new String[0]);
-    }
-
-    @Override
-    @Nullable
-    public StagedApexInfo getStagedApexInfo(String moduleName) {
-        return mPm.mInstallerService.getStagingManager().getStagedApexInfo(moduleName);
+    public StagedApexInfo[] getStagedApexInfos() {
+        return mPm.mInstallerService.getStagingManager().getStagedApexInfos().toArray(
+                new StagedApexInfo[0]);
     }
 }
diff --git a/services/core/java/com/android/server/pm/StagingManager.java b/services/core/java/com/android/server/pm/StagingManager.java
index 74594cce0041..90b09cfe83c1 100644
--- a/services/core/java/com/android/server/pm/StagingManager.java
+++ b/services/core/java/com/android/server/pm/StagingManager.java
@@ -17,7 +17,6 @@
 package com.android.server.pm;
 
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.apex.ApexInfo;
 import android.apex.ApexSessionInfo;
 import android.apex.ApexSessionParams;
@@ -38,7 +37,6 @@ import android.os.SystemProperties;
 import android.os.Trace;
 import android.os.UserHandle;
 import android.text.TextUtils;
-import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.IntArray;
 import android.util.Slog;
@@ -65,9 +63,9 @@ import java.io.File;
 import java.io.FileReader;
 import java.io.FileWriter;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -807,14 +805,13 @@ public class StagingManager {
     }
 
     /**
-     * Returns ApexInfo about APEX contained inside the session as a {@code Map<String, ApexInfo>},
-     * where the key of the map is the module name of the ApexInfo.
+     * Returns ApexInfo about APEX contained inside the session.
      *
-     * Returns an empty map if there is any error.
+     * Returns an empty list if there is any error.
      */
     @VisibleForTesting
     @NonNull
-    Map<String, ApexInfo> getStagedApexInfos(@NonNull StagedSession session) {
+    List<ApexInfo> getStagedApexInfos(@NonNull StagedSession session) {
         Preconditions.checkArgument(session != null, "Session is null");
         Preconditions.checkArgument(!session.hasParentSessionId(),
                 session.sessionId() + " session has parent session");
@@ -824,7 +821,7 @@ public class StagingManager {
         // Even if caller calls this method on ready session, the session could be abandoned
         // right after this method is called.
         if (!session.isSessionReady() || session.isDestroyed()) {
-            return Collections.emptyMap();
+            return Collections.emptyList();
         }
 
         ApexSessionParams params = new ApexSessionParams();
@@ -838,38 +835,17 @@ public class StagingManager {
             }
         }
         params.childSessionIds = childSessionIds.toArray();
-
-        ApexInfo[] infos = mApexManager.getStagedApexInfos(params);
-        Map<String, ApexInfo> result = new ArrayMap<>();
-        for (ApexInfo info : infos) {
-            result.put(info.moduleName, info);
-        }
-        return result;
+        return Arrays.asList(mApexManager.getStagedApexInfos(params));
     }
 
     /**
-     * Returns apex module names of all packages that are staged ready
-     */
-    List<String> getStagedApexModuleNames() {
-        List<String> result = new ArrayList<>();
-        synchronized (mStagedSessions) {
-            for (int i = 0; i < mStagedSessions.size(); i++) {
-                final StagedSession session = mStagedSessions.valueAt(i);
-                if (!session.isSessionReady() || session.isDestroyed()
-                        || session.hasParentSessionId() || !session.containsApexSession()) {
-                    continue;
-                }
-                result.addAll(getStagedApexInfos(session).keySet());
-            }
-        }
-        return result;
-    }
-
-    /**
-     * Returns ApexInfo of the {@code moduleInfo} provided if it is staged, otherwise returns null.
+     * Returns ApexInfo list about APEXes contained inside all staged sessions.
+     *
+     * Returns an empty list if there is any error.
      */
-    @Nullable
-    StagedApexInfo getStagedApexInfo(String moduleName) {
+    @NonNull
+    List<StagedApexInfo> getStagedApexInfos() {
+        List<StagedApexInfo> result = new ArrayList<>();
         synchronized (mStagedSessions) {
             for (int i = 0; i < mStagedSessions.size(); i++) {
                 final StagedSession session = mStagedSessions.valueAt(i);
@@ -877,8 +853,7 @@ public class StagingManager {
                         || session.hasParentSessionId() || !session.containsApexSession()) {
                     continue;
                 }
-                ApexInfo ai = getStagedApexInfos(session).get(moduleName);
-                if (ai != null) {
+                getStagedApexInfos(session).stream().map(ai -> {
                     StagedApexInfo info = new StagedApexInfo();
                     info.moduleName = ai.moduleName;
                     info.diskImagePath = ai.modulePath;
@@ -886,17 +861,19 @@ public class StagingManager {
                     info.versionName = ai.versionName;
                     info.hasClassPathJars = ai.hasClassPathJars;
                     return info;
-                }
+                }).forEach(result::add);
             }
         }
-        return null;
+        return result;
     }
 
     private void notifyStagedApexObservers() {
         synchronized (mStagedApexObservers) {
+            List<StagedApexInfo> stagedApexInfos = getStagedApexInfos();
+            ApexStagedEvent event = new ApexStagedEvent();
+            event.stagedApexInfos =
+                    stagedApexInfos.toArray(new StagedApexInfo[stagedApexInfos.size()]);
             for (IStagedApexObserver observer : mStagedApexObservers) {
-                ApexStagedEvent event = new ApexStagedEvent();
-                event.stagedApexModuleNames = getStagedApexModuleNames().toArray(new String[0]);
                 try {
                     observer.onApexStaged(event);
                 } catch (RemoteException re) {
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java b/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
index 6f9b8dfc023d..52fa331cf7c5 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
@@ -72,7 +72,6 @@ import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
 
@@ -486,7 +485,7 @@ public class StagingManagerTest {
             FakeStagedSession session = new FakeStagedSession(239);
             session.setIsApex(true);
             // Call and verify
-            Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(session);
+            var result = mStagingManager.getStagedApexInfos(session);
             assertThat(result).isEmpty();
         }
         // Invalid session: destroyed
@@ -496,7 +495,7 @@ public class StagingManagerTest {
             session.setIsApex(true);
             session.setDestroyed(true);
             // Call and verify
-            Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(session);
+            var result = mStagingManager.getStagedApexInfos(session);
             assertThat(result).isEmpty();
         }
     }
@@ -520,8 +519,8 @@ public class StagingManagerTest {
         when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
 
         // Call and verify
-        Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(validSession);
-        assertThat(result).containsExactly(fakeApexInfos[0].moduleName, fakeApexInfos[0]);
+        List<ApexInfo> result = mStagingManager.getStagedApexInfos(validSession);
+        assertThat(result).containsExactly(fakeApexInfos[0]);
 
         ArgumentCaptor<ApexSessionParams> argumentCaptor =
                 ArgumentCaptor.forClass(ApexSessionParams.class);
@@ -544,9 +543,8 @@ public class StagingManagerTest {
         when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
 
         // Call and verify
-        Map<String, ApexInfo> result = mStagingManager.getStagedApexInfos(parentSession);
-        assertThat(result).containsExactly(fakeApexInfos[0].moduleName, fakeApexInfos[0],
-                fakeApexInfos[1].moduleName, fakeApexInfos[1]);
+        List<ApexInfo> result = mStagingManager.getStagedApexInfos(parentSession);
+        assertThat(result).containsExactly(fakeApexInfos[0], fakeApexInfos[1]);
 
         ArgumentCaptor<ApexSessionParams> argumentCaptor =
                 ArgumentCaptor.forClass(ApexSessionParams.class);
@@ -557,7 +555,7 @@ public class StagingManagerTest {
     }
 
     @Test
-    public void getStagedApexModuleNames_returnsStagedApexModules() throws Exception {
+    public void getStagedApexInfos_returnsStagedApexModules() throws Exception {
         FakeStagedSession validSession1 = new FakeStagedSession(239);
         validSession1.setIsApex(true);
         validSession1.setSessionReady();
@@ -575,8 +573,8 @@ public class StagingManagerTest {
 
         mockApexManagerGetStagedApexInfoWithSessionId();
 
-        List<String> result = mStagingManager.getStagedApexModuleNames();
-        assertThat(result).containsExactly("239", "123", "124");
+        List<StagedApexInfo> result = mStagingManager.getStagedApexInfos();
+        assertThat(result).containsExactly((Object[]) fakeStagedApexInfos("239", "123", "124"));
         verify(mApexManager, times(2)).getStagedApexInfos(any());
     }
 
@@ -605,26 +603,12 @@ public class StagingManagerTest {
         });
     }
 
-    @Test
-    public void getStagedApexInfo() throws Exception {
-        FakeStagedSession validSession1 = new FakeStagedSession(239);
-        validSession1.setIsApex(true);
-        validSession1.setSessionReady();
-        mStagingManager.createSession(validSession1);
-        ApexInfo[] fakeApexInfos = getFakeApexInfo(Arrays.asList("module1"));
-        when(mApexManager.getStagedApexInfos(any())).thenReturn(fakeApexInfos);
-
-        // Verify null is returned if module name is not found
-        StagedApexInfo result = mStagingManager.getStagedApexInfo("not found");
-        assertThat(result).isNull();
-        verify(mApexManager, times(1)).getStagedApexInfos(any());
-        // Otherwise, the correct object is returned
-        result = mStagingManager.getStagedApexInfo("module1");
-        assertThat(result.moduleName).isEqualTo(fakeApexInfos[0].moduleName);
-        assertThat(result.diskImagePath).isEqualTo(fakeApexInfos[0].modulePath);
-        assertThat(result.versionCode).isEqualTo(fakeApexInfos[0].versionCode);
-        assertThat(result.versionName).isEqualTo(fakeApexInfos[0].versionName);
-        verify(mApexManager, times(2)).getStagedApexInfos(any());
+    private StagedApexInfo[] fakeStagedApexInfos(String... moduleNames) {
+        return Arrays.stream(moduleNames).map(moduleName -> {
+            StagedApexInfo info = new StagedApexInfo();
+            info.moduleName = moduleName;
+            return info;
+        }).toArray(StagedApexInfo[]::new);
     }
 
     @Test
@@ -646,8 +630,8 @@ public class StagingManagerTest {
             ArgumentCaptor<ApexStagedEvent> argumentCaptor = ArgumentCaptor.forClass(
                     ApexStagedEvent.class);
             verify(observer, times(1)).onApexStaged(argumentCaptor.capture());
-            assertThat(argumentCaptor.getValue().stagedApexModuleNames).isEqualTo(
-                    new String[]{"239"});
+            assertThat(argumentCaptor.getValue().stagedApexInfos).isEqualTo(
+                    fakeStagedApexInfos("239"));
         }
 
         // Create another staged session and verify observers are notified of union
@@ -662,8 +646,8 @@ public class StagingManagerTest {
             ArgumentCaptor<ApexStagedEvent> argumentCaptor = ArgumentCaptor.forClass(
                     ApexStagedEvent.class);
             verify(observer, times(1)).onApexStaged(argumentCaptor.capture());
-            assertThat(argumentCaptor.getValue().stagedApexModuleNames).isEqualTo(
-                    new String[]{"239", "240"});
+            assertThat(argumentCaptor.getValue().stagedApexInfos).isEqualTo(
+                    fakeStagedApexInfos("239", "240"));
         }
 
         // Finally, verify that once unregistered, observer is not notified
@@ -699,7 +683,7 @@ public class StagingManagerTest {
         ArgumentCaptor<ApexStagedEvent> argumentCaptor = ArgumentCaptor.forClass(
                 ApexStagedEvent.class);
         verify(observer, times(1)).onApexStaged(argumentCaptor.capture());
-        assertThat(argumentCaptor.getValue().stagedApexModuleNames).hasLength(0);
+        assertThat(argumentCaptor.getValue().stagedApexInfos).hasLength(0);
     }
 
     @Test
diff --git a/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java b/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java
index 0375f66069c3..d9295dd17dd0 100644
--- a/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/app/src/com/android/tests/stagedinstallinternal/StagedInstallInternalTest.java
@@ -515,33 +515,27 @@ public class StagedInstallInternalTest {
                 Install.single(APEX_V2));
     }
 
-    @Test
-    public void testGetStagedModuleNames() throws Exception {
-        // Before staging a session
-        String[] result = getPackageManagerNative().getStagedApexModuleNames();
-        assertThat(result).hasLength(0);
-        // Stage an apex
-        int sessionId = Install.single(APEX_V2).setStaged().commit();
-        result = getPackageManagerNative().getStagedApexModuleNames();
-        assertThat(result).hasLength(1);
-        assertThat(result).isEqualTo(new String[]{SHIM_APEX_PACKAGE_NAME});
-        // Abandon the session
-        InstallUtils.openPackageInstallerSession(sessionId).abandon();
-        result = getPackageManagerNative().getStagedApexModuleNames();
-        assertThat(result).hasLength(0);
+    private StagedApexInfo findStagedApexInfo(StagedApexInfo[] infos, String moduleName) {
+        for (StagedApexInfo info: infos) {
+            if (info.moduleName.equals(moduleName)) {
+                return info;
+            }
+        }
+        return null;
     }
 
     @Test
-    public void testGetStagedApexInfo() throws Exception {
-        // Ask for non-existing module
-        StagedApexInfo result = getPackageManagerNative().getStagedApexInfo("not found");
-        assertThat(result).isNull();
+    public void testGetStagedApexInfos() throws Exception {
+        // Not found before staging
+        StagedApexInfo[] result = getPackageManagerNative().getStagedApexInfos();
+        assertThat(findStagedApexInfo(result, TEST_APEX_PACKAGE_NAME)).isNull();
         // Stage an apex
         int sessionId = Install.single(TEST_APEX_CLASSPATH).setStaged().commit();
         // Query proper module name
-        result = getPackageManagerNative().getStagedApexInfo(TEST_APEX_PACKAGE_NAME);
-        assertThat(result.moduleName).isEqualTo(TEST_APEX_PACKAGE_NAME);
-        assertThat(result.hasClassPathJars).isTrue();
+        result = getPackageManagerNative().getStagedApexInfos();
+        StagedApexInfo found = findStagedApexInfo(result, TEST_APEX_PACKAGE_NAME);
+        assertThat(found).isNotNull();
+        assertThat(found.hasClassPathJars).isTrue();
         InstallUtils.openPackageInstallerSession(sessionId).abandon();
     }
 
@@ -573,14 +567,15 @@ public class StagedInstallInternalTest {
         int sessionId = Install.single(APEX_V2).setStaged().commit();
         ArgumentCaptor<ApexStagedEvent> captor = ArgumentCaptor.forClass(ApexStagedEvent.class);
         verify(observer, timeout(5000)).onApexStaged(captor.capture());
-        assertThat(captor.getValue().stagedApexModuleNames).isEqualTo(
-                new String[] {SHIM_APEX_PACKAGE_NAME});
+        StagedApexInfo found =
+                findStagedApexInfo(captor.getValue().stagedApexInfos, SHIM_APEX_PACKAGE_NAME);
+        assertThat(found).isNotNull();
 
         // Abandon and verify observer is called
         Mockito.clearInvocations(observer);
         InstallUtils.openPackageInstallerSession(sessionId).abandon();
         verify(observer, timeout(5000)).onApexStaged(captor.capture());
-        assertThat(captor.getValue().stagedApexModuleNames).hasLength(0);
+        assertThat(captor.getValue().stagedApexInfos).hasLength(0);
     }
 
     @Test
diff --git a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
index f1fc503c8556..97abcd77e6b9 100644
--- a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
@@ -592,23 +592,15 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
     }
 
     @Test
-    public void testGetStagedModuleNames() throws Exception {
-        assumeTrue("Device does not support updating APEX",
-                mHostUtils.isApexUpdateSupported());
-
-        runPhase("testGetStagedModuleNames");
-    }
-
-    @Test
     @LargeTest
-    public void testGetStagedApexInfo() throws Exception {
+    public void testGetStagedApexInfos() throws Exception {
         assumeTrue("Device does not support updating APEX",
                 mHostUtils.isApexUpdateSupported());
 
         pushTestApex(APEXD_TEST_APEX);
         getDevice().reboot();
 
-        runPhase("testGetStagedApexInfo");
+        runPhase("testGetStagedApexInfos");
     }
 
     @Test