Don't wait for message queues to become idle in wait-for-barrier.

It would take a long time for some of the message queues in
system_server process to become idle and we only need to make
sure any messages currently waiting have been handled as part of
"wait-for-barrier".

Bug: 260158381
Test: TH
Change-Id: Ibcb36ebf767806c23f087ed34ad8792672fbe0d6

2 files changed, 33 insertions, 6 deletions

diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index 7566bab93cc1..efe14f4dec1c 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -18324,7 +18324,7 @@ public class ActivityManagerService extends IActivityManager.Stub
 
     public void waitForBroadcastBarrier(@Nullable PrintWriter pw) {
         enforceCallingPermission(permission.DUMP, "waitForBroadcastBarrier()");
-        BroadcastLoopers.waitForIdle(pw);
+        BroadcastLoopers.waitForBarrier(pw);
         for (BroadcastQueue queue : mBroadcastQueues) {
             queue.waitForBarrier(pw);
         }
diff --git a/services/core/java/com/android/server/am/BroadcastLoopers.java b/services/core/java/com/android/server/am/BroadcastLoopers.java
index b828720c9162..a5535cb13165 100644
--- a/services/core/java/com/android/server/am/BroadcastLoopers.java
+++ b/services/core/java/com/android/server/am/BroadcastLoopers.java
@@ -18,6 +18,7 @@ package com.android.server.am;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
 import android.os.MessageQueue;
@@ -30,6 +31,7 @@ import com.android.internal.annotations.GuardedBy;
 import java.io.PrintWriter;
 import java.util.Objects;
 import java.util.concurrent.CountDownLatch;
+import java.util.function.BiConsumer;
 
 /**
  * Collection of {@link Looper} that are known to be used for broadcast dispatch
@@ -73,19 +75,44 @@ public class BroadcastLoopers {
      * still in the future are ignored for the purposes of the idle test.
      */
     public static void waitForIdle(@Nullable PrintWriter pw) {
+        waitForCondition(pw, (looper, latch) -> {
+            final MessageQueue queue = looper.getQueue();
+            queue.addIdleHandler(() -> {
+                latch.countDown();
+                return false;
+            });
+        });
+    }
+
+    /**
+     * Wait for all registered {@link Looper} instances to handle currently waiting messages.
+     * Note that {@link Message#when} still in the future are ignored for the purposes
+     * of the idle test.
+     */
+    public static void waitForBarrier(@Nullable PrintWriter pw) {
+        waitForCondition(pw, (looper, latch) -> {
+            (new Handler(looper)).post(() -> {
+                latch.countDown();
+            });
+        });
+    }
+
+    /**
+     * Wait for all registered {@link Looper} instances to meet a certain condition.
+     */
+    private static void waitForCondition(@Nullable PrintWriter pw,
+            @NonNull BiConsumer<Looper, CountDownLatch> condition) {
         final CountDownLatch latch;
         synchronized (sLoopers) {
             final int N = sLoopers.size();
             latch = new CountDownLatch(N);
             for (int i = 0; i < N; i++) {
-                final MessageQueue queue = sLoopers.valueAt(i).getQueue();
+                final Looper looper = sLoopers.valueAt(i);
+                final MessageQueue queue = looper.getQueue();
                 if (queue.isIdle()) {
                     latch.countDown();
                 } else {
-                    queue.addIdleHandler(() -> {
-                        latch.countDown();
-                        return false;
-                    });
+                    condition.accept(looper, latch);
                 }
             }
         }