Support Compose Handler calls

- collect and call Handler callbacks on demand
- add logging since callbacks call throw easily

Bug: 139476297
Test: N/A
Change-Id: I1a37721c3871e814139c2eae71b874ecbd2be0b0
(cherry picked from commit 7530bf07e46e3e317a681e9a52e7356ae6f096b1)
diff --git a/bridge/src/android/os/Handler_Delegate.java b/bridge/src/android/os/Handler_Delegate.java
index 2152c8a..ef1a5f3 100644
--- a/bridge/src/android/os/Handler_Delegate.java
+++ b/bridge/src/android/os/Handler_Delegate.java
@@ -16,8 +16,12 @@
 
 package android.os;
 
+import com.android.ide.common.rendering.api.LayoutLog;
+import com.android.layoutlib.bridge.Bridge;
 import com.android.tools.layoutlib.annotations.LayoutlibDelegate;
 
+import java.util.LinkedList;
+import java.util.WeakHashMap;
 
 /**
  * Delegate overriding selected methods of android.os.Handler
@@ -30,6 +34,7 @@
 public class Handler_Delegate {
 
     // -------- Delegate methods
+    private static WeakHashMap<Handler, LinkedList<Runnable>> sRunnablesMap = new WeakHashMap<>();
 
     @LayoutlibDelegate
     /*package*/ static boolean sendMessageAtTime(Handler handler, Message msg, long uptimeMillis) {
@@ -41,7 +46,48 @@
         return true;
     }
 
+    /**
+     * Current implementation of Compose uses {@link Handler#postAtFrontOfQueue} to execute state
+     * updates. We can not intercept postAtFrontOfQueue Compose calls, however we can intecept
+     * internal Handler calls. Since postAtFrontOfQueue is just a wrapper of
+     * sendMessageAtFrontOfQueue we re-define sendMessageAtFrontOfQueue here to catch Compose calls
+     * (we are only interested in them) and execute them.
+     * TODO(b/137794558): Clean/rework this when Compose reworks Handler usage.
+     */
+    @LayoutlibDelegate
+    /*package*/ static boolean sendMessageAtFrontOfQueue(Handler handler, Message msg) {
+        // We will also catch calls from the Choreographer that have no callback.
+        if (msg.callback != null) {
+            LinkedList<Runnable> runnables =
+                    sRunnablesMap.computeIfAbsent(handler, k -> new LinkedList<>());
+            runnables.add(msg.callback);
+        }
+
+        return true;
+    }
+
     // -------- Delegate implementation
+    /**
+     * Executed all the collected callbacks
+     *
+     * @return if there are more callbacks to execute
+     */
+    public static boolean executeCallbacks() {
+        try {
+            while (sRunnablesMap.values().stream().anyMatch(runnables -> !runnables.isEmpty())) {
+                for (LinkedList<Runnable> runnables : sRunnablesMap.values()) {
+                    while (!runnables.isEmpty()) {
+                        Runnable r = runnables.poll();
+                        r.run();
+                    }
+                }
+            }
+        } catch (Throwable t) {
+            Bridge.getLog().error(LayoutLog.TAG_BROKEN, "Failed executing Handler callback", t,
+                null, null);
+        }
+        return false;
+    }
 
     public interface IHandlerCallback {
         void sendMessageAtTime(Handler handler, Message msg, long uptimeMillis);
diff --git a/bridge/src/com/android/layoutlib/bridge/BridgeRenderSession.java b/bridge/src/com/android/layoutlib/bridge/BridgeRenderSession.java
index b7e0018..3a6d87f 100644
--- a/bridge/src/com/android/layoutlib/bridge/BridgeRenderSession.java
+++ b/bridge/src/com/android/layoutlib/bridge/BridgeRenderSession.java
@@ -16,6 +16,7 @@
 
 package com.android.layoutlib.bridge;
 
+import com.android.ide.common.rendering.api.LayoutLog;
 import com.android.ide.common.rendering.api.RenderParams;
 import com.android.ide.common.rendering.api.RenderSession;
 import com.android.ide.common.rendering.api.ResourceReference;
@@ -27,6 +28,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.os.Handler_Delegate;
 import android.view.Choreographer;
 import android.view.MotionEvent;
 
@@ -145,18 +147,26 @@
 
     @Override
     public boolean executeCallbacks(long nanos) {
-        // So far, for Compose, we only have to call doFrame since Compose relies on the frame
-        // callback only. If we want to animate platform widgets as well we will also have to
-        // execute callbacks passes to the Handler (Handler_Delegate in our case) with
-        // sendMessageAtTime. For this purpose, we will have to save those messages with uptimes
-        // and execute appropriate (if uptime has passed) callbacks here.
+        // Currently, Compose relies on Choreographer frame callback and Handler#postAtFrontOfQueue.
+        // Calls to Handler are handled by Handler_Delegate and can be executed by Handler_Delegate#
+        // executeCallbacks. Choreographer frame callback is handled by Choreographer#doFrame.
+        //
+        // If we want to animate platform widgets as well we will also have to execute callbacks
+        // passed to the Handler (Handler_Delegate in our case) with sendMessageAtTime. For this
+        // purpose, we will have to save those messages with uptimes and execute appropriate (if
+        // uptime has passed) callbacks here.
         try {
             Bridge.prepareThread();
+            boolean hasMoreCallbacks = Handler_Delegate.executeCallbacks();
             Choreographer.getInstance().doFrame(nanos, 0);
+            return hasMoreCallbacks;
+        } catch (Throwable t) {
+            Bridge.getLog().error(LayoutLog.TAG_BROKEN, "Failed executing Choreographer#doFrame "
+                    , t, null, null);
+            return false;
         } finally {
             Bridge.cleanupThread();
         }
-        return false;
     }
 
     private static int toMotionEventType(TouchEventType eventType) {
diff --git a/create/src/com/android/tools/layoutlib/create/CreateInfo.java b/create/src/com/android/tools/layoutlib/create/CreateInfo.java
index 5c54d56..f68672c 100644
--- a/create/src/com/android/tools/layoutlib/create/CreateInfo.java
+++ b/create/src/com/android/tools/layoutlib/create/CreateInfo.java
@@ -227,6 +227,7 @@
         "android.graphics.fonts.SystemFonts#mmap",
         "android.os.Binder#getNativeBBinderHolder",
         "android.os.Binder#getNativeFinalizer",
+        "android.os.Handler#sendMessageAtFrontOfQueue",
         "android.os.Handler#sendMessageAtTime",
         "android.os.HandlerThread#run",
         "android.preference.Preference#getView",