Report progress when optimizing packages.

This will be used during boot, to show a dialog "Optimizing app X of Y."

Bug: 260419279
Test: atest ArtServiceTests
Test: -
  1. adb shell setprop pm.dexopt.bg-dexopt.concurrency 4
  2. adb shell pm art optimize-packages bg-dexopt
Ignore-AOSP-First: ART Services.
Change-Id: I3a4250ffeb02d56acc425e0358aad6d0b5725796
diff --git a/libartservice/service/java/com/android/server/art/ArtManagerLocal.java b/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
index 926e0aa..847fb91 100644
--- a/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
+++ b/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
@@ -51,6 +51,7 @@
 import com.android.server.art.model.DeleteResult;
 import com.android.server.art.model.OptimizationStatus;
 import com.android.server.art.model.OptimizeParams;
+import com.android.server.art.model.OptimizeProgress;
 import com.android.server.art.model.OptimizeResult;
 import com.android.server.pm.PackageManagerLocal;
 import com.android.server.pm.pkg.AndroidPackage;
@@ -67,6 +68,7 @@
 import java.util.Objects;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 /**
@@ -313,6 +315,8 @@
      * @param snapshot the snapshot from {@link PackageManagerLocal} to operate on
      * @param reason determines the default list of packages and options
      * @param cancellationSignal provides the ability to cancel this operation
+     * @param processCallbackExecutor the executor to call {@code progressCallback}
+     * @param progressCallback called repeatedly whenever there is an update on the progress
      * @throws IllegalStateException if the operation encounters an error that should never happen
      *         (e.g., an internal logic error), or the callback set by {@link
      *         #setOptimizePackagesCallback(Executor, OptimizePackagesCallback)} provides invalid
@@ -323,7 +327,9 @@
     @NonNull
     public OptimizeResult optimizePackages(@NonNull PackageManagerLocal.FilteredSnapshot snapshot,
             @NonNull @BatchOptimizeReason String reason,
-            @NonNull CancellationSignal cancellationSignal) {
+            @NonNull CancellationSignal cancellationSignal,
+            @Nullable @CallbackExecutor Executor processCallbackExecutor,
+            @Nullable Consumer<OptimizeProgress> progressCallback) {
         List<String> defaultPackages =
                 Collections.unmodifiableList(getDefaultPackages(snapshot, reason));
         OptimizeParams defaultOptimizeParams = new OptimizeParams.Builder(reason).build();
@@ -341,17 +347,15 @@
 
         return mInjector.getDexOptHelper().dexopt(snapshot, params.getPackages(),
                 params.getOptimizeParams(), cancellationSignal,
-                Executors.newFixedThreadPool(ReasonMapping.getConcurrencyForReason(reason)));
+                Executors.newFixedThreadPool(ReasonMapping.getConcurrencyForReason(reason)),
+                processCallbackExecutor, progressCallback);
     }
 
     /**
-     * Overrides the default params for {@link
-     * #optimizePackages(PackageManagerLocal.FilteredSnapshot, String, CancellationSignal). This
-     * method is thread-safe.
+     * Overrides the default params for {@link #optimizePackages}. This method is thread-safe.
      *
-     * This method gives users the opportunity to change the behavior of {@link
-     * #optimizePackages(PackageManagerLocal.FilteredSnapshot, String, CancellationSignal)}, which
-     * is called by ART Service automatically during boot / background dexopt.
+     * This method gives users the opportunity to change the behavior of {@link #optimizePackages},
+     * which is called by ART Service automatically during boot / background dexopt.
      *
      * If this method is not called, the default list of packages and options determined by {@code
      * reason} will be used.
@@ -397,9 +401,7 @@
      * window</i>. For information about <i>maintenance window</i>, see
      * https://developer.android.com/training/monitoring-device-state/doze-standby.
      *
-     * See {@link
-     * #optimizePackages(PackageManagerLocal.FilteredSnapshot, String, CancellationSignal)} for how
-     * to customize the behavior of the job.
+     * See {@link #optimizePackages} for how to customize the behavior of the job.
      *
      * When the job ends (either completed or cancelled), the result is sent to the callbacks added
      * by {@link #addOptimizePackageDoneCallback(Executor, OptimizePackageDoneCallback)} with the
@@ -450,9 +452,7 @@
      * constraints described in {@link #scheduleBackgroundDexoptJob()}, and hence will not be
      * cancelled when they aren't met.
      *
-     * See {@link
-     * #optimizePackages(PackageManagerLocal.FilteredSnapshot, String, CancellationSignal)} for how
-     * to customize the behavior of the job.
+     * See {@link #optimizePackages} for how to customize the behavior of the job.
      *
      * When the job ends (either completed or cancelled), the result is sent to the callbacks added
      * by {@link #addOptimizePackageDoneCallback(Executor, OptimizePackageDoneCallback)} with the
@@ -651,8 +651,7 @@
 
     public interface OptimizePackagesCallback {
         /**
-         * Mutates {@code builder} to override the default params for {@link
-         * #optimizePackages(PackageManagerLocal.FilteredSnapshot, String, CancellationSignal). It
+         * Mutates {@code builder} to override the default params for {@link #optimizePackages}. It
          * must ignore unknown reasons because more reasons may be added in the future.
          *
          * If {@code builder.setPackages} is not called, {@code defaultPackages} will be used as the
@@ -668,9 +667,7 @@
          * https://developer.android.com/training/monitoring-device-state/doze-standby.
          *
          * Changing the reason is not allowed. Doing so will result in {@link IllegalStateException}
-         * when {@link
-         * #optimizePackages(PackageManagerLocal.FilteredSnapshot, String, CancellationSignal)} is
-         * called.
+         * when {@link #optimizePackages} is called.
          */
         void onOverrideBatchOptimizeParams(@NonNull PackageManagerLocal.FilteredSnapshot snapshot,
                 @NonNull @BatchOptimizeReason String reason, @NonNull List<String> defaultPackages,
diff --git a/libartservice/service/java/com/android/server/art/ArtShellCommand.java b/libartservice/service/java/com/android/server/art/ArtShellCommand.java
index eccc7c3..7feea28 100644
--- a/libartservice/service/java/com/android/server/art/ArtShellCommand.java
+++ b/libartservice/service/java/com/android/server/art/ArtShellCommand.java
@@ -53,6 +53,8 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
 
 /**
@@ -141,11 +143,17 @@
                 }
                 case "optimize-packages": {
                     OptimizeResult result;
+                    Executor executor = Executors.newSingleThreadExecutor();
                     try (var signal = new WithCancellationSignal(pw)) {
-                        result = mArtManagerLocal.optimizePackages(
-                                snapshot, getNextArgRequired(), signal.get());
+                        result = mArtManagerLocal.optimizePackages(snapshot, getNextArgRequired(),
+                                signal.get(), executor, progress -> {
+                                    pw.println(String.format("Optimizing packages: %d/%d",
+                                            progress.getDonePackageCount(),
+                                            progress.getTotalPackageCount()));
+                                    pw.flush();
+                                });
                     }
-                    printOptimizeResult(pw, result);
+                    Utils.executeAndWait(executor, () -> printOptimizeResult(pw, result));
                     return 0;
                 }
                 case "cancel": {
diff --git a/libartservice/service/java/com/android/server/art/BackgroundDexOptJob.java b/libartservice/service/java/com/android/server/art/BackgroundDexOptJob.java
index 28677fa..03f8610 100644
--- a/libartservice/service/java/com/android/server/art/BackgroundDexOptJob.java
+++ b/libartservice/service/java/com/android/server/art/BackgroundDexOptJob.java
@@ -195,8 +195,9 @@
         long startTimeMs = SystemClock.uptimeMillis();
         OptimizeResult dexoptResult;
         try (var snapshot = mInjector.getPackageManagerLocal().withFilteredSnapshot()) {
-            dexoptResult = mInjector.getArtManagerLocal().optimizePackages(
-                    snapshot, ReasonMapping.REASON_BG_DEXOPT, cancellationSignal);
+            dexoptResult = mInjector.getArtManagerLocal().optimizePackages(snapshot,
+                    ReasonMapping.REASON_BG_DEXOPT, cancellationSignal,
+                    null /* processCallbackExecutor */, null /* processCallback */);
         }
         return CompletedResult.create(dexoptResult, SystemClock.uptimeMillis() - startTimeMs);
     }
