Improve systemui unlockscreen jank issue

- In VisualStabilityCoordinator, ensure sure setting
  mPipelineRunAllowedis false when the notification is not visible
  after unlock to avoid unnecessary building notification list that
  blocks the UI thread
- Send notifyRendererOfExpensiveFrame to raise CPU work load when
  drawing notification shade window
- Refining WindowRootView#applyMargins performance to layout at once
  after all children's margins has updated

Flag: com.android.systemui.check_lockscreen_gone_transition
Bug: 358301118
Test: atest VisualStabilityCoordinatorTest
Test: run PTS cuj test for testing sysui unlock screen by local command
     ./sysui_unlockscreen.sh -i 5 -w
Change-Id: I61f1dad2dc9485f043c53a8900f1fae2e8ea1187
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index b57629f..391a272 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -1446,3 +1446,13 @@
        purpose: PURPOSE_BUGFIX
    }
 }
+
+flag {
+   name: "check_lockscreen_gone_transition"
+   namespace: "systemui"
+   description: "Run notification pipeline when the lockscreen is not in gone transition for avoiding janky frames during unlocking animation"
+   bug: "358301118"
+   metadata {
+       purpose: PURPOSE_BUGFIX
+   }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
index beba162..ea5c29e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinatorTest.java
@@ -20,8 +20,6 @@
 
 import static junit.framework.Assert.assertFalse;
 
-import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
-
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -31,12 +29,16 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
+
+import android.platform.test.annotations.EnableFlags;
 import android.testing.TestableLooper;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
 import com.android.compose.animation.scene.ObservableTransitionState;
+import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.communal.shared.model.CommunalScenes;
 import com.android.systemui.dump.DumpManager;
@@ -67,9 +69,6 @@
 import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.systemui.util.time.FakeSystemClock;
 
-import kotlinx.coroutines.flow.MutableStateFlow;
-import kotlinx.coroutines.test.TestScope;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -79,6 +78,9 @@
 import org.mockito.MockitoAnnotations;
 import org.mockito.verification.VerificationMode;
 
+import kotlinx.coroutines.flow.MutableStateFlow;
+import kotlinx.coroutines.test.TestScope;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 @TestableLooper.RunWithLooper
@@ -517,6 +519,32 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_CHECK_LOCKSCREEN_GONE_TRANSITION)
+    public void testNotLockscreenInGoneTransition_invalidationCalled() {
+        // GIVEN visual stability is being maintained b/c animation is playing
+        mKosmos.getKeyguardTransitionRepository().sendTransitionStepJava(
+                mTestScope, new TransitionStep(
+                        KeyguardState.LOCKSCREEN,
+                        KeyguardState.GONE,
+                        1f,
+                        TransitionState.RUNNING),  /* validateStep = */ false);
+        mTestScope.getTestScheduler().runCurrent();
+        assertFalse(mNotifStabilityManager.isPipelineRunAllowed());
+
+        // WHEN the animation has stopped playing
+        mKosmos.getKeyguardTransitionRepository().sendTransitionStepJava(
+                mTestScope, new TransitionStep(
+                        KeyguardState.LOCKSCREEN,
+                        KeyguardState.GONE,
+                        1f,
+                        TransitionState.FINISHED),  /* validateStep = */ false);
+        mTestScope.getTestScheduler().runCurrent();
+
+        // invalidate is called, b/c we were previously suppressing the pipeline from running
+        verifyStabilityManagerWasInvalidated(times(1));
+    }
+
+    @Test
     public void testNeverSuppressPipelineRunFromPanelCollapse_noInvalidationCalled() {
         // GIVEN animation is playing
         setPanelCollapsing(true);
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/view/WindowRootView.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/view/WindowRootView.kt
index ea19020..f0f476e 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/view/WindowRootView.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/view/WindowRootView.kt
@@ -25,6 +25,7 @@
 import android.view.WindowInsets
 import android.widget.FrameLayout
 import androidx.core.view.updateMargins
+import com.android.systemui.Flags
 import com.android.systemui.compose.ComposeInitializer
 import com.android.systemui.res.R
 
@@ -103,6 +104,8 @@
 
     private fun applyMargins() {
         val count = childCount
+        val hasFlagsEnabled = Flags.checkLockscreenGoneTransition()
+        var hasChildMarginUpdated = false
         for (i in 0 until count) {
             val child = getChildAt(i)
             if (child.layoutParams is LayoutParams) {
@@ -113,10 +116,17 @@
                             layoutParams.leftMargin != leftInset)
                 ) {
                     layoutParams.updateMargins(left = leftInset, right = rightInset)
-                    child.requestLayout()
+                    hasChildMarginUpdated = true
+                    if (!hasFlagsEnabled) {
+                        child.requestLayout()
+                    }
                 }
             }
         }
+        if (hasFlagsEnabled && hasChildMarginUpdated) {
+            // Request layout at once after all children's margins has updated
+            requestLayout()
+        }
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java
index 7244f8a..7a19838 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java
@@ -58,6 +58,7 @@
 import com.android.internal.statusbar.NotificationVisibility;
 import com.android.internal.widget.LockPatternUtils;
 import com.android.systemui.Dumpable;
