ART services: optimize package - Implement DexOptHelper.

Bug: 229268202
Test: atest ArtServiceTests
Ignore-AOSP-First: ART Services.
Change-Id: Ie7c99b61276699b1121d2f24e9658a2daf9645fc
diff --git a/libartservice/service/java/com/android/server/art/DexOptHelper.java b/libartservice/service/java/com/android/server/art/DexOptHelper.java
index 38dec8c..a5cee0b 100644
--- a/libartservice/service/java/com/android/server/art/DexOptHelper.java
+++ b/libartservice/service/java/com/android/server/art/DexOptHelper.java
@@ -20,7 +20,11 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.apphibernation.AppHibernationManager;
 import android.content.Context;
+import android.os.Binder;
+import android.os.PowerManager;
+import android.os.WorkSource;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.art.model.OptimizeOptions;
@@ -43,6 +47,12 @@
 public class DexOptHelper {
     private static final String TAG = "DexoptHelper";
 
+    /**
+     * Timeout of the wake lock. This is required by AndroidLint, but we set it to a value larger
+     * than artd's {@code kLongTimeoutSec} so that it should normally never triggered.
+     */
+    private static final int WAKE_LOCK_TIMEOUT_MS = 11 * 60 * 1000; // 11 minutes.
+
     @NonNull private final Injector mInjector;
 
     public DexOptHelper(@Nullable Context context) {
@@ -62,7 +72,71 @@
     public OptimizeResult dexopt(@NonNull PackageDataSnapshot snapshot,
             @NonNull PackageState pkgState, @NonNull AndroidPackageApi pkg,
             @NonNull OptimizeOptions options) {
-        throw new UnsupportedOperationException();
+        List<DexFileOptimizeResult> results = new ArrayList<>();
+        Supplier<OptimizeResult> createResult = ()
+                -> new OptimizeResult(pkgState.getPackageName(), options.getCompilerFilter(),
+                        options.getReason(), results);
+
+        if (!canOptimizePackage(pkgState, pkg)) {
+            return createResult.get();
+        }
+
+        long identityToken = Binder.clearCallingIdentity();
+        PowerManager.WakeLock wakeLock = null;
+
+        try {
+            // Acquire a wake lock.
+            // The power manager service may not be ready if this method is called on boot. In this
+            // case, we don't have to acquire a wake lock because there is nothing else we can do.
+            PowerManager powerManager = mInjector.getPowerManager();
+            if (powerManager != null) {
+                wakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
+                wakeLock.setWorkSource(new WorkSource(pkg.getUid()));
+                wakeLock.acquire(WAKE_LOCK_TIMEOUT_MS);
+            }
+
+            if (options.isForPrimaryDex()) {
+                results.addAll(mInjector.getPrimaryDexOptimizer().dexopt(pkgState, pkg, options));
+            }
+
+            if (options.isForSecondaryDex()) {
+                // TODO(jiakaiz): Implement this.
+                throw new UnsupportedOperationException(
+                        "Optimizing secondary dex'es is not implemented yet");
+            }
+
+            if (options.getIncludesDependencies()) {
+                // TODO(jiakaiz): Implement this.
+                throw new UnsupportedOperationException(
+                        "Optimizing dependencies is not implemented yet");
+            }
+        } finally {
+            if (wakeLock != null) {
+                wakeLock.release();
+            }
+            Binder.restoreCallingIdentity(identityToken);
+        }
+
+        return createResult.get();
+    }
+
+    private boolean canOptimizePackage(
+            @NonNull PackageState pkgState, @NonNull AndroidPackageApi pkg) {
+        if (!pkg.isHasCode()) {
+            return false;
+        }
+
+        // We do not dexopt unused packages.
+        // It's possible for this to be called before app hibernation service is ready, especially
+        // on boot. In this case, we ignore the hibernation check here because there is nothing else
+        // we can do.
+        AppHibernationManager ahm = mInjector.getAppHibernationManager();
+        if (ahm != null && ahm.isHibernatingGlobally(pkgState.getPackageName())
+                && ahm.isOatArtifactDeletionEnabled()) {
+            return false;
+        }
+
+        return true;
     }
 
     /**
@@ -78,5 +152,20 @@
         Injector(@Nullable Context context) {
             mContext = context;
         }
+
+        @NonNull
+        PrimaryDexOptimizer getPrimaryDexOptimizer() {
+            return new PrimaryDexOptimizer(mContext);
+        }
+
+        @Nullable
+        public AppHibernationManager getAppHibernationManager() {
+            return mContext != null ? mContext.getSystemService(AppHibernationManager.class) : null;
+        }
+
+        @Nullable
+        public PowerManager getPowerManager() {
+            return mContext != null ? mContext.getSystemService(PowerManager.class) : null;
+        }
     }
 }
diff --git a/libartservice/service/java/com/android/server/art/PrimaryDexOptimizer.java b/libartservice/service/java/com/android/server/art/PrimaryDexOptimizer.java
new file mode 100644
index 0000000..7315040
--- /dev/null
+++ b/libartservice/service/java/com/android/server/art/PrimaryDexOptimizer.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.art;
+
+import static com.android.server.art.model.OptimizeResult.DexFileOptimizeResult;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.content.Context;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.art.model.OptimizeOptions;
+import com.android.server.art.wrapper.AndroidPackageApi;
+import com.android.server.art.wrapper.PackageState;
+
+import java.util.List;
+
+/** @hide */
+public class PrimaryDexOptimizer {
+    private static final String TAG = "PrimaryDexOptimizer";
+
+    @NonNull private final Injector mInjector;
+
+    public PrimaryDexOptimizer(@Nullable Context context) {
+        this(new Injector(context));
+    }
+
+    @VisibleForTesting
+    public PrimaryDexOptimizer(@NonNull Injector injector) {
+        mInjector = injector;
+    }
+
+    /**
+     * DO NOT use this method directly. Use {@link
+     * ArtManagerLocal#optimizePackage(PackageDataSnapshot, String, OptimizeOptions)}.
+     */
+    @NonNull
+    public List<DexFileOptimizeResult> dexopt(@NonNull PackageState pkgState,
+            @NonNull AndroidPackageApi pkg, @NonNull OptimizeOptions options) {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Injector pattern for testing purpose.
+     *
+     * @hide
+     */
+    @VisibleForTesting
+    public static class Injector {
+        // TODO(b/236954191): Make this @NonNull.
+        @Nullable private final Context mContext;
+
+        Injector(@Nullable Context context) {
+            mContext = context;
+        }
+
+        @NonNull
+        public IArtd getArtd() {
+            return Utils.getArtd();
+        }
+    }
+}
diff --git a/libartservice/service/java/com/android/server/art/wrapper/AndroidPackageApi.java b/libartservice/service/java/com/android/server/art/wrapper/AndroidPackageApi.java
index 90adb49..1c3ff6d 100644
--- a/libartservice/service/java/com/android/server/art/wrapper/AndroidPackageApi.java
+++ b/libartservice/service/java/com/android/server/art/wrapper/AndroidPackageApi.java
@@ -114,4 +114,12 @@
             throw new RuntimeException(e);
         }
     }
+
+    public int getUid() {
+        try {
+            return (int) mPkg.getClass().getMethod("getUid").invoke(mPkg);
+        } catch (ReflectiveOperationException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }
diff --git a/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java b/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
new file mode 100644
index 0000000..dffbc1f
--- /dev/null
+++ b/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
@@ -0,0 +1,217 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.art;
+
+import static com.android.server.art.model.OptimizeResult.DexFileOptimizeResult;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.lenient;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.same;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+import android.apphibernation.AppHibernationManager;
+import android.os.PowerManager;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.server.art.model.OptimizeOptions;
+import com.android.server.art.model.OptimizeResult;
+import com.android.server.art.testing.OnSuccessRule;
+import com.android.server.art.wrapper.AndroidPackageApi;
+import com.android.server.art.wrapper.PackageState;
+import com.android.server.pm.snapshot.PackageDataSnapshot;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.InOrder;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import java.util.List;
+
+@SmallTest
+@RunWith(MockitoJUnitRunner.StrictStubs.class)
+public class DexOptHelperTest {
+    private static final String PKG_NAME = "com.example.foo";
+
+    @Mock private DexOptHelper.Injector mInjector;
+    @Mock private PrimaryDexOptimizer mPrimaryDexOptimizer;
+    @Mock private AppHibernationManager mAhm;
+    @Mock private PowerManager mPowerManager;
+    private PackageState mPkgState;
+    private AndroidPackageApi mPkg;
+
+    @Rule
+    public OnSuccessRule onSuccessRule = new OnSuccessRule(() -> {
+        // Don't do this on failure because it will make the failure hard to understand.
+        verifyNoMoreInteractions(mPrimaryDexOptimizer);
+    });
+
+    private final OptimizeOptions mOptions =
+            new OptimizeOptions.Builder("install").setCompilerFilter("speed-profile").build();
+    private final List<DexFileOptimizeResult> mPrimaryResults =
+            List.of(new DexFileOptimizeResult("/data/app/foo/base.apk", "arm64", "verify",
+                            OptimizeResult.OPTIMIZE_PERFORMED),
+                    new DexFileOptimizeResult("/data/app/foo/base.apk", "arm", "verify",
+                            OptimizeResult.OPTIMIZE_FAILED));
+
+    private DexOptHelper mDexOptHelper;
+
+    @Before
+    public void setUp() throws Exception {
+        lenient().when(mInjector.getPrimaryDexOptimizer()).thenReturn(mPrimaryDexOptimizer);
+        lenient().when(mInjector.getAppHibernationManager()).thenReturn(null);
+        lenient().when(mInjector.getPowerManager()).thenReturn(null);
+
+        mPkgState = createPackageState();
+        mPkg = mPkgState.getAndroidPackage();
+
+        mDexOptHelper = new DexOptHelper(mInjector);
+    }
+
+    @Test
+    public void testDexopt() {
+        when(mPrimaryDexOptimizer.dexopt(same(mPkgState), same(mPkg), same(mOptions)))
+                .thenReturn(mPrimaryResults);
+
+        OptimizeResult result =
+                mDexOptHelper.dexopt(mock(PackageDataSnapshot.class), mPkgState, mPkg, mOptions);
+
+        assertThat(result.getPackageName()).isEqualTo(PKG_NAME);
+        assertThat(result.getRequestedCompilerFilter()).isEqualTo("speed-profile");
+        assertThat(result.getReason()).isEqualTo("install");
+        assertThat(result.getFinalStatus()).isEqualTo(OptimizeResult.OPTIMIZE_FAILED);
+        assertThat(result.getDexFileOptimizeResults()).containsExactlyElementsIn(mPrimaryResults);
+    }
+
+    @Test
+    public void testDexoptNoCode() {
+        when(mPkg.isHasCode()).thenReturn(false);
+
+        OptimizeResult result =
+                mDexOptHelper.dexopt(mock(PackageDataSnapshot.class), mPkgState, mPkg, mOptions);
+
+        assertThat(result.getFinalStatus()).isEqualTo(OptimizeResult.OPTIMIZE_SKIPPED);
+        assertThat(result.getDexFileOptimizeResults()).isEmpty();
+    }
+
+    @Test
+    public void testDexoptWithAppHibernationManager() {
+        when(mInjector.getAppHibernationManager()).thenReturn(mAhm);
+        lenient().when(mAhm.isHibernatingGlobally(PKG_NAME)).thenReturn(false);
+        lenient().when(mAhm.isOatArtifactDeletionEnabled()).thenReturn(true);
+
+        when(mPrimaryDexOptimizer.dexopt(same(mPkgState), same(mPkg), same(mOptions)))
+                .thenReturn(mPrimaryResults);
+
+        OptimizeResult result =
+                mDexOptHelper.dexopt(mock(PackageDataSnapshot.class), mPkgState, mPkg, mOptions);
+
+        assertThat(result.getDexFileOptimizeResults()).containsExactlyElementsIn(mPrimaryResults);
+    }
+
+    @Test
+    public void testDexoptIsHibernating() {
+        when(mInjector.getAppHibernationManager()).thenReturn(mAhm);
+        lenient().when(mAhm.isHibernatingGlobally(PKG_NAME)).thenReturn(true);
+        lenient().when(mAhm.isOatArtifactDeletionEnabled()).thenReturn(true);
+
+        OptimizeResult result =
+                mDexOptHelper.dexopt(mock(PackageDataSnapshot.class), mPkgState, mPkg, mOptions);
+
+        assertThat(result.getFinalStatus()).isEqualTo(OptimizeResult.OPTIMIZE_SKIPPED);
+        assertThat(result.getDexFileOptimizeResults()).isEmpty();
+    }
+
+    @Test
+    public void testDexoptIsHibernatingButOatArtifactDeletionDisabled() {
+        when(mInjector.getAppHibernationManager()).thenReturn(mAhm);
+        lenient().when(mAhm.isHibernatingGlobally(PKG_NAME)).thenReturn(true);
+        lenient().when(mAhm.isOatArtifactDeletionEnabled()).thenReturn(false);
+
+        when(mPrimaryDexOptimizer.dexopt(same(mPkgState), same(mPkg), same(mOptions)))
+                .thenReturn(mPrimaryResults);
+
+        OptimizeResult result =
+                mDexOptHelper.dexopt(mock(PackageDataSnapshot.class), mPkgState, mPkg, mOptions);
+
+        assertThat(result.getDexFileOptimizeResults()).containsExactlyElementsIn(mPrimaryResults);
+    }
+
+    @Test
+    public void testDexoptWithPowerManager() {
+        var wakeLock = mock(PowerManager.WakeLock.class);
+        when(mInjector.getPowerManager()).thenReturn(mPowerManager);
+        when(mPowerManager.newWakeLock(eq(PowerManager.PARTIAL_WAKE_LOCK), any()))
+                .thenReturn(wakeLock);
+
+        when(mPrimaryDexOptimizer.dexopt(same(mPkgState), same(mPkg), same(mOptions)))
+                .thenReturn(mPrimaryResults);
+
+        OptimizeResult result =
+                mDexOptHelper.dexopt(mock(PackageDataSnapshot.class), mPkgState, mPkg, mOptions);
+
+        InOrder inOrder = inOrder(mPrimaryDexOptimizer, wakeLock);
+        inOrder.verify(wakeLock).acquire(anyLong());
+        inOrder.verify(mPrimaryDexOptimizer).dexopt(any(), any(), any());
+        inOrder.verify(wakeLock).release();
+    }
+
+    @Test
+    public void testDexoptAlwaysReleasesWakeLock() {
+        var wakeLock = mock(PowerManager.WakeLock.class);
+        when(mInjector.getPowerManager()).thenReturn(mPowerManager);
+        when(mPowerManager.newWakeLock(eq(PowerManager.PARTIAL_WAKE_LOCK), any()))
+                .thenReturn(wakeLock);
+
+        when(mPrimaryDexOptimizer.dexopt(same(mPkgState), same(mPkg), same(mOptions)))
+                .thenThrow(IllegalStateException.class);
+
+        try {
+            OptimizeResult result = mDexOptHelper.dexopt(
+                    mock(PackageDataSnapshot.class), mPkgState, mPkg, mOptions);
+        } catch (Exception e) {
+        }
+
+        verify(wakeLock).release();
+    }
+
+    private AndroidPackageApi createPackage() {
+        AndroidPackageApi pkg = mock(AndroidPackageApi.class);
+        lenient().when(pkg.getUid()).thenReturn(12345);
+        lenient().when(pkg.isHasCode()).thenReturn(true);
+        return pkg;
+    }
+
+    private PackageState createPackageState() {
+        PackageState pkgState = mock(PackageState.class);
+        lenient().when(pkgState.getPackageName()).thenReturn(PKG_NAME);
+        AndroidPackageApi pkg = createPackage();
+        lenient().when(pkgState.getAndroidPackage()).thenReturn(pkg);
+        return pkgState;
+    }
+}
diff --git a/libartservice/service/javatests/com/android/server/art/testing/OnSuccessRule.java b/libartservice/service/javatests/com/android/server/art/testing/OnSuccessRule.java
new file mode 100644
index 0000000..350a8cb
--- /dev/null
+++ b/libartservice/service/javatests/com/android/server/art/testing/OnSuccessRule.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.art.testing;
+
+import org.junit.rules.MethodRule;
+import org.junit.runners.model.FrameworkMethod;
+import org.junit.runners.model.Statement;
+
+/** A JUnit rule that invokes a runnable on success. */
+public class OnSuccessRule implements MethodRule {
+    private RunnableThrowingException mRunnable;
+
+    public OnSuccessRule(RunnableThrowingException runnable) {
+        mRunnable = runnable;
+    }
+
+    @Override
+    public Statement apply(Statement base, FrameworkMethod method, Object target) {
+        return new Statement() {
+            public void evaluate() throws Throwable {
+                base.evaluate();
+                mRunnable.run();
+            }
+        };
+    }
+
+    public interface RunnableThrowingException {
+        void run() throws Exception;
+    }
+}