Support callbacks for results of optimizing package(s).

Bug: 257027956
Test: ArtServiceTests
Ignore-AOSP-First: ART Services.
Change-Id: I2ac70058ef5678641216cc37f68993981157c064
diff --git a/libartservice/service/api/system-server-current.txt b/libartservice/service/api/system-server-current.txt
index 6639926..0fdabff 100644
--- a/libartservice/service/api/system-server-current.txt
+++ b/libartservice/service/api/system-server-current.txt
@@ -4,6 +4,7 @@
   public final class ArtManagerLocal {
     ctor @Deprecated public ArtManagerLocal();
     ctor public ArtManagerLocal(@NonNull android.content.Context);
+    method public void addOptimizePackageDoneCallback(@NonNull java.util.concurrent.Executor, @NonNull com.android.server.art.ArtManagerLocal.OptimizePackageDoneCallback);
     method public void cancelBackgroundDexoptJob();
     method public void clearOptimizePackagesCallback();
     method public void clearScheduleBackgroundDexoptJobCallback();
@@ -15,6 +16,7 @@
     method public void notifyDexContainersLoaded(@NonNull com.android.server.pm.PackageManagerLocal.FilteredSnapshot, @NonNull String, @NonNull java.util.Map<java.lang.String,java.lang.String>);
     method @NonNull public com.android.server.art.model.OptimizeResult optimizePackage(@NonNull com.android.server.pm.PackageManagerLocal.FilteredSnapshot, @NonNull String, @NonNull com.android.server.art.model.OptimizeParams);
     method @NonNull public com.android.server.art.model.OptimizeResult optimizePackage(@NonNull com.android.server.pm.PackageManagerLocal.FilteredSnapshot, @NonNull String, @NonNull com.android.server.art.model.OptimizeParams, @NonNull android.os.CancellationSignal);
+    method public void removeOptimizePackageDoneCallback(@NonNull com.android.server.art.ArtManagerLocal.OptimizePackageDoneCallback);
     method public int scheduleBackgroundDexoptJob();
     method public void setOptimizePackagesCallback(@NonNull java.util.concurrent.Executor, @NonNull com.android.server.art.ArtManagerLocal.OptimizePackagesCallback);
     method public void setScheduleBackgroundDexoptJobCallback(@NonNull java.util.concurrent.Executor, @NonNull com.android.server.art.ArtManagerLocal.ScheduleBackgroundDexoptJobCallback);
@@ -22,6 +24,10 @@
     method public void unscheduleBackgroundDexoptJob();
   }
 
+  public static interface ArtManagerLocal.OptimizePackageDoneCallback {
+    method public void onOptimizePackageDone(@NonNull com.android.server.art.model.OptimizeResult);
+  }
+
   public static interface ArtManagerLocal.OptimizePackagesCallback {
     method public void onOverrideBatchOptimizeParams(@NonNull com.android.server.pm.PackageManagerLocal.FilteredSnapshot, @NonNull String, @NonNull java.util.List<java.lang.String>, @NonNull com.android.server.art.model.BatchOptimizeParams.Builder);
   }