+import com.android.systemui.Flags;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Background;
@@ -802,11 +803,17 @@
 
     private void notifyNotificationStateChanged() {
         if (!Looper.getMainLooper().isCurrentThread()) {
-            mMainExecutor.execute(() -> {
+            if (Flags.checkLockscreenGoneTransition()) {
                 for (NotificationStateChangedListener listener : mNotifStateChangedListeners) {
-                    listener.onNotificationStateChanged();
+                    mMainExecutor.execute(listener::onNotificationStateChanged);
                 }
-            });
+            } else {
+                mMainExecutor.execute(() -> {
+                    for (NotificationStateChangedListener listener : mNotifStateChangedListeners) {
+                        listener.onNotificationStateChanged();
+                    }
+                });
+            }
         } else {
             for (NotificationStateChangedListener listener : mNotifStateChangedListeners) {
                 listener.onNotificationStateChanged();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
index 41419f3..8660cd1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/VisualStabilityCoordinator.java
@@ -22,15 +22,18 @@
 import androidx.annotation.VisibleForTesting;
 
 import com.android.systemui.Dumpable;
+import com.android.systemui.Flags;
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
+import com.android.systemui.keyguard.shared.model.Edge;
 import com.android.systemui.keyguard.shared.model.KeyguardState;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.scene.shared.flag.SceneContainerFlag;
+import com.android.systemui.scene.shared.model.Scenes;
 import com.android.systemui.shade.domain.interactor.ShadeAnimationInteractor;
 import com.android.systemui.shade.domain.interactor.ShadeInteractor;
 import com.android.systemui.statusbar.notification.VisibilityLocationProvider;
@@ -85,6 +88,7 @@
     private boolean mNotifPanelLaunchingActivity;
     private boolean mCommunalShowing = false;
     private boolean mLockscreenShowing = false;
+    private boolean mLockscreenInGoneTransition = false;
 
     private boolean mPipelineRunAllowed;
     private boolean mReorderingAllowed;
@@ -158,6 +162,13 @@
                             KeyguardState.LOCKSCREEN),
                     this::onLockscreenKeyguardStateTransitionValueChanged);
         }
+        if (Flags.checkLockscreenGoneTransition()) {
+            mJavaAdapter.alwaysCollectFlow(mKeyguardTransitionInteractor.isInTransition(
+                            Edge.create(KeyguardState.LOCKSCREEN, Scenes.Gone),
+                            Edge.create(KeyguardState.LOCKSCREEN, KeyguardState.GONE)),
+                    this::onLockscreenInGoneTransitionChanged);
+        }
+
 
         pipeline.setVisualStabilityManager(mNotifStabilityManager);
     }
@@ -239,7 +250,9 @@
     private void updateAllowedStates(String field, boolean value) {
         boolean wasPipelineRunAllowed = mPipelineRunAllowed;
         boolean wasReorderingAllowed = mReorderingAllowed;
-        mPipelineRunAllowed = !isPanelCollapsingOrLaunchingActivity();
+        // No need to run notification pipeline when the lockscreen is in fading animation.
+        mPipelineRunAllowed = !(isPanelCollapsingOrLaunchingActivity()
+                || (Flags.checkLockscreenGoneTransition() && mLockscreenInGoneTransition));
         mReorderingAllowed = isReorderingAllowed();
         if (wasPipelineRunAllowed != mPipelineRunAllowed
                 || wasReorderingAllowed != mReorderingAllowed) {
@@ -330,7 +343,6 @@
                     updateAllowedStates("fullyDozed", fullyDozed);
                 }
             };
-
     final WakefulnessLifecycle.Observer mWakefulnessObserver = new WakefulnessLifecycle.Observer() {
         @Override
         public void onFinishedGoingToSleep() {
@@ -353,6 +365,9 @@
         pw.println("pipelineRunAllowed: " + mPipelineRunAllowed);
         pw.println("  notifPanelCollapsing: " + mNotifPanelCollapsing);
         pw.println("  launchingNotifActivity: " + mNotifPanelLaunchingActivity);
+        if (Flags.checkLockscreenGoneTransition()) {
+            pw.println("  lockscreenInGoneTransition: " + mLockscreenInGoneTransition);
+        }
         pw.println("reorderingAllowed: " + mReorderingAllowed);
         pw.println("  sleepy: " + mSleepy);
         pw.println("  fullyDozed: " + mFullyDozed);
@@ -401,4 +416,15 @@
         mLockscreenShowing = isShowing;
         updateAllowedStates("lockscreenShowing", isShowing);
     }
+
+    private void onLockscreenInGoneTransitionChanged(boolean inGoneTransition) {
+        if (!Flags.checkLockscreenGoneTransition()) {
+            return;
+        }
+        if (inGoneTransition == mLockscreenInGoneTransition) {
+            return;
+        }
+        mLockscreenInGoneTransition = inGoneTransition;
+        updateAllowedStates("lockscreenInGoneTransition", mLockscreenInGoneTransition);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
index 1ea26e5..2c5ba90 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
@@ -55,6 +55,7 @@
 import com.android.keyguard.KeyguardViewController;
 import com.android.keyguard.TrustGrantFlags;
 import com.android.keyguard.ViewMediatorCallback;
+import com.android.systemui.DejankUtils;
 import com.android.systemui.Flags;
 import com.android.systemui.biometrics.domain.interactor.UdfpsOverlayInteractor;
 import com.android.systemui.bouncer.domain.interactor.AlternateBouncerInteractor;
@@ -1202,6 +1203,11 @@
     @Override
     public void hide(long startTime, long fadeoutDuration) {
         Trace.beginSection("StatusBarKeyguardViewManager#hide");
+        if (Flags.checkLockscreenGoneTransition()) {
+            DejankUtils.notifyRendererOfExpensiveFrame(
+                    mNotificationShadeWindowController.getWindowRootView(),
+                    "StatusBarKeyguardViewManager#hide");
+        }
         mKeyguardStateController.notifyKeyguardState(false,
                 mKeyguardStateController.isOccluded());
         launchPendingWakeupAction();