diff --git a/libartservice/service/java/com/android/server/art/DexOptHelper.java b/libartservice/service/java/com/android/server/art/DexOptHelper.java
index df45ab2..5bd70f7 100644
--- a/libartservice/service/java/com/android/server/art/DexOptHelper.java
+++ b/libartservice/service/java/com/android/server/art/DexOptHelper.java
@@ -22,6 +22,7 @@
 import static com.android.server.art.model.OptimizeResult.PackageOptimizeResult;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.apphibernation.AppHibernationManager;
 import android.content.Context;
 import android.os.Binder;
@@ -34,6 +35,7 @@
 import com.android.server.art.model.ArtFlags;
 import com.android.server.art.model.Config;
 import com.android.server.art.model.OptimizeParams;
+import com.android.server.art.model.OptimizeProgress;
 import com.android.server.art.model.OptimizeResult;
 import com.android.server.pm.PackageManagerLocal;
 import com.android.server.pm.pkg.AndroidPackage;
@@ -48,9 +50,11 @@
 import java.util.Objects;
 import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -83,29 +87,43 @@
     }
 
     /**
-     * DO NOT use this method directly. Use {@link
-     * ArtManagerLocal#optimizePackage(PackageManagerLocal.FilteredSnapshot, String,
-     * OptimizeParams)}.
+     * DO NOT use this method directly. Use {@link ArtManagerLocal#optimizePackage} or {@link
+     * ArtManagerLocal#optimizePackages}.
      */
     @NonNull
     public OptimizeResult dexopt(@NonNull PackageManagerLocal.FilteredSnapshot snapshot,
             @NonNull List<String> packageNames, @NonNull OptimizeParams params,
-            @NonNull CancellationSignal cancellationSignal, @NonNull Executor executor) {
-        return dexoptPackages(
-                getPackageStates(snapshot, packageNames,
-                        (params.getFlags() & ArtFlags.FLAG_SHOULD_INCLUDE_DEPENDENCIES) != 0),
-                params, cancellationSignal, executor);
+            @NonNull CancellationSignal cancellationSignal, @NonNull Executor dexoptExecutor) {
+        return dexopt(snapshot, packageNames, params, cancellationSignal, dexoptExecutor,
+                null /* progressCallbackExecutor */, null /* progressCallback */);
     }
 
     /**
-     * DO NOT use this method directly. Use {@link
-     * ArtManagerLocal#optimizePackage(PackageManagerLocal.FilteredSnapshot, String,
-     * OptimizeParams)}.
+     * DO NOT use this method directly. Use {@link ArtManagerLocal#optimizePackage} or {@link
+     * ArtManagerLocal#optimizePackages}.
+     */
+    @NonNull
+    public OptimizeResult dexopt(@NonNull PackageManagerLocal.FilteredSnapshot snapshot,
+            @NonNull List<String> packageNames, @NonNull OptimizeParams params,
+            @NonNull CancellationSignal cancellationSignal, @NonNull Executor dexoptExecutor,
+            @Nullable Executor progressCallbackExecutor,
+            @Nullable Consumer<OptimizeProgress> progressCallback) {
+        return dexoptPackages(
+                getPackageStates(snapshot, packageNames,
+                        (params.getFlags() & ArtFlags.FLAG_SHOULD_INCLUDE_DEPENDENCIES) != 0),
+                params, cancellationSignal, dexoptExecutor, progressCallbackExecutor,
+                progressCallback);
+    }
+
+    /**
+     * DO NOT use this method directly. Use {@link ArtManagerLocal#optimizePackage} or {@link
+     * ArtManagerLocal#optimizePackages}.
      */
     @NonNull
     private OptimizeResult dexoptPackages(@NonNull List<PackageState> pkgStates,
             @NonNull OptimizeParams params, @NonNull CancellationSignal cancellationSignal,
-            @NonNull Executor executor) {
+            @NonNull Executor dexoptExecutor, @Nullable Executor progressCallbackExecutor,
+            @Nullable Consumer<OptimizeProgress> progressCallback) {
         int callingUid = Binder.getCallingUid();
         long identityToken = Binder.clearCallingIdentity();
         PowerManager.WakeLock wakeLock = null;
@@ -117,10 +135,24 @@
             wakeLock.setWorkSource(new WorkSource(callingUid));
             wakeLock.acquire(WAKE_LOCK_TIMEOUT_MS);
 
-            List<Future<PackageOptimizeResult>> futures = new ArrayList<>();
+            List<CompletableFuture<PackageOptimizeResult>> futures = new ArrayList<>();
             for (PackageState pkgState : pkgStates) {
-                futures.add(Utils.execute(
-                        executor, () -> dexoptPackage(pkgState, params, cancellationSignal)));
+                futures.add(CompletableFuture.supplyAsync(
+                        () -> dexoptPackage(pkgState, params, cancellationSignal), dexoptExecutor));
+            }
+
+            if (progressCallback != null) {
+                CompletableFuture.runAsync(() -> {
+                    progressCallback.accept(
+                            OptimizeProgress.create(0 /* donePackageCount */, futures.size()));
+                }, progressCallbackExecutor);
+                AtomicInteger donePackageCount = new AtomicInteger(0);
+                for (CompletableFuture<PackageOptimizeResult> future : futures) {
+                    future.thenRunAsync(() -> {
+                        progressCallback.accept(OptimizeProgress.create(
+                                donePackageCount.incrementAndGet(), futures.size()));
+                    }, progressCallbackExecutor);
+                }
             }
 
             List<PackageOptimizeResult> results =
@@ -129,11 +161,12 @@
             var result =
                     new OptimizeResult(params.getCompilerFilter(), params.getReason(), results);
 
-            for (Callback<OptimizePackageDoneCallback> callback :
+            for (Callback<OptimizePackageDoneCallback> doneCallback :
                     mInjector.getConfig().getOptimizePackageDoneCallbacks()) {
                 // TODO(b/257027956): Consider filtering the packages before calling the callback.
-                Utils.executeAndWait(callback.executor(),
-                        () -> { callback.get().onOptimizePackageDone(result); });
+                CompletableFuture.runAsync(() -> {
+                    doneCallback.get().onOptimizePackageDone(result);
+                }, doneCallback.executor());
             }
 
             return result;
@@ -146,9 +179,8 @@
     }
 
     /**
-     * DO NOT use this method directly. Use {@link
-     * ArtManagerLocal#optimizePackage(PackageManagerLocal.FilteredSnapshot, String,
-     * OptimizeParams)}.
+     * DO NOT use this method directly. Use {@link ArtManagerLocal#optimizePackage} or {@link
+     * ArtManagerLocal#optimizePackages}.
      */
     @NonNull
     private PackageOptimizeResult dexoptPackage(@NonNull PackageState pkgState,
diff --git a/libartservice/service/java/com/android/server/art/ReasonMapping.java b/libartservice/service/java/com/android/server/art/ReasonMapping.java
index 98c82d8..d408759 100644
--- a/libartservice/service/java/com/android/server/art/ReasonMapping.java
+++ b/libartservice/service/java/com/android/server/art/ReasonMapping.java
@@ -69,8 +69,7 @@
             REASON_INSTALL_BULK_DOWNGRADED, REASON_INSTALL_BULK_SECONDARY_DOWNGRADED);
 
     /**
-     * Reasons for
-     * {@link ArtManagerLocal#optimizePackages(PackageManagerLocal.FilteredSnapshot, String)}.
+     * Reasons for {@link ArtManagerLocal#optimizePackages}.
      *
      * @hide
      */
@@ -159,8 +158,8 @@
 
     /**
      * Loads the concurrency from the system property, for batch optimization ({@link
-     * ArtManagerLocal#optimizePackages(PackageManagerLocal.FilteredSnapshot, String)}), or 1 if the
-     * system property is not found or cannot be parsed.
+     * ArtManagerLocal#optimizePackages}), or 1 if the system property is not found or cannot be
+     * parsed.
      *
      * @hide
      */
diff --git a/libartservice/service/java/com/android/server/art/Utils.java b/libartservice/service/java/com/android/server/art/Utils.java
index 5083428..7752ba7 100644
--- a/libartservice/service/java/com/android/server/art/Utils.java
+++ b/libartservice/service/java/com/android/server/art/Utils.java
@@ -38,12 +38,10 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
-import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
-import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
-import java.util.concurrent.FutureTask;
 import java.util.stream.Collectors;
 
 /** @hide */
@@ -218,13 +216,7 @@
     }
 
     public static void executeAndWait(@NonNull Executor executor, @NonNull Runnable runnable) {
-        getFuture(execute(executor, Executors.callable(runnable)));
-    }
-
-    public static <T> Future<T> execute(@NonNull Executor executor, @NonNull Callable<T> callable) {
-        var future = new FutureTask<T>(callable);
-        executor.execute(future);
-        return future;
+        getFuture(CompletableFuture.runAsync(runnable, executor));
     }
 
     public static <T> T getFuture(Future<T> future) {
diff --git a/libartservice/service/java/com/android/server/art/model/OptimizeProgress.java b/libartservice/service/java/com/android/server/art/model/OptimizeProgress.java
new file mode 100644
index 0000000..39310a5
--- /dev/null
+++ b/libartservice/service/java/com/android/server/art/model/OptimizeProgress.java
@@ -0,0 +1,50 @@
+/*
+ * 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.model;
+
+import static com.android.server.art.model.OptimizeResult.PackageOptimizeResult;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+
+import com.android.internal.annotations.Immutable;
+
+import com.google.auto.value.AutoValue;
+
+/** @hide */
+@Immutable
+@AutoValue
+public abstract class OptimizeProgress {
+    /** @hide */
+    protected OptimizeProgress() {}
+
+    /** @hide */
+    public static @NonNull OptimizeProgress create(int donePackageCount, int totalPackageCount) {
+        return new AutoValue_OptimizeProgress(donePackageCount, totalPackageCount);
+    }
+
+    /**
+     * The number of packages, for which optimization has been done, regardless of the results
+     * (performed, failed, skipped, etc.). Can be 0, which means the optimization was just started.
+     */
+    public abstract int getDonePackageCount();
+
+    /**
+     * The total number of packages to optimize. Stays constant during the operation.
+     */
+    public abstract int getTotalPackageCount();
+}
diff --git a/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java b/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java
index 85a7b87..a50c801 100644
--- a/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java
+++ b/libartservice/service/javatests/com/android/server/art/ArtManagerLocalTest.java
@@ -345,10 +345,11 @@
 
         // It should use the default package list and params.
         when(mDexOptHelper.dexopt(any(), deepEq(List.of(PKG_NAME, PKG_NAME_SYS_UI)), any(),
-                     same(cancellationSignal), any()))
+                     same(cancellationSignal), any(), any(), any()))
                 .thenReturn(result);
 