diff --git a/libartservice/service/java/com/android/server/art/ArtManagerLocal.java b/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
index ed14e19..90f027a 100644
--- a/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
+++ b/libartservice/service/java/com/android/server/art/ArtManagerLocal.java
@@ -489,6 +489,28 @@
     }
 
     /**
+     * Adds a global listener that listens to any result of optimizing package(s), no matter run
+     * manually or automatically. Calling this method multiple times with different callbacks is
+     * allowed. Callbacks are executed in the same order as the one in which they were added. This
+     * method is thread-safe.
+     *
+     * @throws IllegalStateException if the same callback instance is already added
+     */
+    public void addOptimizePackageDoneCallback(@NonNull @CallbackExecutor Executor executor,
+            @NonNull OptimizePackageDoneCallback callback) {
+        mInjector.getConfig().addOptimizePackageDoneCallback(executor, callback);
+    }
+
+    /**
+     * Removes the listener added by {@link #addOptimizePackageDoneCallback(Executor,
+     * OptimizePackageDoneCallback)}. Does nothing if the callback was not added. This method is
+     * thread-safe.
+     */
+    public void removeOptimizePackageDoneCallback(@NonNull OptimizePackageDoneCallback callback) {
+        mInjector.getConfig().removeOptimizePackageDoneCallback(callback);
+    }
+
+    /**
      * Should be used by {@link BackgroundDexOptJobService} ONLY.
      *
      * @hide
@@ -543,6 +565,10 @@
         void onOverrideJobInfo(@NonNull JobInfo.Builder builder);
     }
 
+    public interface OptimizePackageDoneCallback {
+        void onOptimizePackageDone(@NonNull OptimizeResult result);
+    }
+
     /**
      * Injector pattern for testing purpose.
      *
@@ -586,7 +612,7 @@
 
         @NonNull
         public DexOptHelper getDexOptHelper() {
-            return new DexOptHelper(getContext());
+            return new DexOptHelper(getContext(), getConfig());
         }
 
         @NonNull
diff --git a/libartservice/service/java/com/android/server/art/DexOptHelper.java b/libartservice/service/java/com/android/server/art/DexOptHelper.java
index 43be7d3..25bc325 100644
--- a/libartservice/service/java/com/android/server/art/DexOptHelper.java
+++ b/libartservice/service/java/com/android/server/art/DexOptHelper.java
@@ -16,6 +16,8 @@
 
 package com.android.server.art;
 
+import static com.android.server.art.ArtManagerLocal.OptimizePackageDoneCallback;
+import static com.android.server.art.model.Config.Callback;
 import static com.android.server.art.model.OptimizeResult.DexContainerFileOptimizeResult;
 import static com.android.server.art.model.OptimizeResult.PackageOptimizeResult;
 
@@ -30,6 +32,7 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 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.OptimizeResult;
 import com.android.server.pm.PackageManagerLocal;
@@ -70,8 +73,8 @@
 
     @NonNull private final Injector mInjector;
 
-    public DexOptHelper(@NonNull Context context) {
-        this(new Injector(context));
+    public DexOptHelper(@NonNull Context context, @NonNull Config config) {
+        this(new Injector(context, config));
     }
 
     @VisibleForTesting
@@ -123,7 +126,17 @@
             List<PackageOptimizeResult> results =
                     futures.stream().map(Utils::getFuture).collect(Collectors.toList());
 
-            return new OptimizeResult(params.getCompilerFilter(), params.getReason(), results);
+            var result =
+                    new OptimizeResult(params.getCompilerFilter(), params.getReason(), results);
+
+            for (Callback<OptimizePackageDoneCallback> callback :
+                    mInjector.getConfig().getOptimizePackageDoneCallbacks()) {
+                // TODO(b/257027956): Consider filtering the packages before calling the callback.
+                Utils.executeAndWait(callback.executor(),
+                        () -> { callback.get().onOptimizePackageDone(result); });
+            }
+
+            return result;
         } finally {
             if (wakeLock != null) {
                 wakeLock.release();
@@ -241,9 +254,11 @@
     @VisibleForTesting
     public static class Injector {
         @NonNull private final Context mContext;
+        @NonNull private final Config mConfig;
 
-        Injector(@NonNull Context context) {
+        Injector(@NonNull Context context, @NonNull Config config) {
             mContext = context;
+            mConfig = config;
         }
 
         @NonNull
@@ -269,5 +284,10 @@
         public PowerManager getPowerManager() {
             return Objects.requireNonNull(mContext.getSystemService(PowerManager.class));
         }
+
+        @NonNull
+        public Config getConfig() {
+            return mConfig;
+        }
     }
 }
diff --git a/libartservice/service/java/com/android/server/art/model/Config.java b/libartservice/service/java/com/android/server/art/model/Config.java
index a8ed18b..8b84bcc 100644
--- a/libartservice/service/java/com/android/server/art/model/Config.java
+++ b/libartservice/service/java/com/android/server/art/model/Config.java
@@ -16,6 +16,7 @@
 
 package com.android.server.art.model;
 
+import static com.android.server.art.ArtManagerLocal.OptimizePackageDoneCallback;
 import static com.android.server.art.ArtManagerLocal.OptimizePackagesCallback;
 import static com.android.server.art.ArtManagerLocal.ScheduleBackgroundDexoptJobCallback;
 
@@ -27,6 +28,9 @@
 
 import com.google.auto.value.AutoValue;
 
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.concurrent.Executor;
 
 /**
@@ -50,6 +54,14 @@
     private Callback<ScheduleBackgroundDexoptJobCallback> mScheduleBackgroundDexoptJobCallback =
             null;
 
+    /**
+     * @see ArtManagerLocal#addOptimizePackageDoneCallback(Executor, OptimizePackageDoneCallback)
+     */
+    @GuardedBy("this")
+    @NonNull
+    private LinkedHashMap<OptimizePackageDoneCallback, Callback<OptimizePackageDoneCallback>>
+            mOptimizePackageDoneCallbacks = new LinkedHashMap<>();
+
     public synchronized void setOptimizePackagesCallback(
             @NonNull Executor executor, @NonNull OptimizePackagesCallback callback) {
         mOptimizePackagesCallback = Callback.<OptimizePackagesCallback>create(callback, executor);
@@ -80,6 +92,26 @@
         return mScheduleBackgroundDexoptJobCallback;
     }
 
