Merge "Remove deadlock path between InteractionJankMonitor and FrameTracker" into sc-v2-dev am: 9de009c3a5

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/15828030

Change-Id: I516149ab133ca36bcc915bb9cbc65191342a8a01

core/java/com/android/internal/jank/FrameTracker.java

@@ -99,6 +99,7 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
     private final ViewRootImpl.SurfaceChangedCallback mSurfaceChangedCallback;
     private final Handler mHandler;
     private final ChoreographerWrapper mChoreographer;
+    private final Object mLock = InteractionJankMonitor.getInstance().getLock();
 
     @VisibleForTesting
     public final boolean mSurfaceOnly;
@@ -181,7 +182,7 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
             mSurfaceChangedCallback = new ViewRootImpl.SurfaceChangedCallback() {
                 @Override
                 public void surfaceCreated(SurfaceControl.Transaction t) {
-                    synchronized (FrameTracker.this) {
+                    synchronized (mLock) {
                         if (mSurfaceControl == null) {
                             mSurfaceControl = mViewRoot.getSurfaceControl();
                             if (mBeginVsyncId != INVALID_ID) {
@@ -203,12 +204,12 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
                     // Wait a while to give the system a chance for the remaining
                     // frames to arrive, then force finish the session.
                     mHandler.postDelayed(() -> {
-                        synchronized (FrameTracker.this) {
+                        synchronized (mLock) {
                             if (DEBUG) {
                                 Log.d(TAG, "surfaceDestroyed: " + mSession.getName()
                                         + ", finalized=" + mMetricsFinalized
                                         + ", info=" + mJankInfos.size()
-                                        + ", vsync=" + mBeginVsyncId + "-" + mEndVsyncId);
+                                        + ", vsync=" + mBeginVsyncId);
                             }
                             if (!mMetricsFinalized) {
                                 end(REASON_END_SURFACE_DESTROYED);
@@ -227,20 +228,20 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
     /**
      * Begin a trace session of the CUJ.
      */
-    public synchronized void begin() {
+    public void begin() {
-        mBeginVsyncId = mChoreographer.getVsyncId() + 1;
+        synchronized (mLock) {
-        if (DEBUG) {
+            mBeginVsyncId = mChoreographer.getVsyncId() + 1;
-            Log.d(TAG, "begin: " + mSession.getName() + ", begin=" + mBeginVsyncId);
+            if (DEBUG) {
-        }
+                Log.d(TAG, "begin: " + mSession.getName() + ", begin=" + mBeginVsyncId);
-        if (mSurfaceControl != null) {
+            }
-            postTraceStartMarker();
+            if (mSurfaceControl != null) {
-            mSurfaceControlWrapper.addJankStatsListener(this, mSurfaceControl);
+                postTraceStartMarker();
-        }
+                mSurfaceControlWrapper.addJankStatsListener(this, mSurfaceControl);
-        if (!mSurfaceOnly) {
+            }
-            mRendererWrapper.addObserver(mObserver);
+            if (!mSurfaceOnly) {
-        }
+                mRendererWrapper.addObserver(mObserver);
-        if (mListener != null) {
+            }
-            mListener.onCujEvents(mSession, ACTION_SESSION_BEGIN);
+            notifyCujEvent(ACTION_SESSION_BEGIN);
         }
     }
 
@@ -250,7 +251,7 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
     @VisibleForTesting
     public void postTraceStartMarker() {
         mChoreographer.mChoreographer.postCallback(Choreographer.CALLBACK_INPUT, () -> {
-            synchronized (FrameTracker.this) {
+            synchronized (mLock) {
                 if (mCancelled || mEndVsyncId != INVALID_ID) {
                     return;
                 }
@@ -263,88 +264,98 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
     /**
      * End the trace session of the CUJ.
      */
-    public synchronized void end(@Reasons int reason) {
+    public boolean end(@Reasons int reason) {
-        if (mEndVsyncId != INVALID_ID) return;
+        synchronized (mLock) {
-        mEndVsyncId = mChoreographer.getVsyncId();
+            if (mCancelled || mEndVsyncId != INVALID_ID) return false;
+            mEndVsyncId = mChoreographer.getVsyncId();
+            // Cancel the session if:
+            // 1. The session begins and ends at the same vsync id.
+            // 2. The session never begun.
+            if (mBeginVsyncId == INVALID_ID) {
+                return cancel(REASON_CANCEL_NOT_BEGUN);
+            } else if (mEndVsyncId <= mBeginVsyncId) {
+                return cancel(REASON_CANCEL_SAME_VSYNC);
+            } else {
+                if (DEBUG) {
+                    Log.d(TAG, "end: " + mSession.getName()
+                            + ", end=" + mEndVsyncId + ", reason=" + reason);
+                }
+                Trace.endAsyncSection(mSession.getName(), (int) mBeginVsyncId);
+                mSession.setReason(reason);
 
-        // Cancel the session if:
+                // We don't remove observer here,
-        // 1. The session begins and ends at the same vsync id.
+                // will remove it when all the frame metrics in this duration are called back.
-        // 2. The session never begun.
+                // See onFrameMetricsAvailable for the logic of removing the observer.
-        if (mBeginVsyncId == INVALID_ID) {
+                // Waiting at most 10 seconds for all callbacks to finish.
-            cancel(REASON_CANCEL_NOT_BEGUN);
+                mWaitForFinishTimedOut = () -> {
-        } else if (mEndVsyncId <= mBeginVsyncId) {
+                    Log.e(TAG, "force finish cuj because of time out:" + mSession.getName());
-            cancel(REASON_CANCEL_SAME_VSYNC);
+                    finish(mJankInfos.size() - 1);
-        } else {
+                };
-            if (DEBUG) {
+                mHandler.postDelayed(mWaitForFinishTimedOut, TimeUnit.SECONDS.toMillis(10));
-                Log.d(TAG, "end: " + mSession.getName()
+                notifyCujEvent(ACTION_SESSION_END);
-                        + ", end=" + mEndVsyncId + ", reason=" + reason);
+                return true;
             }
-            Trace.endAsyncSection(mSession.getName(), (int) mBeginVsyncId);
-            mSession.setReason(reason);
-            if (mListener != null) {
-                mListener.onCujEvents(mSession, ACTION_SESSION_END);
-            }
-
-            // We don't remove observer here,
-            // will remove it when all the frame metrics in this duration are called back.
-            // See onFrameMetricsAvailable for the logic of removing the observer.
-            // Waiting at most 10 seconds for all callbacks to finish.
-            mWaitForFinishTimedOut = () -> {
-                Log.e(TAG, "force finish cuj because of time out:" + mSession.getName());
-                finish(mJankInfos.size() - 1);
-            };
-            mHandler.postDelayed(mWaitForFinishTimedOut, TimeUnit.SECONDS.toMillis(10));
         }
     }
 
     /**
      * Cancel the trace session of the CUJ.
      */
-    public synchronized void cancel(@Reasons int reason) {
+    public boolean cancel(@Reasons int reason) {
-        mCancelled = true;
+        synchronized (mLock) {
+            final boolean cancelFromEnd =
+                    reason == REASON_CANCEL_NOT_BEGUN || reason == REASON_CANCEL_SAME_VSYNC;
+            if (mCancelled || (mEndVsyncId != INVALID_ID && !cancelFromEnd)) return false;
+            mCancelled = true;
+            // We don't need to end the trace section if it never begun.
+            if (mTracingStarted) {
+                Trace.endAsyncSection(mSession.getName(), (int) mBeginVsyncId);
+            }
 
-        // We don't need to end the trace section if it never begun.
+            // Always remove the observers in cancel call to avoid leakage.
-        if (mTracingStarted) {
+            removeObservers();
-            Trace.endAsyncSection(mSession.getName(), (int) mBeginVsyncId);
-        }
 
-        // Always remove the observers in cancel call to avoid leakage.
+            if (DEBUG) {
-        removeObservers();
+                Log.d(TAG, "cancel: " + mSession.getName() + ", begin=" + mBeginVsyncId
+                        + ", end=" + mEndVsyncId + ", reason=" + reason);
+            }
 
-        if (DEBUG) {
+            mSession.setReason(reason);
-            Log.d(TAG, "cancel: " + mSession.getName()
+            // Notify the listener the session has been cancelled.
-                    + ", begin=" + mBeginVsyncId + ", end=" + mEndVsyncId + ", reason=" + reason);
+            // We don't notify the listeners if the session never begun.
-        }
+            notifyCujEvent(ACTION_SESSION_CANCEL);
-
+            return true;
-        mSession.setReason(reason);
-        // Notify the listener the session has been cancelled.
-        // We don't notify the listeners if the session never begun.
-        if (mListener != null) {
-            mListener.onCujEvents(mSession, ACTION_SESSION_CANCEL);
         }
     }
 
-    @Override
+    private void notifyCujEvent(String action) {
-    public synchronized void onJankDataAvailable(SurfaceControl.JankData[] jankData) {
+        if (mListener == null) return;
-        if (mCancelled) {
+        mListener.onCujEvents(mSession, action);
-            return;
+    }
-        }
 
-        for (SurfaceControl.JankData jankStat : jankData) {
+    @Override
-            if (!isInRange(jankStat.frameVsyncId)) {
+    public void onJankDataAvailable(SurfaceControl.JankData[] jankData) {
-                continue;
+        synchronized (mLock) {
+            if (mCancelled) {
+                return;
             }
-            JankInfo info = findJankInfo(jankStat.frameVsyncId);
+
-            if (info != null) {
+            for (SurfaceControl.JankData jankStat : jankData) {
-                info.surfaceControlCallbackFired = true;
+                if (!isInRange(jankStat.frameVsyncId)) {
-                info.jankType = jankStat.jankType;
+                    continue;
-            } else {
+                }
-                mJankInfos.put((int) jankStat.frameVsyncId,
+                JankInfo info = findJankInfo(jankStat.frameVsyncId);
-                        JankInfo.createFromSurfaceControlCallback(
+                if (info != null) {
-                                jankStat.frameVsyncId, jankStat.jankType));
+                    info.surfaceControlCallbackFired = true;
+                    info.jankType = jankStat.jankType;
+                } else {
+                    mJankInfos.put((int) jankStat.frameVsyncId,
+                            JankInfo.createFromSurfaceControlCallback(
+                                    jankStat.frameVsyncId, jankStat.jankType));
+                }
             }
+            processJankInfos();
         }
-        processJankInfos();
     }
 
     private @Nullable JankInfo findJankInfo(long frameVsyncId) {
@@ -359,31 +370,34 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
     }
 
     @Override
-    public synchronized void onFrameMetricsAvailable(int dropCountSinceLastInvocation) {
+    public void onFrameMetricsAvailable(int dropCountSinceLastInvocation) {
-        if (mCancelled) {
+        synchronized (mLock) {
-            return;
+            if (mCancelled) {
-        }
+                return;
+            }
 
-        // Since this callback might come a little bit late after the end() call.
+            // Since this callback might come a little bit late after the end() call.
-        // We should keep tracking the begin / end timestamp.
+            // We should keep tracking the begin / end timestamp that we can compare with
-        // Then compare with vsync timestamp to check if the frame is in the duration of the CUJ.
+            // vsync timestamp to check if the frame is in the duration of the CUJ.
-        long totalDurationNanos = mMetricsWrapper.getMetric(FrameMetrics.TOTAL_DURATION);
+            long totalDurationNanos = mMetricsWrapper.getMetric(FrameMetrics.TOTAL_DURATION);
-        boolean isFirstFrame = mMetricsWrapper.getMetric(FrameMetrics.FIRST_DRAW_FRAME) == 1;
+            boolean isFirstFrame = mMetricsWrapper.getMetric(FrameMetrics.FIRST_DRAW_FRAME) == 1;
-        long frameVsyncId = mMetricsWrapper.getTiming()[FrameMetrics.Index.FRAME_TIMELINE_VSYNC_ID];
+            long frameVsyncId =
+                    mMetricsWrapper.getTiming()[FrameMetrics.Index.FRAME_TIMELINE_VSYNC_ID];
 
-        if (!isInRange(frameVsyncId)) {
+            if (!isInRange(frameVsyncId)) {
-            return;
+                return;
+            }
+            JankInfo info = findJankInfo(frameVsyncId);
+            if (info != null) {
+                info.hwuiCallbackFired = true;
+                info.totalDurationNanos = totalDurationNanos;
+                info.isFirstFrame = isFirstFrame;
+            } else {
+                mJankInfos.put((int) frameVsyncId, JankInfo.createFromHwuiCallback(
+                        frameVsyncId, totalDurationNanos, isFirstFrame));
+            }
+            processJankInfos();
         }
-        JankInfo info = findJankInfo(frameVsyncId);
-        if (info != null) {
-            info.hwuiCallbackFired = true;
-            info.totalDurationNanos = totalDurationNanos;
-            info.isFirstFrame = isFirstFrame;
-        } else {
-            mJankInfos.put((int) frameVsyncId, JankInfo.createFromHwuiCallback(
-                    frameVsyncId, totalDurationNanos, isFirstFrame));
-        }
-        processJankInfos();
     }
 
     /**
@@ -497,11 +511,7 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
                 (int) (maxFrameTimeNanos / NANOS_IN_MILLISECOND));
 
         // Trigger perfetto if necessary.
-        boolean overMissedFramesThreshold = mTraceThresholdMissedFrames != -1
+        if (shouldTriggerPerfetto(missedFramesCount, (int) maxFrameTimeNanos)) {
-                && missedFramesCount >= mTraceThresholdMissedFrames;
-        boolean overFrameTimeThreshold = !mSurfaceOnly && mTraceThresholdFrameTimeMillis != -1
-                && maxFrameTimeNanos >= mTraceThresholdFrameTimeMillis * NANOS_IN_MILLISECOND;
-        if (overMissedFramesThreshold || overFrameTimeThreshold) {
             triggerPerfetto();
         }
         if (mSession.logToStatsd()) {
@@ -513,9 +523,7 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
                     maxFrameTimeNanos, /* will be 0 if mSurfaceOnly == true */
                     missedSfFramesCount,
                     missedAppFramesCount);
-            if (mListener != null) {
+            notifyCujEvent(ACTION_METRICS_LOGGED);
-                mListener.onCujEvents(mSession, ACTION_METRICS_LOGGED);
-            }
         }
         if (DEBUG) {
             Log.i(TAG, "finish: CUJ=" + mSession.getName()
@@ -528,6 +536,14 @@ public class FrameTracker extends SurfaceControl.OnJankDataListener
         }
     }
 
+    private boolean shouldTriggerPerfetto(int missedFramesCount, int maxFrameTimeNanos) {
+        boolean overMissedFramesThreshold = mTraceThresholdMissedFrames != -1
+                && missedFramesCount >= mTraceThresholdMissedFrames;
+        boolean overFrameTimeThreshold = !mSurfaceOnly && mTraceThresholdFrameTimeMillis != -1
+                && maxFrameTimeNanos >= mTraceThresholdFrameTimeMillis * NANOS_IN_MILLISECOND;
+        return overMissedFramesThreshold || overFrameTimeThreshold;
+    }
+
     /**
      * Remove all the registered listeners, observers and callbacks.
      */

core/java/com/android/internal/jank/InteractionJankMonitor.java

@@ -230,6 +230,7 @@ public class InteractionJankMonitor {
     private final SparseArray<FrameTracker> mRunningTrackers;
     private final SparseArray<Runnable> mTimeoutActions;
     private final HandlerThread mWorker;
+    private final Object mLock = new Object();
 
     private boolean mEnabled = DEFAULT_ENABLED;
     private int mSamplingInterval = DEFAULT_SAMPLING_INTERVAL;
@@ -325,6 +326,10 @@ public class InteractionJankMonitor {
                 mPropertiesChangedListener);
     }
 
+    Object getLock() {
+        return mLock;
+    }
+
     /**
      * Creates a {@link FrameTracker} instance.
      *
@@ -344,7 +349,7 @@ public class InteractionJankMonitor {
         final ChoreographerWrapper choreographer =
                 new ChoreographerWrapper(Choreographer.getInstance());
 
-        synchronized (this) {
+        synchronized (mLock) {
             FrameTrackerListener eventsListener =
                     (s, act) -> handleCujEvents(config.getContext(), act, s);
             return new FrameTracker(session, mWorker.getThreadHandler(),
@@ -372,11 +377,16 @@ public class InteractionJankMonitor {
         final boolean badEnd = action.equals(ACTION_SESSION_END)
                 && session.getReason() != REASON_END_NORMAL;
         final boolean badCancel = action.equals(ACTION_SESSION_CANCEL)
-                && session.getReason() != REASON_CANCEL_NORMAL;
+                && !(session.getReason() == REASON_CANCEL_NORMAL
+                || session.getReason() == REASON_CANCEL_TIMEOUT);
         return badEnd || badCancel;
     }
 
-    private void notifyEvents(Context context, String action, Session session) {
+    /**
+     * Notifies who may interest in some CUJ events.
+     */
+    @VisibleForTesting
+    public void notifyEvents(Context context, String action, Session session) {
         if (action.equals(ACTION_SESSION_CANCEL)
                 && session.getReason() == REASON_CANCEL_NOT_BEGUN) {
             return;
@@ -389,7 +399,7 @@ public class InteractionJankMonitor {
     }
 
     private void removeTimeout(@CujType int cujType) {
-        synchronized (this) {
+        synchronized (mLock) {
             Runnable timeout = mTimeoutActions.get(cujType);
             if (timeout != null) {
                 mWorker.getThreadHandler().removeCallbacks(timeout);
@@ -432,17 +442,9 @@ public class InteractionJankMonitor {
     }
 
     private boolean beginInternal(@NonNull Configuration conf) {
-        synchronized (this) {
+        synchronized (mLock) {
             int cujType = conf.mCujType;
-            boolean shouldSample = ThreadLocalRandom.current().nextInt() % mSamplingInterval == 0;
+            if (!shouldMonitor(cujType)) return false;
-            if (!mEnabled || !shouldSample) {
-                if (DEBUG) {
-                    Log.d(TAG, "Skip monitoring cuj: " + getNameOfCuj(cujType)
-                            + ", enable=" + mEnabled + ", debuggable=" + DEFAULT_ENABLED
-                            + ", sample=" + shouldSample + ", interval=" + mSamplingInterval);
-                }
-                return false;
-            }
             FrameTracker tracker = getTracker(cujType);
             // Skip subsequent calls if we already have an ongoing tracing.
             if (tracker != null) return false;
@@ -459,6 +461,24 @@ public class InteractionJankMonitor {
         }
     }
 
+    /**
+     * Check if the monitoring is enabled and if it should be sampled.
+     */
+    @SuppressWarnings("RandomModInteger")
+    @VisibleForTesting
+    public boolean shouldMonitor(@CujType int cujType) {
+        boolean shouldSample = ThreadLocalRandom.current().nextInt() % mSamplingInterval == 0;
+        if (!mEnabled || !shouldSample) {
+            if (DEBUG) {
+                Log.d(TAG, "Skip monitoring cuj: " + getNameOfCuj(cujType)
+                        + ", enable=" + mEnabled + ", debuggable=" + DEFAULT_ENABLED
+                        + ", sample=" + shouldSample + ", interval=" + mSamplingInterval);
+            }
+            return false;
+        }
+        return true;
+    }
+
     /**
      * Schedules a timeout action.
      * @param cuj cuj type
@@ -478,14 +498,16 @@ public class InteractionJankMonitor {
      * @return boolean true if the tracker is ended successfully, false otherwise.
      */
     public boolean end(@CujType int cujType) {
-        synchronized (this) {
+        synchronized (mLock) {
             // remove the timeout action first.
             removeTimeout(cujType);
             FrameTracker tracker = getTracker(cujType);
             // Skip this call since we haven't started a trace yet.
             if (tracker == null) return false;
-            tracker.end(REASON_END_NORMAL);
+            // if the end call doesn't return true, another thread is handling end of the cuj.
-            removeTracker(cujType);
+            if (tracker.end(REASON_END_NORMAL)) {
+                removeTracker(cujType);
+            }
             return true;
         }
     }
@@ -499,33 +521,37 @@ public class InteractionJankMonitor {
         return cancel(cujType, REASON_CANCEL_NORMAL);
     }
 
-    boolean cancel(@CujType int cujType, @Reasons int reason) {
+    /**
-        synchronized (this) {
+     * Cancels the trace session.
+     *
+     * @return boolean true if the tracker is cancelled successfully, false otherwise.
+     */
+    @VisibleForTesting
+    public boolean cancel(@CujType int cujType, @Reasons int reason) {
+        synchronized (mLock) {
             // remove the timeout action first.
             removeTimeout(cujType);
             FrameTracker tracker = getTracker(cujType);
             // Skip this call since we haven't started a trace yet.
             if (tracker == null) return false;
-            tracker.cancel(reason);
+            // if the cancel call doesn't return true, another thread is handling cancel of the cuj.
-            removeTracker(cujType);
+            if (tracker.cancel(reason)) {
+                removeTracker(cujType);
+            }
             return true;
         }
     }
 
     private FrameTracker getTracker(@CujType int cuj) {
-        synchronized (this) {
+        return mRunningTrackers.get(cuj);
-            return mRunningTrackers.get(cuj);
-        }
     }
 
     private void removeTracker(@CujType int cuj) {
-        synchronized (this) {
+        mRunningTrackers.remove(cuj);
-            mRunningTrackers.remove(cuj);
-        }
     }
 
     private void updateProperties(DeviceConfig.Properties properties) {
-        synchronized (this) {
+        synchronized (mLock) {
             mSamplingInterval = properties.getInt(SETTINGS_SAMPLING_INTERVAL_KEY,
                     DEFAULT_SAMPLING_INTERVAL);
             mEnabled = properties.getBoolean(SETTINGS_ENABLED_KEY, DEFAULT_ENABLED);
@@ -547,10 +573,8 @@ public class InteractionJankMonitor {
      */
     @VisibleForTesting
     public void trigger(Session session) {
-        synchronized (this) {
+        mWorker.getThreadHandler().post(
-            mWorker.getThreadHandler().post(
+                () -> PerfettoTrigger.trigger(session.getPerfettoTrigger()));
-                    () -> PerfettoTrigger.trigger(session.getPerfettoTrigger()));
-        }
     }
 
     /**

core/tests/coretests/src/com/android/internal/jank/InteractionJankMonitorTest.java

@@ -16,6 +16,8 @@
 
 package com.android.internal.jank;
 
+import static com.android.internal.jank.FrameTracker.REASON_CANCEL_TIMEOUT;
+import static com.android.internal.jank.FrameTracker.REASON_END_NORMAL;
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE;
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_TO_STATSD_INTERACTION_TYPE;
 
@@ -34,6 +36,7 @@ import static org.mockito.Mockito.when;
 
 import android.os.Handler;
 import android.os.HandlerThread;
+import android.os.SystemClock;
 import android.provider.DeviceConfig;
 import android.view.View;
 import android.view.ViewAttachTestActivity;
@@ -82,36 +85,23 @@ public class InteractionJankMonitorTest {
 
         Handler handler = spy(new Handler(mActivity.getMainLooper()));
         doReturn(true).when(handler).sendMessageAtTime(any(), anyLong());
-        mWorker = spy(new HandlerThread("Interaction-jank-monitor-test"));
+        mWorker = mock(HandlerThread.class);
-        doNothing().when(mWorker).start();
         doReturn(handler).when(mWorker).getThreadHandler();
     }
 
     @Test
     public void testBeginEnd() {
-        // Should return false if the view is not attached.
+        InteractionJankMonitor monitor = createMockedInteractionJankMonitor();
-        InteractionJankMonitor monitor = spy(new InteractionJankMonitor(mWorker));
+        FrameTracker tracker = createMockedFrameTracker(null);
-        verify(mWorker).start();
-
-        Session session = new Session(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX);
-        Configuration config = mock(Configuration.class);
-        when(config.isSurfaceOnly()).thenReturn(false);
-        FrameTracker tracker = spy(new FrameTracker(session, mWorker.getThreadHandler(),
-                new ThreadedRendererWrapper(mView.getThreadedRenderer()),
-                new ViewRootWrapper(mView.getViewRootImpl()),
-                new SurfaceControlWrapper(), mock(ChoreographerWrapper.class),
-                new FrameMetricsWrapper(),
-                /* traceThresholdMissedFrames= */ 1, /* traceThresholdFrameTimeMillis= */ -1,
-                /* FrameTrackerListener */ null, config));
         doReturn(tracker).when(monitor).createFrameTracker(any(), any());
-        doNothing().when(tracker).triggerPerfetto();
+        doNothing().when(tracker).begin();
-        doNothing().when(tracker).postTraceStartMarker();
+        doReturn(true).when(tracker).end(anyInt());
 
         // Simulate a trace session and see if begin / end are invoked.
-        assertThat(monitor.begin(mView, session.getCuj())).isTrue();
+        assertThat(monitor.begin(mView, CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
         verify(tracker).begin();
-        assertThat(monitor.end(session.getCuj())).isTrue();
+        assertThat(monitor.end(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
-        verify(tracker).end(FrameTracker.REASON_END_NORMAL);
+        verify(tracker).end(REASON_END_NORMAL);
     }
 
     @Test
@@ -140,33 +130,23 @@ public class InteractionJankMonitorTest {
     }
 
     @Test
-    public void testBeginCancel() {
+    public void testBeginTimeout() {
-        InteractionJankMonitor monitor = spy(new InteractionJankMonitor(mWorker));
-
         ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-
+        InteractionJankMonitor monitor = createMockedInteractionJankMonitor();
-        Session session = new Session(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX);
+        FrameTracker tracker = createMockedFrameTracker(null);
-        Configuration config = mock(Configuration.class);
-        when(config.isSurfaceOnly()).thenReturn(false);
-        FrameTracker tracker = spy(new FrameTracker(session, mWorker.getThreadHandler(),
-                new ThreadedRendererWrapper(mView.getThreadedRenderer()),
-                new ViewRootWrapper(mView.getViewRootImpl()),
-                new SurfaceControlWrapper(), mock(FrameTracker.ChoreographerWrapper.class),
-                new FrameMetricsWrapper(),
-                /* traceThresholdMissedFrames= */ 1, /* traceThresholdFrameTimeMillis= */ -1,
-                /* FrameTrackerListener */ null, config));
         doReturn(tracker).when(monitor).createFrameTracker(any(), any());
-        doNothing().when(tracker).triggerPerfetto();
+        doNothing().when(tracker).begin();
-        doNothing().when(tracker).postTraceStartMarker();
+        doReturn(true).when(tracker).cancel(anyInt());
 
-        assertThat(monitor.begin(mView, session.getCuj())).isTrue();
+        assertThat(monitor.begin(mView, CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE)).isTrue();
         verify(tracker).begin();
         verify(monitor).scheduleTimeoutAction(anyInt(), anyLong(), captor.capture());
         Runnable runnable = captor.getValue();
         assertThat(runnable).isNotNull();
         mWorker.getThreadHandler().removeCallbacks(runnable);
         runnable.run();
-        verify(tracker).cancel(FrameTracker.REASON_CANCEL_TIMEOUT);
+        verify(monitor).cancel(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, REASON_CANCEL_TIMEOUT);
+        verify(tracker).cancel(REASON_CANCEL_TIMEOUT);
     }
 
     @Test
@@ -192,4 +172,43 @@ public class InteractionJankMonitorTest {
                     .isTrue();
         }
     }
+
+    private InteractionJankMonitor createMockedInteractionJankMonitor() {
+        InteractionJankMonitor monitor = spy(new InteractionJankMonitor(mWorker));
+        doReturn(true).when(monitor).shouldMonitor(anyInt());
+        doNothing().when(monitor).notifyEvents(any(), any(), any());
+        return monitor;
+    }
+
+    private FrameTracker createMockedFrameTracker(FrameTracker.FrameTrackerListener listener) {
+        Session session = spy(new Session(CUJ_NOTIFICATION_SHADE_EXPAND_COLLAPSE, CUJ_POSTFIX));
+        doReturn(false).when(session).logToStatsd();
+
+        ThreadedRendererWrapper threadedRenderer = mock(ThreadedRendererWrapper.class);
+        doNothing().when(threadedRenderer).addObserver(any());
+        doNothing().when(threadedRenderer).removeObserver(any());
+
+        ViewRootWrapper viewRoot = spy(new ViewRootWrapper(mView.getViewRootImpl()));
+        doNothing().when(viewRoot).addSurfaceChangedCallback(any());
+
+        SurfaceControlWrapper surfaceControl = mock(SurfaceControlWrapper.class);
+        doNothing().when(surfaceControl).addJankStatsListener(any(), any());
+        doNothing().when(surfaceControl).removeJankStatsListener(any());
+
+        final ChoreographerWrapper choreographer = mock(ChoreographerWrapper.class);
+        doReturn(SystemClock.elapsedRealtime()).when(choreographer).getVsyncId();
+
+        Configuration configuration = mock(Configuration.class);
+        when(configuration.isSurfaceOnly()).thenReturn(false);
+
+        FrameTracker tracker = spy(new FrameTracker(session, mWorker.getThreadHandler(),
+                threadedRenderer, viewRoot, surfaceControl, choreographer,
+                new FrameMetricsWrapper(), /* traceThresholdMissedFrames= */ 1,
+                /* traceThresholdFrameTimeMillis= */ -1, listener, configuration));
+
+        doNothing().when(tracker).postTraceStartMarker();
+        doNothing().when(tracker).triggerPerfetto();
+
+        return tracker;
+    }
 }