-        assertThat(mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal))
+        assertThat(mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal,
+                           null /* processCallbackExecutor */, null /* processCallback */))
                 .isSameInstanceAs(result);
     }
 
@@ -367,10 +368,11 @@
 
         // It should use the overridden package list and params.
         when(mDexOptHelper.dexopt(any(), deepEq(List.of(PKG_NAME)), same(params),
-                     same(cancellationSignal), any()))
+                     same(cancellationSignal), any(), any(), any()))
                 .thenReturn(result);
 
-        assertThat(mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal))
+        assertThat(mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal,
+                           null /* processCallbackExecutor */, null /* processCallback */))
                 .isSameInstanceAs(result);
     }
 
@@ -388,10 +390,11 @@
 
         // It should use the default package list and params.
         when(mDexOptHelper.dexopt(any(), deepEq(List.of(PKG_NAME, PKG_NAME_SYS_UI)),
-                     not(same(params)), same(cancellationSignal), any()))
+                     not(same(params)), same(cancellationSignal), any(), any(), any()))
                 .thenReturn(result);
 
-        assertThat(mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal))
+        assertThat(mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal,
+                           null /* processCallbackExecutor */, null /* processCallback */))
                 .isSameInstanceAs(result);
     }
 
@@ -405,7 +408,8 @@
                     builder.setOptimizeParams(params);
                 });
 
