Use Choreographer from ViewRootImpl

This change adds a method to retrieve the choreographer instance that is
tied to the ViewRootImpl. It also updates JankTracker to use that
instance instead of using the global choreographer instance. This should
ensure the vsync ids we are using to keep track of active states line up
with the ids that are being returned by OnJankDataListener.

With this change the instantiation of StateTracker and JankDataProcessor
are now delayed until after the OnWindowAttachListener is invoked. This
was required as the DecorView still returns null when calling
getViewRootImpl while JankTracker is instantiated within the Activity.
StateTracker has a dependency on Choreographer so an active instance of
ViewRootImpl is required to obtain that.

Change-Id: I41728d4488bfb2df7d3c355ff8713cc673d13fc8
Bug: 377960907
Test: atest CoreAppJankTestCases
Flag: android.app.jank.viewroot_choreographer

5 files changed, 98 insertions, 10 deletions

diff --git a/core/java/android/app/Activity.java b/core/java/android/app/Activity.java
index b4f653354e07..d5df48a2ea22 100644
--- a/core/java/android/app/Activity.java
+++ b/core/java/android/app/Activity.java
@@ -10060,9 +10060,11 @@ public class Activity extends ContextThemeWrapper
                     }
                 });
                 if (mJankTracker == null) {
-                    // TODO b/377960907 use the Choreographer attached to the ViewRootImpl instead.
-                    mJankTracker = new JankTracker(Choreographer.getInstance(),
-                            decorView);
+                    if (android.app.jank.Flags.viewrootChoreographer()) {
+                        mJankTracker = new JankTracker(decorView);
+                    } else {
+                        mJankTracker = new JankTracker(Choreographer.getInstance(), decorView);
+                    }
                 }
                 // TODO b/377674765 confirm this is the string we want logged.
                 mJankTracker.setActivityName(getComponentName().getClassName());
diff --git a/core/java/android/app/jank/JankTracker.java b/core/java/android/app/jank/JankTracker.java
index 9c85b09f6be3..e3f67811757c 100644
--- a/core/java/android/app/jank/JankTracker.java
+++ b/core/java/android/app/jank/JankTracker.java
@@ -25,6 +25,7 @@ import android.view.AttachedSurfaceControl;
 import android.view.Choreographer;
 import android.view.SurfaceControl;
 import android.view.View;
+import android.view.ViewRootImpl;
 import android.view.ViewTreeObserver;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -100,7 +101,7 @@ public class JankTracker {
                         public void run() {
                             mDecorView.getViewTreeObserver()
                                     .removeOnWindowAttachListener(mOnWindowAttachListener);
-                            registerForJankData();
+                            initializeJankTrackingComponents();
                         }
                     }, REGISTRATION_DELAY_MS);
                 }
@@ -115,6 +116,7 @@ public class JankTracker {
                 }
             };
 
