[PIP][jank] Add jank CUJs for expand transitions

As part of improving jank cuj coverage, add cujs for expanding pip
in PIP2. This also adds PipInteractionHandler which will handle all
other jank CUJ tracking.

Flag: EXEMPT jank
Bug: b/352738878
Test: atest WMShellUnitTests:PipInteractionHandlerTest
Test: Manually check events are being logged
Change-Id: Iad2f4874f199479cfccd35b561c2d223672bc5ea

6 files changed, 237 insertions, 5 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
index 2bd5c72eb34b..6f0919e1d045 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/pip/Pip2Module.java
@@ -20,6 +20,7 @@ import android.annotation.NonNull;
 import android.content.Context;
 import android.os.Handler;
 
+import com.android.internal.jank.InteractionJankMonitor;
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
@@ -45,6 +46,7 @@ import com.android.wm.shell.desktopmode.DesktopUserRepositories;
 import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler;
 import com.android.wm.shell.pip2.phone.PhonePipMenuController;
 import com.android.wm.shell.pip2.phone.PipController;
+import com.android.wm.shell.pip2.phone.PipInteractionHandler;
 import com.android.wm.shell.pip2.phone.PipMotionHelper;
 import com.android.wm.shell.pip2.phone.PipScheduler;
 import com.android.wm.shell.pip2.phone.PipTaskListener;
@@ -88,12 +90,13 @@ public abstract class Pip2Module {
             @NonNull PipUiStateChangeController pipUiStateChangeController,
             DisplayController displayController,
             Optional<SplitScreenController> splitScreenControllerOptional,
-            PipDesktopState pipDesktopState) {
+            PipDesktopState pipDesktopState,
+            PipInteractionHandler pipInteractionHandler) {
         return new PipTransition(context, shellInit, shellTaskOrganizer, transitions,
                 pipBoundsState, null, pipBoundsAlgorithm, pipTaskListener,
                 pipScheduler, pipStackListenerController, pipDisplayLayoutState,
                 pipUiStateChangeController, displayController, splitScreenControllerOptional,
-                pipDesktopState);
+                pipDesktopState, pipInteractionHandler);
     }
 
     @WMSingleton