-        mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal);
+        mArtManagerLocal.optimizePackages(mSnapshot, "bg-dexopt", cancellationSignal,
+                null /* processCallbackExecutor */, null /* processCallback */);
     }
 
     @Test
diff --git a/libartservice/service/javatests/com/android/server/art/BackgroundDexOptJobTest.java b/libartservice/service/javatests/com/android/server/art/BackgroundDexOptJobTest.java
index 3bf229c..16480ac 100644
--- a/libartservice/service/javatests/com/android/server/art/BackgroundDexOptJobTest.java
+++ b/libartservice/service/javatests/com/android/server/art/BackgroundDexOptJobTest.java
@@ -115,7 +115,7 @@
     @Test
     public void testStart() {
         when(mArtManagerLocal.optimizePackages(
-                     same(mSnapshot), eq(ReasonMapping.REASON_BG_DEXOPT), any()))
+                     same(mSnapshot), eq(ReasonMapping.REASON_BG_DEXOPT), any(), any(), any()))
                 .thenReturn(mOptimizeResult);
 
         Result result = Utils.getFuture(mBackgroundDexOptJob.start());
@@ -126,10 +126,11 @@
     @Test
     public void testStartAlreadyRunning() {
         Semaphore optimizeDone = new Semaphore(0);
-        when(mArtManagerLocal.optimizePackages(any(), any(), any())).thenAnswer(invocation -> {
-            assertThat(optimizeDone.tryAcquire(TIMEOUT_SEC, TimeUnit.SECONDS)).isTrue();
-            return mOptimizeResult;
-        });
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
+                .thenAnswer(invocation -> {
+                    assertThat(optimizeDone.tryAcquire(TIMEOUT_SEC, TimeUnit.SECONDS)).isTrue();
+                    return mOptimizeResult;
+                });
 
         Future<Result> future1 = mBackgroundDexOptJob.start();
         Future<Result> future2 = mBackgroundDexOptJob.start();
@@ -138,12 +139,13 @@
         optimizeDone.release();
         Utils.getFuture(future1);
 
-        verify(mArtManagerLocal, times(1)).optimizePackages(any(), any(), any());
+        verify(mArtManagerLocal, times(1)).optimizePackages(any(), any(), any(), any(), any());
     }
 
     @Test
     public void testStartAnother() {
-        when(mArtManagerLocal.optimizePackages(any(), any(), any())).thenReturn(mOptimizeResult);
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
+                .thenReturn(mOptimizeResult);
 
         Future<Result> future1 = mBackgroundDexOptJob.start();
         Utils.getFuture(future1);
@@ -154,7 +156,7 @@
 
     @Test
     public void testStartFatalError() {
-        when(mArtManagerLocal.optimizePackages(any(), any(), any()))
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
                 .thenThrow(IllegalStateException.class);
 
         Result result = Utils.getFuture(mBackgroundDexOptJob.start());
@@ -167,7 +169,8 @@
                 .when(SystemProperties.getBoolean(eq("pm.dexopt.disable_bg_dexopt"), anyBoolean()))
                 .thenReturn(true);
 
-        when(mArtManagerLocal.optimizePackages(any(), any(), any())).thenReturn(mOptimizeResult);
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
+                .thenReturn(mOptimizeResult);
 
         // The `start` method should ignore the system property. The system property is for
         // `schedule`.
@@ -177,12 +180,14 @@
     @Test
     public void testCancel() {
         Semaphore optimizeCancelled = new Semaphore(0);
-        when(mArtManagerLocal.optimizePackages(any(), any(), any())).thenAnswer(invocation -> {
-            assertThat(optimizeCancelled.tryAcquire(TIMEOUT_SEC, TimeUnit.SECONDS)).isTrue();
-            var cancellationSignal = invocation.<CancellationSignal>getArgument(2);
-            assertThat(cancellationSignal.isCanceled()).isTrue();
-            return mOptimizeResult;
-        });
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
+                .thenAnswer(invocation -> {
+                    assertThat(optimizeCancelled.tryAcquire(TIMEOUT_SEC, TimeUnit.SECONDS))
+                            .isTrue();
+                    var cancellationSignal = invocation.<CancellationSignal>getArgument(2);
+                    assertThat(cancellationSignal.isCanceled()).isTrue();
+                    return mOptimizeResult;
+                });
 
         Future<Result> future = mBackgroundDexOptJob.start();
         mBackgroundDexOptJob.cancel();
@@ -261,7 +266,8 @@
     @Test
     public void testWantsRescheduleFalsePerformed() throws Exception {
         when(mOptimizeResult.getFinalStatus()).thenReturn(OptimizeResult.OPTIMIZE_PERFORMED);
-        when(mArtManagerLocal.optimizePackages(any(), any(), any())).thenReturn(mOptimizeResult);
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
+                .thenReturn(mOptimizeResult);
 
         mBackgroundDexOptJob.onStartJob(mJobService, mJobParameters);
         assertThat(mJobFinishedCalled.tryAcquire(TIMEOUT_SEC, TimeUnit.SECONDS)).isTrue();
@@ -271,7 +277,7 @@
 
     @Test
     public void testWantsRescheduleFalseFatalError() throws Exception {
-        when(mArtManagerLocal.optimizePackages(any(), any(), any()))
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
                 .thenThrow(RuntimeException.class);
 
         mBackgroundDexOptJob.onStartJob(mJobService, mJobParameters);
@@ -283,7 +289,8 @@
     @Test
     public void testWantsRescheduleTrue() throws Exception {
         when(mOptimizeResult.getFinalStatus()).thenReturn(OptimizeResult.OPTIMIZE_CANCELLED);
-        when(mArtManagerLocal.optimizePackages(any(), any(), any())).thenReturn(mOptimizeResult);
+        when(mArtManagerLocal.optimizePackages(any(), any(), any(), any(), any()))
+                .thenReturn(mOptimizeResult);
 
         mBackgroundDexOptJob.onStartJob(mJobService, mJobParameters);
         assertThat(mJobFinishedCalled.tryAcquire(TIMEOUT_SEC, TimeUnit.SECONDS)).isTrue();
diff --git a/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java b/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
index ac79d26..c090bd5 100644
--- a/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
+++ b/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
@@ -25,6 +25,7 @@
 import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.argThat;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.lenient;
@@ -45,6 +46,7 @@
 import com.android.server.art.model.ArtFlags;
 import com.android.server.art.model.Config;
 import com.android.server.art.model.OptimizeParams;
+import com.android.server.art.model.OptimizeProgress;
 import com.android.server.art.model.OptimizeResult;
 import com.android.server.pm.PackageManagerLocal;
 import com.android.server.pm.pkg.AndroidPackage;
@@ -61,8 +63,10 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 @SmallTest
@@ -496,6 +500,40 @@
         mConfig.addOptimizePackageDoneCallback(Runnable::run, callback);
     }
 
+    @Test
+    public void testProgressCallback() throws Exception {
+        mParams = new OptimizeParams.Builder("install")
+                          .setCompilerFilter("speed-profile")
+                          .setFlags(ArtFlags.FLAG_FOR_SECONDARY_DEX,
+                                  ArtFlags.FLAG_FOR_SECONDARY_DEX
+                                          | ArtFlags.FLAG_SHOULD_INCLUDE_DEPENDENCIES)
+                          .build();
+
+        // Delay the executor to verify that the commands passed to the executor are not bound to
+        // changing variables.
+        var progressCallbackExecutor = new DelayedExecutor();
+        Consumer<OptimizeProgress> progressCallback = mock(Consumer.class);
+
+        mDexOptHelper.dexopt(mSnapshot, mRequestedPackages, mParams, mCancellationSignal, mExecutor,
+                progressCallbackExecutor, progressCallback);
+
+        progressCallbackExecutor.runAll();
+
+        InOrder inOrder = inOrder(progressCallback);
+        inOrder.verify(progressCallback)
+                .accept(eq(OptimizeProgress.create(
+                        0 /* donePackageCount */, 3 /* totalPackageCount */)));
+        inOrder.verify(progressCallback)
+                .accept(eq(OptimizeProgress.create(
+                        1 /* donePackageCount */, 3 /* totalPackageCount */)));
+        inOrder.verify(progressCallback)
+                .accept(eq(OptimizeProgress.create(
+                        2 /* donePackageCount */, 3 /* totalPackageCount */)));
+        inOrder.verify(progressCallback)
+                .accept(eq(OptimizeProgress.create(
+                        3 /* donePackageCount */, 3 /* totalPackageCount */)));
+    }
+
     private AndroidPackage createPackage(boolean multiSplit) {
         AndroidPackage pkg = mock(AndroidPackage.class);
 
@@ -625,4 +663,20 @@
                                                    .flatMap(r -> r.stream())
                                                    .collect(Collectors.toList()));
     }
+
+    /** An executor that delays execution until `runAll` is called. */
+    private static class DelayedExecutor implements Executor {
+        private List<Runnable> mCommands = new ArrayList<>();
+
+        public void execute(Runnable command) {
+            mCommands.add(command);
+        }
+
+        public void runAll() {
+            for (Runnable command : mCommands) {
+                command.run();
+            }
+            mCommands.clear();
+        }
+    }
 }