+    public synchronized void addOptimizePackageDoneCallback(
+            @NonNull Executor executor, @NonNull OptimizePackageDoneCallback callback) {
+        if (mOptimizePackageDoneCallbacks.putIfAbsent(
+                    callback, Callback.<OptimizePackageDoneCallback>create(callback, executor))
+                != null) {
+            throw new IllegalStateException("callback already added");
+        }
+    }
+
+    public synchronized void removeOptimizePackageDoneCallback(
+            @NonNull OptimizePackageDoneCallback callback) {
+        mOptimizePackageDoneCallbacks.remove(callback);
+    }
+
+    @NonNull
+    public synchronized List<Callback<OptimizePackageDoneCallback>>
+    getOptimizePackageDoneCallbacks() {
+        return new ArrayList<>(mOptimizePackageDoneCallbacks.values());
+    }
+
     @AutoValue
     public static abstract class Callback<T> {
         public abstract @NonNull T get();
diff --git a/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java b/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
index f71f46d..e12cd5d 100644
--- a/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
+++ b/libartservice/service/javatests/com/android/server/art/DexOptHelperTest.java
@@ -16,6 +16,7 @@
 
 package com.android.server.art;
 
+import static com.android.server.art.ArtManagerLocal.OptimizePackageDoneCallback;
 import static com.android.server.art.model.OptimizeResult.DexContainerFileOptimizeResult;
 import static com.android.server.art.model.OptimizeResult.PackageOptimizeResult;
 
@@ -41,6 +42,7 @@
 import androidx.test.filters.SmallTest;
 
 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.OptimizeResult;
 import com.android.server.pm.PackageManagerLocal;
@@ -56,6 +58,7 @@
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -95,15 +98,13 @@
     private ExecutorService mExecutor = Executors.newSingleThreadExecutor();
     private List<DexContainerFileOptimizeResult> mPrimaryResults;
     private List<DexContainerFileOptimizeResult> mSecondaryResults;
+    private Config mConfig;
     private OptimizeParams mParams;
     private List<String> mRequestedPackages;
     private DexOptHelper mDexOptHelper;
 
     @Before
     public void setUp() throws Exception {
-        lenient().when(mInjector.getAppHibernationManager()).thenReturn(mAhm);
-        lenient().when(mInjector.getPowerManager()).thenReturn(mPowerManager);
-
         lenient()
                 .when(mPowerManager.newWakeLock(eq(PowerManager.PARTIAL_WAKE_LOCK), any()))
                 .thenReturn(mWakeLock);
@@ -112,6 +113,7 @@
         lenient().when(mAhm.isOatArtifactDeletionEnabled()).thenReturn(true);
 
         mCancellationSignal = new CancellationSignal();
+        mConfig = new Config();
 
         preparePackagesAndLibraries();
 
@@ -137,6 +139,10 @@
                                           | ArtFlags.FLAG_SHOULD_INCLUDE_DEPENDENCIES)
                           .build();
 
+        lenient().when(mInjector.getAppHibernationManager()).thenReturn(mAhm);
+        lenient().when(mInjector.getPowerManager()).thenReturn(mPowerManager);
+        lenient().when(mInjector.getConfig()).thenReturn(mConfig);
+
         mDexOptHelper = new DexOptHelper(mInjector);
     }
 
@@ -420,6 +426,47 @@
         verifyNoDexopt();
     }
 
+    @Test
+    public void testCallbacks() throws Exception {
+        List<OptimizeResult> list1 = new ArrayList<>();
+        mConfig.addOptimizePackageDoneCallback(Runnable::run, result -> list1.add(result));
+
+        List<OptimizeResult> list2 = new ArrayList<>();
+        mConfig.addOptimizePackageDoneCallback(Runnable::run, result -> list2.add(result));
+
+        OptimizeResult result = mDexOptHelper.dexopt(
+                mSnapshot, mRequestedPackages, mParams, mCancellationSignal, mExecutor);
+
+        assertThat(list1).containsExactly(result);
+        assertThat(list2).containsExactly(result);
+    }
+
+    @Test
+    public void testCallbackRemoved() throws Exception {
+        List<OptimizeResult> list1 = new ArrayList<>();
+        OptimizePackageDoneCallback callback1 = result -> list1.add(result);
+        mConfig.addOptimizePackageDoneCallback(Runnable::run, callback1);
+
+        List<OptimizeResult> list2 = new ArrayList<>();
+        mConfig.addOptimizePackageDoneCallback(Runnable::run, result -> list2.add(result));
+
+        mConfig.removeOptimizePackageDoneCallback(callback1);
+
+        OptimizeResult result = mDexOptHelper.dexopt(
+                mSnapshot, mRequestedPackages, mParams, mCancellationSignal, mExecutor);
+
+        assertThat(list1).isEmpty();
+        assertThat(list2).containsExactly(result);
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void testCallbackAlreadyAdded() throws Exception {
+        List<OptimizeResult> list = new ArrayList<>();
+        OptimizePackageDoneCallback callback = result -> list.add(result);
+        mConfig.addOptimizePackageDoneCallback(Runnable::run, callback);
+        mConfig.addOptimizePackageDoneCallback(Runnable::run, callback);
+    }
+
     private AndroidPackage createPackage() {
         AndroidPackage pkg = mock(AndroidPackage.class);
         var baseSplit = mock(AndroidPackageSplit.class);