@@ -249,4 +252,14 @@ public abstract class Pip2Module {
 
     @BindsOptionalOf
     abstract DragToDesktopTransitionHandler optionalDragToDesktopTransitionHandler();
+
+    @WMSingleton
+    @Provides
+    static PipInteractionHandler providePipInteractionHandler(
+            Context context,
+            @ShellMainThread Handler mainHandler
+    ) {
+        return new PipInteractionHandler(context, mainHandler,
+                InteractionJankMonitor.getInstance());
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipInteractionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipInteractionHandler.java
new file mode 100644
index 000000000000..321952480094
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipInteractionHandler.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.pip2.phone;
+
+import static com.android.internal.jank.Cuj.CUJ_PIP_TRANSITION;
+
+import android.annotation.IntDef;
+import android.content.Context;
+import android.os.Handler;
+import android.view.SurfaceControl;
+
+import com.android.internal.jank.InteractionJankMonitor;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/**
+ * Helps track PIP CUJ interactions
+ */
+public class PipInteractionHandler {
+    @IntDef(prefix = {"INTERACTION_"}, value = {
+            INTERACTION_EXIT_PIP,
+            INTERACTION_EXIT_PIP_TO_SPLIT
+    })
+
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface Interaction {}
+
+    public static final int INTERACTION_EXIT_PIP = 0;
+    public static final int INTERACTION_EXIT_PIP_TO_SPLIT = 1;
+
+    private final Context mContext;
+    private final Handler mHandler;
+    private final InteractionJankMonitor mInteractionJankMonitor;
+
+    public PipInteractionHandler(Context context, Handler handler,
+            InteractionJankMonitor interactionJankMonitor) {
+        mContext = context;
+        mHandler = handler;
+        mInteractionJankMonitor = interactionJankMonitor;
+    }
+
+    /**
+     * Begin tracking PIP CUJ.
+     *
+     * @param leash PIP leash.
+     * @param interaction Tag for interaction.
+     */
+    public void begin(SurfaceControl leash, @Interaction int interaction) {
+        mInteractionJankMonitor.begin(leash, mContext, mHandler, CUJ_PIP_TRANSITION,
+                pipInteractionToString(interaction));
+    }
+
+    /**
+     * End tracking CUJ.
+     */
+    public void end() {
+        mInteractionJankMonitor.end(CUJ_PIP_TRANSITION);
+    }
+
+    /**
+     * Converts an interaction to a string representation used for tagging.
+     *
+     * @param interaction Interaction to track.
+     * @return String representation of the interaction.
+     */
+    public static String pipInteractionToString(@Interaction int interaction) {
+        return switch (interaction) {
+            case INTERACTION_EXIT_PIP -> "EXIT_PIP";
+            case INTERACTION_EXIT_PIP_TO_SPLIT -> "EXIT_PIP_TO_SPLIT";
+            default -> "";
+        };
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
index b51a58e604bf..9bb2e38e1526 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
@@ -113,6 +113,7 @@ public class PipTransition extends PipTransitionController implements
     private final DisplayController mDisplayController;
     private final PipSurfaceTransactionHelper mPipSurfaceTransactionHelper;
     private final PipDesktopState mPipDesktopState;
+    private final PipInteractionHandler mPipInteractionHandler;
 
     //
     // Transition caches
@@ -154,7 +155,8 @@ public class PipTransition extends PipTransitionController implements
             PipUiStateChangeController pipUiStateChangeController,
             DisplayController displayController,
             Optional<SplitScreenController> splitScreenControllerOptional,
-            PipDesktopState pipDesktopState) {
+            PipDesktopState pipDesktopState,
+            PipInteractionHandler pipInteractionHandler) {
         super(shellInit, shellTaskOrganizer, transitions, pipBoundsState, pipMenuController,
                 pipBoundsAlgorithm);
 
@@ -168,9 +170,11 @@ public class PipTransition extends PipTransitionController implements
         mDisplayController = displayController;
         mPipSurfaceTransactionHelper = new PipSurfaceTransactionHelper(mContext);
         mPipDesktopState = pipDesktopState;
+        mPipInteractionHandler = pipInteractionHandler;
 
         mExpandHandler = new PipExpandHandler(mContext, pipBoundsState, pipBoundsAlgorithm,
-                pipTransitionState, pipDisplayLayoutState, splitScreenControllerOptional);
+                pipTransitionState, pipDisplayLayoutState, pipInteractionHandler,
+                splitScreenControllerOptional);
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandler.java
index db4942b2fb95..3274f4ae354a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandler.java
@@ -45,6 +45,7 @@ import com.android.wm.shell.common.pip.PipBoundsAlgorithm;
 import com.android.wm.shell.common.pip.PipBoundsState;
 import com.android.wm.shell.common.pip.PipDisplayLayoutState;
 import com.android.wm.shell.pip2.animation.PipExpandAnimator;
+import com.android.wm.shell.pip2.phone.PipInteractionHandler;
 import com.android.wm.shell.pip2.phone.PipTransitionState;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.splitscreen.SplitScreenController;
@@ -58,6 +59,7 @@ public class PipExpandHandler implements Transitions.TransitionHandler {
     private final PipBoundsAlgorithm mPipBoundsAlgorithm;
     private final PipTransitionState mPipTransitionState;
     private final PipDisplayLayoutState mPipDisplayLayoutState;
+    private final PipInteractionHandler mPipInteractionHandler;
     private final Optional<SplitScreenController> mSplitScreenControllerOptional;
 
     @Nullable
@@ -72,12 +74,14 @@ public class PipExpandHandler implements Transitions.TransitionHandler {
             PipBoundsAlgorithm pipBoundsAlgorithm,
             PipTransitionState pipTransitionState,
             PipDisplayLayoutState pipDisplayLayoutState,
+            PipInteractionHandler pipInteractionHandler,
             Optional<SplitScreenController> splitScreenControllerOptional) {
         mContext = context;
         mPipBoundsState = pipBoundsState;
         mPipBoundsAlgorithm = pipBoundsAlgorithm;
         mPipTransitionState = pipTransitionState;
         mPipDisplayLayoutState = pipDisplayLayoutState;
+        mPipInteractionHandler = pipInteractionHandler;
         mSplitScreenControllerOptional = splitScreenControllerOptional;
 
         mPipExpandAnimatorSupplier = PipExpandAnimator::new;
@@ -183,6 +187,8 @@ public class PipExpandHandler implements Transitions.TransitionHandler {
         PipExpandAnimator animator = mPipExpandAnimatorSupplier.get(mContext, pipLeash,
                 startTransaction, finishTransaction, endBounds, startBounds, endBounds,
                 sourceRectHint, delta);
+        animator.setAnimationStartCallback(() -> mPipInteractionHandler.begin(pipLeash,
+                PipInteractionHandler.INTERACTION_EXIT_PIP));
         animator.setAnimationEndCallback(() -> {
             if (parentBeforePip != null) {
                 // TODO b/377362511: Animate local leash instead to also handle letterbox case.
@@ -190,6 +196,7 @@ public class PipExpandHandler implements Transitions.TransitionHandler {
                 finishTransaction.setCrop(pipLeash, null);
             }
             finishTransition();
+            mPipInteractionHandler.end();
         });
         cacheAndStartTransitionAnimator(animator);
         saveReentryState();
@@ -248,6 +255,8 @@ public class PipExpandHandler implements Transitions.TransitionHandler {
             splitController.finishEnterSplitScreen(finishTransaction);
         });
 
+        animator.setAnimationStartCallback(() -> mPipInteractionHandler.begin(pipLeash,
+                PipInteractionHandler.INTERACTION_EXIT_PIP_TO_SPLIT));
         animator.setAnimationEndCallback(() -> {
             if (parentBeforePip == null) {
                 // After PipExpandAnimator is done modifying finishTransaction, we need to make
@@ -256,6 +265,7 @@ public class PipExpandHandler implements Transitions.TransitionHandler {
                 finishTransaction.setPosition(pipLeash, 0, 0);
             }
             finishTransition();
+            mPipInteractionHandler.end();
         });
         cacheAndStartTransitionAnimator(animator);
         saveReentryState();
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/PipInteractionHandlerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/PipInteractionHandlerTest.java
new file mode 100644
index 000000000000..9c0127ea2414
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/PipInteractionHandlerTest.java
@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.pip2.phone;
+
+import static com.android.internal.jank.Cuj.CUJ_PIP_TRANSITION;
+import static com.android.wm.shell.pip2.phone.PipInteractionHandler.INTERACTION_EXIT_PIP;
+import static com.android.wm.shell.pip2.phone.PipInteractionHandler.INTERACTION_EXIT_PIP_TO_SPLIT;
+
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.kotlin.VerificationKt.times;
+
+import android.content.Context;
+import android.os.Handler;
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
+import android.view.SurfaceControl;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.internal.jank.InteractionJankMonitor;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Unit test against {@link PipInteractionHandler}.
+ */
+@SmallTest
+@TestableLooper.RunWithLooper
+@RunWith(AndroidTestingRunner.class)
+public class PipInteractionHandlerTest {
+    @Mock private Context mMockContext;
+    @Mock private Handler mMockHandler;
+    @Mock private InteractionJankMonitor mMockInteractionJankMonitor;
+
+    private SurfaceControl mTestLeash;
+
+    private PipInteractionHandler mPipInteractionHandler;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mPipInteractionHandler = new PipInteractionHandler(mMockContext, mMockHandler,
+                mMockInteractionJankMonitor);
+        mTestLeash = new SurfaceControl.Builder()
+                .setContainerLayer()
+                .setName("PipInteractionHandlerTest")
+                .setCallsite("PipInteractionHandlerTest")
+                .build();
+    }
+
+    @Test
+    public void begin_expand_startsTracking() {
+        mPipInteractionHandler.begin(mTestLeash, INTERACTION_EXIT_PIP);
+
+        verify(mMockInteractionJankMonitor, times(1)).begin(eq(mTestLeash),
+                eq(mMockContext), eq(mMockHandler), eq(CUJ_PIP_TRANSITION),
+                eq(PipInteractionHandler.pipInteractionToString(INTERACTION_EXIT_PIP)));
+    }
+
+    @Test
+    public void begin_expandToSplit_startsTracking() {
+        mPipInteractionHandler.begin(mTestLeash, INTERACTION_EXIT_PIP_TO_SPLIT);
+
+        verify(mMockInteractionJankMonitor, times(1)).begin(eq(mTestLeash),
+                eq(mMockContext), eq(mMockHandler), eq(CUJ_PIP_TRANSITION),
+                eq(PipInteractionHandler.pipInteractionToString(INTERACTION_EXIT_PIP_TO_SPLIT)));
+    }
+
+    @Test
+    public void end_stopsTracking() {
+        mPipInteractionHandler.end();
+
+        verify(mMockInteractionJankMonitor, times(1)).end(CUJ_PIP_TRANSITION);
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandlerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandlerTest.java
index 2a22842eda1a..cc66f00525b5 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandlerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip2/phone/transition/PipExpandHandlerTest.java
@@ -23,6 +23,7 @@ import static com.android.wm.shell.transition.Transitions.TRANSIT_EXIT_PIP;
 import static com.android.wm.shell.transition.Transitions.TRANSIT_EXIT_PIP_TO_SPLIT;
 
 import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -48,11 +49,13 @@ import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
 import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.wm.shell.common.pip.PipBoundsAlgorithm;
 import com.android.wm.shell.common.pip.PipBoundsState;
 import com.android.wm.shell.common.pip.PipDisplayLayoutState;
 import com.android.wm.shell.pip2.animation.PipExpandAnimator;
+import com.android.wm.shell.pip2.phone.PipInteractionHandler;
 import com.android.wm.shell.pip2.phone.PipTransitionState;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.transition.TransitionInfoBuilder;
@@ -61,6 +64,8 @@ import com.android.wm.shell.util.StubTransaction;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -79,6 +84,7 @@ public class PipExpandHandlerTest {
     @Mock private PipBoundsAlgorithm mMockPipBoundsAlgorithm;
     @Mock private PipTransitionState mMockPipTransitionState;
     @Mock private PipDisplayLayoutState mMockPipDisplayLayoutState;
+    @Mock private PipInteractionHandler mMockPipInteractionHandler;
     @Mock private SplitScreenController mMockSplitScreenController;
 
     @Mock private IBinder mMockTransitionToken;
@@ -89,6 +95,8 @@ public class PipExpandHandlerTest {
 
     @Mock private PipExpandAnimator mMockPipExpandAnimator;
 
+    @Captor private ArgumentCaptor<Runnable> mAnimatorCallbackArgumentCaptor;
+
     @Surface.Rotation
     private static final int DISPLAY_ROTATION = Surface.ROTATION_0;
 
@@ -108,7 +116,7 @@ public class PipExpandHandlerTest {
 
         mPipExpandHandler = new PipExpandHandler(mMockContext, mMockPipBoundsState,
                 mMockPipBoundsAlgorithm, mMockPipTransitionState, mMockPipDisplayLayoutState,
-                Optional.of(mMockSplitScreenController));
+                mMockPipInteractionHandler, Optional.of(mMockSplitScreenController));
         mPipExpandHandler.setPipExpandAnimatorSupplier((context, leash, startTransaction,
                 finishTransaction, baseBounds, startBounds, endBounds,
                 sourceRectHint, rotation) -> mMockPipExpandAnimator);
@@ -138,6 +146,13 @@ public class PipExpandHandlerTest {
 
         verify(mMockPipExpandAnimator, times(1)).start();
         verify(mMockPipBoundsState, times(1)).saveReentryState(SNAP_FRACTION);
+
+        verify(mMockPipExpandAnimator, times(1))
+                .setAnimationStartCallback(mAnimatorCallbackArgumentCaptor.capture());
+        InstrumentationRegistry.getInstrumentation()
+                .runOnMainSync(mAnimatorCallbackArgumentCaptor.getValue());
+        verify(mMockPipInteractionHandler, times(1)).begin(any(),
+                eq(PipInteractionHandler.INTERACTION_EXIT_PIP));
     }
 
     @Test
@@ -158,6 +173,13 @@ public class PipExpandHandlerTest {
         verify(mMockSplitScreenController, times(1)).finishEnterSplitScreen(eq(mFinishT));
         verify(mMockPipExpandAnimator, times(1)).start();
         verify(mMockPipBoundsState, times(1)).saveReentryState(SNAP_FRACTION);
+
+        verify(mMockPipExpandAnimator, times(1))
+                .setAnimationStartCallback(mAnimatorCallbackArgumentCaptor.capture());
+        InstrumentationRegistry.getInstrumentation()
+                .runOnMainSync(mAnimatorCallbackArgumentCaptor.getValue());
+        verify(mMockPipInteractionHandler, times(1)).begin(any(),
+                eq(PipInteractionHandler.INTERACTION_EXIT_PIP_TO_SPLIT));
     }
 
     private TransitionInfo getExpandFromPipTransitionInfo(@WindowManager.TransitionType int type,