+    // TODO remove this once the viewroot_choreographer bugfix has been rolled out. b/399724640
     public JankTracker(Choreographer choreographer, View decorView) {
         mStateTracker = new StateTracker(choreographer);
         mJankDataProcessor = new JankDataProcessor(mStateTracker);
@@ -124,6 +126,19 @@ public class JankTracker {
     }
 
     /**
+     * Using this constructor delays the instantiation of the StateTracker and JankDataProcessor
+     * until after the OnWindowAttachListener is fired and the instance of Choreographer attached to
+     * the ViewRootImpl can be passed to StateTracker. This should ensures the vsync ids we are
+     * using to keep track of active states line up with the ids that are being returned by
+     * OnJankDataListener.
+     */
+    public JankTracker(View decorView) {
+        mDecorView = decorView;
+        mHandlerThread.start();
+        registerWindowListeners();
+    }
+
+    /**
      * Merges app jank stats reported by components outside the platform to the current pending
      * stats
      */
@@ -131,6 +146,9 @@ public class JankTracker {
         getHandler().post(new Runnable() {
             @Override
             public void run() {
+                if (mJankDataProcessor == null) {
+                    return;
+                }
                 mJankDataProcessor.mergeJankStats(appJankStats, mActivityName);
             }
         });
@@ -192,8 +210,7 @@ public class JankTracker {
     public void enableAppJankTracking() {
         // Add the activity as a state, this will ensure we track frames to the activity without the
         // need for a decorated widget to be used.
-        // TODO b/376116199 replace "NONE" with UNSPECIFIED once the API changes are merged.
-        mStateTracker.putState("NONE", mActivityName, "NONE");
+        addActivityToStateTracking();
         mTrackingEnabled = true;
         registerForJankData();
     }
@@ -203,27 +220,33 @@ public class JankTracker {
      */
     public void disableAppJankTracking() {
         mTrackingEnabled = false;
-        // TODO b/376116199 replace "NONE" with UNSPECIFIED once the API changes are merged.
-        mStateTracker.removeState("NONE", mActivityName, "NONE");
+        removeActivityFromStateTracking();
         unregisterForJankData();
     }
 
     /**
-     * Retrieve all pending widget states, this is intended for testing purposes only.
+     * Retrieve all pending widget states, this is intended for testing purposes only. If
+     * this is called before StateTracker has been created the method will just return without
+     * copying any data to the stateDataList parameter.
      *
      * @param stateDataList the ArrayList that will be populated with the pending states.
      */
     @VisibleForTesting
     public void getAllUiStates(@NonNull ArrayList<StateTracker.StateData> stateDataList) {
+        if (mStateTracker == null) return;
         mStateTracker.retrieveAllStates(stateDataList);
     }
 
     /**
      * Retrieve all pending jank stats before they are logged, this is intended for testing
-     * purposes only.
+     * purposes only. If this method is called before JankDataProcessor is created it will return
+     * an empty HashMap.
      */
     @VisibleForTesting
     public HashMap<String, JankDataProcessor.PendingJankStat> getPendingJankStats() {
+        if (mJankDataProcessor == null) {
+            return new HashMap<>();
+        }
         return mJankDataProcessor.getPendingJankStats();
     }
 
@@ -233,8 +256,10 @@ public class JankTracker {
      */
     @VisibleForTesting
     public void forceListenerRegistration() {
+        addActivityToStateTracking();
         mSurfaceControl = mDecorView.getRootSurfaceControl();
         registerJankDataListener();
+        mListenersRegistered = true;
     }
 
     private void unregisterForJankData() {
@@ -270,6 +295,10 @@ public class JankTracker {
      */
     @VisibleForTesting
     public boolean shouldTrack() {
+        if (DEBUG) {
+            Log.d(DEBUG_KEY, String.format("mTrackingEnabled: %s | mListenersRegistered: %s",
+                    mTrackingEnabled, mListenersRegistered));
+        }
         return mTrackingEnabled && mListenersRegistered;
     }
 
@@ -313,4 +342,36 @@ public class JankTracker {
         }
         return mHandler;
     }
+
+    private void addActivityToStateTracking() {
+        if (mStateTracker == null) return;
+
+        mStateTracker.putState(AppJankStats.WIDGET_CATEGORY_UNSPECIFIED, mActivityName,
+                AppJankStats.WIDGET_STATE_UNSPECIFIED);
+    }
+
+    private void removeActivityFromStateTracking() {
+        if (mStateTracker == null) return;
+
+        mStateTracker.removeState(AppJankStats.WIDGET_CATEGORY_UNSPECIFIED, mActivityName,
+                AppJankStats.WIDGET_STATE_UNSPECIFIED);
+    }
+
+    private void initializeJankTrackingComponents() {
+        ViewRootImpl viewRoot = mDecorView.getViewRootImpl();
+        if (viewRoot == null || viewRoot.getChoreographer() == null) {
+            return;
+        }
+
+        if (mStateTracker == null) {
+            mStateTracker = new StateTracker(viewRoot.getChoreographer());
+        }
+
+        if (mJankDataProcessor == null) {
+            mJankDataProcessor = new JankDataProcessor(mStateTracker);
+        }
+
+        addActivityToStateTracking();
+        registerForJankData();
+    }
 }
diff --git a/core/java/android/app/jank/flags.aconfig b/core/java/android/app/jank/flags.aconfig
index a62df1b3cbf7..de98b88d2aca 100644
--- a/core/java/android/app/jank/flags.aconfig
+++ b/core/java/android/app/jank/flags.aconfig
@@ -14,4 +14,14 @@ flag {
   namespace: "system_performance"
   description: "Controls whether the system will log frame metrics related to app jank"
   bug: "366265225"
+}
+
+flag {
+  name: "viewroot_choreographer"
+  namespace: "system_performance"
+  description: "when enabled janktracker will get the instance of choreographer from viewrootimpl"
+  bug: "377960907"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
 }
\ No newline at end of file
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 9d0773f0a606..98e6cdde664a 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -13654,4 +13654,11 @@ public final class ViewRootImpl implements ViewParent,
             ThreadedRenderer.preInitBufferAllocator();
         }
     }
+
+    /**
+     * @hide
+     */
+    public Choreographer getChoreographer() {
+        return mChoreographer;
+    }
 }
diff --git a/tests/AppJankTest/src/android/app/jank/tests/IntegrationTests.java b/tests/AppJankTest/src/android/app/jank/tests/IntegrationTests.java
index 3498974b348e..229c7bfb53e9 100644
--- a/tests/AppJankTest/src/android/app/jank/tests/IntegrationTests.java
+++ b/tests/AppJankTest/src/android/app/jank/tests/IntegrationTests.java
@@ -103,6 +103,10 @@ public class IntegrationTests {
     @RequiresFlagsEnabled(Flags.FLAG_DETAILED_APP_JANK_METRICS_API)
     public void reportJankStats_confirmPendingStatsIncreases() {
         Activity jankTrackerActivity = mJankTrackerActivityRule.launchActivity(null);
+        mDevice.wait(Until.findObject(
+                By.text(jankTrackerActivity.getString(R.string.continue_test))),
+                WAIT_FOR_TIMEOUT_MS);
+
         EditText editText = jankTrackerActivity.findViewById(R.id.edit_text);
         JankTracker jankTracker = editText.getJankTracker();
 
@@ -135,6 +139,10 @@ public class IntegrationTests {
     public void simulateWidgetStateChanges_confirmStateChangesAreTracked() {
         JankTrackerActivity jankTrackerActivity =
                 mJankTrackerActivityRule.launchActivity(null);
+        mDevice.wait(Until.findObject(
+                        By.text(jankTrackerActivity.getString(R.string.continue_test))),
+                WAIT_FOR_TIMEOUT_MS);
+
         TestWidget testWidget = jankTrackerActivity.findViewById(R.id.jank_tracker_widget);
         JankTracker jankTracker = testWidget.getJankTracker();
         jankTracker.forceListenerRegistration();