Reland "Remove MessageQueue reflection from TestLooper"

This reverts commit b652cb193092176cdc88d1268d5c9430d7a9aafa.

Reason for revert: fixed underlying issue

Change-Id: I21bd4453b42be42bc4dcc70a38242a7fa9bc8e1e

1 files changed, 146 insertions, 22 deletions

diff --git a/tests/utils/testutils/java/android/os/test/TestLooper.java b/tests/utils/testutils/java/android/os/test/TestLooper.java
index 83d22d923c78..4d379e45a81a 100644
--- a/tests/utils/testutils/java/android/os/test/TestLooper.java
+++ b/tests/utils/testutils/java/android/os/test/TestLooper.java
@@ -18,18 +18,24 @@ package android.os.test;
 
 import static org.junit.Assert.assertTrue;
 
+import android.os.Build;
 import android.os.Handler;
 import android.os.HandlerExecutor;
 import android.os.Looper;
 import android.os.Message;
 import android.os.MessageQueue;
 import android.os.SystemClock;
+import android.os.TestLooperManager;
 import android.util.Log;
 
+import androidx.test.platform.app.InstrumentationRegistry;
+
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.ArrayDeque;
+import java.util.Queue;
 import java.util.concurrent.Executor;
 
 /**
@@ -44,7 +50,9 @@ import java.util.concurrent.Executor;
  *     The Robolectric class also allows advancing time.
  */
 public class TestLooper {
-    protected final Looper mLooper;
+    private final Looper mLooper;
+    private final TestLooperManager mTestLooperManager;
+    private final Clock mClock;
 
     private static final Constructor<Looper> LOOPER_CONSTRUCTOR;
     private static final Field THREAD_LOCAL_LOOPER_FIELD;
@@ -54,24 +62,46 @@ public class TestLooper {
     private static final Method MESSAGE_MARK_IN_USE_METHOD;
     private static final String TAG = "TestLooper";
 
-    private final Clock mClock;
-
     private AutoDispatchThread mAutoDispatchThread;
 
+    /**
+     * Baklava introduces new {@link TestLooperManager} APIs that we can use instead of reflection.
+     */
+    private static boolean isAtLeastBaklava() {
+        Method[] methods = TestLooperManager.class.getMethods();
+        for (Method method : methods) {
+            if (method.getName().equals("peekWhen")) {
+                return true;
+            }
+        }
+        return false;
+        // TODO(shayba): delete the above, uncomment the below.
+        // SDK_INT has not yet ramped to Baklava in all 25Q2 builds.
+        // return Build.VERSION.SDK_INT >= Build.VERSION_CODES.BAKLAVA;
+    }
+
     static {
         try {
             LOOPER_CONSTRUCTOR = Looper.class.getDeclaredConstructor(Boolean.TYPE);
             LOOPER_CONSTRUCTOR.setAccessible(true);
             THREAD_LOCAL_LOOPER_FIELD = Looper.class.getDeclaredField("sThreadLocal");
             THREAD_LOCAL_LOOPER_FIELD.setAccessible(true);
-            MESSAGE_QUEUE_MESSAGES_FIELD = MessageQueue.class.getDeclaredField("mMessages");
-            MESSAGE_QUEUE_MESSAGES_FIELD.setAccessible(true);
-            MESSAGE_NEXT_FIELD = Message.class.getDeclaredField("next");
-            MESSAGE_NEXT_FIELD.setAccessible(true);
-            MESSAGE_WHEN_FIELD = Message.class.getDeclaredField("when");
-            MESSAGE_WHEN_FIELD.setAccessible(true);
-            MESSAGE_MARK_IN_USE_METHOD = Message.class.getDeclaredMethod("markInUse");
-            MESSAGE_MARK_IN_USE_METHOD.setAccessible(true);
+
+            if (isAtLeastBaklava()) {
+                MESSAGE_QUEUE_MESSAGES_FIELD = null;
+                MESSAGE_NEXT_FIELD = null;
+                MESSAGE_WHEN_FIELD = null;
+                MESSAGE_MARK_IN_USE_METHOD = null;
+            } else {
+                MESSAGE_QUEUE_MESSAGES_FIELD = MessageQueue.class.getDeclaredField("mMessages");
+                MESSAGE_QUEUE_MESSAGES_FIELD.setAccessible(true);
+                MESSAGE_NEXT_FIELD = Message.class.getDeclaredField("next");
+                MESSAGE_NEXT_FIELD.setAccessible(true);
+                MESSAGE_WHEN_FIELD = Message.class.getDeclaredField("when");
+                MESSAGE_WHEN_FIELD.setAccessible(true);
+                MESSAGE_MARK_IN_USE_METHOD = Message.class.getDeclaredMethod("markInUse");
+                MESSAGE_MARK_IN_USE_METHOD.setAccessible(true);
+            }
         } catch (NoSuchFieldException | NoSuchMethodException e) {
             throw new RuntimeException("Failed to initialize TestLooper", e);
         }
@@ -106,6 +136,13 @@ public class TestLooper {
             throw new RuntimeException("Reflection error constructing or accessing looper", e);
         }
 
+        if (isAtLeastBaklava()) {
+            mTestLooperManager =
+                InstrumentationRegistry.getInstrumentation().acquireLooperManager(mLooper);
+        } else {
+            mTestLooperManager = null;
+        }
+
         mClock = clock;
     }
 
@@ -117,19 +154,61 @@ public class TestLooper {
         return new HandlerExecutor(new Handler(getLooper()));
     }
 
-    private Message getMessageLinkedList() {
+    private Message getMessageLinkedListLegacy() {
         try {
             MessageQueue queue = mLooper.getQueue();
             return (Message) MESSAGE_QUEUE_MESSAGES_FIELD.get(queue);
         } catch (IllegalAccessException e) {
             throw new RuntimeException("Access failed in TestLooper: get - MessageQueue.mMessages",
-                    e);
+                e);
         }
     }
 
     public void moveTimeForward(long milliSeconds) {
+        if (isAtLeastBaklava()) {
+            moveTimeForwardBaklava(milliSeconds);
+        } else {
+            moveTimeForwardLegacy(milliSeconds);
+        }
+    }
+
+    private void moveTimeForwardBaklava(long milliSeconds) {
+        // Drain all Messages from the queue.
+        Queue<Message> messages = new ArrayDeque<>();
+        while (true) {
+            Message message = mTestLooperManager.poll();
+            if (message == null) {
+                break;
+            }
+            messages.add(message);
+        }
+
+        // Repost all Messages back to the queue with a new time.
+        while (true) {
+            Message message = messages.poll();
+            if (message == null) {
+                break;
+            }
+
+            // Ugly trick to reset the Message's "in use" flag.
+            // This is needed because the Message cannot be re-enqueued if it's
+            // marked in use.
+            message.copyFrom(message);
+
+            // Adjust the Message's delivery time.
+            long newWhen = message.getWhen() - milliSeconds;
+            if (newWhen < 0) {
+                newWhen = 0;
+            }
+
+            // Send the Message back to its Handler to be re-enqueued.
+            message.getTarget().sendMessageAtTime(message, newWhen);
+        }
+    }
+
+    private void moveTimeForwardLegacy(long milliSeconds) {
         try {
-            Message msg = getMessageLinkedList();
+            Message msg = getMessageLinkedListLegacy();
             while (msg != null) {
                 long updatedWhen = msg.getWhen() - milliSeconds;
                 if (updatedWhen < 0) {
@@ -147,12 +226,12 @@ public class TestLooper {
         return mClock.uptimeMillis();
     }
 
-    private Message messageQueueNext() {
+    private Message messageQueueNextLegacy() {
         try {
             long now = currentTime();
 
             Message prevMsg = null;
-            Message msg = getMessageLinkedList();
+            Message msg = getMessageLinkedListLegacy();
             if (msg != null && msg.getTarget() == null) {
                 // Stalled by a barrier. Find the next asynchronous message in
                 // the queue.
@@ -185,18 +264,46 @@ public class TestLooper {
     /**
      * @return true if there are pending messages in the message queue
      */
-    public synchronized boolean isIdle() {
-        Message messageList = getMessageLinkedList();
+    public boolean isIdle() {
+        if (isAtLeastBaklava()) {
+            return isIdleBaklava();
+        } else {
+            return isIdleLegacy();
+        }
+    }
+
+    private boolean isIdleBaklava() {
+        Long when = mTestLooperManager.peekWhen();
+        return when != null && currentTime() >= when;
+    }
 
+    private synchronized boolean isIdleLegacy() {
+        Message messageList = getMessageLinkedListLegacy();
         return messageList != null && currentTime() >= messageList.getWhen();
     }
 
     /**
      * @return the next message in the Looper's message queue or null if there is none
      */
-    public synchronized Message nextMessage() {
+    public Message nextMessage() {
+        if (isAtLeastBaklava()) {
+            return nextMessageBaklava();
+        } else {
+            return nextMessageLegacy();
+        }
+    }
+
+    private Message nextMessageBaklava() {
+        if (isIdle()) {
+            return mTestLooperManager.poll();
+        } else {
+            return null;
+        }
+    }
+
+    private synchronized Message nextMessageLegacy() {
         if (isIdle()) {
-            return messageQueueNext();
+            return messageQueueNextLegacy();
         } else {
             return null;
         }
@@ -206,9 +313,26 @@ public class TestLooper {
      * Dispatch the next message in the queue
      * Asserts that there is a message in the queue
      */
-    public synchronized void dispatchNext() {
+    public void dispatchNext() {
+        if (isAtLeastBaklava()) {
+            dispatchNextBaklava();
+        } else {
+            dispatchNextLegacy();
+        }
+    }
+
+    private void dispatchNextBaklava() {
+        assertTrue(isIdle());
+        Message msg = mTestLooperManager.poll();
+        if (msg == null) {
+            return;
+        }
+        msg.getTarget().dispatchMessage(msg);
+    }
+
+    private synchronized void dispatchNextLegacy() {
         assertTrue(isIdle());
-        Message msg = messageQueueNext();
+        Message msg = messageQueueNextLegacy();
         if (msg == null) {
             return;
         }