[pm] shell command to wait for handler(s)

Adding a shell command to wait for the Package Manager handler to finish
handling all messages that are currently waiting in the message queue.
This gives a better indicator for tests to know when to wait for
incoming broadcasts.

Added a similar command for the background handler as well.

BUG: 280604852
Test: atest android.content.pm.cts.PackageManagerShellCommandMultiUserTest
Change-Id: I43c281a4f17edc3982a7885a39b1182be57c4504
Merged-In: I43c281a4f17edc3982a7885a39b1182be57c4504

3 files changed, 83 insertions, 0 deletions

diff --git a/core/java/android/content/pm/IPackageManager.aidl b/core/java/android/content/pm/IPackageManager.aidl
index d802b46f1fde..47a5db8ea22a 100644
--- a/core/java/android/content/pm/IPackageManager.aidl
+++ b/core/java/android/content/pm/IPackageManager.aidl
@@ -791,4 +791,6 @@ interface IPackageManager {
     void setKeepUninstalledPackages(in List<String> packageList);
 
     boolean[] canPackageQuery(String sourcePackageName, in String[] targetPackageNames, int userId);
+
+    boolean waitForHandler(long timeoutMillis, boolean forBackgroundHandler);
 }
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index ae520c00d977..bba8043af5be 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -279,6 +279,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
@@ -6295,6 +6296,36 @@ public class PackageManagerService implements PackageSender, TestUtilityService
             }
         }
 
+        /**
+         * Wait for the handler to finish handling all pending messages.
+         * @param timeoutMillis Maximum time in milliseconds to wait.
+         * @param forBackgroundHandler Whether to wait for the background handler instead.
+         * @return True if all the waiting messages in the handler has been handled.
+         *         False if timeout.
+         */
+        @Override
+        public boolean waitForHandler(long timeoutMillis, boolean forBackgroundHandler) {
+            final CountDownLatch latch = new CountDownLatch(1);
+            if (forBackgroundHandler) {
+                mBackgroundHandler.post(latch::countDown);
+            } else {
+                mHandler.post(latch::countDown);
+            }
+            final long endTimeMillis = System.currentTimeMillis() + timeoutMillis;
+            while (latch.getCount() > 0) {
+                try {
+                    final long remainingTimeMillis = endTimeMillis - System.currentTimeMillis();
+                    if (remainingTimeMillis <= 0) {
+                        return false;
+                    }
+                    return latch.await(remainingTimeMillis, TimeUnit.MILLISECONDS);
+                } catch (InterruptedException e) {
+                    // ignore and retry
+                }
+            }
+            return true;
+        }
+
         @Override
         public boolean onTransact(int code, Parcel data, Parcel reply, int flags)
                 throws RemoteException {
diff --git a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
index 05bfec48464b..e1f010f62232 100644
--- a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
+++ b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
@@ -354,6 +354,10 @@ class PackageManagerShellCommand extends ShellCommand {
                     return runSetSilentUpdatesPolicy();
                 case "get-app-metadata":
                     return runGetAppMetadata();
+                case "wait-for-handler":
+                    return runWaitForHandler(/* forBackgroundHandler= */ false);
+                case "wait-for-background-handler":
+                    return runWaitForHandler(/* forBackgroundHandler= */ true);
                 default: {
                     if (ART_SERVICE_COMMANDS.contains(cmd)) {
                         if (DexOptHelper.useArtService()) {
@@ -3601,6 +3605,40 @@ class PackageManagerShellCommand extends ShellCommand {
         return 1;
     }
 
+    private int runWaitForHandler(boolean forBackgroundHandler) {
+        final PrintWriter pw = getOutPrintWriter();
+        long timeoutMillis = 60000; // default timeout is 60 seconds
+        String opt;
+        while ((opt = getNextOption()) != null) {
+            switch (opt) {
+                case "--timeout":
+                    timeoutMillis = Long.parseLong(getNextArgRequired());
+                    break;
+                default:
+                    pw.println("Error: Unknown option: " + opt);
+                    return -1;
+            }
+        }
+        if (timeoutMillis <= 0) {
+            pw.println("Error: --timeout value must be positive: " + timeoutMillis);
+            return -1;
+        }
+        final boolean success;
+        try {
+            success = mInterface.waitForHandler(timeoutMillis, forBackgroundHandler);
+        } catch (RemoteException e) {
+            pw.println("Failure [" + e.getClass().getName() + " - " + e.getMessage() + "]");
+            return -1;
+        }
+        if (success) {
+            pw.println("Success");
+            return 0;
+        } else {
+            pw.println("Timeout. PackageManager handlers are still busy.");
+            return -1;
+        }
+    }
+
     private int runArtServiceCommand() {
         try (var in = ParcelFileDescriptor.dup(getInFileDescriptor());
                 var out = ParcelFileDescriptor.dup(getOutFileDescriptor());
@@ -4427,6 +4465,18 @@ class PackageManagerShellCommand extends ShellCommand {
         pw.println("      --reset: restore the installer and throttle time to the default, and");
         pw.println("        clear tracks of silent updates in the system.");
         pw.println("");
+        pw.println("  wait-for-handler --timeout <MILLIS>");
+        pw.println("    Wait for a given amount of time till the package manager handler finishes");
+        pw.println("    handling all pending messages.");
+        pw.println("      --timeout: wait for a given number of milliseconds. If the handler(s)");
+        pw.println("        fail to finish before the timeout, the command returns error.");
+        pw.println("");
+        pw.println("  wait-for-background-handler --timeout <MILLIS>");
+        pw.println("    Wait for a given amount of time till the package manager's background");
+        pw.println("    handler finishes handling all pending messages.");
+        pw.println("      --timeout: wait for a given number of milliseconds. If the handler(s)");
+        pw.println("        fail to finish before the timeout, the command returns error.");
+        pw.println("");
         if (DexOptHelper.useArtService()) {
             printArtServiceHelp();
         } else {