Optimize DexUseManagerLocal.findOwningPackage - Step 1.

This CLs contains no semantic change. Optimizations are:

1. Using UnfilteredSnapshot to defer the shouldFilterApplication check
2. Check if the dex is secondary dex of the loading package before
   checking if it's primary dex of other packages
3. Adding an LRU cache for recent packages

Bug: 328673771
Test: atest ArtServiceTests
Change-Id: I7974b0240a81e6a62ad15a6cb5825a6278ce05d1

4 files changed, 250 insertions, 56 deletions

diff --git a/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java b/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java
index 9ef804919a..6c05c372ad 100644
--- a/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java
+++ b/libartservice/service/java/com/android/server/art/DexUseManagerLocal.java
@@ -16,8 +16,10 @@
 
 package com.android.server.art;
 
+import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
 import android.content.BroadcastReceiver;
 import android.content.Context;
@@ -30,6 +32,7 @@ import android.os.Process;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
 import android.os.UserHandle;
+import android.util.LruCache;
 
 import androidx.annotation.RequiresApi;
 
@@ -59,9 +62,9 @@ import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.nio.file.StandardCopyOption;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -72,6 +75,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.SequencedMap;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.Executors;
@@ -117,6 +121,13 @@ public class DexUseManagerLocal {
     @NonNull private final Injector mInjector;
     @NonNull private final Debouncer mDebouncer;
 
+    // The cache is motivated by the fact that only a handful of packages are commonly used by other
+    // packages. The cache size is arbitrarily decided.
+    /** A map from recently used dex files to their package names. */
+    @NonNull
+    private final LruCache<String, String> mRecentDexFilesToPackageNames =
+            new LruCache<>(50 /* maxSize */);
+
     private final Object mLock = new Object();
     @GuardedBy("mLock") @NonNull private DexUse mDexUse; // Initialized by `load`.
     @GuardedBy("mLock") private int mRevision = 0;
@@ -406,52 +417,145 @@ public class DexUseManagerLocal {
         for (var entry : classLoaderContextByDexContainerFile.entrySet()) {
             String dexPath = Utils.assertNonEmpty(entry.getKey());
             String classLoaderContext = Utils.assertNonEmpty(entry.getValue());
-            String owningPackageName = findOwningPackage(snapshot, loadingPackageName,
-                    (pkgState) -> isOwningPackageForPrimaryDex(pkgState, dexPath));
-            if (owningPackageName != null) {
-                addPrimaryDexUse(owningPackageName, dexPath, loadingPackageName, isolatedProcess,
-                        lastUsedAtMs);
+            FindResult findResult = findOwningPackage(snapshot, loadingPackageName, dexPath);
+            if (findResult == null) {
                 continue;
             }
-            Path path = Paths.get(dexPath);
-            synchronized (mLock) {
-                owningPackageName = findOwningPackage(snapshot, loadingPackageName,
-                        (pkgState) -> isOwningPackageForSecondaryDexLocked(pkgState, path));
-            }
-            if (owningPackageName != null) {
-                PackageState loadingPkgState =
-                        Utils.getPackageStateOrThrow(snapshot, loadingPackageName);
-                // An app is always launched with its primary ABI.
-                Utils.Abi abi = Utils.getPrimaryAbi(loadingPkgState);
-                addSecondaryDexUse(owningPackageName, dexPath, loadingPackageName, isolatedProcess,
-                        classLoaderContext, abi.name(), lastUsedAtMs);
-                continue;
+
+            switch (findResult.type()) {
+                case TYPE_PRIMARY:
+                    addPrimaryDexUse(findResult.owningPackageName(), dexPath, loadingPackageName,
+                            isolatedProcess, lastUsedAtMs);
+                    break;
+                case TYPE_SECONDARY:
+                    PackageState loadingPkgState =
+                            Utils.getPackageStateOrThrow(snapshot, loadingPackageName);
+                    // An app is always launched with its primary ABI.
+                    Utils.Abi abi = Utils.getPrimaryAbi(loadingPkgState);
+                    addSecondaryDexUse(findResult.owningPackageName(), dexPath, loadingPackageName,
+                            isolatedProcess, classLoaderContext, abi.name(), lastUsedAtMs);
+                    break;
+                default:
+                    // Intentionally ignore.
             }
-            // It is expected that a dex file isn't owned by any package. For example, the dex
-            // file could be a shared library jar.
         }
     }
 
     @Nullable
-    private static String findOwningPackage(@NonNull PackageManagerLocal.FilteredSnapshot snapshot,
-            @NonNull String loadingPackageName,
-            @NonNull Function<PackageState, Boolean> predicate) {
+    private FindResult findOwningPackage(@NonNull PackageManagerLocal.FilteredSnapshot snapshot,
+            @NonNull String loadingPackageName, @NonNull String dexPath) {
         // Most likely, the package is loading its own dex file, so we check this first as an
         // optimization.
         PackageState loadingPkgState = Utils.getPackageStateOrThrow(snapshot, loadingPackageName);
-        if (predicate.apply(loadingPkgState)) {
-            return loadingPkgState.getPackageName();
+        FindResult result = checkForPackage(loadingPkgState, dexPath);
+        if (result != null) {
+            return result;
         }
 
-        for (PackageState pkgState : snapshot.getPackageStates().values()) {
-            if (predicate.apply(pkgState)) {
-                return pkgState.getPackageName();
+        // Check all packages.
+        result = checkForAllPackages(dexPath);
+        if (result != null && result.type() != TYPE_DONT_RECORD
+                && snapshot.getPackageState(result.owningPackageName()) != null) {
+            return result;
+        }
+
+        // It is expected that there is no result. For example, the app could be loading a dex file
+        // from a non-canonical location, or it could be sending a bogus dex filename.
+        return null;
+    }
+
+    /**
+     * Returns the owner of the given dex file, found among all packages. The return value is
+     * unfiltered and must be checked whether it's visible to the calling app.
+     */
+    @SuppressLint("NewApi") // Using new Libcore APIs from the same module.
+    @Nullable
+    private FindResult checkForAllPackages(@NonNull String dexPath) {
+        // Iterating over the filtered snapshot is slow because it involves a
+        // `shouldFilterApplication` call on every iteration, which is considerably more expensive
+        // than a `checkForPackage` call. Therefore, we iterate over the unfiltered snapshot
+        // instead.
+        // There may be some inconsistencies between the filtered snapshot and the unfiltered
+        // snapshot, as they are not created atomically, but this is fine. If a package is in the
+        // filtered snapshot but not in the unfiltered snapshot, it means the package got removed,
+        // so we don't need to record it.
+        try (PackageManagerLocal.UnfilteredSnapshot unfilteredSnapshot =
+                        mInjector.getPackageManagerLocal().withUnfilteredSnapshot()) {
+            Map<String, PackageState> packageStates = unfilteredSnapshot.getPackageStates();
+            Set<String> visitedPackages = new HashSet<>();
+
+            Function<String, FindResult> visitPackage = (packageName) -> {
+                if (visitedPackages.contains(packageName)) {
+                    return null;
+                }
+                visitedPackages.add(packageName);
+                PackageState pkgState = packageStates.get(packageName);
+                if (pkgState == null) {
+                    mRecentDexFilesToPackageNames.remove(dexPath);
+                    return null;
+                }
+                FindResult result = checkForPackage(pkgState, dexPath);
+                if (result != null) {
+                    mRecentDexFilesToPackageNames.put(dexPath, packageName);
+                    return result;
+                }
+                return null;
+            };
+
+            String cachedPackageName = mRecentDexFilesToPackageNames.get(dexPath);
+            if (cachedPackageName != null) {
+                FindResult result = visitPackage.apply(cachedPackageName);
+                if (result != null) {
+                    return result;
+                }
+            }
+
+            var recentDexFilesToPackageNamesSnapshot =
+                    (SequencedMap<String, String>) mRecentDexFilesToPackageNames.snapshot();
+
+            // Check recent packages first.
+            for (String packageName :
+                    recentDexFilesToPackageNamesSnapshot.sequencedValues().reversed()) {
+                FindResult result = visitPackage.apply(packageName);
+                if (result != null) {
+                    return result;
+                }
+            }
+
+            // Check remaining packages.
+            for (PackageState pkgState : packageStates.values()) {
+                if (visitedPackages.contains(pkgState.getPackageName())) {
+                    continue;
+                }
+                FindResult result = checkForPackage(pkgState, dexPath);
+                if (result != null) {
+                    mRecentDexFilesToPackageNames.put(dexPath, pkgState.getPackageName());
+                    return result;
+                }
             }
         }
 
         return null;
     }
 
+    @Nullable
+    private FindResult checkForPackage(@NonNull PackageState pkgState, @NonNull String dexPath) {
+        if (isOwningPackageForPrimaryDex(pkgState, dexPath)) {
+            return new FindResult(TYPE_PRIMARY, pkgState.getPackageName());
+        }
+        synchronized (mLock) {
+            if (isOwningPackageForSecondaryDexLocked(pkgState, dexPath)) {
+                return new FindResult(TYPE_SECONDARY, pkgState.getPackageName());
+            }
+        }
+        String packageCodeDir = getPackageCodeDir(pkgState);
+        if (packageCodeDir != null && Utils.pathStartsWith(dexPath, packageCodeDir)) {
+            // TODO(b/351761207): Support secondary dex files in package dir.
+            return new FindResult(TYPE_DONT_RECORD, null);
+        }
+        return null;
+    }
+
     private static boolean isOwningPackageForPrimaryDex(
             @NonNull PackageState pkgState, @NonNull String dexPath) {
         AndroidPackage pkg = pkgState.getAndroidPackage();
@@ -469,17 +573,33 @@ public class DexUseManagerLocal {
 
     @GuardedBy("mLock")
     private boolean isOwningPackageForSecondaryDexLocked(
-            @NonNull PackageState pkgState, @NonNull Path dexPath) {
+            @NonNull PackageState pkgState, @NonNull String dexPath) {
         UserHandle userHandle = Binder.getCallingUserHandle();
-        List<Path> locations = mSecondaryDexLocationManager.getLocations(pkgState, userHandle);
+        List<String> locations = mSecondaryDexLocationManager.getLocations(pkgState, userHandle);
         for (int i = 0; i < locations.size(); i++) {
-            if (dexPath.startsWith(locations.get(i))) {
+            if (Utils.pathStartsWith(dexPath, locations.get(i))) {
                 return true;
             }
         }
         return false;
     }
 
+    @Nullable
+    private String getPackageCodeDir(@NonNull PackageState pkgState) {
+        AndroidPackage pkg = pkgState.getAndroidPackage();
+        if (pkg == null) {
+            return null;
+        }
+        List<AndroidPackageSplit> splits = pkg.getSplits();
+        if (splits.size() == 0) {
+            return null;
+        }
+        String path = splits.get(0).getPath();
+        int pos = path.lastIndexOf('/');
+        Utils.check(pos >= 0);
+        return path.substring(0, pos + 1);
+    }
+
     private void addPrimaryDexUse(@NonNull String owningPackageName, @NonNull String dexPath,
             @NonNull String loadingPackageName, boolean isolatedProcess, long lastUsedAtMs) {
         synchronized (mLock) {
@@ -1109,7 +1229,7 @@ public class DexUseManagerLocal {
     static class SecondaryDexLocationManager {
         private @NonNull Map<CacheKey, CacheValue> mCache = new HashMap<>();
 
-        public @NonNull List<Path> getLocations(
+        public @NonNull List<String> getLocations(
                 @NonNull PackageState pkgState, @NonNull UserHandle userHandle) {
             AndroidPackage pkg = pkgState.getAndroidPackage();
             if (pkg == null) {
@@ -1129,11 +1249,12 @@ public class DexUseManagerLocal {
                     storageUuid, userHandle, packageName);
             File deDir = Environment.getDataDePackageDirectoryForUser(
                     storageUuid, userHandle, packageName);
-            List<Path> locations = List.of(ceDir.toPath(), deDir.toPath());
-            mCache.put(cacheKey, CacheValue.create(locations, storageUuid));
+            List<String> locations = List.of(ceDir.getAbsolutePath(), deDir.getAbsolutePath());
+            mCache.put(cacheKey, new CacheValue(locations, storageUuid));
             return locations;
         }
 
+        // TODO(b/351994199): Don't replace this with record because the latter is too slow.
         @Immutable
         @AutoValue
         abstract static class CacheKey {
@@ -1147,19 +1268,28 @@ public class DexUseManagerLocal {
             abstract @NonNull UserHandle userHandle();
         }
 
-        @Immutable
-        @AutoValue
-        abstract static class CacheValue {
-            static CacheValue create(@NonNull List<Path> locations, @NonNull UUID storageUuid) {
-                return new AutoValue_DexUseManagerLocal_SecondaryDexLocationManager_CacheValue(
-                        locations, storageUuid);
-            }
+        private record CacheValue(@NonNull List<String> locations, @NonNull UUID storageUuid) {}
+    }
 
-            abstract @NonNull List<Path> locations();
+    /** Result found but don't record it. */
+    private static final int TYPE_DONT_RECORD = 0;
+    /** Primary dex file. */
+    private static final int TYPE_PRIMARY = 1;
+    /** Secondary dex file. */
+    private static final int TYPE_SECONDARY = 2;
 
-            abstract @NonNull UUID storageUuid();
-        }
-    }
+    /** @hide */
+    // clang-format off
+    @IntDef(prefix = "TYPE_", value = {
+        TYPE_DONT_RECORD,
+        TYPE_PRIMARY,
+        TYPE_SECONDARY,
+    })
+    // clang-format on
+    @Retention(RetentionPolicy.SOURCE)
+    private @interface DexType {}
+
+    private record FindResult(@DexType int type, @Nullable String owningPackageName) {}
 
     /**
      * Injector pattern for testing purpose.
@@ -1215,7 +1345,7 @@ public class DexUseManagerLocal {
         }
 
         @NonNull
-        private PackageManagerLocal getPackageManagerLocal() {
+        public PackageManagerLocal getPackageManagerLocal() {
             return Objects.requireNonNull(
                     LocalManagerRegistry.getManager(PackageManagerLocal.class));
         }
diff --git a/libartservice/service/java/com/android/server/art/Utils.java b/libartservice/service/java/com/android/server/art/Utils.java
index 09ce19c5ac..600671fcb2 100644
--- a/libartservice/service/java/com/android/server/art/Utils.java
+++ b/libartservice/service/java/com/android/server/art/Utils.java
@@ -483,6 +483,24 @@ public final class Utils {
         }
     }
 
+    public static boolean pathStartsWith(@NonNull String path, @NonNull String prefix) {
+        check(!prefix.isEmpty() && !path.isEmpty() && prefix.charAt(0) == '/'
+                && path.charAt(0) == '/');
+        int prefixLen = prefix.length();
+        if (prefix.charAt(prefixLen - 1) == '/') {
+            prefixLen--;
+        }
+        if (path.length() < prefixLen) {
+            return false;
+        }
+        for (int i = 0; i < prefixLen; i++) {
+            if (path.charAt(i) != prefix.charAt(i)) {
+                return false;
+            }
+        }
+        return path.length() == prefixLen || path.charAt(prefixLen) == '/';
+    }
+
     @AutoValue
     public abstract static class Abi {
         static @NonNull Abi create(
diff --git a/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java b/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java
index ffc702f338..95d111774d 100644
--- a/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java
+++ b/libartservice/service/javatests/com/android/server/art/DexUseManagerTest.java
@@ -77,10 +77,15 @@ import java.util.UUID;
 @RunWith(AndroidJUnit4.class)
 public class DexUseManagerTest {
     private static final String LOADING_PKG_NAME = "com.example.loadingpackage";
+
     private static final String OWNING_PKG_NAME = "com.example.owningpackage";
     private static final String BASE_APK = "/somewhere/app/" + OWNING_PKG_NAME + "/base.apk";
     private static final String SPLIT_APK = "/somewhere/app/" + OWNING_PKG_NAME + "/split_0.apk";
 
+    private static final String INVISIBLE_PKG_NAME = "com.example.invisiblepackage";
+    private static final String INVISIBLE_BASE_APK =
+            "/somewhere/app/" + INVISIBLE_PKG_NAME + "/base.apk";
+
     @Rule
     public StaticMockitoRule mockitoRule = new StaticMockitoRule(
             SystemProperties.class, Constants.class, Process.class, ArtJni.class);
@@ -99,6 +104,8 @@ public class DexUseManagerTest {
     @Mock private IArtd mArtd;
     @Mock private Context mContext;
     @Mock private ArtManagerLocal mArtManagerLocal;
+    @Mock private PackageManagerLocal mPackageManagerLocal;
+    @Mock private PackageManagerLocal.UnfilteredSnapshot mUnfilteredSnapshot;
     private DexUseManagerLocal mDexUseManager;
     private String mCeDir;
     private String mDeDir;
@@ -126,18 +133,21 @@ public class DexUseManagerTest {
         // Put the null package in front of other packages to verify that it's properly skipped.
         PackageState nullPkgState =
                 createPackageState("com.example.null", "arm64-v8a", false /* hasPackage */);
-        addPackage("com.example.null", nullPkgState);
+        addPackage("com.example.null", nullPkgState, true /* isVisible */);
         PackageState loadingPkgState =
                 createPackageState(LOADING_PKG_NAME, "armeabi-v7a", true /* hasPackage */);
-        addPackage(LOADING_PKG_NAME, loadingPkgState);
+        addPackage(LOADING_PKG_NAME, loadingPkgState, true /* isVisible */);
         PackageState owningPkgState =
                 createPackageState(OWNING_PKG_NAME, "arm64-v8a", true /* hasPackage */);
-        addPackage(OWNING_PKG_NAME, owningPkgState);
+        addPackage(OWNING_PKG_NAME, owningPkgState, true /* isVisible */);
         PackageState platformPkgState =
                 createPackageState(Utils.PLATFORM_PACKAGE_NAME, "arm64-v8a", true /* hasPackage */);
-        addPackage(Utils.PLATFORM_PACKAGE_NAME, platformPkgState);
+        addPackage(Utils.PLATFORM_PACKAGE_NAME, platformPkgState, true /* isVisible */);
+        PackageState invisiblePkgState =
+                createPackageState(INVISIBLE_PKG_NAME, "arm64-v8a", true /* hasPackage */);
+        addPackage(INVISIBLE_PKG_NAME, invisiblePkgState, false /* isVisible */);
 
-        lenient().when(mSnapshot.getPackageStates()).thenReturn(mPackageStates);
+        lenient().when(mUnfilteredSnapshot.getPackageStates()).thenReturn(mPackageStates);
 
         mBroadcastReceiverCaptor = ArgumentCaptor.forClass(BroadcastReceiver.class);
         lenient()
@@ -160,6 +170,10 @@ public class DexUseManagerTest {
         lenient().when(ArtJni.validateDexPath(any())).thenReturn(null);
         lenient().when(ArtJni.validateClassLoaderContext(any(), any())).thenReturn(null);
 
+        lenient()
+                .when(mPackageManagerLocal.withUnfilteredSnapshot())
+                .thenReturn(mUnfilteredSnapshot);
+
         lenient().when(mInjector.getArtd()).thenReturn(mArtd);
         lenient().when(mInjector.getCurrentTimeMillis()).thenReturn(0l);
         lenient().when(mInjector.getFilename()).thenReturn(mTempFile.getPath());
@@ -170,6 +184,7 @@ public class DexUseManagerTest {
         lenient().when(mInjector.getAllPackageNames()).thenReturn(mPackageStates.keySet());
         lenient().when(mInjector.isPreReboot()).thenReturn(false);
         lenient().when(mInjector.getArtManagerLocal()).thenReturn(mArtManagerLocal);
+        lenient().when(mInjector.getPackageManagerLocal()).thenReturn(mPackageManagerLocal);
 
         mDexUseManager = new DexUseManagerLocal(mInjector);
         mDexUseManager.systemReady();
@@ -232,6 +247,18 @@ public class DexUseManagerTest {
                 .isFalse();
     }
 
+    @Test
+    public void testPrimaryDexInvisible() {
+        mDexUseManager.notifyDexContainersLoaded(
+                mSnapshot, LOADING_PKG_NAME, Map.of(INVISIBLE_BASE_APK, "CLC"));
+
+        assertThat(mDexUseManager.getPrimaryDexLoaders(INVISIBLE_PKG_NAME, INVISIBLE_BASE_APK))
+                .isEmpty();
+        assertThat(
+                mDexUseManager.isPrimaryDexUsedByOtherApps(INVISIBLE_PKG_NAME, INVISIBLE_BASE_APK))
+                .isFalse();
+    }
+
     /** Checks that it ignores and dedups things correctly. */
     @Test
     public void testPrimaryDexMultipleEntries() throws Exception {
@@ -650,7 +677,7 @@ public class DexUseManagerTest {
     public void testCleanupDeletedPackage() throws Exception {
         PackageState pkgState = createPackageState(
                 "com.example.deletedpackage", "arm64-v8a", true /* hasPackage */);
-        addPackage("com.example.deletedpackage", pkgState);
+        addPackage("com.example.deletedpackage", pkgState, true /* isVisible */);
         lenient()
                 .when(mArtd.getDexFileVisibility(
                         "/somewhere/app/com.example.deletedpackage/base.apk"))
@@ -879,8 +906,10 @@ public class DexUseManagerTest {
         return pkgState;
     }
 
-    private void addPackage(String packageName, PackageState pkgState) {
-        lenient().when(mSnapshot.getPackageState(packageName)).thenReturn(pkgState);
+    private void addPackage(String packageName, PackageState pkgState, boolean isVisible) {
+        if (isVisible) {
+            lenient().when(mSnapshot.getPackageState(packageName)).thenReturn(pkgState);
+        }
         mPackageStates.put(packageName, pkgState);
     }
 
diff --git a/libartservice/service/javatests/com/android/server/art/UtilsTest.java b/libartservice/service/javatests/com/android/server/art/UtilsTest.java
index 95fd747fbb..b3220ffdff 100644
--- a/libartservice/service/javatests/com/android/server/art/UtilsTest.java
+++ b/libartservice/service/javatests/com/android/server/art/UtilsTest.java
@@ -203,4 +203,21 @@ public class UtilsTest {
         Executor executor = ForkJoinPool.commonPool();
         Utils.executeAndWait(executor, () -> { throw new IllegalArgumentException(); });
     }
+
+    @Test
+    public void testPathStartsWith() {
+        assertThat(Utils.pathStartsWith("/a/b", "/a")).isTrue();
+        assertThat(Utils.pathStartsWith("/a/b", "/a/")).isTrue();
+
+        assertThat(Utils.pathStartsWith("/a/c", "/a/b")).isFalse();
+        assertThat(Utils.pathStartsWith("/ab", "/a")).isFalse();
+
+        assertThat(Utils.pathStartsWith("/a", "/a")).isTrue();
+        assertThat(Utils.pathStartsWith("/a/", "/a")).isTrue();
+        assertThat(Utils.pathStartsWith("/a", "/a/")).isTrue();
+
+        assertThat(Utils.pathStartsWith("/a", "/")).isTrue();
+        assertThat(Utils.pathStartsWith("/", "/")).isTrue();
+        assertThat(Utils.pathStartsWith("/", "/a")).isFalse();
+    }
 }