Merge "Move active unlock chipbar flag to a ResourceFlag" into tm-qpr-dev
diff --git a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
index 5d9f335..a6f47d4 100644
--- a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
+++ b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
@@ -109,6 +109,7 @@
 import android.util.SparseArray;
 import android.util.SparseBooleanArray;
 import android.util.SparseLongArray;
+import android.util.SparseSetArray;
 import android.util.TimeUtils;
 import android.view.Display;
 import android.widget.Toast;
@@ -452,6 +453,12 @@
     private final Map<String, String> mAppStandbyProperties = new ArrayMap<>();
 
     /**
+     * Set of apps that were restored via backup & restore, per user, that need their
+     * standby buckets to be adjusted when installed.
+     */
+    private final SparseSetArray<String> mAppsToRestoreToRare = new SparseSetArray<>();
+
+    /**
      * List of app-ids of system packages, populated on boot, when system services are ready.
      */
     private final ArrayList<Integer> mSystemPackagesAppIds = new ArrayList<>();
@@ -968,13 +975,17 @@
                                     + standbyBucketToString(newBucket));
                         }
                     } else {
-                        newBucket = getBucketForLocked(packageName, userId,
-                                elapsedRealtime);
-                        if (DEBUG) {
-                            Slog.d(TAG, "Evaluated AOSP newBucket = "
-                                    + standbyBucketToString(newBucket));
+                        // Don't update the standby state for apps that were restored
+                        if (!(oldMainReason == REASON_MAIN_DEFAULT
+                                && (app.bucketingReason & REASON_SUB_MASK)
+                                        == REASON_SUB_DEFAULT_APP_RESTORED)) {
+                            newBucket = getBucketForLocked(packageName, userId, elapsedRealtime);
+                            if (DEBUG) {
+                                Slog.d(TAG, "Evaluated AOSP newBucket = "
+                                        + standbyBucketToString(newBucket));
+                            }
+                            reason = REASON_MAIN_TIMEOUT;
                         }
-                        reason = REASON_MAIN_TIMEOUT;
                     }
                 }
 
@@ -1610,18 +1621,29 @@
         final int reason = REASON_MAIN_DEFAULT | REASON_SUB_DEFAULT_APP_RESTORED;
         final long nowElapsed = mInjector.elapsedRealtime();
         for (String packageName : restoredApps) {
-            // If the package is not installed, don't allow the bucket to be set.
+            // If the package is not installed, don't allow the bucket to be set. Instead, add it
+            // to a list of all packages whose buckets need to be adjusted when installed.
             if (!mInjector.isPackageInstalled(packageName, 0, userId)) {
-                Slog.e(TAG, "Tried to restore bucket for uninstalled app: " + packageName);
+                Slog.i(TAG, "Tried to restore bucket for uninstalled app: " + packageName);
+                mAppsToRestoreToRare.add(userId, packageName);
                 continue;
             }
 
-            final int standbyBucket = getAppStandbyBucket(packageName, userId, nowElapsed, false);
-            // Only update the standby bucket to RARE if the app is still in the NEVER bucket.
-            if (standbyBucket == STANDBY_BUCKET_NEVER) {
-                setAppStandbyBucket(packageName, userId, STANDBY_BUCKET_RARE, reason,
-                        nowElapsed, false);
-            }
+            restoreAppToRare(packageName, userId, nowElapsed, reason);
+        }
+        // Clear out the list of restored apps that need to have their standby buckets adjusted
+        // if they still haven't been installed eight hours after restore.
+        // Note: if the device reboots within these first 8 hours, this list will be lost since it's
+        // not persisted - this is the expected behavior for now and may be updated in the future.
+        mHandler.postDelayed(() -> mAppsToRestoreToRare.remove(userId), 8 * ONE_HOUR);
+    }
+
+    /** Adjust the standby bucket of the given package for the user to RARE. */
+    private void restoreAppToRare(String pkgName, int userId, long nowElapsed, int reason) {
+        final int standbyBucket = getAppStandbyBucket(pkgName, userId, nowElapsed, false);
+        // Only update the standby bucket to RARE if the app is still in the NEVER bucket.
+        if (standbyBucket == STANDBY_BUCKET_NEVER) {
+            setAppStandbyBucket(pkgName, userId, STANDBY_BUCKET_RARE, reason, nowElapsed, false);
         }
     }
 
@@ -2116,15 +2138,24 @@
                 }
                 // component-level enable/disable can affect bucketing, so we always
                 // reevaluate that for any PACKAGE_CHANGED
-                mHandler.obtainMessage(MSG_CHECK_PACKAGE_IDLE_STATE, userId, -1, pkgName)
-                    .sendToTarget();
+                if (Intent.ACTION_PACKAGE_CHANGED.equals(action)) {
+                    mHandler.obtainMessage(MSG_CHECK_PACKAGE_IDLE_STATE, userId, -1, pkgName)
+                            .sendToTarget();
+                }
             }
             if ((Intent.ACTION_PACKAGE_REMOVED.equals(action) ||
                     Intent.ACTION_PACKAGE_ADDED.equals(action))) {
                 if (intent.getBooleanExtra(Intent.EXTRA_REPLACING, false)) {
                     maybeUnrestrictBuggyApp(pkgName, userId);
-                } else {
+                } else if (!Intent.ACTION_PACKAGE_ADDED.equals(action)) {
                     clearAppIdleForPackage(pkgName, userId);
+                } else {
+                    // Package was just added and it's not being replaced.
+                    if (mAppsToRestoreToRare.contains(userId, pkgName)) {
+                        restoreAppToRare(pkgName, userId, mInjector.elapsedRealtime(),
+                                REASON_MAIN_DEFAULT | REASON_SUB_DEFAULT_APP_RESTORED);
+                        mAppsToRestoreToRare.remove(userId, pkgName);
+                    }
                 }
             }
         }
diff --git a/core/java/android/app/trust/ITrustListener.aidl b/core/java/android/app/trust/ITrustListener.aidl
index 6b9d2c73..e4ac0119 100644
--- a/core/java/android/app/trust/ITrustListener.aidl
+++ b/core/java/android/app/trust/ITrustListener.aidl
@@ -24,8 +24,8 @@
  * {@hide}
  */
 oneway interface ITrustListener {
-    void onTrustChanged(boolean enabled, int userId, int flags,
+    void onTrustChanged(boolean enabled, boolean newlyUnlocked, int userId, int flags,
         in List<String> trustGrantedMessages);
     void onTrustManagedChanged(boolean managed, int userId);
     void onTrustError(in CharSequence message);
-}
\ No newline at end of file
+}
diff --git a/core/java/android/app/trust/TrustManager.java b/core/java/android/app/trust/TrustManager.java
index 9e825b72..62f755d 100644
--- a/core/java/android/app/trust/TrustManager.java
+++ b/core/java/android/app/trust/TrustManager.java
@@ -22,6 +22,7 @@
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.hardware.biometrics.BiometricSourceType;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
@@ -45,6 +46,7 @@
 
     private static final String TAG = "TrustManager";
     private static final String DATA_FLAGS = "initiatedByUser";
+    private static final String DATA_NEWLY_UNLOCKED = "newlyUnlocked";
     private static final String DATA_MESSAGE = "message";
     private static final String DATA_GRANTED_MESSAGES = "grantedMessages";
 
@@ -171,13 +173,14 @@
         try {
             ITrustListener.Stub iTrustListener = new ITrustListener.Stub() {
                 @Override
-                public void onTrustChanged(boolean enabled, int userId, int flags,
-                        List<String> trustGrantedMessages) {
+                public void onTrustChanged(boolean enabled, boolean newlyUnlocked, int userId,
+                        int flags, List<String> trustGrantedMessages) {
                     Message m = mHandler.obtainMessage(MSG_TRUST_CHANGED, (enabled ? 1 : 0), userId,
                             trustListener);
                     if (flags != 0) {
                         m.getData().putInt(DATA_FLAGS, flags);
                     }
+                    m.getData().putInt(DATA_NEWLY_UNLOCKED, newlyUnlocked ? 1 : 0);
                     m.getData().putCharSequenceArrayList(
                             DATA_GRANTED_MESSAGES, (ArrayList) trustGrantedMessages);
                     m.sendToTarget();
@@ -265,9 +268,14 @@
         public void handleMessage(Message msg) {
             switch(msg.what) {
                 case MSG_TRUST_CHANGED:
-                    int flags = msg.peekData() != null ? msg.peekData().getInt(DATA_FLAGS) : 0;
-                    ((TrustListener) msg.obj).onTrustChanged(msg.arg1 != 0, msg.arg2, flags,
-                            msg.getData().getStringArrayList(DATA_GRANTED_MESSAGES));
+                    Bundle data = msg.peekData();
+                    int flags = data != null ? data.getInt(DATA_FLAGS) : 0;
+                    boolean enabled = msg.arg1 != 0;
+                    int newlyUnlockedInt =
+                            data != null ? data.getInt(DATA_NEWLY_UNLOCKED) : 0;
+                    boolean newlyUnlocked = newlyUnlockedInt != 0;
+                    ((TrustListener) msg.obj).onTrustChanged(enabled, newlyUnlocked, msg.arg2,
+                            flags, msg.getData().getStringArrayList(DATA_GRANTED_MESSAGES));
                     break;
                 case MSG_TRUST_MANAGED_CHANGED:
                     ((TrustListener)msg.obj).onTrustManagedChanged(msg.arg1 != 0, msg.arg2);
@@ -284,6 +292,8 @@
         /**
          * Reports that the trust state has changed.
          * @param enabled If true, the system believes the environment to be trusted.
+         * @param newlyUnlocked If true, the system believes the device is newly unlocked due
+         *        to the trust changing.
          * @param userId The user, for which the trust changed.
          * @param flags Flags specified by the trust agent when granting trust. See
          *     {@link android.service.trust.TrustAgentService#grantTrust(CharSequence, long, int)
@@ -291,7 +301,7 @@
          * @param trustGrantedMessages Messages to display to the user when trust has been granted
          *        by one or more trust agents.
          */
-        void onTrustChanged(boolean enabled, int userId, int flags,
+        void onTrustChanged(boolean enabled, boolean newlyUnlocked, int userId, int flags,
                 List<String> trustGrantedMessages);
 
         /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index c743582..09f5cf1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -61,6 +61,7 @@
 import com.android.wm.shell.desktopmode.DesktopModeController;
 import com.android.wm.shell.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
+import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.displayareahelper.DisplayAreaHelper;
 import com.android.wm.shell.displayareahelper.DisplayAreaHelperController;
 import com.android.wm.shell.draganddrop.DragAndDropController;
@@ -677,7 +678,11 @@
     @WMSingleton
     @Provides
     static Optional<DesktopMode> provideDesktopMode(
-            Optional<DesktopModeController> desktopModeController) {
+            Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController) {
+        if (DesktopModeStatus.isProto2Enabled()) {
+            return desktopTasksController.map(DesktopTasksController::asDesktopMode);
+        }
         return desktopModeController.map(DesktopModeController::asDesktopMode);
     }
 
@@ -700,6 +705,23 @@
 
     @BindsOptionalOf
     @DynamicOverride
+    abstract DesktopTasksController optionalDesktopTasksController();
+
+    @WMSingleton
+    @Provides
+    static Optional<DesktopTasksController> providesDesktopTasksController(
+            @DynamicOverride Optional<Lazy<DesktopTasksController>> desktopTasksController) {
+        // Use optional-of-lazy for the dependency that this provider relies on.
+        // Lazy ensures that this provider will not be the cause the dependency is created
+        // when it will not be returned due to the condition below.
+        if (DesktopModeStatus.isProto2Enabled()) {
+            return desktopTasksController.map(Lazy::get);
+        }
+        return Optional.empty();
+    }
+
+    @BindsOptionalOf
+    @DynamicOverride
     abstract DesktopModeTaskRepository optionalDesktopModeTaskRepository();
 
     @WMSingleton
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 6be8305..701a3a4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -50,6 +50,7 @@
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.desktopmode.DesktopModeController;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
+import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.freeform.FreeformComponents;
 import com.android.wm.shell.freeform.FreeformTaskListener;
@@ -189,7 +190,8 @@
             ShellTaskOrganizer taskOrganizer,
             DisplayController displayController,
             SyncTransactionQueue syncQueue,
-            Optional<DesktopModeController> desktopModeController) {
+            Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController) {
         return new CaptionWindowDecorViewModel(
                     context,
                     mainHandler,
@@ -197,7 +199,8 @@
                     taskOrganizer,
                     displayController,
                     syncQueue,
-                    desktopModeController);
+                    desktopModeController,
+                    desktopTasksController);
     }
 
     //
@@ -616,6 +619,22 @@
     @WMSingleton
     @Provides
     @DynamicOverride
+    static DesktopTasksController provideDesktopTasksController(
+            Context context,
+            ShellInit shellInit,
+            ShellController shellController,
+            ShellTaskOrganizer shellTaskOrganizer,
+            Transitions transitions,
+            @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository,
+            @ShellMainThread ShellExecutor mainExecutor
+    ) {
+        return new DesktopTasksController(context, shellInit, shellController, shellTaskOrganizer,
+                transitions, desktopModeTaskRepository, mainExecutor);
+    }
+
+    @WMSingleton
+    @Provides
+    @DynamicOverride
     static DesktopModeTaskRepository provideDesktopModeTaskRepository() {
         return new DesktopModeTaskRepository();
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java
index 67f4a19..055949f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeStatus.java
@@ -70,9 +70,13 @@
      * @return {@code true} if active
      */
     public static boolean isActive(Context context) {
-        if (!IS_SUPPORTED) {
+        if (!isAnyEnabled()) {
             return false;
         }
+        if (isProto2Enabled()) {
+            // Desktop mode is always active in prototype 2
+            return true;
+        }
         try {
             int result = Settings.System.getIntForUser(context.getContentResolver(),
                     Settings.System.DESKTOP_MODE, UserHandle.USER_CURRENT);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
new file mode 100644
index 0000000..b075b14
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -0,0 +1,231 @@
+/*
+ * Copyright (C) 2022 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.desktopmode
+
+import android.app.ActivityManager
+import android.app.WindowConfiguration
+import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
+import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
+import android.app.WindowConfiguration.WindowingMode
+import android.content.Context
+import android.view.WindowManager
+import android.window.WindowContainerTransaction
+import androidx.annotation.BinderThread
+import com.android.internal.protolog.common.ProtoLog
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.common.ExecutorUtils
+import com.android.wm.shell.common.ExternalInterfaceBinder
+import com.android.wm.shell.common.RemoteCallable
+import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.common.annotations.ExternalThread
+import com.android.wm.shell.common.annotations.ShellMainThread
+import com.android.wm.shell.desktopmode.DesktopModeTaskRepository.VisibleTasksListener
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.sysui.ShellController
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.sysui.ShellSharedConstants
+import com.android.wm.shell.transition.Transitions
+import java.util.concurrent.Executor
+import java.util.function.Consumer
+
+/** Handles moving tasks in and out of desktop */
+class DesktopTasksController(
+    private val context: Context,
+    shellInit: ShellInit,
+    private val shellController: ShellController,
+    private val shellTaskOrganizer: ShellTaskOrganizer,
+    private val transitions: Transitions,
+    private val desktopModeTaskRepository: DesktopModeTaskRepository,
+    @ShellMainThread private val mainExecutor: ShellExecutor
+) : RemoteCallable<DesktopTasksController> {
+
+    private val desktopMode: DesktopModeImpl
+
+    init {
+        desktopMode = DesktopModeImpl()
+        if (DesktopModeStatus.isProto2Enabled()) {
+            shellInit.addInitCallback({ onInit() }, this)
+        }
+    }
+
+    private fun onInit() {
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "Initialize DesktopTasksController")
+        shellController.addExternalInterface(
+            ShellSharedConstants.KEY_EXTRA_SHELL_DESKTOP_MODE,
+            { createExternalInterface() },
+            this
+        )
+    }
+
+    /** Show all tasks, that are part of the desktop, on top of launcher */
+    fun showDesktopApps() {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "showDesktopApps")
+        val wct = WindowContainerTransaction()
+
+        bringDesktopAppsToFront(wct)
+
+        // Execute transaction if there are pending operations
+        if (!wct.isEmpty) {
+            if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+                transitions.startTransition(WindowManager.TRANSIT_TO_FRONT, wct, null /* handler */)
+            } else {
+                shellTaskOrganizer.applyTransaction(wct)
+            }
+        }
+    }
+
+    /** Move a task with given `taskId` to desktop */
+    fun moveToDesktop(taskId: Int) {
+        shellTaskOrganizer.getRunningTaskInfo(taskId)?.let { task -> moveToDesktop(task) }
+    }
+
+    /** Move a task to desktop */
+    fun moveToDesktop(task: ActivityManager.RunningTaskInfo) {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "moveToDesktop: %d", task.taskId)
+
+        val wct = WindowContainerTransaction()
+        // Bring other apps to front first
+        bringDesktopAppsToFront(wct)
+
+        wct.setWindowingMode(task.getToken(), WindowConfiguration.WINDOWING_MODE_FREEFORM)
+        wct.reorder(task.getToken(), true /* onTop */)
+
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            transitions.startTransition(WindowManager.TRANSIT_CHANGE, wct, null /* handler */)
+        } else {
+            shellTaskOrganizer.applyTransaction(wct)
+        }
+    }
+
+    /** Move a task with given `taskId` to fullscreen */
+    fun moveToFullscreen(taskId: Int) {
+        shellTaskOrganizer.getRunningTaskInfo(taskId)?.let { task -> moveToFullscreen(task) }
+    }
+
+    /** Move a task to fullscreen */
+    fun moveToFullscreen(task: ActivityManager.RunningTaskInfo) {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "moveToFullscreen: %d", task.taskId)
+
+        val wct = WindowContainerTransaction()
+        wct.setWindowingMode(task.getToken(), WindowConfiguration.WINDOWING_MODE_FULLSCREEN)
+        wct.setBounds(task.getToken(), null)
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            transitions.startTransition(WindowManager.TRANSIT_CHANGE, wct, null /* handler */)
+        } else {
+            shellTaskOrganizer.applyTransaction(wct)
+        }
+    }
+
+    /**
+     * Get windowing move for a given `taskId`
+     *
+     * @return [WindowingMode] for the task or [WINDOWING_MODE_UNDEFINED] if task is not found
+     */
+    @WindowingMode
+    fun getTaskWindowingMode(taskId: Int): Int {
+        return shellTaskOrganizer.getRunningTaskInfo(taskId)?.windowingMode
+            ?: WINDOWING_MODE_UNDEFINED
+    }
+
+    private fun bringDesktopAppsToFront(wct: WindowContainerTransaction) {
+        val activeTasks = desktopModeTaskRepository.getActiveTasks()
+
+        // Skip if all tasks are already visible
+        if (activeTasks.isNotEmpty() && activeTasks.all(desktopModeTaskRepository::isVisibleTask)) {
+            ProtoLog.d(
+                WM_SHELL_DESKTOP_MODE,
+                "bringDesktopAppsToFront: active tasks are already in front, skipping."
+            )
+            return
+        }
+
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "bringDesktopAppsToFront")
+
+        // First move home to front and then other tasks on top of it
+        moveHomeTaskToFront(wct)
+
+        val allTasksInZOrder = desktopModeTaskRepository.getFreeformTasksInZOrder()
+        activeTasks
+            // Sort descending as the top task is at index 0. It should be ordered to top last
+            .sortedByDescending { taskId -> allTasksInZOrder.indexOf(taskId) }
+            .mapNotNull { taskId -> shellTaskOrganizer.getRunningTaskInfo(taskId) }
+            .forEach { task -> wct.reorder(task.token, true /* onTop */) }
+    }
+
+    private fun moveHomeTaskToFront(wct: WindowContainerTransaction) {
+        shellTaskOrganizer
+            .getRunningTasks(context.displayId)
+            .firstOrNull { task -> task.activityType == ACTIVITY_TYPE_HOME }
+            ?.let { homeTask -> wct.reorder(homeTask.getToken(), true /* onTop */) }
+    }
+
+    override fun getContext(): Context {
+        return context
+    }
+
+    override fun getRemoteCallExecutor(): ShellExecutor {
+        return mainExecutor
+    }
+
+    /** Creates a new instance of the external interface to pass to another process. */
+    private fun createExternalInterface(): ExternalInterfaceBinder {
+        return IDesktopModeImpl(this)
+    }
+
+    /** Get connection interface between sysui and shell */
+    fun asDesktopMode(): DesktopMode {
+        return desktopMode
+    }
+
+    /**
+     * Adds a listener to find out about changes in the visibility of freeform tasks.
+     *
+     * @param listener the listener to add.
+     * @param callbackExecutor the executor to call the listener on.
+     */
+    fun addListener(listener: VisibleTasksListener, callbackExecutor: Executor) {
+        desktopModeTaskRepository.addVisibleTasksListener(listener, callbackExecutor)
+    }
+
+    /** The interface for calls from outside the shell, within the host process. */
+    @ExternalThread
+    private inner class DesktopModeImpl : DesktopMode {
+        override fun addListener(listener: VisibleTasksListener, callbackExecutor: Executor) {
+            mainExecutor.execute {
+                this@DesktopTasksController.addListener(listener, callbackExecutor)
+            }
+        }
+    }
+
+    /** The interface for calls from outside the host process. */
+    @BinderThread
+    private class IDesktopModeImpl(private var controller: DesktopTasksController?) :
+        IDesktopMode.Stub(), ExternalInterfaceBinder {
+        /** Invalidates this instance, preventing future calls from updating the controller. */
+        override fun invalidate() {
+            controller = null
+        }
+
+        override fun showDesktopApps() {
+            ExecutorUtils.executeRemoteCallWithTaskPermission(
+                controller,
+                "showDesktopApps",
+                Consumer(DesktopTasksController::showDesktopApps)
+            )
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
index 42e2b3f..299284f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModel.java
@@ -53,6 +53,7 @@
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.DesktopModeController;
 import com.android.wm.shell.desktopmode.DesktopModeStatus;
+import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
 import com.android.wm.shell.transition.Transitions;
 
@@ -75,6 +76,7 @@
     private final SyncTransactionQueue mSyncQueue;
     private FreeformTaskTransitionStarter mTransitionStarter;
     private Optional<DesktopModeController> mDesktopModeController;
+    private Optional<DesktopTasksController> mDesktopTasksController;
     private boolean mTransitionDragActive;
 
     private SparseArray<EventReceiver> mEventReceiversByDisplay = new SparseArray<>();
@@ -90,7 +92,8 @@
             ShellTaskOrganizer taskOrganizer,
             DisplayController displayController,
             SyncTransactionQueue syncQueue,
-            Optional<DesktopModeController> desktopModeController) {
+            Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController) {
         this(
                 context,
                 mainHandler,
@@ -99,6 +102,7 @@
                 displayController,
                 syncQueue,
                 desktopModeController,
+                desktopTasksController,
                 new CaptionWindowDecoration.Factory(),
                 new InputMonitorFactory());
     }
@@ -112,6 +116,7 @@
             DisplayController displayController,
             SyncTransactionQueue syncQueue,
             Optional<DesktopModeController> desktopModeController,
+            Optional<DesktopTasksController> desktopTasksController,
             CaptionWindowDecoration.Factory captionWindowDecorFactory,
             InputMonitorFactory inputMonitorFactory) {
         mContext = context;
@@ -122,6 +127,7 @@
         mDisplayController = displayController;
         mSyncQueue = syncQueue;
         mDesktopModeController = desktopModeController;
+        mDesktopTasksController = desktopTasksController;
 
         mCaptionWindowDecorFactory = captionWindowDecorFactory;
         mInputMonitorFactory = inputMonitorFactory;
@@ -242,11 +248,13 @@
                 decoration.createHandleMenu();
             } else if (id == R.id.desktop_button) {
                 mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(true));
+                mDesktopTasksController.ifPresent(c -> c.moveToDesktop(mTaskId));
                 decoration.closeHandleMenu();
             } else if (id == R.id.fullscreen_button) {
                 mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(false));
+                mDesktopTasksController.ifPresent(c -> c.moveToFullscreen(mTaskId));
                 decoration.closeHandleMenu();
-                decoration.setButtonVisibility();
+                decoration.setButtonVisibility(false);
             }
         }
 
@@ -299,8 +307,13 @@
          */
         private void handleEventForMove(MotionEvent e) {
             RunningTaskInfo taskInfo = mTaskOrganizer.getRunningTaskInfo(mTaskId);
-            if (mDesktopModeController.isPresent()
-                    && mDesktopModeController.get().getDisplayAreaWindowingMode(taskInfo.displayId)
+            if (DesktopModeStatus.isProto2Enabled()
+                    && taskInfo.getWindowingMode() == WINDOWING_MODE_FULLSCREEN) {
+                return;
+            }
+            if (DesktopModeStatus.isProto1Enabled() && mDesktopModeController.isPresent()
+                    && mDesktopModeController.get().getDisplayAreaWindowingMode(
+                    taskInfo.displayId)
                     == WINDOWING_MODE_FULLSCREEN) {
                 return;
             }
@@ -324,9 +337,20 @@
                             .stableInsets().top;
                     mDragResizeCallback.onDragResizeEnd(
                             e.getRawX(dragPointerIdx), e.getRawY(dragPointerIdx));
-                    if (e.getRawY(dragPointerIdx) <= statusBarHeight
-                            && DesktopModeStatus.isActive(mContext)) {
-                        mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(false));
+                    if (e.getRawY(dragPointerIdx) <= statusBarHeight) {
+                        if (DesktopModeStatus.isProto2Enabled()) {
+                            if (taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM) {
+                                // Switch a single task to fullscreen
+                                mDesktopTasksController.ifPresent(
+                                        c -> c.moveToFullscreen(taskInfo));
+                            }
+                        } else if (DesktopModeStatus.isProto1Enabled()) {
+                            if (DesktopModeStatus.isActive(mContext)) {
+                                // Turn off desktop mode
+                                mDesktopModeController.ifPresent(
+                                        c -> c.setDesktopModeActive(false));
+                            }
+                        }
                     }
                     break;
                 }
@@ -408,13 +432,27 @@
      * @param ev the {@link MotionEvent} received by {@link EventReceiver}
      */
     private void handleReceivedMotionEvent(MotionEvent ev, InputMonitor inputMonitor) {
-        if (!DesktopModeStatus.isActive(mContext)) {
-            handleCaptionThroughStatusBar(ev);
+        if (DesktopModeStatus.isProto2Enabled()) {
+            CaptionWindowDecoration focusedDecor = getFocusedDecor();
+            if (focusedDecor == null
+                    || focusedDecor.mTaskInfo.getWindowingMode() != WINDOWING_MODE_FREEFORM) {
+                handleCaptionThroughStatusBar(ev);
+            }
+        } else if (DesktopModeStatus.isProto1Enabled()) {
+            if (!DesktopModeStatus.isActive(mContext)) {
+                handleCaptionThroughStatusBar(ev);
+            }
         }
         handleEventOutsideFocusedCaption(ev);
         // Prevent status bar from reacting to a caption drag.
-        if (mTransitionDragActive && !DesktopModeStatus.isActive(mContext)) {
-            inputMonitor.pilferPointers();
+        if (DesktopModeStatus.isProto2Enabled()) {
+            if (mTransitionDragActive) {
+                inputMonitor.pilferPointers();
+            }
+        } else if (DesktopModeStatus.isProto1Enabled()) {
+            if (mTransitionDragActive && !DesktopModeStatus.isActive(mContext)) {
+                inputMonitor.pilferPointers();
+            }
         }
     }
 
@@ -443,9 +481,20 @@
             case MotionEvent.ACTION_DOWN: {
                 // Begin drag through status bar if applicable.
                 CaptionWindowDecoration focusedDecor = getFocusedDecor();
-                if (focusedDecor != null && !DesktopModeStatus.isActive(mContext)
-                        && focusedDecor.checkTouchEventInHandle(ev)) {
-                    mTransitionDragActive = true;
+                if (focusedDecor != null) {
+                    boolean dragFromStatusBarAllowed = false;
+                    if (DesktopModeStatus.isProto2Enabled()) {
+                        // In proto2 any full screen task can be dragged to freeform
+                        dragFromStatusBarAllowed = focusedDecor.mTaskInfo.getWindowingMode()
+                                == WINDOWING_MODE_FULLSCREEN;
+                    } else if (DesktopModeStatus.isProto1Enabled()) {
+                        // In proto1 task can be dragged to freeform when not in desktop mode
+                        dragFromStatusBarAllowed = !DesktopModeStatus.isActive(mContext);
+                    }
+
+                    if (dragFromStatusBarAllowed && focusedDecor.checkTouchEventInHandle(ev)) {
+                        mTransitionDragActive = true;
+                    }
                 }
                 break;
             }
@@ -460,7 +509,13 @@
                     int statusBarHeight = mDisplayController
                             .getDisplayLayout(focusedDecor.mTaskInfo.displayId).stableInsets().top;
                     if (ev.getY() > statusBarHeight) {
-                        mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(true));
+                        if (DesktopModeStatus.isProto2Enabled()) {
+                            mDesktopTasksController.ifPresent(
+                                    c -> c.moveToDesktop(focusedDecor.mTaskInfo));
+                        } else if (DesktopModeStatus.isProto1Enabled()) {
+                            mDesktopModeController.ifPresent(c -> c.setDesktopModeActive(true));
+                        }
+
                         return;
                     }
                 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
index 037ca20..f7c7a87 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
@@ -16,8 +16,9 @@
 
 package com.android.wm.shell.windowdecor;
 
+import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
+
 import android.app.ActivityManager;
-import android.app.WindowConfiguration;
 import android.content.Context;
 import android.content.res.ColorStateList;
 import android.content.res.Resources;
@@ -117,7 +118,7 @@
                 ? R.dimen.freeform_decor_shadow_focused_thickness
                 : R.dimen.freeform_decor_shadow_unfocused_thickness;
         final boolean isFreeform =
-                taskInfo.getWindowingMode() == WindowConfiguration.WINDOWING_MODE_FREEFORM;
+                taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM;
         final boolean isDragResizeable = isFreeform && taskInfo.isResizeable;
 
         WindowDecorLinearLayout oldRootView = mResult.mRootView;
@@ -167,11 +168,17 @@
         // If this task is not focused, do not show caption.
         setCaptionVisibility(mTaskInfo.isFocused);
 
-        // Only handle should show if Desktop Mode is inactive.
-        boolean desktopCurrentStatus = DesktopModeStatus.isActive(mContext);
-        if (mDesktopActive != desktopCurrentStatus && mTaskInfo.isFocused) {
-            mDesktopActive = desktopCurrentStatus;
-            setButtonVisibility();
+        if (mTaskInfo.isFocused) {
+            if (DesktopModeStatus.isProto2Enabled()) {
+                updateButtonVisibility();
+            } else if (DesktopModeStatus.isProto1Enabled()) {
+                // Only handle should show if Desktop Mode is inactive.
+                boolean desktopCurrentStatus = DesktopModeStatus.isActive(mContext);
+                if (mDesktopActive != desktopCurrentStatus) {
+                    mDesktopActive = desktopCurrentStatus;
+                    setButtonVisibility(mDesktopActive);
+                }
+            }
         }
 
         if (!isDragResizeable) {
@@ -214,7 +221,7 @@
         View handle = caption.findViewById(R.id.caption_handle);
         handle.setOnTouchListener(mOnCaptionTouchListener);
         handle.setOnClickListener(mOnCaptionButtonClickListener);
-        setButtonVisibility();
+        updateButtonVisibility();
     }
 
     private void setupHandleMenu() {
@@ -244,14 +251,25 @@
     /**
      * Sets the visibility of buttons and color of caption based on desktop mode status
      */
-    void setButtonVisibility() {
-        mDesktopActive = DesktopModeStatus.isActive(mContext);
-        int v = mDesktopActive ? View.VISIBLE : View.GONE;
+    void updateButtonVisibility() {
+        if (DesktopModeStatus.isProto2Enabled()) {
+            setButtonVisibility(mTaskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM);
+        } else if (DesktopModeStatus.isProto1Enabled()) {
+            mDesktopActive = DesktopModeStatus.isActive(mContext);
+            setButtonVisibility(mDesktopActive);
+        }
+    }
+
+    /**
+     * Show or hide buttons
+     */
+    void setButtonVisibility(boolean visible) {
+        int visibility = visible ? View.VISIBLE : View.GONE;
         View caption = mResult.mRootView.findViewById(R.id.caption);
         View back = caption.findViewById(R.id.back_button);
         View close = caption.findViewById(R.id.close_window);
-        back.setVisibility(v);
-        close.setVisibility(v);
+        back.setVisibility(visibility);
+        close.setVisibility(visibility);
         int buttonTintColorRes =
                 mDesktopActive ? R.color.decor_button_dark_color
                         : R.color.decor_button_light_color;
@@ -260,7 +278,7 @@
         View handle = caption.findViewById(R.id.caption_handle);
         VectorDrawable handleBackground = (VectorDrawable) handle.getBackground();
         handleBackground.setTintList(buttonTintColor);
-        caption.getBackground().setTint(v == View.VISIBLE ? Color.WHITE : Color.TRANSPARENT);
+        caption.getBackground().setTint(visible ? Color.WHITE : Color.TRANSPARENT);
     }
 
     boolean isHandleMenuActive() {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java
index 707c049..08af3d3 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeControllerTest.java
@@ -16,8 +16,6 @@
 
 package com.android.wm.shell.desktopmode;
 
-import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
-import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
@@ -29,6 +27,9 @@
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+import static com.android.wm.shell.desktopmode.DesktopTestHelpers.createFreeformTask;
+import static com.android.wm.shell.desktopmode.DesktopTestHelpers.createFullscreenTask;
+import static com.android.wm.shell.desktopmode.DesktopTestHelpers.createHomeTask;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -48,7 +49,6 @@
 import android.testing.AndroidTestingRunner;
 import android.window.DisplayAreaInfo;
 import android.window.TransitionRequestInfo;
-import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 import android.window.WindowContainerTransaction.Change;
 import android.window.WindowContainerTransaction.HierarchyOp;
@@ -59,7 +59,6 @@
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
-import com.android.wm.shell.TestRunningTaskInfoBuilder;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.sysui.ShellController;
@@ -355,7 +354,7 @@
     @Test
     public void testHandleTransitionRequest_taskOpen_returnsWct() {
         RunningTaskInfo trigger = new RunningTaskInfo();
-        trigger.token = new MockToken().mToken;
+        trigger.token = new MockToken().token();
         trigger.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FREEFORM);
         WindowContainerTransaction wct = mController.handleRequest(
                 mock(IBinder.class),
@@ -366,7 +365,7 @@
     @Test
     public void testHandleTransitionRequest_taskToFront_returnsWct() {
         RunningTaskInfo trigger = new RunningTaskInfo();
-        trigger.token = new MockToken().mToken;
+        trigger.token = new MockToken().token();
         trigger.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FREEFORM);
         WindowContainerTransaction wct = mController.handleRequest(
                 mock(IBinder.class),
@@ -381,40 +380,13 @@
     }
 
     private DisplayAreaInfo createMockDisplayArea() {
-        DisplayAreaInfo displayAreaInfo = new DisplayAreaInfo(new MockToken().mToken,
+        DisplayAreaInfo displayAreaInfo = new DisplayAreaInfo(new MockToken().token(),
                 mContext.getDisplayId(), 0);
         when(mRootTaskDisplayAreaOrganizer.getDisplayAreaInfo(mContext.getDisplayId()))
                 .thenReturn(displayAreaInfo);
         return displayAreaInfo;
     }
 
-    private RunningTaskInfo createFreeformTask() {
-        return new TestRunningTaskInfoBuilder()
-                .setToken(new MockToken().token())
-                .setActivityType(ACTIVITY_TYPE_STANDARD)
-                .setWindowingMode(WINDOWING_MODE_FREEFORM)
-                .setLastActiveTime(100)
-                .build();
-    }
-
-    private RunningTaskInfo createFullscreenTask() {
-        return new TestRunningTaskInfoBuilder()
-                .setToken(new MockToken().token())
-                .setActivityType(ACTIVITY_TYPE_STANDARD)
-                .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
-                .setLastActiveTime(100)
-                .build();
-    }
-
-    private RunningTaskInfo createHomeTask() {
-        return new TestRunningTaskInfoBuilder()
-                .setToken(new MockToken().token())
-                .setActivityType(ACTIVITY_TYPE_HOME)
-                .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
-                .setLastActiveTime(100)
-                .build();
-    }
-
     private WindowContainerTransaction getDesktopModeSwitchTransaction() {
         ArgumentCaptor<WindowContainerTransaction> arg = ArgumentCaptor.forClass(
                 WindowContainerTransaction.class);
@@ -442,18 +414,4 @@
         assertThat(change.getConfiguration().windowConfiguration.getBounds().isEmpty()).isTrue();
     }
 
-    private static class MockToken {
-        private final WindowContainerToken mToken;
-        private final IBinder mBinder;
-
-        MockToken() {
-            mToken = mock(WindowContainerToken.class);
-            mBinder = mock(IBinder.class);
-            when(mToken.asBinder()).thenReturn(mBinder);
-        }
-
-        WindowContainerToken token() {
-            return mToken;
-        }
-    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
new file mode 100644
index 0000000..de2473b
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -0,0 +1,281 @@
+/*
+ * Copyright (C) 2022 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.desktopmode
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
+import android.testing.AndroidTestingRunner
+import android.window.WindowContainerTransaction
+import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER
+import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.dx.mockito.inline.extended.ExtendedMockito.never
+import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestShellExecutor
+import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFreeformTask
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFullscreenTask
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createHomeTask
+import com.android.wm.shell.sysui.ShellController
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.transition.Transitions
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.isNull
+import org.mockito.Mock
+import org.mockito.Mockito
+import org.mockito.Mockito.any
+import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.clearInvocations
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.`when` as whenever
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class DesktopTasksControllerTest : ShellTestCase() {
+
+    @Mock lateinit var testExecutor: ShellExecutor
+    @Mock lateinit var shellController: ShellController
+    @Mock lateinit var shellTaskOrganizer: ShellTaskOrganizer
+    @Mock lateinit var transitions: Transitions
+
+    lateinit var mockitoSession: StaticMockitoSession
+    lateinit var controller: DesktopTasksController
+    lateinit var shellInit: ShellInit
+    lateinit var desktopModeTaskRepository: DesktopModeTaskRepository
+
+    // Mock running tasks are registered here so we can get the list from mock shell task organizer
+    private val runningTasks = mutableListOf<RunningTaskInfo>()
+
+    @Before
+    fun setUp() {
+        mockitoSession = mockitoSession().mockStatic(DesktopModeStatus::class.java).startMocking()
+        whenever(DesktopModeStatus.isProto2Enabled()).thenReturn(true)
+
+        shellInit = Mockito.spy(ShellInit(testExecutor))
+        desktopModeTaskRepository = DesktopModeTaskRepository()
+
+        whenever(shellTaskOrganizer.getRunningTasks(anyInt())).thenAnswer { runningTasks }
+
+        controller = createController()
+
+        shellInit.init()
+    }
+
+    private fun createController(): DesktopTasksController {
+        return DesktopTasksController(
+            context,
+            shellInit,
+            shellController,
+            shellTaskOrganizer,
+            transitions,
+            desktopModeTaskRepository,
+            TestShellExecutor()
+        )
+    }
+
+    @After
+    fun tearDown() {
+        mockitoSession.finishMocking()
+
+        runningTasks.clear()
+    }
+
+    @Test
+    fun instantiate_addInitCallback() {
+        verify(shellInit).addInitCallback(any(), any<DesktopTasksController>())
+    }
+
+    @Test
+    fun instantiate_flagOff_doNotAddInitCallback() {
+        whenever(DesktopModeStatus.isProto2Enabled()).thenReturn(false)
+        clearInvocations(shellInit)
+
+        createController()
+
+        verify(shellInit, never()).addInitCallback(any(), any<DesktopTasksController>())
+    }
+
+    @Test
+    fun showDesktopApps_allAppsInvisible_bringsToFront() {
+        val homeTask = setUpHomeTask()
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskHidden(task1)
+        markTaskHidden(task2)
+
+        controller.showDesktopApps()
+
+        val wct = getLatestWct()
+        assertThat(wct.hierarchyOps).hasSize(3)
+        // Expect order to be from bottom: home, task1, task2
+        wct.assertReorderAt(index = 0, homeTask)
+        wct.assertReorderAt(index = 1, task1)
+        wct.assertReorderAt(index = 2, task2)
+    }
+
+    @Test
+    fun showDesktopApps_appsAlreadyVisible_doesNothing() {
+        setUpHomeTask()
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskVisible(task1)
+        markTaskVisible(task2)
+
+        controller.showDesktopApps()
+
+        verifyWCTNotExecuted()
+    }
+
+    @Test
+    fun showDesktopApps_someAppsInvisible_reordersAll() {
+        val homeTask = setUpHomeTask()
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskHidden(task1)
+        markTaskVisible(task2)
+
+        controller.showDesktopApps()
+
+        val wct = getLatestWct()
+        assertThat(wct.hierarchyOps).hasSize(3)
+        // Expect order to be from bottom: home, task1, task2
+        wct.assertReorderAt(index = 0, homeTask)
+        wct.assertReorderAt(index = 1, task1)
+        wct.assertReorderAt(index = 2, task2)
+    }
+
+    @Test
+    fun showDesktopApps_noActiveTasks_reorderHomeToTop() {
+        val homeTask = setUpHomeTask()
+
+        controller.showDesktopApps()
+
+        val wct = getLatestWct()
+        assertThat(wct.hierarchyOps).hasSize(1)
+        wct.assertReorderAt(index = 0, homeTask)
+    }
+
+    @Test
+    fun moveToDesktop() {
+        val task = setUpFullscreenTask()
+        controller.moveToDesktop(task)
+        val wct = getLatestWct()
+        assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
+            .isEqualTo(WINDOWING_MODE_FREEFORM)
+    }
+
+    @Test
+    fun moveToDesktop_nonExistentTask_doesNothing() {
+        controller.moveToDesktop(999)
+        verifyWCTNotExecuted()
+    }
+
+    @Test
+    fun moveToFullscreen() {
+        val task = setUpFreeformTask()
+        controller.moveToFullscreen(task)
+        val wct = getLatestWct()
+        assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
+            .isEqualTo(WINDOWING_MODE_FULLSCREEN)
+    }
+
+    @Test
+    fun moveToFullscreen_nonExistentTask_doesNothing() {
+        controller.moveToFullscreen(999)
+        verifyWCTNotExecuted()
+    }
+
+    @Test
+    fun getTaskWindowingMode() {
+        val fullscreenTask = setUpFullscreenTask()
+        val freeformTask = setUpFreeformTask()
+
+        assertThat(controller.getTaskWindowingMode(fullscreenTask.taskId))
+            .isEqualTo(WINDOWING_MODE_FULLSCREEN)
+        assertThat(controller.getTaskWindowingMode(freeformTask.taskId))
+            .isEqualTo(WINDOWING_MODE_FREEFORM)
+        assertThat(controller.getTaskWindowingMode(999)).isEqualTo(WINDOWING_MODE_UNDEFINED)
+    }
+
+    private fun setUpFreeformTask(): RunningTaskInfo {
+        val task = createFreeformTask()
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        desktopModeTaskRepository.addActiveTask(task.taskId)
+        desktopModeTaskRepository.addOrMoveFreeformTaskToTop(task.taskId)
+        runningTasks.add(task)
+        return task
+    }
+
+    private fun setUpHomeTask(): RunningTaskInfo {
+        val task = createHomeTask()
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        runningTasks.add(task)
+        return task
+    }
+
+    private fun setUpFullscreenTask(): RunningTaskInfo {
+        val task = createFullscreenTask()
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        runningTasks.add(task)
+        return task
+    }
+
+    private fun markTaskVisible(task: RunningTaskInfo) {
+        desktopModeTaskRepository.updateVisibleFreeformTasks(task.taskId, visible = true)
+    }
+
+    private fun markTaskHidden(task: RunningTaskInfo) {
+        desktopModeTaskRepository.updateVisibleFreeformTasks(task.taskId, visible = false)
+    }
+
+    private fun getLatestWct(): WindowContainerTransaction {
+        val arg = ArgumentCaptor.forClass(WindowContainerTransaction::class.java)
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            verify(transitions).startTransition(anyInt(), arg.capture(), isNull())
+        } else {
+            verify(shellTaskOrganizer).applyTransaction(arg.capture())
+        }
+        return arg.value
+    }
+
+    private fun verifyWCTNotExecuted() {
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            verify(transitions, never()).startTransition(anyInt(), any(), isNull())
+        } else {
+            verify(shellTaskOrganizer, never()).applyTransaction(any())
+        }
+    }
+}
+
+private fun WindowContainerTransaction.assertReorderAt(index: Int, task: RunningTaskInfo) {
+    assertWithMessage("WCT does not have a hierarchy operation at index $index")
+        .that(hierarchyOps.size)
+        .isGreaterThan(index)
+    val op = hierarchyOps[index]
+    assertThat(op.type).isEqualTo(HIERARCHY_OP_TYPE_REORDER)
+    assertThat(op.container).isEqualTo(task.token.asBinder())
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTestHelpers.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTestHelpers.kt
new file mode 100644
index 0000000..dc91d75
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTestHelpers.kt
@@ -0,0 +1,61 @@
+/*
+ * Copyright (C) 2022 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.desktopmode
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+
+class DesktopTestHelpers {
+    companion object {
+        /** Create a task that has windowing mode set to [WINDOWING_MODE_FREEFORM] */
+        @JvmStatic
+        fun createFreeformTask(): RunningTaskInfo {
+            return TestRunningTaskInfoBuilder()
+                    .setToken(MockToken().token())
+                    .setActivityType(ACTIVITY_TYPE_STANDARD)
+                    .setWindowingMode(WINDOWING_MODE_FREEFORM)
+                    .setLastActiveTime(100)
+                    .build()
+        }
+
+        /** Create a task that has windowing mode set to [WINDOWING_MODE_FULLSCREEN] */
+        @JvmStatic
+        fun createFullscreenTask(): RunningTaskInfo {
+            return TestRunningTaskInfoBuilder()
+                    .setToken(MockToken().token())
+                    .setActivityType(ACTIVITY_TYPE_STANDARD)
+                    .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
+                    .setLastActiveTime(100)
+                    .build()
+        }
+
+        /** Create a new home task */
+        @JvmStatic
+        fun createHomeTask(): RunningTaskInfo {
+            return TestRunningTaskInfoBuilder()
+                    .setToken(MockToken().token())
+                    .setActivityType(ACTIVITY_TYPE_HOME)
+                    .setWindowingMode(WINDOWING_MODE_FULLSCREEN)
+                    .setLastActiveTime(100)
+                    .build()
+        }
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/MockToken.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/MockToken.java
new file mode 100644
index 0000000..09d474d
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/MockToken.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2022 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.desktopmode;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.os.IBinder;
+import android.window.WindowContainerToken;
+
+/**
+ * {@link WindowContainerToken} wrapper that supports a mock binder
+ */
+class MockToken {
+    private final WindowContainerToken mToken;
+
+    MockToken() {
+        mToken = mock(WindowContainerToken.class);
+        IBinder binder = mock(IBinder.class);
+        when(mToken.asBinder()).thenReturn(binder);
+    }
+
+    WindowContainerToken token() {
+        return mToken;
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java
index 0dbf30d..4875832 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/CaptionWindowDecorViewModelTests.java
@@ -48,6 +48,7 @@
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.DesktopModeController;
+import com.android.wm.shell.desktopmode.DesktopTasksController;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -74,6 +75,7 @@
     @Mock private DisplayController mDisplayController;
     @Mock private SyncTransactionQueue mSyncQueue;
     @Mock private DesktopModeController mDesktopModeController;
+    @Mock private DesktopTasksController mDesktopTasksController;
     @Mock private InputMonitor mInputMonitor;
     @Mock private InputManager mInputManager;
 
@@ -95,6 +97,7 @@
                 mDisplayController,
                 mSyncQueue,
                 Optional.of(mDesktopModeController),
+                Optional.of(mDesktopTasksController),
                 mCaptionWindowDecorFactory,
                 mMockInputMonitorFactory
             );
diff --git a/packages/SettingsLib/res/values/strings.xml b/packages/SettingsLib/res/values/strings.xml
index d1ac7d0..cf37205 100644
--- a/packages/SettingsLib/res/values/strings.xml
+++ b/packages/SettingsLib/res/values/strings.xml
@@ -1125,9 +1125,9 @@
     <!-- [CHAR_LIMIT=40] Label for battery level chart when charging with duration -->
     <string name="power_charging_duration"><xliff:g id="level">%1$s</xliff:g> - <xliff:g id="time">%2$s</xliff:g> left until full</string>
     <!-- [CHAR_LIMIT=80] Label for battery level chart when charge been limited -->
-    <string name="power_charging_limited"><xliff:g id="level">%1$s</xliff:g> - Charging paused</string>
+    <string name="power_charging_limited"><xliff:g id="level">%1$s</xliff:g> - Charging optimized</string>
     <!-- [CHAR_LIMIT=80] Label for battery charging future pause -->
-    <string name="power_charging_future_paused"><xliff:g id="level">%1$s</xliff:g> - Charging to <xliff:g id="dock_defender_threshold">%2$s</xliff:g></string>
+    <string name="power_charging_future_paused"><xliff:g id="level">%1$s</xliff:g> - Charging optimized</string>
 
     <!-- Battery Info screen. Value for a status item.  Used for diagnostic info screens, precise translation isn't needed -->
     <string name="battery_info_status_unknown">Unknown</string>
diff --git a/packages/SystemUI/AndroidManifest.xml b/packages/SystemUI/AndroidManifest.xml
index 4f08a30..2f5b42f 100644
--- a/packages/SystemUI/AndroidManifest.xml
+++ b/packages/SystemUI/AndroidManifest.xml
@@ -886,7 +886,7 @@
                   android:showForAllUsers="true"
                   android:finishOnTaskLaunch="true"
                   android:launchMode="singleInstance"
-                  android:configChanges="screenSize|smallestScreenSize|screenLayout|keyboard|keyboardHidden"
+                  android:configChanges="screenSize|smallestScreenSize|screenLayout|orientation|keyboard|keyboardHidden"
                   android:visibleToInstantApps="true">
         </activity>
 
diff --git a/packages/SystemUI/res-keyguard/values/strings.xml b/packages/SystemUI/res-keyguard/values/strings.xml
index da485a9..f522167 100644
--- a/packages/SystemUI/res-keyguard/values/strings.xml
+++ b/packages/SystemUI/res-keyguard/values/strings.xml
@@ -53,7 +53,7 @@
     <string name="keyguard_plugged_in_charging_slowly"><xliff:g id="percentage">%s</xliff:g> • Charging slowly</string>
 
     <!-- When the lock screen is showing and the phone plugged in, and the defend mode is triggered, say that charging is temporarily limited.  -->
-    <string name="keyguard_plugged_in_charging_limited"><xliff:g id="percentage">%s</xliff:g> • Charging paused to protect battery</string>
+    <string name="keyguard_plugged_in_charging_limited"><xliff:g id="percentage">%s</xliff:g> • Charging optimized to protect battery</string>
 
     <!-- On the keyguard screen, when pattern lock is disabled, only tell them to press menu to unlock.  This is shown in small font at the bottom. -->
     <string name="keyguard_instructions_when_pattern_disabled">Press Menu to unlock.</string>
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/rotation/RotationButtonController.java b/packages/SystemUI/shared/src/com/android/systemui/shared/rotation/RotationButtonController.java
index 6087655..8ee893c 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/rotation/RotationButtonController.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/rotation/RotationButtonController.java
@@ -71,7 +71,7 @@
  */
 public class RotationButtonController {
 
-    private static final String TAG = "StatusBar/RotationButtonController";
+    private static final String TAG = "RotationButtonController";
     private static final int BUTTON_FADE_IN_OUT_DURATION_MS = 100;
     private static final int NAVBAR_HIDDEN_PENDING_ICON_TIMEOUT_MS = 20000;
     private static final Interpolator LINEAR_INTERPOLATOR = new LinearInterpolator();
@@ -377,6 +377,7 @@
         }
 
         // Prepare to show the navbar icon by updating the icon style to change anim params
+        Log.i(TAG, "onRotationProposal(rotation=" + rotation + ")");
         mLastRotationSuggestion = rotation; // Remember rotation for click
         final boolean rotationCCW = Utilities.isRotationAnimationCCW(windowRotation, rotation);
         if (windowRotation == Surface.ROTATION_0 || windowRotation == Surface.ROTATION_180) {
@@ -499,6 +500,7 @@
         mUiEventLogger.log(RotationButtonEvent.ROTATION_SUGGESTION_ACCEPTED);
         incrementNumAcceptedRotationSuggestionsIfNeeded();
         setRotationLockedAtAngle(mLastRotationSuggestion);
+        Log.i(TAG, "onRotateSuggestionClick() mLastRotationSuggestion=" + mLastRotationSuggestion);
         v.performHapticFeedback(HapticFeedbackConstants.VIRTUAL_KEY);
     }
 
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
index ec4b780..8283ce8 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
@@ -480,7 +480,7 @@
     }
 
     @Override
-    public void onTrustChanged(boolean enabled, int userId, int flags,
+    public void onTrustChanged(boolean enabled, boolean newlyUnlocked, int userId, int flags,
             List<String> trustGrantedMessages) {
         Assert.isMainThread();
         boolean wasTrusted = mUserHasTrust.get(userId, false);
diff --git a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
index 4b7e9a5..98ac2c0 100644
--- a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
+++ b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
@@ -16,31 +16,25 @@
 
 package com.android.keyguard.mediator
 
+import android.annotation.BinderThread
 import android.os.Trace
-
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.keyguard.ScreenLifecycle
-import com.android.systemui.util.concurrency.Execution
-import com.android.systemui.util.concurrency.PendingTasksContainer
 import com.android.systemui.unfold.SysUIUnfoldComponent
+import com.android.systemui.util.concurrency.PendingTasksContainer
 import com.android.systemui.util.kotlin.getOrNull
-
 import java.util.Optional
-
 import javax.inject.Inject
 
 /**
  * Coordinates screen on/turning on animations for the KeyguardViewMediator. Specifically for
  * screen on events, this will invoke the onDrawn Runnable after all tasks have completed. This
- * should route back to the KeyguardService, which informs the system_server that keyguard has
- * drawn.
+ * should route back to the [com.android.systemui.keyguard.KeyguardService], which informs
+ * the system_server that keyguard has drawn.
  */
 @SysUISingleton
 class ScreenOnCoordinator @Inject constructor(
-    screenLifecycle: ScreenLifecycle,
-    unfoldComponent: Optional<SysUIUnfoldComponent>,
-    private val execution: Execution
-) : ScreenLifecycle.Observer {
+    unfoldComponent: Optional<SysUIUnfoldComponent>
+) {
 
     private val unfoldLightRevealAnimation = unfoldComponent.map(
         SysUIUnfoldComponent::getUnfoldLightRevealOverlayAnimation).getOrNull()
@@ -48,15 +42,12 @@
         SysUIUnfoldComponent::getFoldAodAnimationController).getOrNull()
     private val pendingTasks = PendingTasksContainer()
 
-    init {
-        screenLifecycle.addObserver(this)
-    }
-
     /**
      * When turning on, registers tasks that may need to run before invoking [onDrawn].
+     * This is called on a binder thread from [com.android.systemui.keyguard.KeyguardService].
      */
-    override fun onScreenTurningOn(onDrawn: Runnable) {
-        execution.assertIsMainThread()
+    @BinderThread
+    fun onScreenTurningOn(onDrawn: Runnable) {
         Trace.beginSection("ScreenOnCoordinator#onScreenTurningOn")
 
         pendingTasks.reset()
@@ -68,11 +59,13 @@
         Trace.endSection()
     }
 
-    override fun onScreenTurnedOn() {
-        execution.assertIsMainThread()
-
+    /**
+     * Called when screen is fully turned on and screen on blocker is removed.
+     * This is called on a binder thread from [com.android.systemui.keyguard.KeyguardService].
+     */
+    @BinderThread
+    fun onScreenTurnedOn() {
         foldAodAnimationController?.onScreenTurnedOn()
-
         pendingTasks.reset()
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
index 9f4e7e2..ff0c821 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
@@ -451,6 +451,11 @@
     // TODO(b/261538825): Tracking Bug
     @JvmField
     val OUTPUT_SWITCHER_ADVANCED_LAYOUT = unreleasedFlag(2500, "output_switcher_advanced_layout")
+    @JvmField
+    val OUTPUT_SWITCHER_ROUTES_PROCESSING =
+        unreleasedFlag(2501, "output_switcher_routes_processing")
+    @JvmField
+    val OUTPUT_SWITCHER_DEVICE_STATUS = unreleasedFlag(2502, "output_switcher_device_status")
 
     // TODO(b259590361): Tracking bug
     val EXPERIMENTAL_FLAG = unreleasedFlag(2, "exp_flag_release")
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
index 822b1cf..757afb6 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
@@ -18,11 +18,8 @@
 
 import android.os.Handler;
 import android.os.Message;
-import android.os.RemoteException;
 import android.os.Trace;
-import android.util.Log;
 
-import com.android.internal.policy.IKeyguardDrawnCallback;
 import com.android.systemui.dagger.SysUISingleton;
 
 import javax.inject.Inject;
@@ -80,33 +77,10 @@
     private Handler mHandler = new Handler() {
         @Override
         public void handleMessage(Message msg) {
-            final Object obj = msg.obj;
             switch (msg.what) {
                 case SCREEN_TURNING_ON:
                     Trace.beginSection("KeyguardLifecyclesDispatcher#SCREEN_TURNING_ON");
-                    final String onDrawWaitingTraceTag =
-                            "Waiting for KeyguardDrawnCallback#onDrawn";
-                    int traceCookie = System.identityHashCode(msg);
-                    Trace.beginAsyncSection(onDrawWaitingTraceTag, traceCookie);
-                    // Ensure the drawn callback is only ever called once
-                    mScreenLifecycle.dispatchScreenTurningOn(new Runnable() {
-                            boolean mInvoked;
-                            @Override
-                            public void run() {
-                                if (obj == null) return;
-                                if (!mInvoked) {
-                                    mInvoked = true;
-                                    try {
-                                        Trace.endAsyncSection(onDrawWaitingTraceTag, traceCookie);
-                                        ((IKeyguardDrawnCallback) obj).onDrawn();
-                                    } catch (RemoteException e) {
-                                        Log.w(TAG, "Exception calling onDrawn():", e);
-                                    }
-                                } else {
-                                    Log.w(TAG, "KeyguardDrawnCallback#onDrawn() invoked > 1 times");
-                                }
-                            }
-                        });
+                    mScreenLifecycle.dispatchScreenTurningOn();
                     Trace.endSection();
                     break;
                 case SCREEN_TURNED_ON:
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index c332a0d..f4a1227 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -78,6 +78,7 @@
 import com.android.internal.policy.IKeyguardExitCallback;
 import com.android.internal.policy.IKeyguardService;
 import com.android.internal.policy.IKeyguardStateCallback;
+import com.android.keyguard.mediator.ScreenOnCoordinator;
 import com.android.systemui.SystemUIApplication;
 import com.android.wm.shell.transition.ShellTransitions;
 import com.android.wm.shell.transition.Transitions;
@@ -120,6 +121,7 @@
 
     private final KeyguardViewMediator mKeyguardViewMediator;
     private final KeyguardLifecyclesDispatcher mKeyguardLifecyclesDispatcher;
+    private final ScreenOnCoordinator mScreenOnCoordinator;
     private final ShellTransitions mShellTransitions;
 
     private static int newModeToLegacyMode(int newMode) {
@@ -283,10 +285,12 @@
     @Inject
     public KeyguardService(KeyguardViewMediator keyguardViewMediator,
                            KeyguardLifecyclesDispatcher keyguardLifecyclesDispatcher,
+                           ScreenOnCoordinator screenOnCoordinator,
                            ShellTransitions shellTransitions) {
         super();
         mKeyguardViewMediator = keyguardViewMediator;
         mKeyguardLifecyclesDispatcher = keyguardLifecyclesDispatcher;
+        mScreenOnCoordinator = screenOnCoordinator;
         mShellTransitions = shellTransitions;
     }
 
@@ -583,6 +587,31 @@
             checkPermission();
             mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNING_ON,
                     callback);
+
+            final String onDrawWaitingTraceTag = "Waiting for KeyguardDrawnCallback#onDrawn";
+            final int traceCookie = System.identityHashCode(callback);
+            Trace.beginAsyncSection(onDrawWaitingTraceTag, traceCookie);
+
+            // Ensure the drawn callback is only ever called once
+            mScreenOnCoordinator.onScreenTurningOn(new Runnable() {
+                boolean mInvoked;
+                @Override
+                public void run() {
+                    if (callback == null) return;
+                    if (!mInvoked) {
+                        mInvoked = true;
+                        try {
+                            Trace.endAsyncSection(onDrawWaitingTraceTag, traceCookie);
+                            callback.onDrawn();
+                        } catch (RemoteException e) {
+                            Log.w(TAG, "Exception calling onDrawn():", e);
+                        }
+                    } else {
+                        Log.w(TAG, "KeyguardDrawnCallback#onDrawn() invoked > 1 times");
+                    }
+                }
+            });
+
             Trace.endSection();
         }
 
@@ -591,6 +620,7 @@
             Trace.beginSection("KeyguardService.mBinder#onScreenTurnedOn");
             checkPermission();
             mKeyguardLifecyclesDispatcher.dispatch(KeyguardLifecyclesDispatcher.SCREEN_TURNED_ON);
+            mScreenOnCoordinator.onScreenTurnedOn();
             Trace.endSection();
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 96ec43d..d231870 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -723,7 +723,7 @@
 
         @Override
         public void keyguardDone(boolean strongAuth, int targetUserId) {
-            if (targetUserId != mUserTracker.getUserId()) {
+            if (targetUserId != KeyguardUpdateMonitor.getCurrentUser()) {
                 return;
             }
             if (DEBUG) Log.d(TAG, "keyguardDone");
@@ -746,7 +746,7 @@
         public void keyguardDonePending(boolean strongAuth, int targetUserId) {
             Trace.beginSection("KeyguardViewMediator.mViewMediatorCallback#keyguardDonePending");
             if (DEBUG) Log.d(TAG, "keyguardDonePending");
-            if (targetUserId != mUserTracker.getUserId()) {
+            if (targetUserId != KeyguardUpdateMonitor.getCurrentUser()) {
                 Trace.endSection();
                 return;
             }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt b/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt
index 0a55294..0b04fb4 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt
@@ -46,7 +46,7 @@
         listeners.forEach(ScreenListener::onScreenTurningOff)
     }
 
-    override fun onScreenTurningOn(ignored: Runnable) {
+    override fun onScreenTurningOn() {
         listeners.forEach(ScreenListener::onScreenTurningOn)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java b/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
index b348121..8535eda 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ScreenLifecycle.java
@@ -18,8 +18,6 @@
 
 import android.os.Trace;
 
-import androidx.annotation.NonNull;
-
 import com.android.systemui.Dumpable;
 import com.android.systemui.dump.DumpManager;
 
@@ -50,14 +48,9 @@
         return mScreenState;
     }
 
-    /**
-     * Dispatch screen turning on events to the registered observers
-     *
-     * @param onDrawn Invoke to notify the caller that the event has been processed
-     */
-    public void dispatchScreenTurningOn(@NonNull Runnable onDrawn) {
+    public void dispatchScreenTurningOn() {
         setScreenState(SCREEN_TURNING_ON);
-        dispatch(Observer::onScreenTurningOn, onDrawn);
+        dispatch(Observer::onScreenTurningOn);
     }
 
     public void dispatchScreenTurnedOn() {
@@ -87,12 +80,7 @@
     }
 
     public interface Observer {
-        /**
-         * Receive the screen turning on event
-         *
-         * @param onDrawn Invoke to notify the caller that the event has been processed
-         */
-        default void onScreenTurningOn(@NonNull Runnable onDrawn) {}
+        default void onScreenTurningOn() {}
         default void onScreenTurnedOn() {}
         default void onScreenTurningOff() {}
         default void onScreenTurnedOff() {}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt
index 3218f96..5cb7d70 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/LockscreenBouncerTransitionInteractor.kt
@@ -22,7 +22,7 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
-import com.android.systemui.keyguard.shared.model.StatusBarState.SHADE_LOCKED
+import com.android.systemui.keyguard.shared.model.StatusBarState.KEYGUARD
 import com.android.systemui.keyguard.shared.model.TransitionInfo
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.WakefulnessState
@@ -50,27 +50,22 @@
 
     override fun start() {
         listenForDraggingUpToBouncer()
-        listenForBouncerHiding()
+        listenForBouncer()
     }
 
-    private fun listenForBouncerHiding() {
+    private fun listenForBouncer() {
         scope.launch {
             keyguardInteractor.isBouncerShowing
                 .sample(
                     combine(
                         keyguardInteractor.wakefulnessModel,
                         keyguardTransitionInteractor.startedKeyguardTransitionStep,
-                    ) { wakefulnessModel, transitionStep ->
-                        Pair(wakefulnessModel, transitionStep)
-                    }
-                ) { bouncerShowing, wakefulnessAndTransition ->
-                    Triple(
-                        bouncerShowing,
-                        wakefulnessAndTransition.first,
-                        wakefulnessAndTransition.second
-                    )
-                }
-                .collect { (isBouncerShowing, wakefulnessState, lastStartedTransitionStep) ->
+                        ::Pair
+                    ),
+                    ::toTriple
+                )
+                .collect { triple ->
+                    val (isBouncerShowing, wakefulnessState, lastStartedTransitionStep) = triple
                     if (
                         !isBouncerShowing && lastStartedTransitionStep.to == KeyguardState.BOUNCER
                     ) {
@@ -91,7 +86,19 @@
                                 animator = getAnimator(),
                             )
                         )
+                    } else if (
+                        isBouncerShowing && lastStartedTransitionStep.to == KeyguardState.LOCKSCREEN
+                    ) {
+                        keyguardTransitionRepository.startTransition(
+                            TransitionInfo(
+                                ownerName = name,
+                                from = KeyguardState.LOCKSCREEN,
+                                to = KeyguardState.BOUNCER,
+                                animator = getAnimator(),
+                            )
+                        )
                     }
+                    Unit
                 }
         }
     }
@@ -104,24 +111,20 @@
                     combine(
                         keyguardTransitionInteractor.finishedKeyguardState,
                         keyguardInteractor.statusBarState,
-                    ) { finishedKeyguardState, statusBarState ->
-                        Pair(finishedKeyguardState, statusBarState)
-                    }
-                ) { shadeModel, keyguardStateAndStatusBarState ->
-                    Triple(
-                        shadeModel,
-                        keyguardStateAndStatusBarState.first,
-                        keyguardStateAndStatusBarState.second
-                    )
-                }
-                .collect { (shadeModel, keyguardState, statusBarState) ->
+                        ::Pair
+                    ),
+                    ::toTriple
+                )
+                .collect { triple ->
+                    val (shadeModel, keyguardState, statusBarState) = triple
+
                     val id = transitionId
                     if (id != null) {
                         // An existing `id` means a transition is started, and calls to
                         // `updateTransition` will control it until FINISHED
                         keyguardTransitionRepository.updateTransition(
                             id,
-                            shadeModel.expansionAmount,
+                            1f - shadeModel.expansionAmount,
                             if (
                                 shadeModel.expansionAmount == 0f || shadeModel.expansionAmount == 1f
                             ) {
@@ -137,7 +140,7 @@
                         if (
                             keyguardState == KeyguardState.LOCKSCREEN &&
                                 shadeModel.isUserDragging &&
-                                statusBarState != SHADE_LOCKED
+                                statusBarState == KEYGUARD
                         ) {
                             transitionId =
                                 keyguardTransitionRepository.startTransition(
diff --git a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
index ec2e340..48a68be 100644
--- a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
@@ -123,7 +123,7 @@
     @SysUISingleton
     @QSLog
     public static LogBuffer provideQuickSettingsLogBuffer(LogBufferFactory factory) {
-        return factory.create("QSLog", 500 /* maxSize */, false /* systrace */);
+        return factory.create("QSLog", 700 /* maxSize */, false /* systrace */);
     }
 
     /** Provides a logging buffer for {@link com.android.systemui.broadcast.BroadcastDispatcher} */
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
index e2f55f0..8914552 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
@@ -464,6 +464,8 @@
 
     @Override
     public void dump(@NonNull PrintWriter pw, @NonNull String[] args) {
+        pw.println("mIsTablet=" + mIsTablet);
+        pw.println("mNavMode=" + mNavMode);
         for (int i = 0; i < mNavigationBars.size(); i++) {
             if (i > 0) {
                 pw.println();
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/OWNERS b/packages/SystemUI/src/com/android/systemui/notetask/OWNERS
new file mode 100644
index 0000000..7ccb316
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/notetask/OWNERS
@@ -0,0 +1,8 @@
+# Bug component: 1254381
+azappone@google.com
+achalke@google.com
+juliacr@google.com
+madym@google.com
+mgalhardo@google.com
+petrcermak@google.com
+vanjan@google.com
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java b/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java
index 8ceee1a..7523d6e 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/PagedTileLayout.java
@@ -11,7 +11,6 @@
 import android.content.res.Configuration;
 import android.os.Bundle;
 import android.util.AttributeSet;
-import android.util.Log;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
@@ -31,6 +30,7 @@
 import com.android.systemui.plugins.qs.QSTile;
 import com.android.systemui.qs.QSPanel.QSTileLayout;
 import com.android.systemui.qs.QSPanelControllerBase.TileRecord;
+import com.android.systemui.qs.logging.QSLogger;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -38,11 +38,9 @@
 
 public class PagedTileLayout extends ViewPager implements QSTileLayout {
 
-    private static final boolean DEBUG = false;
     private static final String CURRENT_PAGE = "current_page";
     private static final int NO_PAGE = -1;
 
-    private static final String TAG = "PagedTileLayout";
     private static final int REVEAL_SCROLL_DURATION_MILLIS = 750;
     private static final float BOUNCE_ANIMATION_TENSION = 1.3f;
     private static final long BOUNCE_ANIMATION_DURATION = 450L;
@@ -55,6 +53,7 @@
     private final ArrayList<TileRecord> mTiles = new ArrayList<>();
     private final ArrayList<TileLayout> mPages = new ArrayList<>();
 
+    private QSLogger mLogger;
     @Nullable
     private PageIndicator mPageIndicator;
     private float mPageIndicatorPosition;
@@ -146,9 +145,15 @@
         }
         if (mLayoutOrientation != newConfig.orientation) {
             mLayoutOrientation = newConfig.orientation;
-            mDistributeTiles = true;
+            forceTilesRedistribution("orientation changed to " + mLayoutOrientation);
             setCurrentItem(0, false);
             mPageToRestore = 0;
+        } else {
+            // logging in case we missed redistribution because orientation was not changed
+            // while configuration changed, can be removed after b/255208946 is fixed
+            mLogger.d(
+                    "Orientation didn't change, tiles might be not redistributed, new config",
+                    newConfig);
         }
     }
 
@@ -226,7 +231,7 @@
             // Keep on drawing until the animation has finished.
             postInvalidateOnAnimation();
         } catch (NullPointerException e) {
-            Log.e(TAG, "FakeDragBy called before begin", e);
+            mLogger.logException("FakeDragBy called before begin", e);
             // If we were trying to fake drag, it means we just added a new tile to the last
             // page, so animate there.
             final int lastPageNumber = mPages.size() - 1;
@@ -246,7 +251,7 @@
             super.endFakeDrag();
         } catch (NullPointerException e) {
             // Not sure what's going on. Let's log it
-            Log.e(TAG, "endFakeDrag called without velocityTracker", e);
+            mLogger.logException("endFakeDrag called without velocityTracker", e);
         }
     }
 
@@ -304,14 +309,14 @@
     @Override
     public void addTile(TileRecord tile) {
         mTiles.add(tile);
-        mDistributeTiles = true;
+        forceTilesRedistribution("adding new tile");
         requestLayout();
     }
 
     @Override
     public void removeTile(TileRecord tile) {
         if (mTiles.remove(tile)) {
-            mDistributeTiles = true;
+            forceTilesRedistribution("removing tile");
             requestLayout();
         }
     }
@@ -367,19 +372,11 @@
         final int tilesPerPageCount = mPages.get(0).maxTiles();
         int index = 0;
         final int totalTilesCount = mTiles.size();
-        if (DEBUG) {
-            Log.d(TAG, "Distributing tiles: "
-                    + "[tilesPerPageCount=" + tilesPerPageCount + "]"
-                    + "[totalTilesCount=" + totalTilesCount + "]"
-            );
-        }
+        mLogger.logTileDistributionInProgress(tilesPerPageCount, totalTilesCount);
         for (int i = 0; i < totalTilesCount; i++) {
             TileRecord tile = mTiles.get(i);
             if (mPages.get(index).mRecords.size() == tilesPerPageCount) index++;
-            if (DEBUG) {
-                Log.d(TAG, "Adding " + tile.tile.getClass().getSimpleName() + " to "
-                        + index);
-            }
+            mLogger.logTileDistributed(tile.tile.getClass().getSimpleName(), index);
             mPages.get(index).addTile(tile);
         }
     }
@@ -394,11 +391,11 @@
             return;
         }
         while (mPages.size() < numPages) {
-            if (DEBUG) Log.d(TAG, "Adding page");
+            mLogger.d("Adding new page");
             mPages.add(createTileLayout());
         }
         while (mPages.size() > numPages) {
-            if (DEBUG) Log.d(TAG, "Removing page");
+            mLogger.d("Removing page");
             mPages.remove(mPages.size() - 1);
         }
         mPageIndicator.setNumPages(mPages.size());
@@ -417,8 +414,12 @@
             changed |= mPages.get(i).updateResources();
         }
         if (changed) {
-            mDistributeTiles = true;
+            forceTilesRedistribution("resources in pages changed");
             requestLayout();
+        } else {
+            // logging in case we missed redistribution because number of column in updateResources
+            // was not changed, can be removed after b/255208946 is fixed
+            mLogger.d("resource in pages didn't change, tiles might be not redistributed");
         }
         return changed;
     }
@@ -430,7 +431,7 @@
         for (int i = 0; i < mPages.size(); i++) {
             if (mPages.get(i).setMinRows(minRows)) {
                 changed = true;
-                mDistributeTiles = true;
+                forceTilesRedistribution("minRows changed in page");
             }
         }
         return changed;
@@ -443,7 +444,7 @@
         for (int i = 0; i < mPages.size(); i++) {
             if (mPages.get(i).setMaxColumns(maxColumns)) {
                 changed = true;
-                mDistributeTiles = true;
+                forceTilesRedistribution("maxColumns in pages changed");
             }
         }
         return changed;
@@ -710,14 +711,14 @@
     private final PagerAdapter mAdapter = new PagerAdapter() {
         @Override
         public void destroyItem(ViewGroup container, int position, Object object) {
-            if (DEBUG) Log.d(TAG, "Destantiating " + position);
+            mLogger.d("Destantiating page at", position);
             container.removeView((View) object);
             updateListening();
         }
 
         @Override
         public Object instantiateItem(ViewGroup container, int position) {
-            if (DEBUG) Log.d(TAG, "Instantiating " + position);
+            mLogger.d("Instantiating page at", position);
             if (isLayoutRtl()) {
                 position = mPages.size() - 1 - position;
             }
@@ -745,10 +746,15 @@
      * Force all tiles to be redistributed across pages.
      * Should be called when one of the following changes: rows, columns, number of tiles.
      */
-    public void forceTilesRedistribution() {
+    public void forceTilesRedistribution(String reason) {
+        mLogger.d("forcing tile redistribution across pages, reason", reason);
         mDistributeTiles = true;
     }
 
+    public void setLogger(QSLogger qsLogger) {
+        mLogger = qsLogger;
+    }
+
     public interface PageListener {
         int INVALID_PAGE = -1;
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
index 6517ff3..7067c220 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanel.java
@@ -43,6 +43,7 @@
 import com.android.internal.widget.RemeasuringLinearLayout;
 import com.android.systemui.R;
 import com.android.systemui.plugins.qs.QSTile;
+import com.android.systemui.qs.logging.QSLogger;
 import com.android.systemui.settings.brightness.BrightnessSliderController;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.tuner.TunerService.Tunable;
@@ -106,6 +107,7 @@
     private ViewGroup mMediaHostView;
     private boolean mShouldMoveMediaOnExpansion = true;
     private boolean mUsingCombinedHeaders = false;
+    private QSLogger mQsLogger;
 
     public QSPanel(Context context, AttributeSet attrs) {
         super(context, attrs);
@@ -122,7 +124,8 @@
 
     }
 
-    void initialize() {
+    void initialize(QSLogger qsLogger) {
+        mQsLogger = qsLogger;
         mTileLayout = getOrCreateTileLayout();
 
         if (mUsingMediaPlayer) {
@@ -206,6 +209,7 @@
         if (mTileLayout == null) {
             mTileLayout = (QSTileLayout) LayoutInflater.from(mContext)
                     .inflate(R.layout.qs_paged_tile_layout, this, false);
+            mTileLayout.setLogger(mQsLogger);
             mTileLayout.setSquishinessFraction(mSquishinessFraction);
         }
         return mTileLayout;
@@ -735,6 +739,8 @@
         default void setExpansion(float expansion, float proposedTranslation) {}
 
         int getNumVisibleTiles();
+
+        default void setLogger(QSLogger qsLogger) { }
     }
 
     interface OnConfigurationChangedListener {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java
index b2ca6b7..cabe1da 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanelController.java
@@ -122,9 +122,8 @@
         }
         switchTileLayout(true);
         mBrightnessMirrorHandler.onQsPanelAttached();
-
-        ((PagedTileLayout) mView.getOrCreateTileLayout())
-                .setOnTouchListener(mTileLayoutTouchListener);
+        PagedTileLayout pagedTileLayout= ((PagedTileLayout) mView.getOrCreateTileLayout());
+        pagedTileLayout.setOnTouchListener(mTileLayoutTouchListener);
     }
 
     @Override
@@ -150,7 +149,8 @@
 
     @Override
     protected void onSplitShadeChanged() {
-        ((PagedTileLayout) mView.getOrCreateTileLayout()).forceTilesRedistribution();
+        ((PagedTileLayout) mView.getOrCreateTileLayout())
+                .forceTilesRedistribution("Split shade state changed");
     }
 
     /** */
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java b/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java
index 60d2c17..7bb672c 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSPanelControllerBase.java
@@ -70,7 +70,7 @@
     protected final MediaHost mMediaHost;
     protected final MetricsLogger mMetricsLogger;
     private final UiEventLogger mUiEventLogger;
-    private final QSLogger mQSLogger;
+    protected final QSLogger mQSLogger;
     private final DumpManager mDumpManager;
     protected final ArrayList<TileRecord> mRecords = new ArrayList<>();
     protected boolean mShouldUseSplitNotificationShade;
@@ -152,7 +152,7 @@
 
     @Override
     protected void onInit() {
-        mView.initialize();
+        mView.initialize(mQSLogger);
         mQSLogger.logAllTilesChangeListening(mView.isListening(), mView.getDumpableTag(), "");
     }
 
@@ -430,6 +430,7 @@
             pw.println("  horizontal layout: " + mUsingHorizontalLayout);
             pw.println("  last orientation: " + mLastOrientation);
         }
+        pw.println("  mShouldUseSplitNotificationShade: " + mShouldUseSplitNotificationShade);
     }
 
     public QSPanel.QSTileLayout getTileLayout() {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/logging/QSLogger.kt b/packages/SystemUI/src/com/android/systemui/qs/logging/QSLogger.kt
index 9f6317f..d682853f5 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/logging/QSLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/logging/QSLogger.kt
@@ -21,10 +21,13 @@
 import com.android.systemui.plugins.log.LogBuffer
 import com.android.systemui.plugins.log.LogLevel
 import com.android.systemui.plugins.log.LogLevel.DEBUG
+import com.android.systemui.plugins.log.LogLevel.ERROR
 import com.android.systemui.plugins.log.LogLevel.VERBOSE
+import com.android.systemui.plugins.log.LogLevel.WARNING
 import com.android.systemui.plugins.log.LogMessage
 import com.android.systemui.plugins.qs.QSTile
 import com.android.systemui.statusbar.StatusBarState
+import com.google.errorprone.annotations.CompileTimeConstant
 import javax.inject.Inject
 
 private const val TAG = "QSLog"
@@ -33,6 +36,26 @@
     @QSLog private val buffer: LogBuffer
 ) {
 
+    fun d(@CompileTimeConstant msg: String) = buffer.log(TAG, DEBUG, msg)
+
+    fun e(@CompileTimeConstant msg: String) = buffer.log(TAG, ERROR, msg)
+
+    fun v(@CompileTimeConstant msg: String) = buffer.log(TAG, VERBOSE, msg)
+
+    fun w(@CompileTimeConstant msg: String) = buffer.log(TAG, WARNING, msg)
+
+    fun logException(@CompileTimeConstant logMsg: String, ex: Exception) {
+        buffer.log(TAG, ERROR, {}, { logMsg }, exception = ex)
+    }
+
+    fun v(@CompileTimeConstant msg: String, arg: Any) {
+        buffer.log(TAG, VERBOSE, { str1 = arg.toString() }, { "$msg: $str1" })
+    }
+
+    fun d(@CompileTimeConstant msg: String, arg: Any) {
+        buffer.log(TAG, DEBUG, { str1 = arg.toString() }, { "$msg: $str1" })
+    }
+
     fun logTileAdded(tileSpec: String) {
         log(DEBUG, {
             str1 = tileSpec
@@ -236,6 +259,24 @@
         })
     }
 
+    fun logTileDistributionInProgress(tilesPerPageCount: Int, totalTilesCount: Int) {
+        log(DEBUG, {
+            int1 = tilesPerPageCount
+            int2 = totalTilesCount
+        }, {
+            "Distributing tiles: [tilesPerPageCount=$int1] [totalTilesCount=$int2]"
+        })
+    }
+
+    fun logTileDistributed(tileName: String, pageIndex: Int) {
+        log(DEBUG, {
+            str1 = tileName
+            int1 = pageIndex
+        }, {
+            "Adding $str1 to page number $int1"
+        })
+    }
+
     private fun toStateString(state: Int): String {
         return when (state) {
             Tile.STATE_ACTIVE -> "active"
diff --git a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
index 00d129a..4d005be 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
@@ -889,7 +889,7 @@
          * Notifies the Launcher that screen is starting to turn on.
          */
         @Override
-        public void onScreenTurningOn(@NonNull Runnable ignored) {
+        public void onScreenTurningOn() {
             try {
                 if (mOverviewProxy != null) {
                     mOverviewProxy.onScreenTurningOn();
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeEventsModule.java b/packages/SystemUI/src/com/android/systemui/shade/ShadeEventsModule.java
index 959c339..f87a1ed 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeEventsModule.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeEventsModule.java
@@ -16,14 +16,17 @@
 
 package com.android.systemui.shade;
 
-import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.shade.data.repository.ShadeRepository;
+import com.android.systemui.shade.data.repository.ShadeRepositoryImpl;
 
 import dagger.Binds;
 import dagger.Module;
 
-/** Provides a {@link ShadeStateEvents} in {@link SysUISingleton} scope. */
+/** Provides Shade-related events and information. */
 @Module
 public abstract class ShadeEventsModule {
     @Binds
     abstract ShadeStateEvents bindShadeEvents(ShadeExpansionStateManager impl);
+
+    @Binds abstract ShadeRepository shadeRepository(ShadeRepositoryImpl impl);
 }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
index 09019a6..bf7a2be 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
@@ -27,11 +27,17 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
 
+interface ShadeRepository {
+    /** ShadeModel information regarding shade expansion events */
+    val shadeModel: Flow<ShadeModel>
+}
+
 /** Business logic for shade interactions */
 @SysUISingleton
-class ShadeRepository @Inject constructor(shadeExpansionStateManager: ShadeExpansionStateManager) {
-
-    val shadeModel: Flow<ShadeModel> =
+class ShadeRepositoryImpl
+@Inject
+constructor(shadeExpansionStateManager: ShadeExpansionStateManager) : ShadeRepository {
+    override val shadeModel: Flow<ShadeModel> =
         conflatedCallbackFlow {
                 val callback =
                     object : ShadeExpansionListener {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
index 0eb0000..dc9b416 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
@@ -306,9 +306,12 @@
  */
 class RoundableState(
     internal val targetView: View,
-    roundable: Roundable,
-    internal val maxRadius: Float,
+    private val roundable: Roundable,
+    maxRadius: Float,
 ) {
+    internal var maxRadius = maxRadius
+        private set
+
     /** Animatable for top roundness */
     private val topAnimatable = topAnimatable(roundable)
 
@@ -356,6 +359,13 @@
         PropertyAnimator.setProperty(targetView, bottomAnimatable, value, DURATION, animated)
     }
 
+    fun setMaxRadius(radius: Float) {
+        if (maxRadius != radius) {
+            maxRadius = radius
+            roundable.applyRoundnessAndInvalidate()
+        }
+    }
+
     fun debugString() = buildString {
         append("TargetView: ${targetView.hashCode()} ")
         append("Top: $topRoundness ")
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
index 0380fff..1fcf17f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemory.kt
@@ -17,10 +17,14 @@
 
 package com.android.systemui.statusbar.notification.logging
 
+import android.app.Notification
+
 /** Describes usage of a notification. */
 data class NotificationMemoryUsage(
     val packageName: String,
+    val uid: Int,
     val notificationKey: String,
+    val notification: Notification,
     val objectUsage: NotificationObjectUsage,
     val viewUsage: List<NotificationViewUsage>
 )
@@ -34,7 +38,8 @@
     val smallIcon: Int,
     val largeIcon: Int,
     val extras: Int,
-    val style: String?,
+    /** Style type, integer from [android.stats.sysui.NotificationEnums] */
+    val style: Int,
     val styleIcon: Int,
     val bigPicture: Int,
     val extender: Int,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryDumper.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryDumper.kt
new file mode 100644
index 0000000..ffd931c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryDumper.kt
@@ -0,0 +1,173 @@
+/*
+ *
+ * Copyright (C) 2022 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.systemui.statusbar.notification.logging
+
+import android.stats.sysui.NotificationEnums
+import com.android.systemui.Dumpable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dump.DumpManager
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import java.io.PrintWriter
+import javax.inject.Inject
+
+/** Dumps current notification memory use to bug reports for easier debugging. */
+@SysUISingleton
+class NotificationMemoryDumper
+@Inject
+constructor(val dumpManager: DumpManager, val notificationPipeline: NotifPipeline) : Dumpable {
+
+    fun init() {
+        dumpManager.registerNormalDumpable(javaClass.simpleName, this)
+    }
+
+    override fun dump(pw: PrintWriter, args: Array<out String>) {
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(notificationPipeline.allNotifs)
+                .sortedWith(compareBy({ it.packageName }, { it.notificationKey }))
+        dumpNotificationObjects(pw, memoryUse)
+        dumpNotificationViewUsage(pw, memoryUse)
+    }
+
+    /** Renders a table of notification object usage into passed [PrintWriter]. */
+    private fun dumpNotificationObjects(pw: PrintWriter, memoryUse: List<NotificationMemoryUsage>) {
+        pw.println("Notification Object Usage")
+        pw.println("-----------")
+        pw.println(
+            "Package".padEnd(35) +
+                "\t\tSmall\tLarge\t${"Style".padEnd(15)}\t\tStyle\tBig\tExtend.\tExtras\tCustom"
+        )
+        pw.println("".padEnd(35) + "\t\tIcon\tIcon\t${"".padEnd(15)}\t\tIcon\tPicture\t \t \tView")
+        pw.println()
+
+        memoryUse.forEach { use ->
+            pw.println(
+                use.packageName.padEnd(35) +
+                    "\t\t" +
+                    "${use.objectUsage.smallIcon}\t${use.objectUsage.largeIcon}\t" +
+                    (styleEnumToString(use.objectUsage.style).take(15) ?: "").padEnd(15) +
+                    "\t\t${use.objectUsage.styleIcon}\t" +
+                    "${use.objectUsage.bigPicture}\t${use.objectUsage.extender}\t" +
+                    "${use.objectUsage.extras}\t${use.objectUsage.hasCustomView}\t" +
+                    use.notificationKey
+            )
+        }
+
+        // Calculate totals for easily glanceable summary.
+        data class Totals(
+            var smallIcon: Int = 0,
+            var largeIcon: Int = 0,
+            var styleIcon: Int = 0,
+            var bigPicture: Int = 0,
+            var extender: Int = 0,
+            var extras: Int = 0,
+        )
+
+        val totals =
+            memoryUse.fold(Totals()) { t, usage ->
+                t.smallIcon += usage.objectUsage.smallIcon
+                t.largeIcon += usage.objectUsage.largeIcon
+                t.styleIcon += usage.objectUsage.styleIcon
+                t.bigPicture += usage.objectUsage.bigPicture
+                t.extender += usage.objectUsage.extender
+                t.extras += usage.objectUsage.extras
+                t
+            }
+
+        pw.println()
+        pw.println("TOTALS")
+        pw.println(
+            "".padEnd(35) +
+                "\t\t" +
+                "${toKb(totals.smallIcon)}\t${toKb(totals.largeIcon)}\t" +
+                "".padEnd(15) +
+                "\t\t${toKb(totals.styleIcon)}\t" +
+                "${toKb(totals.bigPicture)}\t${toKb(totals.extender)}\t" +
+                toKb(totals.extras)
+        )
+        pw.println()
+    }
+
+    /** Renders a table of notification view usage into passed [PrintWriter] */
+    private fun dumpNotificationViewUsage(
+        pw: PrintWriter,
+        memoryUse: List<NotificationMemoryUsage>,
+    ) {
+
+        data class Totals(
+            var smallIcon: Int = 0,
+            var largeIcon: Int = 0,
+            var style: Int = 0,
+            var customViews: Int = 0,
+            var softwareBitmapsPenalty: Int = 0,
+        )
+
+        val totals = Totals()
+        pw.println("Notification View Usage")
+        pw.println("-----------")
+        pw.println("View Type".padEnd(24) + "\tSmall\tLarge\tStyle\tCustom\tSoftware")
+        pw.println("".padEnd(24) + "\tIcon\tIcon\tUse\tView\tBitmaps")
+        pw.println()
+        memoryUse
+            .filter { it.viewUsage.isNotEmpty() }
+            .forEach { use ->
+                pw.println(use.packageName + " " + use.notificationKey)
+                use.viewUsage.forEach { view ->
+                    pw.println(
+                        "  ${view.viewType.toString().padEnd(24)}\t${view.smallIcon}" +
+                            "\t${view.largeIcon}\t${view.style}" +
+                            "\t${view.customViews}\t${view.softwareBitmapsPenalty}"
+                    )
+
+                    if (view.viewType == ViewType.TOTAL) {
+                        totals.smallIcon += view.smallIcon
+                        totals.largeIcon += view.largeIcon
+                        totals.style += view.style
+                        totals.customViews += view.customViews
+                        totals.softwareBitmapsPenalty += view.softwareBitmapsPenalty
+                    }
+                }
+            }
+        pw.println()
+        pw.println("TOTALS")
+        pw.println(
+            "  ${"".padEnd(24)}\t${toKb(totals.smallIcon)}" +
+                "\t${toKb(totals.largeIcon)}\t${toKb(totals.style)}" +
+                "\t${toKb(totals.customViews)}\t${toKb(totals.softwareBitmapsPenalty)}"
+        )
+        pw.println()
+    }
+
+    private fun styleEnumToString(styleEnum: Int): String =
+        when (styleEnum) {
+            NotificationEnums.STYLE_UNSPECIFIED -> "Unspecified"
+            NotificationEnums.STYLE_NONE -> "None"
+            NotificationEnums.STYLE_BIG_PICTURE -> "BigPicture"
+            NotificationEnums.STYLE_BIG_TEXT -> "BigText"
+            NotificationEnums.STYLE_CALL -> "Call"
+            NotificationEnums.STYLE_DECORATED_CUSTOM_VIEW -> "DCustomView"
+            NotificationEnums.STYLE_INBOX -> "Inbox"
+            NotificationEnums.STYLE_MEDIA -> "Media"
+            NotificationEnums.STYLE_MESSAGING -> "Messaging"
+            NotificationEnums.STYLE_RANKER_GROUP -> "RankerGroup"
+            else -> "Unknown"
+        }
+
+    private fun toKb(bytes: Int): String {
+        return (bytes / 1024).toString() + " KB"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryLogger.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryLogger.kt
new file mode 100644
index 0000000..ec8501a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryLogger.kt
@@ -0,0 +1,194 @@
+/*
+ *
+ * Copyright (C) 2022 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.systemui.statusbar.notification.logging
+
+import android.app.StatsManager
+import android.util.StatsEvent
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.shared.system.SysUiStatsLog
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.util.traceSection
+import java.util.concurrent.Executor
+import javax.inject.Inject
+import kotlin.math.roundToInt
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.runBlocking
+
+/** Periodically logs current state of notification memory consumption. */
+@SysUISingleton
+class NotificationMemoryLogger
+@Inject
+constructor(
+    private val notificationPipeline: NotifPipeline,
+    private val statsManager: StatsManager,
+    @Main private val mainDispatcher: CoroutineDispatcher,
+    @Background private val backgroundExecutor: Executor
+) : StatsManager.StatsPullAtomCallback {
+
+    /**
+     * This class is used to accumulate and aggregate data - the fields mirror values in statd Atom
+     * with ONE IMPORTANT difference - the values are in bytes, not KB!
+     */
+    internal data class NotificationMemoryUseAtomBuilder(val uid: Int, val style: Int) {
+        var count: Int = 0
+        var countWithInflatedViews: Int = 0
+        var smallIconObject: Int = 0
+        var smallIconBitmapCount: Int = 0
+        var largeIconObject: Int = 0
+        var largeIconBitmapCount: Int = 0
+        var bigPictureObject: Int = 0
+        var bigPictureBitmapCount: Int = 0
+        var extras: Int = 0
+        var extenders: Int = 0
+        var smallIconViews: Int = 0
+        var largeIconViews: Int = 0
+        var systemIconViews: Int = 0
+        var styleViews: Int = 0
+        var customViews: Int = 0
+        var softwareBitmaps: Int = 0
+        var seenCount = 0
+    }
+
+    fun init() {
+        statsManager.setPullAtomCallback(
+            SysUiStatsLog.NOTIFICATION_MEMORY_USE,
+            null,
+            backgroundExecutor,
+            this
+        )
+    }
+
+    /** Called by statsd to pull data. */
+    override fun onPullAtom(atomTag: Int, data: MutableList<StatsEvent>): Int =
+        traceSection("NML#onPullAtom") {
+            if (atomTag != SysUiStatsLog.NOTIFICATION_MEMORY_USE) {
+                return StatsManager.PULL_SKIP
+            }
+
+            // Notifications can only be retrieved on the main thread, so switch to that thread.
+            val notifications = getAllNotificationsOnMainThread()
+            val notificationMemoryUse =
+                NotificationMemoryMeter.notificationMemoryUse(notifications)
+                    .sortedWith(
+                        compareBy(
+                            { it.packageName },
+                            { it.objectUsage.style },
+                            { it.notificationKey }
+                        )
+                    )
+            val usageData = aggregateMemoryUsageData(notificationMemoryUse)
+            usageData.forEach { (_, use) ->
+                data.add(
+                    SysUiStatsLog.buildStatsEvent(
+                        SysUiStatsLog.NOTIFICATION_MEMORY_USE,
+                        use.uid,
+                        use.style,
+                        use.count,
+                        use.countWithInflatedViews,
+                        toKb(use.smallIconObject),
+                        use.smallIconBitmapCount,
+                        toKb(use.largeIconObject),
+                        use.largeIconBitmapCount,
+                        toKb(use.bigPictureObject),
+                        use.bigPictureBitmapCount,
+                        toKb(use.extras),
+                        toKb(use.extenders),
+                        toKb(use.smallIconViews),
+                        toKb(use.largeIconViews),
+                        toKb(use.systemIconViews),
+                        toKb(use.styleViews),
+                        toKb(use.customViews),
+                        toKb(use.softwareBitmaps),
+                        use.seenCount
+                    )
+                )
+            }
+
+            return StatsManager.PULL_SUCCESS
+        }
+
+    private fun getAllNotificationsOnMainThread() =
+        runBlocking(mainDispatcher) {
+            traceSection("NML#getNotifications") { notificationPipeline.allNotifs }
+        }
+
+    /** Aggregates memory usage data by package and style, returning sums. */
+    private fun aggregateMemoryUsageData(
+        notificationMemoryUse: List<NotificationMemoryUsage>
+    ): Map<Pair<String, Int>, NotificationMemoryUseAtomBuilder> {
+        return notificationMemoryUse
+            .groupingBy { Pair(it.packageName, it.objectUsage.style) }
+            .aggregate {
+                _,
+                accumulator: NotificationMemoryUseAtomBuilder?,
+                element: NotificationMemoryUsage,
+                first ->
+                val use =
+                    if (first) {
+                        NotificationMemoryUseAtomBuilder(element.uid, element.objectUsage.style)
+                    } else {
+                        accumulator!!
+                    }
+
+                use.count++
+                // If the views of the notification weren't inflated, the list of memory usage
+                // parameters will be empty.
+                if (element.viewUsage.isNotEmpty()) {
+                    use.countWithInflatedViews++
+                }
+
+                use.smallIconObject += element.objectUsage.smallIcon
+                if (element.objectUsage.smallIcon > 0) {
+                    use.smallIconBitmapCount++
+                }
+
+                use.largeIconObject += element.objectUsage.largeIcon
+                if (element.objectUsage.largeIcon > 0) {
+                    use.largeIconBitmapCount++
+                }
+
+                use.bigPictureObject += element.objectUsage.bigPicture
+                if (element.objectUsage.bigPicture > 0) {
+                    use.bigPictureBitmapCount++
+                }
+
+                use.extras += element.objectUsage.extras
+                use.extenders += element.objectUsage.extender
+
+                // Use totals count which are more accurate when aggregated
+                // in this manner.
+                element.viewUsage
+                    .firstOrNull { vu -> vu.viewType == ViewType.TOTAL }
+                    ?.let {
+                        use.smallIconViews += it.smallIcon
+                        use.largeIconViews += it.largeIcon
+                        use.systemIconViews += it.systemIcons
+                        use.styleViews += it.style
+                        use.customViews += it.style
+                        use.softwareBitmaps += it.softwareBitmapsPenalty
+                    }
+
+                return@aggregate use
+            }
+    }
+
+    /** Rounds the passed value to the nearest KB - e.g. 700B rounds to 1KB. */
+    private fun toKb(value: Int): Int = (value.toFloat() / 1024f).roundToInt()
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt
index 7d39e18..41fb91e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeter.kt
@@ -1,12 +1,20 @@
 package com.android.systemui.statusbar.notification.logging
 
 import android.app.Notification
+import android.app.Notification.BigPictureStyle
+import android.app.Notification.BigTextStyle
+import android.app.Notification.CallStyle
+import android.app.Notification.DecoratedCustomViewStyle
+import android.app.Notification.InboxStyle
+import android.app.Notification.MediaStyle
+import android.app.Notification.MessagingStyle
 import android.app.Person
 import android.graphics.Bitmap
 import android.graphics.drawable.Icon
 import android.os.Bundle
 import android.os.Parcel
 import android.os.Parcelable
+import android.stats.sysui.NotificationEnums
 import androidx.annotation.WorkerThread
 import com.android.systemui.statusbar.notification.NotificationUtils
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
@@ -19,6 +27,7 @@
     private const val TV_EXTENSIONS = "android.tv.EXTENSIONS"
     private const val WEARABLE_EXTENSIONS = "android.wearable.EXTENSIONS"
     private const val WEARABLE_EXTENSIONS_BACKGROUND = "background"
+    private const val AUTOGROUP_KEY = "ranker_group"
 
     /** Returns a list of memory use entries for currently shown notifications. */
     @WorkerThread
@@ -29,12 +38,15 @@
             .asSequence()
             .map { entry ->
                 val packageName = entry.sbn.packageName
+                val uid = entry.sbn.uid
                 val notificationObjectUsage =
                     notificationMemoryUse(entry.sbn.notification, hashSetOf())
                 val notificationViewUsage = NotificationMemoryViewWalker.getViewUsage(entry.row)
                 NotificationMemoryUsage(
                     packageName,
+                    uid,
                     NotificationUtils.logKey(entry.sbn.key),
+                    entry.sbn.notification,
                     notificationObjectUsage,
                     notificationViewUsage
                 )
@@ -49,7 +61,9 @@
     ): NotificationMemoryUsage {
         return NotificationMemoryUsage(
             entry.sbn.packageName,
+            entry.sbn.uid,
             NotificationUtils.logKey(entry.sbn.key),
+            entry.sbn.notification,
             notificationMemoryUse(entry.sbn.notification, seenBitmaps),
             NotificationMemoryViewWalker.getViewUsage(entry.row)
         )
@@ -116,7 +130,13 @@
         val wearExtenderBackground =
             computeParcelableUse(wearExtender, WEARABLE_EXTENSIONS_BACKGROUND, seenBitmaps)
 
-        val style = notification.notificationStyle
+        val style =
+            if (notification.group == AUTOGROUP_KEY) {
+                NotificationEnums.STYLE_RANKER_GROUP
+            } else {
+                styleEnum(notification.notificationStyle)
+            }
+
         val hasCustomView = notification.contentView != null || notification.bigContentView != null
         val extrasSize = computeBundleSize(extras)
 
@@ -124,7 +144,7 @@
             smallIcon = smallIconUse,
             largeIcon = largeIconUse,
             extras = extrasSize,
-            style = style?.simpleName,
+            style = style,
             styleIcon =
                 bigPictureIconUse +
                     peopleUse +
@@ -144,6 +164,25 @@
     }
 
     /**
+     * Returns logging style enum based on current style class.
+     *
+     * @return style value in [NotificationEnums]
+     */
+    private fun styleEnum(style: Class<out Notification.Style>?): Int =
+        when (style?.name) {
+            null -> NotificationEnums.STYLE_NONE
+            BigTextStyle::class.java.name -> NotificationEnums.STYLE_BIG_TEXT
+            BigPictureStyle::class.java.name -> NotificationEnums.STYLE_BIG_PICTURE
+            InboxStyle::class.java.name -> NotificationEnums.STYLE_INBOX
+            MediaStyle::class.java.name -> NotificationEnums.STYLE_MEDIA
+            DecoratedCustomViewStyle::class.java.name ->
+                NotificationEnums.STYLE_DECORATED_CUSTOM_VIEW
+            MessagingStyle::class.java.name -> NotificationEnums.STYLE_MESSAGING
+            CallStyle::class.java.name -> NotificationEnums.STYLE_CALL
+            else -> NotificationEnums.STYLE_UNSPECIFIED
+        }
+
+    /**
      * Calculates size of the bundle data (excluding FDs and other shared objects like ashmem
      * bitmaps). Can be slow.
      */
@@ -176,7 +215,7 @@
      *
      * @return memory usage in bytes or 0 if the icon is Uri/Resource based
      */
-    private fun computeIconUse(icon: Icon?, seenBitmaps: HashSet<Int>) =
+    private fun computeIconUse(icon: Icon?, seenBitmaps: HashSet<Int>): Int =
         when (icon?.type) {
             Icon.TYPE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
             Icon.TYPE_ADAPTIVE_BITMAP -> computeBitmapUse(icon.bitmap, seenBitmaps)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
index c09cc43..f38c1e5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMonitor.kt
@@ -18,11 +18,10 @@
 package com.android.systemui.statusbar.notification.logging
 
 import android.util.Log
-import com.android.systemui.Dumpable
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dump.DumpManager
-import com.android.systemui.statusbar.notification.collection.NotifPipeline
-import java.io.PrintWriter
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
+import dagger.Lazy
 import javax.inject.Inject
 
 /** This class monitors and logs current Notification memory use. */
@@ -30,9 +29,10 @@
 class NotificationMemoryMonitor
 @Inject
 constructor(
-    val notificationPipeline: NotifPipeline,
-    val dumpManager: DumpManager,
-) : Dumpable {
+    private val featureFlags: FeatureFlags,
+    private val notificationMemoryDumper: NotificationMemoryDumper,
+    private val notificationMemoryLogger: Lazy<NotificationMemoryLogger>,
+) {
 
     companion object {
         private const val TAG = "NotificationMemory"
@@ -40,127 +40,10 @@
 
     fun init() {
         Log.d(TAG, "NotificationMemoryMonitor initialized.")
-        dumpManager.registerDumpable(javaClass.simpleName, this)
-    }
-
-    override fun dump(pw: PrintWriter, args: Array<out String>) {
-        val memoryUse =
-            NotificationMemoryMeter.notificationMemoryUse(notificationPipeline.allNotifs)
-                .sortedWith(compareBy({ it.packageName }, { it.notificationKey }))
-        dumpNotificationObjects(pw, memoryUse)
-        dumpNotificationViewUsage(pw, memoryUse)
-    }
-
-    /** Renders a table of notification object usage into passed [PrintWriter]. */
-    private fun dumpNotificationObjects(pw: PrintWriter, memoryUse: List<NotificationMemoryUsage>) {
-        pw.println("Notification Object Usage")
-        pw.println("-----------")
-        pw.println(
-            "Package".padEnd(35) +
-                "\t\tSmall\tLarge\t${"Style".padEnd(15)}\t\tStyle\tBig\tExtend.\tExtras\tCustom"
-        )
-        pw.println("".padEnd(35) + "\t\tIcon\tIcon\t${"".padEnd(15)}\t\tIcon\tPicture\t \t \tView")
-        pw.println()
-
-        memoryUse.forEach { use ->
-            pw.println(
-                use.packageName.padEnd(35) +
-                    "\t\t" +
-                    "${use.objectUsage.smallIcon}\t${use.objectUsage.largeIcon}\t" +
-                    (use.objectUsage.style?.take(15) ?: "").padEnd(15) +
-                    "\t\t${use.objectUsage.styleIcon}\t" +
-                    "${use.objectUsage.bigPicture}\t${use.objectUsage.extender}\t" +
-                    "${use.objectUsage.extras}\t${use.objectUsage.hasCustomView}\t" +
-                    use.notificationKey
-            )
+        notificationMemoryDumper.init()
+        if (featureFlags.isEnabled(Flags.NOTIFICATION_MEMORY_LOGGING_ENABLED)) {
+            Log.d(TAG, "Notification memory logging enabled.")
+            notificationMemoryLogger.get().init()
         }
-
-        // Calculate totals for easily glanceable summary.
-        data class Totals(
-            var smallIcon: Int = 0,
-            var largeIcon: Int = 0,
-            var styleIcon: Int = 0,
-            var bigPicture: Int = 0,
-            var extender: Int = 0,
-            var extras: Int = 0,
-        )
-
-        val totals =
-            memoryUse.fold(Totals()) { t, usage ->
-                t.smallIcon += usage.objectUsage.smallIcon
-                t.largeIcon += usage.objectUsage.largeIcon
-                t.styleIcon += usage.objectUsage.styleIcon
-                t.bigPicture += usage.objectUsage.bigPicture
-                t.extender += usage.objectUsage.extender
-                t.extras += usage.objectUsage.extras
-                t
-            }
-
-        pw.println()
-        pw.println("TOTALS")
-        pw.println(
-            "".padEnd(35) +
-                "\t\t" +
-                "${toKb(totals.smallIcon)}\t${toKb(totals.largeIcon)}\t" +
-                "".padEnd(15) +
-                "\t\t${toKb(totals.styleIcon)}\t" +
-                "${toKb(totals.bigPicture)}\t${toKb(totals.extender)}\t" +
-                toKb(totals.extras)
-        )
-        pw.println()
-    }
-
-    /** Renders a table of notification view usage into passed [PrintWriter] */
-    private fun dumpNotificationViewUsage(
-        pw: PrintWriter,
-        memoryUse: List<NotificationMemoryUsage>,
-    ) {
-
-        data class Totals(
-            var smallIcon: Int = 0,
-            var largeIcon: Int = 0,
-            var style: Int = 0,
-            var customViews: Int = 0,
-            var softwareBitmapsPenalty: Int = 0,
-        )
-
-        val totals = Totals()
-        pw.println("Notification View Usage")
-        pw.println("-----------")
-        pw.println("View Type".padEnd(24) + "\tSmall\tLarge\tStyle\tCustom\tSoftware")
-        pw.println("".padEnd(24) + "\tIcon\tIcon\tUse\tView\tBitmaps")
-        pw.println()
-        memoryUse
-            .filter { it.viewUsage.isNotEmpty() }
-            .forEach { use ->
-                pw.println(use.packageName + " " + use.notificationKey)
-                use.viewUsage.forEach { view ->
-                    pw.println(
-                        "  ${view.viewType.toString().padEnd(24)}\t${view.smallIcon}" +
-                            "\t${view.largeIcon}\t${view.style}" +
-                            "\t${view.customViews}\t${view.softwareBitmapsPenalty}"
-                    )
-
-                    if (view.viewType == ViewType.TOTAL) {
-                        totals.smallIcon += view.smallIcon
-                        totals.largeIcon += view.largeIcon
-                        totals.style += view.style
-                        totals.customViews += view.customViews
-                        totals.softwareBitmapsPenalty += view.softwareBitmapsPenalty
-                    }
-                }
-            }
-        pw.println()
-        pw.println("TOTALS")
-        pw.println(
-            "  ${"".padEnd(24)}\t${toKb(totals.smallIcon)}" +
-                "\t${toKb(totals.largeIcon)}\t${toKb(totals.style)}" +
-                "\t${toKb(totals.customViews)}\t${toKb(totals.softwareBitmapsPenalty)}"
-        )
-        pw.println()
-    }
-
-    private fun toKb(bytes: Int): String {
-        return (bytes / 1024).toString() + " KB"
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
index a0bee15..2d04211 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
@@ -50,7 +50,11 @@
 
     /**
      * Returns memory usage of public and private views contained in passed
-     * [ExpandableNotificationRow]
+     * [ExpandableNotificationRow]. Each entry will correspond to one of the [ViewType] values with
+     * [ViewType.TOTAL] totalling all memory use. If a type of view is missing, the corresponding
+     * entry will not appear in resulting list.
+     *
+     * This will return an empty list if the ExpandableNotificationRow has no views inflated.
      */
     fun getViewUsage(row: ExpandableNotificationRow?): List<NotificationViewUsage> {
         if (row == null) {
@@ -58,42 +62,72 @@
         }
 
         // The ordering here is significant since it determines deduplication of seen drawables.
-        return listOf(
-            getViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, row.privateLayout?.expandedChild),
-            getViewUsage(ViewType.PRIVATE_CONTRACTED_VIEW, row.privateLayout?.contractedChild),
-            getViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, row.privateLayout?.headsUpChild),
-            getViewUsage(ViewType.PUBLIC_VIEW, row.publicLayout),
-            getTotalUsage(row)
-        )
+        val perViewUsages =
+            listOf(
+                    getViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, row.privateLayout?.expandedChild),
+                    getViewUsage(
+                        ViewType.PRIVATE_CONTRACTED_VIEW,
+                        row.privateLayout?.contractedChild
+                    ),
+                    getViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, row.privateLayout?.headsUpChild),
+                    getViewUsage(
+                        ViewType.PUBLIC_VIEW,
+                        row.publicLayout?.expandedChild,
+                        row.publicLayout?.contractedChild,
+                        row.publicLayout?.headsUpChild
+                    ),
+                )
+                .filterNotNull()
+
+        return if (perViewUsages.isNotEmpty()) {
+            // Attach summed totals field only if there was any view actually measured.
+            // This reduces bug report noise and makes checks for collapsed views easier.
+            val totals = getTotalUsage(row)
+            if (totals == null) {
+                perViewUsages
+            } else {
+                perViewUsages + totals
+            }
+        } else {
+            listOf()
+        }
     }
 
     /**
      * Calculate total usage of all views - we need to do a separate traversal to make sure we don't
      * double count fields.
      */
-    private fun getTotalUsage(row: ExpandableNotificationRow): NotificationViewUsage {
-        val totalUsage = UsageBuilder()
+    private fun getTotalUsage(row: ExpandableNotificationRow): NotificationViewUsage? {
         val seenObjects = hashSetOf<Int>()
-
-        row.publicLayout?.let { computeViewHierarchyUse(it, totalUsage, seenObjects) }
-        row.privateLayout?.let { child ->
-            for (view in listOf(child.expandedChild, child.contractedChild, child.headsUpChild)) {
-                (view as? ViewGroup)?.let { v ->
-                    computeViewHierarchyUse(v, totalUsage, seenObjects)
-                }
-            }
-        }
-        return totalUsage.build(ViewType.TOTAL)
+        return getViewUsage(
+            ViewType.TOTAL,
+            row.privateLayout?.expandedChild,
+            row.privateLayout?.contractedChild,
+            row.privateLayout?.headsUpChild,
+            row.publicLayout?.expandedChild,
+            row.publicLayout?.contractedChild,
+            row.publicLayout?.headsUpChild,
+            seenObjects = seenObjects
+        )
     }
 
     private fun getViewUsage(
         type: ViewType,
-        rootView: View?,
+        vararg rootViews: View?,
         seenObjects: HashSet<Int> = hashSetOf()
-    ): NotificationViewUsage {
-        val usageBuilder = UsageBuilder()
-        (rootView as? ViewGroup)?.let { computeViewHierarchyUse(it, usageBuilder, seenObjects) }
-        return usageBuilder.build(type)
+    ): NotificationViewUsage? {
+        val usageBuilder = lazy { UsageBuilder() }
+        rootViews.forEach { rootView ->
+            (rootView as? ViewGroup)?.let { rootViewGroup ->
+                computeViewHierarchyUse(rootViewGroup, usageBuilder.value, seenObjects)
+            }
+        }
+
+        return if (usageBuilder.isInitialized()) {
+            usageBuilder.value.build(type)
+        } else {
+            null
+        }
     }
 
     private fun computeViewHierarchyUse(
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
index 0213b96..2041245 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
@@ -214,7 +214,11 @@
         } else {
             maxRadius = res.getDimensionPixelSize(R.dimen.notification_corner_radius);
         }
-        mRoundableState = new RoundableState(this, this, maxRadius);
+        if (mRoundableState == null) {
+            mRoundableState = new RoundableState(this, this, maxRadius);
+        } else {
+            mRoundableState.setMaxRadius(maxRadius);
+        }
         setClipToOutline(mAlwaysRoundBothCorners);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
index 005cd1b..93da0b8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
@@ -3574,7 +3574,7 @@
 
     final ScreenLifecycle.Observer mScreenObserver = new ScreenLifecycle.Observer() {
         @Override
-        public void onScreenTurningOn(Runnable onDrawn) {
+        public void onScreenTurningOn() {
             mFalsingCollector.onScreenTurningOn();
             mNotificationPanelViewController.onScreenTurningOn();
         }
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt b/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt
index 6216acd..101bd44 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/FoldAodAnimationController.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.unfold
 
+import android.annotation.BinderThread
 import android.content.Context
 import android.hardware.devicestate.DeviceStateManager
 import android.os.PowerManager
@@ -41,7 +42,6 @@
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
-import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.launch
 
 /**
@@ -52,7 +52,7 @@
 class FoldAodAnimationController
 @Inject
 constructor(
-    @Main private val executor: DelayableExecutor,
+    @Main private val mainExecutor: DelayableExecutor,
     private val context: Context,
     private val deviceStateManager: DeviceStateManager,
     private val wakefulnessLifecycle: WakefulnessLifecycle,
@@ -89,7 +89,7 @@
     override fun initialize(centralSurfaces: CentralSurfaces, lightRevealScrim: LightRevealScrim) {
         this.centralSurfaces = centralSurfaces
 
-        deviceStateManager.registerCallback(executor, FoldListener())
+        deviceStateManager.registerCallback(mainExecutor, FoldListener())
         wakefulnessLifecycle.addObserver(this)
 
         // TODO(b/254878364): remove this call to NPVC.getView()
@@ -139,7 +139,8 @@
      * @param onReady callback when the animation is ready
      * @see [com.android.systemui.keyguard.KeyguardViewMediator]
      */
-    fun onScreenTurningOn(onReady: Runnable) {
+    @BinderThread
+    fun onScreenTurningOn(onReady: Runnable) = mainExecutor.execute {
         if (shouldPlayAnimation) {
             // The device was not dozing and going to sleep after folding, play the animation
 
@@ -179,12 +180,13 @@
         }
     }
 
-    fun onScreenTurnedOn() {
+    @BinderThread
+    fun onScreenTurnedOn() = mainExecutor.execute {
         if (shouldPlayAnimation) {
             cancelAnimation?.run()
 
             // Post starting the animation to the next frame to avoid junk due to inset changes
-            cancelAnimation = executor.executeDelayed(startAnimationRunnable, /* delayMillis= */ 0)
+            cancelAnimation = mainExecutor.executeDelayed(startAnimationRunnable, /* delayMillis= */ 0)
             shouldPlayAnimation = false
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
index b2ec27c..9cca950 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
@@ -15,6 +15,7 @@
  */
 package com.android.systemui.unfold
 
+import android.annotation.BinderThread
 import android.content.ContentResolver
 import android.content.Context
 import android.graphics.PixelFormat
@@ -34,9 +35,7 @@
 import android.view.SurfaceSession
 import android.view.WindowManager
 import android.view.WindowlessWindowManager
-import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
-import com.android.systemui.dagger.qualifiers.UiBackground
 import com.android.systemui.statusbar.LightRevealEffect
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.statusbar.LinearLightRevealEffect
@@ -45,7 +44,7 @@
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
 import com.android.systemui.unfold.updates.RotationChangeProvider
 import com.android.systemui.unfold.util.ScaleAwareTransitionProgressProvider.Companion.areAnimationsEnabled
-import com.android.systemui.util.Assert.isMainThread
+import com.android.systemui.util.concurrency.ThreadFactory
 import com.android.systemui.util.traceSection
 import com.android.wm.shell.displayareahelper.DisplayAreaHelper
 import java.util.Optional
@@ -64,14 +63,16 @@
     private val unfoldTransitionProgressProvider: UnfoldTransitionProgressProvider,
     private val displayAreaHelper: Optional<DisplayAreaHelper>,
     @Main private val executor: Executor,
-    @UiBackground private val backgroundExecutor: Executor,
-    @Background private val bgHandler: Handler,
+    private val threadFactory: ThreadFactory,
     private val rotationChangeProvider: RotationChangeProvider,
 ) {
 
     private val transitionListener = TransitionListener()
     private val rotationWatcher = RotationWatcher()
 
+    private lateinit var bgHandler: Handler
+    private lateinit var bgExecutor: Executor
+
     private lateinit var wwm: WindowlessWindowManager
     private lateinit var unfoldedDisplayInfo: DisplayInfo
     private lateinit var overlayContainer: SurfaceControl
@@ -84,7 +85,12 @@
     private var currentRotation: Int = context.display!!.rotation
 
     fun init() {
-        deviceStateManager.registerCallback(executor, FoldListener())
+        // This method will be called only on devices where this animation is enabled,
+        // so normally this thread won't be created
+        bgHandler = threadFactory.buildHandlerOnNewThread(TAG)
+        bgExecutor = threadFactory.buildDelayableExecutorOnHandler(bgHandler)
+
+        deviceStateManager.registerCallback(bgExecutor, FoldListener())
         unfoldTransitionProgressProvider.addCallback(transitionListener)
         rotationChangeProvider.addCallback(rotationWatcher)
 
@@ -122,20 +128,23 @@
      * @param onOverlayReady callback when the overlay is drawn and visible on the screen
      * @see [com.android.systemui.keyguard.KeyguardViewMediator]
      */
+    @BinderThread
     fun onScreenTurningOn(onOverlayReady: Runnable) {
-        Trace.beginSection("UnfoldLightRevealOverlayAnimation#onScreenTurningOn")
-        try {
-            // Add the view only if we are unfolding and this is the first screen on
-            if (!isFolded && !isUnfoldHandled && contentResolver.areAnimationsEnabled()) {
-                executeInBackground { addOverlay(onOverlayReady, reason = UNFOLD) }
-                isUnfoldHandled = true
-            } else {
-                // No unfold transition, immediately report that overlay is ready
-                executeInBackground { ensureOverlayRemoved() }
-                onOverlayReady.run()
+        executeInBackground {
+            Trace.beginSection("$TAG#onScreenTurningOn")
+            try {
+                // Add the view only if we are unfolding and this is the first screen on
+                if (!isFolded && !isUnfoldHandled && contentResolver.areAnimationsEnabled()) {
+                    addOverlay(onOverlayReady, reason = UNFOLD)
+                    isUnfoldHandled = true
+                } else {
+                    // No unfold transition, immediately report that overlay is ready
+                    ensureOverlayRemoved()
+                    onOverlayReady.run()
+                }
+            } finally {
+                Trace.endSection()
             }
-        } finally {
-            Trace.endSection()
         }
     }
 
@@ -154,17 +163,18 @@
             LightRevealScrim(context, null).apply {
                 revealEffect = createLightRevealEffect()
                 isScrimOpaqueChangedListener = Consumer {}
-                revealAmount = when (reason) {
-                    FOLD -> TRANSPARENT
-                    UNFOLD -> BLACK
-                }
+                revealAmount =
+                    when (reason) {
+                        FOLD -> TRANSPARENT
+                        UNFOLD -> BLACK
+                    }
             }
 
         val params = getLayoutParams()
         newRoot.setView(newView, params)
 
         if (onOverlayReady != null) {
-            Trace.beginAsyncSection("UnfoldLightRevealOverlayAnimation#relayout", 0)
+            Trace.beginAsyncSection("$TAG#relayout", 0)
 
             newRoot.relayout(params) { transaction ->
                 val vsyncId = Choreographer.getSfInstance().vsyncId
@@ -179,8 +189,8 @@
 
                 transaction
                     .setFrameTimelineVsync(vsyncId + 1)
-                    .addTransactionCommittedListener(backgroundExecutor) {
-                        Trace.endAsyncSection("UnfoldLightRevealOverlayAnimation#relayout", 0)
+                    .addTransactionCommittedListener(bgExecutor) {
+                        Trace.endAsyncSection("$TAG#relayout", 0)
                         onOverlayReady.run()
                     }
                     .apply()
@@ -233,7 +243,8 @@
     }
 
     private fun getUnfoldedDisplayInfo(): DisplayInfo =
-        displayManager.getDisplays(DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED)
+        displayManager
+            .getDisplays(DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED)
             .asSequence()
             .map { DisplayInfo().apply { it.getDisplayInfo(this) } }
             .filter { it.type == Display.TYPE_INTERNAL }
@@ -261,10 +272,10 @@
 
     private inner class RotationWatcher : RotationChangeProvider.RotationListener {
         override fun onRotationChanged(newRotation: Int) {
-            traceSection("UnfoldLightRevealOverlayAnimation#onRotationChanged") {
-                if (currentRotation != newRotation) {
-                    currentRotation = newRotation
-                    executeInBackground {
+            executeInBackground {
+                traceSection("$TAG#onRotationChanged") {
+                    if (currentRotation != newRotation) {
+                        currentRotation = newRotation
                         scrimView?.revealEffect = createLightRevealEffect()
                         root?.relayout(getLayoutParams())
                     }
@@ -274,7 +285,10 @@
     }
 
     private fun executeInBackground(f: () -> Unit) {
-        ensureInMainThread()
+        check(Looper.myLooper() != bgHandler.looper) {
+            "Trying to execute using background handler while already running" +
+                " in the background handler"
+        }
         // The UiBackground executor is not used as it doesn't have a prepared looper.
         bgHandler.post(f)
     }
@@ -283,25 +297,25 @@
         check(Looper.myLooper() == bgHandler.looper) { "Not being executed in the background!" }
     }
 
-    private fun ensureInMainThread() {
-        isMainThread()
-    }
-
     private inner class FoldListener :
         FoldStateListener(
             context,
             Consumer { isFolded ->
                 if (isFolded) {
-                    executeInBackground { ensureOverlayRemoved() }
+                    ensureOverlayRemoved()
                     isUnfoldHandled = false
                 }
                 this.isFolded = isFolded
             }
         )
 
-    private enum class AddOverlayReason { FOLD, UNFOLD }
+    private enum class AddOverlayReason {
+        FOLD,
+        UNFOLD
+    }
 
     private companion object {
+        const val TAG = "UnfoldLightRevealOverlayAnimation"
         const val ROTATION_ANIMATION_OVERLAY_Z_INDEX = Integer.MAX_VALUE
 
         // Put the unfold overlay below the rotation animation screenshot to hide the moment
diff --git a/packages/SystemUI/src/com/android/systemui/user/UserSwitcherActivity.kt b/packages/SystemUI/src/com/android/systemui/user/UserSwitcherActivity.kt
index 7da2d47..52b7fb6 100644
--- a/packages/SystemUI/src/com/android/systemui/user/UserSwitcherActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/user/UserSwitcherActivity.kt
@@ -39,9 +39,9 @@
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         setContentView(R.layout.user_switcher_fullscreen)
-        window.decorView.getWindowInsetsController().apply {
-            setSystemBarsBehavior(BEHAVIOR_SHOW_TRANSIENT_BARS_BY_SWIPE)
-            hide(Type.systemBars())
+        window.decorView.windowInsetsController?.let { controller ->
+            controller.systemBarsBehavior = BEHAVIOR_SHOW_TRANSIENT_BARS_BY_SWIPE
+            controller.hide(Type.systemBars())
         }
         val viewModel =
             ViewModelProvider(this, viewModelFactory.get())[UserSwitcherViewModel::class.java]
diff --git a/packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt b/packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt
index 6cd384f..ceebcb7 100644
--- a/packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt
@@ -25,8 +25,11 @@
  */
 class PendingTasksContainer {
 
-    private var pendingTasksCount: AtomicInteger = AtomicInteger(0)
-    private var completionCallback: AtomicReference<Runnable> = AtomicReference()
+    @Volatile
+    private var pendingTasksCount = AtomicInteger(0)
+
+    @Volatile
+    private var completionCallback = AtomicReference<Runnable>()
 
     /**
      * Registers a task that we should wait for
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
index 849ff08..2aef726 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
@@ -823,7 +823,7 @@
         mKeyguardUpdateMonitor.dispatchStartedWakingUp(PowerManager.WAKE_REASON_POWER_BUTTON);
         mTestableLooper.processAllMessages();
         lockscreenBypassIsAllowed();
-        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */,
+        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */, true /* newlyUnlocked */,
                 KeyguardUpdateMonitor.getCurrentUser(), 0 /* flags */,
                 new ArrayList<>());
         keyguardIsVisible();
@@ -834,7 +834,7 @@
     public void testIgnoresAuth_whenTrustAgentOnKeyguard_withoutBypass() {
         mKeyguardUpdateMonitor.dispatchStartedWakingUp(PowerManager.WAKE_REASON_POWER_BUTTON);
         mTestableLooper.processAllMessages();
-        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */,
+        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */, true /* newlyUnlocked */,
                 KeyguardUpdateMonitor.getCurrentUser(), 0 /* flags */, new ArrayList<>());
         keyguardIsVisible();
         verify(mFaceManager, never()).authenticate(any(), any(), any(), any(), anyInt(),
@@ -1049,8 +1049,8 @@
     @Test
     public void testGetUserCanSkipBouncer_whenTrust() {
         int user = KeyguardUpdateMonitor.getCurrentUser();
-        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */, user, 0 /* flags */,
-                new ArrayList<>());
+        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */, true /* newlyUnlocked */,
+                user, 0 /* flags */, new ArrayList<>());
         assertThat(mKeyguardUpdateMonitor.getUserCanSkipBouncer(user)).isTrue();
     }
 
@@ -1314,7 +1314,7 @@
         when(mStrongAuthTracker.hasUserAuthenticatedSinceBoot()).thenReturn(true);
 
         // WHEN trust is enabled (ie: via smartlock)
-        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */,
+        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */, true /* newlyUnlocked */,
                 KeyguardUpdateMonitor.getCurrentUser(), 0 /* flags */, new ArrayList<>());
 
         // THEN we shouldn't listen for udfps
@@ -1418,7 +1418,7 @@
     @Test
     public void testShowTrustGrantedMessage_onTrustGranted() {
         // WHEN trust is enabled (ie: via some trust agent) with a trustGranted string
-        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */,
+        mKeyguardUpdateMonitor.onTrustChanged(true /* enabled */, true /* newlyUnlocked */,
                 KeyguardUpdateMonitor.getCurrentUser(), 0 /* flags */,
                 Arrays.asList("Unlocked by wearable"));
 
@@ -1870,6 +1870,7 @@
         // WHEN onTrustChanged with TRUST_DISMISS_KEYGUARD flag
         mKeyguardUpdateMonitor.onTrustChanged(
                 true /* enabled */,
+                true /* newlyUnlocked */,
                 getCurrentUser() /* userId */,
                 TrustAgentService.FLAG_GRANT_TRUST_DISMISS_KEYGUARD /* flags */,
                 null /* trustGrantedMessages */);
@@ -1893,6 +1894,7 @@
         // WHEN onTrustChanged with TRUST_DISMISS_KEYGUARD flag
         mKeyguardUpdateMonitor.onTrustChanged(
                 true /* enabled */,
+                true /* newlyUnlocked */,
                 getCurrentUser() /* userId */,
                 TrustAgentService.FLAG_GRANT_TRUST_DISMISS_KEYGUARD /* flags */,
                 null /* trustGrantedMessages */);
@@ -1917,6 +1919,7 @@
         // WHEN onTrustChanged for a different user
         mKeyguardUpdateMonitor.onTrustChanged(
                 true /* enabled */,
+                true /* newlyUnlocked */,
                 546 /* userId, not the current userId */,
                 0 /* flags */,
                 null /* trustGrantedMessages */);
@@ -1941,6 +1944,7 @@
         // flags (temporary & rewable is active unlock)
         mKeyguardUpdateMonitor.onTrustChanged(
                 true /* enabled */,
+                true /* newlyUnlocked */,
                 getCurrentUser() /* userId */,
                 TrustAgentService.FLAG_GRANT_TRUST_DISMISS_KEYGUARD
                         | TrustAgentService.FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE /* flags */,
@@ -1968,6 +1972,7 @@
         // WHEN onTrustChanged with INITIATED_BY_USER flag
         mKeyguardUpdateMonitor.onTrustChanged(
                 true /* enabled */,
+                true /* newlyUnlocked */,
                 getCurrentUser() /* userId, not the current userId */,
                 TrustAgentService.FLAG_GRANT_TRUST_INITIATED_BY_USER /* flags */,
                 null /* trustGrantedMessages */);
@@ -1992,6 +1997,7 @@
         // WHEN onTrustChanged with INITIATED_BY_USER flag
         mKeyguardUpdateMonitor.onTrustChanged(
                 true /* enabled */,
+                true /* newlyUnlocked */,
                 getCurrentUser() /* userId, not the current userId */,
                 TrustAgentService.FLAG_GRANT_TRUST_INITIATED_BY_USER
                         | TrustAgentService.FLAG_GRANT_TRUST_TEMPORARY_AND_RENEWABLE /* flags */,
@@ -2310,6 +2316,7 @@
     private void currentUserDoesNotHaveTrust() {
         mKeyguardUpdateMonitor.onTrustChanged(
                 false,
+                false,
                 KeyguardUpdateMonitor.getCurrentUser(),
                 -1,
                 new ArrayList<>()
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
index 5734c3d..34e78eb 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
@@ -52,8 +52,6 @@
     private lateinit var foldAodAnimationController: FoldAodAnimationController
     @Mock
     private lateinit var unfoldAnimation: UnfoldLightRevealOverlayAnimation
-    @Mock
-    private lateinit var screenLifecycle: ScreenLifecycle
     @Captor
     private lateinit var readyCaptor: ArgumentCaptor<Runnable>
 
@@ -69,13 +67,8 @@
                 .thenReturn(foldAodAnimationController)
 
         screenOnCoordinator = ScreenOnCoordinator(
-            screenLifecycle,
             Optional.of(unfoldComponent),
-            FakeExecution()
         )
-
-        // Make sure screen events are registered to observe
-        verify(screenLifecycle).addObserver(screenOnCoordinator)
     }
 
     @Test
@@ -93,9 +86,7 @@
     fun testUnfoldTransitionDisabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback() {
         // Recreate with empty unfoldComponent
         screenOnCoordinator = ScreenOnCoordinator(
-            screenLifecycle,
             Optional.empty(),
-            FakeExecution()
         )
         screenOnCoordinator.onScreenTurningOn(runnable)
 
@@ -105,11 +96,11 @@
 
     private fun onUnfoldOverlayReady() {
         verify(unfoldAnimation).onScreenTurningOn(capture(readyCaptor))
-        readyCaptor.getValue().run()
+        readyCaptor.value.run()
     }
 
     private fun onFoldAodReady() {
         verify(foldAodAnimationController).onScreenTurningOn(capture(readyCaptor))
-        readyCaptor.getValue().run()
+        readyCaptor.value.run()
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java
index f46d58d..70a0415 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ScreenLifecycleTest.java
@@ -58,15 +58,15 @@
     @Test
     public void screenTurningOn() throws Exception {
         Runnable onDrawn = () -> {};
-        mScreen.dispatchScreenTurningOn(onDrawn);
+        mScreen.dispatchScreenTurningOn();
 
         assertEquals(ScreenLifecycle.SCREEN_TURNING_ON, mScreen.getScreenState());
-        verify(mScreenObserverMock).onScreenTurningOn(onDrawn);
+        verify(mScreenObserverMock).onScreenTurningOn();
     }
 
     @Test
     public void screenTurnedOn() throws Exception {
-        mScreen.dispatchScreenTurningOn(null);
+        mScreen.dispatchScreenTurningOn();
         mScreen.dispatchScreenTurnedOn();
 
         assertEquals(ScreenLifecycle.SCREEN_ON, mScreen.getScreenState());
@@ -75,7 +75,7 @@
 
     @Test
     public void screenTurningOff() throws Exception {
-        mScreen.dispatchScreenTurningOn(null);
+        mScreen.dispatchScreenTurningOn();
         mScreen.dispatchScreenTurnedOn();
         mScreen.dispatchScreenTurningOff();
 
@@ -85,7 +85,7 @@
 
     @Test
     public void screenTurnedOff() throws Exception {
-        mScreen.dispatchScreenTurningOn(null);
+        mScreen.dispatchScreenTurningOn();
         mScreen.dispatchScreenTurnedOn();
         mScreen.dispatchScreenTurningOff();
         mScreen.dispatchScreenTurnedOff();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
index ce9c1da..5d2f0eb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepositoryTest.kt
@@ -16,13 +16,10 @@
 
 package com.android.systemui.keyguard.data.repository
 
-import android.animation.AnimationHandler.AnimationFrameCallbackProvider
 import android.animation.ValueAnimator
 import android.util.Log
 import android.util.Log.TerribleFailure
 import android.util.Log.TerribleFailureHandler
-import android.view.Choreographer.FrameCallback
-import androidx.test.filters.FlakyTest
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.Interpolators
@@ -32,22 +29,17 @@
 import com.android.systemui.keyguard.shared.model.TransitionInfo
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.keyguard.util.KeyguardTransitionRunner
 import com.google.common.truth.Truth.assertThat
 import java.math.BigDecimal
 import java.math.RoundingMode
 import java.util.UUID
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.yield
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 import org.junit.After
-import org.junit.Assert.fail
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -60,12 +52,14 @@
     private lateinit var underTest: KeyguardTransitionRepository
     private lateinit var oldWtfHandler: TerribleFailureHandler
     private lateinit var wtfHandler: WtfHandler
+    private lateinit var runner: KeyguardTransitionRunner
 
     @Before
     fun setUp() {
         underTest = KeyguardTransitionRepositoryImpl()
         wtfHandler = WtfHandler()
         oldWtfHandler = Log.setWtfHandler(wtfHandler)
+        runner = KeyguardTransitionRunner(underTest)
     }
 
     @After
@@ -75,56 +69,37 @@
 
     @Test
     fun `startTransition runs animator to completion`() =
-        runBlocking(IMMEDIATE) {
-            val (animator, provider) = setupAnimator(this)
-
+        TestScope().runTest {
             val steps = mutableListOf<TransitionStep>()
             val job = underTest.transition(AOD, LOCKSCREEN).onEach { steps.add(it) }.launchIn(this)
 
-            underTest.startTransition(TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, animator))
-
-            val startTime = System.currentTimeMillis()
-            while (animator.isRunning()) {
-                yield()
-                if (System.currentTimeMillis() - startTime > MAX_TEST_DURATION) {
-                    fail("Failed test due to excessive runtime of: $MAX_TEST_DURATION")
-                }
-            }
+            runner.startTransition(
+                this,
+                TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
+                maxFrames = 100
+            )
 
             assertSteps(steps, listWithStep(BigDecimal(.1)), AOD, LOCKSCREEN)
-
             job.cancel()
-            provider.stop()
         }
 
     @Test
-    @FlakyTest(bugId = 260213291)
-    fun `starting second transition will cancel the first transition`() {
-        runBlocking(IMMEDIATE) {
-            val (animator, provider) = setupAnimator(this)
-
+    fun `starting second transition will cancel the first transition`() =
+        TestScope().runTest {
             val steps = mutableListOf<TransitionStep>()
             val job = underTest.transition(AOD, LOCKSCREEN).onEach { steps.add(it) }.launchIn(this)
-
-            underTest.startTransition(TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, animator))
-            // 3 yields(), alternating with the animator, results in a value 0.1, which can be
-            // canceled and tested against
-            yield()
-            yield()
-            yield()
+            runner.startTransition(
+                this,
+                TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, getAnimator()),
+                maxFrames = 3,
+            )
 
             // Now start 2nd transition, which will interrupt the first
             val job2 = underTest.transition(LOCKSCREEN, AOD).onEach { steps.add(it) }.launchIn(this)
-            val (animator2, provider2) = setupAnimator(this)
-            underTest.startTransition(TransitionInfo(OWNER_NAME, LOCKSCREEN, AOD, animator2))
-
-            val startTime = System.currentTimeMillis()
-            while (animator2.isRunning()) {
-                yield()
-                if (System.currentTimeMillis() - startTime > MAX_TEST_DURATION) {
-                    fail("Failed test due to excessive runtime of: $MAX_TEST_DURATION")
-                }
-            }
+            runner.startTransition(
+                this,
+                TransitionInfo(OWNER_NAME, LOCKSCREEN, AOD, getAnimator()),
+            )
 
             val firstTransitionSteps = listWithStep(step = BigDecimal(.1), stop = BigDecimal(.1))
             assertSteps(steps.subList(0, 4), firstTransitionSteps, AOD, LOCKSCREEN)
@@ -134,31 +109,25 @@
 
             job.cancel()
             job2.cancel()
-            provider.stop()
-            provider2.stop()
         }
-    }
 
     @Test
     fun `Null animator enables manual control with updateTransition`() =
-        runBlocking(IMMEDIATE) {
+        TestScope().runTest {
             val steps = mutableListOf<TransitionStep>()
             val job = underTest.transition(AOD, LOCKSCREEN).onEach { steps.add(it) }.launchIn(this)
 
             val uuid =
                 underTest.startTransition(
-                    TransitionInfo(
-                        ownerName = OWNER_NAME,
-                        from = AOD,
-                        to = LOCKSCREEN,
-                        animator = null,
-                    )
+                    TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, animator = null)
                 )
+            runCurrent()
 
             checkNotNull(uuid).let {
                 underTest.updateTransition(it, 0.5f, TransitionState.RUNNING)
                 underTest.updateTransition(it, 1f, TransitionState.FINISHED)
             }
+            runCurrent()
 
             assertThat(steps.size).isEqualTo(3)
             assertThat(steps[0])
@@ -256,57 +225,11 @@
         assertThat(wtfHandler.failed).isFalse()
     }
 
-    private fun setupAnimator(
-        scope: CoroutineScope
-    ): Pair<ValueAnimator, TestFrameCallbackProvider> {
-        val animator =
-            ValueAnimator().apply {
-                setInterpolator(Interpolators.LINEAR)
-                setDuration(ANIMATION_DURATION)
-            }
-
-        val provider = TestFrameCallbackProvider(animator, scope)
-        provider.start()
-
-        return Pair(animator, provider)
-    }
-
-    /** Gives direct control over ValueAnimator. See [AnimationHandler] */
-    private class TestFrameCallbackProvider(
-        private val animator: ValueAnimator,
-        private val scope: CoroutineScope,
-    ) : AnimationFrameCallbackProvider {
-
-        private var frameCount = 1L
-        private var frames = MutableStateFlow(Pair<Long, FrameCallback?>(0L, null))
-        private var job: Job? = null
-
-        fun start() {
-            animator.getAnimationHandler().setProvider(this)
-
-            job =
-                scope.launch {
-                    frames.collect {
-                        // Delay is required for AnimationHandler to properly register a callback
-                        yield()
-                        val (frameNumber, callback) = it
-                        callback?.doFrame(frameNumber)
-                    }
-                }
+    private fun getAnimator(): ValueAnimator {
+        return ValueAnimator().apply {
+            setInterpolator(Interpolators.LINEAR)
+            setDuration(10)
         }
-
-        fun stop() {
-            job?.cancel()
-            animator.getAnimationHandler().setProvider(null)
-        }
-
-        override fun postFrameCallback(cb: FrameCallback) {
-            frames.value = Pair(frameCount++, cb)
-        }
-        override fun postCommitCallback(runnable: Runnable) {}
-        override fun getFrameTime() = frameCount
-        override fun getFrameDelay() = 1L
-        override fun setFrameDelay(delay: Long) {}
     }
 
     private class WtfHandler : TerribleFailureHandler {
@@ -317,9 +240,6 @@
     }
 
     companion object {
-        private const val MAX_TEST_DURATION = 100L
-        private const val ANIMATION_DURATION = 10L
-        private const val OWNER_NAME = "Test"
-        private val IMMEDIATE = Dispatchers.Main.immediate
+        private const val OWNER_NAME = "KeyguardTransitionRunner"
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
new file mode 100644
index 0000000..a6cf840
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
@@ -0,0 +1,140 @@
+/*
+ * Copyright (C) 2022 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.systemui.keyguard.domain.interactor
+
+import android.animation.ValueAnimator
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.animation.Interpolators
+import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepositoryImpl
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.TransitionInfo
+import com.android.systemui.keyguard.shared.model.WakeSleepReason
+import com.android.systemui.keyguard.shared.model.WakefulnessModel
+import com.android.systemui.keyguard.shared.model.WakefulnessState
+import com.android.systemui.keyguard.util.KeyguardTransitionRunner
+import com.android.systemui.shade.data.repository.FakeShadeRepository
+import com.android.systemui.shade.data.repository.ShadeRepository
+import com.android.systemui.util.mockito.withArgCaptor
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.cancelChildren
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+/**
+ * Class for testing user journeys through the interactors. They will all be activated during setup,
+ * to ensure the expected transitions are still triggered.
+ */
+@SmallTest
+@RunWith(JUnit4::class)
+class KeyguardTransitionScenariosTest : SysuiTestCase() {
+    private lateinit var testScope: TestScope
+
+    private lateinit var keyguardRepository: FakeKeyguardRepository
+    private lateinit var shadeRepository: ShadeRepository
+
+    // Used to issue real transition steps for test input
+    private lateinit var runner: KeyguardTransitionRunner
+    private lateinit var transitionRepository: KeyguardTransitionRepository
+
+    // Used to verify transition requests for test output
+    @Mock private lateinit var mockTransitionRepository: KeyguardTransitionRepository
+
+    private lateinit var lockscreenBouncerTransitionInteractor:
+        LockscreenBouncerTransitionInteractor
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        testScope = TestScope()
+
+        keyguardRepository = FakeKeyguardRepository()
+        shadeRepository = FakeShadeRepository()
+
+        /* Used to issue full transition steps, to better simulate a real device */
+        transitionRepository = KeyguardTransitionRepositoryImpl()
+        runner = KeyguardTransitionRunner(transitionRepository)
+
+        lockscreenBouncerTransitionInteractor =
+            LockscreenBouncerTransitionInteractor(
+                scope = testScope,
+                keyguardInteractor = KeyguardInteractor(keyguardRepository),
+                shadeRepository = shadeRepository,
+                keyguardTransitionRepository = mockTransitionRepository,
+                keyguardTransitionInteractor = KeyguardTransitionInteractor(transitionRepository),
+            )
+        lockscreenBouncerTransitionInteractor.start()
+    }
+
+    @Test
+    fun `LOCKSCREEN to BOUNCER via bouncer showing call`() =
+        testScope.runTest {
+            // GIVEN a device that has at least woken up
+            keyguardRepository.setWakefulnessModel(startingToWake())
+            runCurrent()
+
+            // GIVEN a transition has run to LOCKSCREEN
+            runner.startTransition(
+                testScope,
+                TransitionInfo(
+                    ownerName = "",
+                    from = KeyguardState.OFF,
+                    to = KeyguardState.LOCKSCREEN,
+                    animator =
+                        ValueAnimator().apply {
+                            duration = 10
+                            interpolator = Interpolators.LINEAR
+                        },
+                )
+            )
+            runCurrent()
+
+            // WHEN the bouncer is set to show
+            keyguardRepository.setBouncerShowing(true)
+            runCurrent()
+
+            val info =
+                withArgCaptor<TransitionInfo> {
+                    verify(mockTransitionRepository).startTransition(capture())
+                }
+            // THEN a transition to BOUNCER should occur
+            assertThat(info.ownerName).isEqualTo("LockscreenBouncerTransitionInteractor")
+            assertThat(info.from).isEqualTo(KeyguardState.LOCKSCREEN)
+            assertThat(info.to).isEqualTo(KeyguardState.BOUNCER)
+            assertThat(info.animator).isNotNull()
+
+            coroutineContext.cancelChildren()
+        }
+
+    private fun startingToWake() =
+        WakefulnessModel(
+            WakefulnessState.STARTING_TO_WAKE,
+            true,
+            WakeSleepReason.OTHER,
+            WakeSleepReason.OTHER
+        )
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt
new file mode 100644
index 0000000..c88f84a
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/util/KeyguardTransitionRunner.kt
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2022 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.systemui.keyguard.util
+
+import android.animation.AnimationHandler.AnimationFrameCallbackProvider
+import android.animation.ValueAnimator
+import android.view.Choreographer.FrameCallback
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.TransitionInfo
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import org.junit.Assert.fail
+
+/**
+ * Gives direct control over ValueAnimator, in order to make transition tests deterministic. See
+ * [AnimationHandler]. Animators are required to be run on the main thread, so dispatch accordingly.
+ */
+class KeyguardTransitionRunner(
+    val repository: KeyguardTransitionRepository,
+) : AnimationFrameCallbackProvider {
+
+    private var frameCount = 1L
+    private var frames = MutableStateFlow(Pair<Long, FrameCallback?>(0L, null))
+    private var job: Job? = null
+    private var isTerminated = false
+
+    /**
+     * For transitions being directed by an animator. Will control the number of frames being
+     * generated so the values are deterministic.
+     */
+    suspend fun startTransition(scope: CoroutineScope, info: TransitionInfo, maxFrames: Int = 100) {
+        // AnimationHandler uses ThreadLocal storage, and ValueAnimators MUST start from main
+        // thread
+        withContext(Dispatchers.Main) {
+            info.animator!!.getAnimationHandler().setProvider(this@KeyguardTransitionRunner)
+        }
+
+        job =
+            scope.launch {
+                frames.collect {
+                    val (frameNumber, callback) = it
+
+                    isTerminated = frameNumber >= maxFrames
+                    if (!isTerminated) {
+                        withContext(Dispatchers.Main) { callback?.doFrame(frameNumber) }
+                    }
+                }
+            }
+        withContext(Dispatchers.Main) { repository.startTransition(info) }
+
+        waitUntilComplete(info.animator!!)
+    }
+
+    suspend private fun waitUntilComplete(animator: ValueAnimator) {
+        withContext(Dispatchers.Main) {
+            val startTime = System.currentTimeMillis()
+            while (!isTerminated && animator.isRunning()) {
+                delay(1)
+                if (System.currentTimeMillis() - startTime > MAX_TEST_DURATION) {
+                    fail("Failed test due to excessive runtime of: $MAX_TEST_DURATION")
+                }
+            }
+
+            animator.getAnimationHandler().setProvider(null)
+        }
+
+        job?.cancel()
+    }
+
+    override fun postFrameCallback(cb: FrameCallback) {
+        frames.value = Pair(frameCount++, cb)
+    }
+    override fun postCommitCallback(runnable: Runnable) {}
+    override fun getFrameTime() = frameCount
+    override fun getFrameDelay() = 1L
+    override fun setFrameDelay(delay: Long) {}
+
+    companion object {
+        private const val MAX_TEST_DURATION = 100L
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/OWNERS b/packages/SystemUI/tests/src/com/android/systemui/notetask/OWNERS
new file mode 100644
index 0000000..7ccb316
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/OWNERS
@@ -0,0 +1,8 @@
+# Bug component: 1254381
+azappone@google.com
+achalke@google.com
+juliacr@google.com
+madym@google.com
+mgalhardo@google.com
+petrcermak@google.com
+vanjan@google.com
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerBaseTest.java
index caf8321..5058373 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerBaseTest.java
@@ -226,7 +226,8 @@
                 + "    " + mockTileViewString + "\n"
                 + "  media bounds: null\n"
                 + "  horizontal layout: false\n"
-                + "  last orientation: 0\n";
+                + "  last orientation: 0\n"
+                + "  mShouldUseSplitNotificationShade: false\n";
         assertEquals(expected, w.getBuffer().toString());
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.kt
index 5e082f6..6cf642c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelControllerTest.kt
@@ -135,10 +135,10 @@
     fun configurationChange_onlySplitShadeConfigChanges_tileAreRedistributed() {
         testableResources.addOverride(R.bool.config_use_split_notification_shade, false)
         controller.mOnConfigurationChangedListener.onConfigurationChange(configuration)
-        verify(pagedTileLayout, never()).forceTilesRedistribution()
+        verify(pagedTileLayout, never()).forceTilesRedistribution(any())
 
         testableResources.addOverride(R.bool.config_use_split_notification_shade, true)
         controller.mOnConfigurationChangedListener.onConfigurationChange(configuration)
-        verify(pagedTileLayout).forceTilesRedistribution()
+        verify(pagedTileLayout).forceTilesRedistribution("Split shade state changed")
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelTest.kt
index 7c930b1..d52b296 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QSPanelTest.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.plugins.qs.QSTile
+import com.android.systemui.qs.logging.QSLogger
 import com.android.systemui.qs.tileimpl.QSIconViewImpl
 import com.android.systemui.qs.tileimpl.QSTileViewImpl
 import com.google.common.truth.Truth.assertThat
@@ -34,6 +35,7 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mock
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
@@ -42,6 +44,9 @@
 @RunWithLooper
 @SmallTest
 class QSPanelTest : SysuiTestCase() {
+
+    @Mock private lateinit var qsLogger: QSLogger
+
     private lateinit var testableLooper: TestableLooper
     private lateinit var qsPanel: QSPanel
 
@@ -57,7 +62,7 @@
             qsPanel = QSPanel(context, null)
             qsPanel.mUsingMediaPlayer = true
 
-            qsPanel.initialize()
+            qsPanel.initialize(qsLogger)
             // QSPanel inflates a footer inside of it, mocking it here
             footer = LinearLayout(context).apply { id = R.id.qs_footer }
             qsPanel.addView(footer, MATCH_PARENT, 100)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelTest.kt
index a6a584d..3fba393 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/QuickQSPanelTest.kt
@@ -7,10 +7,12 @@
 import android.widget.FrameLayout
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.qs.logging.QSLogger
 import com.google.common.truth.Truth
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mock
 import org.mockito.Mockito
 import org.mockito.MockitoAnnotations
 
@@ -19,6 +21,8 @@
 @SmallTest
 class QuickQSPanelTest : SysuiTestCase() {
 
+    @Mock private lateinit var qsLogger: QSLogger
+
     private lateinit var testableLooper: TestableLooper
     private lateinit var quickQSPanel: QuickQSPanel
 
@@ -32,7 +36,7 @@
 
         testableLooper.runWithLooper {
             quickQSPanel = QuickQSPanel(mContext, null)
-            quickQSPanel.initialize()
+            quickQSPanel.initialize(qsLogger)
 
             quickQSPanel.onFinishInflate()
             // Provides a parent with non-zero size for QSPanel
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryLoggerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryLoggerTest.kt
new file mode 100644
index 0000000..33b94e3
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryLoggerTest.kt
@@ -0,0 +1,127 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.notification.logging
+
+import android.app.Notification
+import android.app.StatsManager
+import android.graphics.Bitmap
+import android.graphics.drawable.Icon
+import android.testing.AndroidTestingRunner
+import android.util.StatsEvent
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.shared.system.SysUiStatsLog
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
+import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.time.FakeSystemClock
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.Dispatchers
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class NotificationMemoryLoggerTest : SysuiTestCase() {
+
+    private val bgExecutor = FakeExecutor(FakeSystemClock())
+    private val immediate = Dispatchers.Main.immediate
+
+    @Mock private lateinit var statsManager: StatsManager
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+    }
+
+    @Test
+    fun onInit_registersCallback() {
+        val logger = createLoggerWithNotifications(listOf())
+        logger.init()
+        verify(statsManager)
+            .setPullAtomCallback(SysUiStatsLog.NOTIFICATION_MEMORY_USE, null, bgExecutor, logger)
+    }
+
+    @Test
+    fun onPullAtom_wrongAtomId_returnsSkip() {
+        val logger = createLoggerWithNotifications(listOf())
+        val data: MutableList<StatsEvent> = mutableListOf()
+        assertThat(logger.onPullAtom(111, data)).isEqualTo(StatsManager.PULL_SKIP)
+        assertThat(data).isEmpty()
+    }
+
+    @Test
+    fun onPullAtom_emptyNotifications_returnsZeros() {
+        val logger = createLoggerWithNotifications(listOf())
+        val data: MutableList<StatsEvent> = mutableListOf()
+        assertThat(logger.onPullAtom(SysUiStatsLog.NOTIFICATION_MEMORY_USE, data))
+            .isEqualTo(StatsManager.PULL_SUCCESS)
+        assertThat(data).isEmpty()
+    }
+
+    @Test
+    fun onPullAtom_notificationPassed_populatesData() {
+        val icon = Icon.createWithBitmap(Bitmap.createBitmap(300, 300, Bitmap.Config.ARGB_8888))
+        val notification =
+            Notification.Builder(context).setSmallIcon(icon).setContentTitle("title").build()
+        val logger = createLoggerWithNotifications(listOf(notification))
+        val data: MutableList<StatsEvent> = mutableListOf()
+
+        assertThat(logger.onPullAtom(SysUiStatsLog.NOTIFICATION_MEMORY_USE, data))
+            .isEqualTo(StatsManager.PULL_SUCCESS)
+        assertThat(data).hasSize(1)
+    }
+
+    @Test
+    fun onPullAtom_multipleNotificationsPassed_populatesData() {
+        val icon = Icon.createWithBitmap(Bitmap.createBitmap(300, 300, Bitmap.Config.ARGB_8888))
+        val notification =
+            Notification.Builder(context).setSmallIcon(icon).setContentTitle("title").build()
+        val iconTwo = Icon.createWithBitmap(Bitmap.createBitmap(300, 300, Bitmap.Config.ARGB_8888))
+
+        val notificationTwo =
+            Notification.Builder(context)
+                .setStyle(Notification.BigTextStyle().bigText("text"))
+                .setSmallIcon(iconTwo)
+                .setContentTitle("titleTwo")
+                .build()
+        val logger = createLoggerWithNotifications(listOf(notification, notificationTwo))
+        val data: MutableList<StatsEvent> = mutableListOf()
+
+        assertThat(logger.onPullAtom(SysUiStatsLog.NOTIFICATION_MEMORY_USE, data))
+            .isEqualTo(StatsManager.PULL_SUCCESS)
+        assertThat(data).hasSize(2)
+    }
+
+    private fun createLoggerWithNotifications(
+        notifications: List<Notification>
+    ): NotificationMemoryLogger {
+        val pipeline: NotifPipeline = mock()
+        val notifications =
+            notifications.map { notification ->
+                NotificationEntryBuilder().setTag("test").setNotification(notification).build()
+            }
+        whenever(pipeline.allNotifs).thenReturn(notifications)
+        return NotificationMemoryLogger(pipeline, statsManager, immediate, bgExecutor)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeterTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeterTest.kt
index f69839b..072a497 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeterTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryMeterTest.kt
@@ -23,6 +23,7 @@
 import android.content.Intent
 import android.graphics.Bitmap
 import android.graphics.drawable.Icon
+import android.stats.sysui.NotificationEnums
 import android.testing.AndroidTestingRunner
 import android.widget.RemoteViews
 import androidx.test.filters.SmallTest
@@ -50,7 +51,27 @@
             extras = 3316,
             bigPicture = 0,
             extender = 0,
-            style = null,
+            style = NotificationEnums.STYLE_NONE,
+            styleIcon = 0,
+            hasCustomView = false,
+        )
+    }
+
+    @Test
+    fun currentNotificationMemoryUse_rankerGroupNotification() {
+        val notification = createBasicNotification().build()
+        val memoryUse =
+            NotificationMemoryMeter.notificationMemoryUse(
+                createNotificationEntry(createBasicNotification().setGroup("ranker_group").build())
+            )
+        assertNotificationObjectSizes(
+            memoryUse,
+            smallIcon = notification.smallIcon.bitmap.allocationByteCount,
+            largeIcon = notification.getLargeIcon().bitmap.allocationByteCount,
+            extras = 3316,
+            bigPicture = 0,
+            extender = 0,
+            style = NotificationEnums.STYLE_RANKER_GROUP,
             styleIcon = 0,
             hasCustomView = false,
         )
@@ -69,7 +90,7 @@
             extras = 3316,
             bigPicture = 0,
             extender = 0,
-            style = null,
+            style = NotificationEnums.STYLE_NONE,
             styleIcon = 0,
             hasCustomView = false,
         )
@@ -92,7 +113,7 @@
             extras = 3384,
             bigPicture = 0,
             extender = 0,
-            style = null,
+            style = NotificationEnums.STYLE_NONE,
             styleIcon = 0,
             hasCustomView = true,
         )
@@ -112,7 +133,7 @@
             extras = 3212,
             bigPicture = 0,
             extender = 0,
-            style = null,
+            style = NotificationEnums.STYLE_NONE,
             styleIcon = 0,
             hasCustomView = false,
         )
@@ -141,7 +162,7 @@
             extras = 4092,
             bigPicture = bigPicture.bitmap.allocationByteCount,
             extender = 0,
-            style = "BigPictureStyle",
+            style = NotificationEnums.STYLE_BIG_PICTURE,
             styleIcon = bigPictureIcon.bitmap.allocationByteCount,
             hasCustomView = false,
         )
@@ -167,7 +188,7 @@
             extras = 4084,
             bigPicture = 0,
             extender = 0,
-            style = "CallStyle",
+            style = NotificationEnums.STYLE_CALL,
             styleIcon = personIcon.bitmap.allocationByteCount,
             hasCustomView = false,
         )
@@ -203,7 +224,7 @@
             extras = 5024,
             bigPicture = 0,
             extender = 0,
-            style = "MessagingStyle",
+            style = NotificationEnums.STYLE_MESSAGING,
             styleIcon =
                 personIcon.bitmap.allocationByteCount +
                     historicPersonIcon.bitmap.allocationByteCount,
@@ -225,7 +246,7 @@
             extras = 3612,
             bigPicture = 0,
             extender = 556656,
-            style = null,
+            style = NotificationEnums.STYLE_NONE,
             styleIcon = 0,
             hasCustomView = false,
         )
@@ -246,7 +267,7 @@
             extras = 3820,
             bigPicture = 0,
             extender = 388 + wearBackground.allocationByteCount,
-            style = null,
+            style = NotificationEnums.STYLE_NONE,
             styleIcon = 0,
             hasCustomView = false,
         )
@@ -272,7 +293,7 @@
         extras: Int,
         bigPicture: Int,
         extender: Int,
-        style: String?,
+        style: Int,
         styleIcon: Int,
         hasCustomView: Boolean,
     ) {
@@ -282,11 +303,7 @@
         assertThat(memoryUse.objectUsage.smallIcon).isEqualTo(smallIcon)
         assertThat(memoryUse.objectUsage.largeIcon).isEqualTo(largeIcon)
         assertThat(memoryUse.objectUsage.bigPicture).isEqualTo(bigPicture)
-        if (style == null) {
-            assertThat(memoryUse.objectUsage.style).isNull()
-        } else {
-            assertThat(memoryUse.objectUsage.style).isEqualTo(style)
-        }
+        assertThat(memoryUse.objectUsage.style).isEqualTo(style)
         assertThat(memoryUse.objectUsage.styleIcon).isEqualTo(styleIcon)
         assertThat(memoryUse.objectUsage.hasCustomView).isEqualTo(hasCustomView)
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt
index 3a16fb3..a0f5048 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalkerTest.kt
@@ -8,6 +8,7 @@
 import android.widget.RemoteViews
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.notification.row.NotificationRowContentBinder
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper
 import com.android.systemui.tests.R
 import com.google.common.truth.Truth.assertThat
@@ -39,16 +40,84 @@
     fun testViewWalker_plainNotification() {
         val row = testHelper.createRow()
         val result = NotificationMemoryViewWalker.getViewUsage(row)
-        assertThat(result).hasSize(5)
-        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
-        assertThat(result)
-            .contains(NotificationViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result).hasSize(3)
         assertThat(result)
             .contains(NotificationViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, 0, 0, 0, 0, 0, 0))
         assertThat(result)
             .contains(NotificationViewUsage(ViewType.PRIVATE_CONTRACTED_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result).contains(NotificationViewUsage(ViewType.TOTAL, 0, 0, 0, 0, 0, 0))
+    }
+
+    @Test
+    fun testViewWalker_plainNotification_withPublicView() {
+        val icon = Icon.createWithBitmap(Bitmap.createBitmap(20, 20, Bitmap.Config.ARGB_8888))
+        val publicIcon = Icon.createWithBitmap(Bitmap.createBitmap(40, 40, Bitmap.Config.ARGB_8888))
+        testHelper.setDefaultInflationFlags(NotificationRowContentBinder.FLAG_CONTENT_VIEW_ALL)
+        val row =
+            testHelper.createRow(
+                Notification.Builder(mContext)
+                    .setContentText("Test")
+                    .setContentTitle("title")
+                    .setSmallIcon(icon)
+                    .setPublicVersion(
+                        Notification.Builder(mContext)
+                            .setContentText("Public Test")
+                            .setContentTitle("title")
+                            .setSmallIcon(publicIcon)
+                            .build()
+                    )
+                    .build()
+            )
+        val result = NotificationMemoryViewWalker.getViewUsage(row)
+        assertThat(result).hasSize(4)
         assertThat(result)
-            .contains(NotificationViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, 0, 0, 0, 0, 0, 0))
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_EXPANDED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    0,
+                    icon.bitmap.allocationByteCount
+                )
+            )
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PRIVATE_CONTRACTED_VIEW,
+                    icon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    0,
+                    icon.bitmap.allocationByteCount
+                )
+            )
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.PUBLIC_VIEW,
+                    publicIcon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    0,
+                    publicIcon.bitmap.allocationByteCount
+                )
+            )
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.TOTAL,
+                    icon.bitmap.allocationByteCount + publicIcon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    0,
+                    icon.bitmap.allocationByteCount + publicIcon.bitmap.allocationByteCount
+                )
+            )
     }
 
     @Test
@@ -67,7 +136,7 @@
                     .build()
             )
         val result = NotificationMemoryViewWalker.getViewUsage(row)
-        assertThat(result).hasSize(5)
+        assertThat(result).hasSize(3)
         assertThat(result)
             .contains(
                 NotificationViewUsage(
@@ -95,8 +164,20 @@
                     icon.bitmap.allocationByteCount + largeIcon.bitmap.allocationByteCount
                 )
             )
-        // Due to deduplication, this should all be 0.
-        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.TOTAL,
+                    icon.bitmap.allocationByteCount,
+                    largeIcon.bitmap.allocationByteCount,
+                    0,
+                    bigPicture.allocationByteCount,
+                    0,
+                    bigPicture.allocationByteCount +
+                        icon.bitmap.allocationByteCount +
+                        largeIcon.bitmap.allocationByteCount
+                )
+            )
     }
 
     @Test
@@ -117,7 +198,7 @@
                     .build()
             )
         val result = NotificationMemoryViewWalker.getViewUsage(row)
-        assertThat(result).hasSize(5)
+        assertThat(result).hasSize(3)
         assertThat(result)
             .contains(
                 NotificationViewUsage(
@@ -142,7 +223,17 @@
                     bitmap.allocationByteCount + icon.bitmap.allocationByteCount
                 )
             )
-        // Due to deduplication, this should all be 0.
-        assertThat(result).contains(NotificationViewUsage(ViewType.PUBLIC_VIEW, 0, 0, 0, 0, 0, 0))
+        assertThat(result)
+            .contains(
+                NotificationViewUsage(
+                    ViewType.TOTAL,
+                    icon.bitmap.allocationByteCount,
+                    0,
+                    0,
+                    0,
+                    bitmap.allocationByteCount,
+                    bitmap.allocationByteCount + icon.bitmap.allocationByteCount
+                )
+            )
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt
index 5f57695..3f61af0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt
@@ -26,6 +26,7 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.statusbar.notification.FakeShadowView
 import com.android.systemui.statusbar.notification.NotificationUtils
+import com.android.systemui.statusbar.notification.SourceType
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
@@ -83,4 +84,17 @@
         mView.updateBackgroundColors()
         assertThat(mView.currentBackgroundTint).isEqualTo(mNormalColor)
     }
+
+    @Test
+    fun roundnessShouldBeTheSame_after_onDensityOrFontScaleChanged() {
+        val roundableState = mView.roundableState
+        assertThat(mView.topRoundness).isEqualTo(0f)
+        mView.requestTopRoundness(1f, SourceType.from(""))
+        assertThat(mView.topRoundness).isEqualTo(1f)
+
+        mView.onDensityOrFontScaleChanged()
+
+        assertThat(mView.topRoundness).isEqualTo(1f)
+        assertThat(mView.roundableState.hashCode()).isEqualTo(roundableState.hashCode())
+    }
 }
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/stylus/OWNERS b/packages/SystemUI/tests/src/com/android/systemui/stylus/OWNERS
new file mode 100644
index 0000000..7ccb316
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/stylus/OWNERS
@@ -0,0 +1,8 @@
+# Bug component: 1254381
+azappone@google.com
+achalke@google.com
+juliacr@google.com
+madym@google.com
+mgalhardo@google.com
+petrcermak@google.com
+vanjan@google.com
\ No newline at end of file
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
index 5c2a915..5501949 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
@@ -120,6 +120,14 @@
         _dozeAmount.value = dozeAmount
     }
 
+    fun setWakefulnessModel(model: WakefulnessModel) {
+        _wakefulnessModel.value = model
+    }
+
+    fun setBouncerShowing(isShowing: Boolean) {
+        _isBouncerShowing.value = isShowing
+    }
+
     fun setBiometricUnlockState(state: BiometricUnlockModel) {
         _biometricUnlockState.tryEmit(state)
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
new file mode 100644
index 0000000..2c0a8fd
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2022 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.systemui.shade.data.repository
+
+import com.android.systemui.shade.domain.model.ShadeModel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+
+/** Fake implementation of [KeyguardRepository] */
+class FakeShadeRepository : ShadeRepository {
+
+    private val _shadeModel = MutableStateFlow(ShadeModel())
+    override val shadeModel: Flow<ShadeModel> = _shadeModel
+
+    fun setShadeModel(model: ShadeModel) {
+        _shadeModel.value = model
+    }
+}
diff --git a/services/core/java/com/android/server/GestureLauncherService.java b/services/core/java/com/android/server/GestureLauncherService.java
index e529010..daba4c0 100644
--- a/services/core/java/com/android/server/GestureLauncherService.java
+++ b/services/core/java/com/android/server/GestureLauncherService.java
@@ -79,7 +79,7 @@
      * completed faster than this, we assume it's not performed by human and the
      * event gets ignored.
      */
-    @VisibleForTesting static final int EMERGENCY_GESTURE_TAP_DETECTION_MIN_TIME_MS = 160;
+    @VisibleForTesting static final int EMERGENCY_GESTURE_TAP_DETECTION_MIN_TIME_MS = 200;
 
     /**
      * Interval in milliseconds in which the power button must be depressed in succession to be
diff --git a/services/core/java/com/android/server/biometrics/sensors/fingerprint/hidl/Fingerprint21UdfpsMock.java b/services/core/java/com/android/server/biometrics/sensors/fingerprint/hidl/Fingerprint21UdfpsMock.java
index bea0f4f..846c2d9 100644
--- a/services/core/java/com/android/server/biometrics/sensors/fingerprint/hidl/Fingerprint21UdfpsMock.java
+++ b/services/core/java/com/android/server/biometrics/sensors/fingerprint/hidl/Fingerprint21UdfpsMock.java
@@ -417,7 +417,7 @@
     }
 
     @Override
-    public void onTrustChanged(boolean enabled, int userId, int flags,
+    public void onTrustChanged(boolean enabled, boolean newlyUnlocked, int userId, int flags,
             List<String> trustGrantedMessages) {
         mUserHasTrust.put(userId, enabled);
     }
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 85a2a5d..dbe049a 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -1049,12 +1049,6 @@
             return;
         }
 
-        // Make sure the device locks. Unfortunately, this has the side-effect of briefly revealing
-        // the lock screen before the dream appears. Note that this locking behavior needs to
-        // happen regardless of whether we end up dreaming (below) or not.
-        // TODO(b/261662912): Find a better way to lock the device that doesn't result in jank.
-        lockNow(null);
-
         // Don't dream if the user isn't user zero.
         // TODO(b/261907079): Move this check to DreamManagerService#canStartDreamingInternal().
         if (ActivityManager.getCurrentUser() != UserHandle.USER_SYSTEM) {
@@ -1068,6 +1062,12 @@
             return;
         }
 
+        // Make sure the device locks. Unfortunately, this has the side-effect of briefly revealing
+        // the lock screen before the dream appears. Note that locking is a side-effect of the no
+        // dream action that is executed if we early return above.
+        // TODO(b/261662912): Find a better way to lock the device that doesn't result in jank.
+        lockNow(null);
+
         dreamManagerInternal.requestDream();
     }
 
diff --git a/services/core/java/com/android/server/trust/TrustManagerService.java b/services/core/java/com/android/server/trust/TrustManagerService.java
index 6c191eb3..afe8d3e 100644
--- a/services/core/java/com/android/server/trust/TrustManagerService.java
+++ b/services/core/java/com/android/server/trust/TrustManagerService.java
@@ -563,7 +563,12 @@
             changed = mUserIsTrusted.get(userId) != trusted;
             mUserIsTrusted.put(userId, trusted);
         }
-        dispatchOnTrustChanged(trusted, userId, flags, getTrustGrantedMessages(userId));
+        dispatchOnTrustChanged(
+                trusted,
+                false /* newlyUnlocked */,
+                userId,
+                flags,
+                getTrustGrantedMessages(userId));
         if (changed) {
             refreshDeviceLockedForUser(userId);
             if (!trusted) {
@@ -628,7 +633,9 @@
         if (DEBUG) Slog.d(TAG, "pendingTrustState: " + pendingTrustState);
 
         boolean isNowTrusted = pendingTrustState == TrustState.TRUSTED;
-        dispatchOnTrustChanged(isNowTrusted, userId, flags, getTrustGrantedMessages(userId));
+        boolean newlyUnlocked = !alreadyUnlocked && isNowTrusted;
+        dispatchOnTrustChanged(
+                isNowTrusted, newlyUnlocked, userId, flags, getTrustGrantedMessages(userId));
         if (isNowTrusted != wasTrusted) {
             refreshDeviceLockedForUser(userId);
             if (!isNowTrusted) {
@@ -643,8 +650,7 @@
             }
         }
 
-        boolean wasLocked = !alreadyUnlocked;
-        boolean shouldSendCallback = wasLocked && pendingTrustState == TrustState.TRUSTED;
+        boolean shouldSendCallback = newlyUnlocked;
         if (shouldSendCallback) {
             if (resultCallback != null) {
                 if (DEBUG) Slog.d(TAG, "calling back with UNLOCKED_BY_GRANT");
@@ -1387,16 +1393,17 @@
         }
     }
 
-    private void dispatchOnTrustChanged(boolean enabled, int userId, int flags,
-            @NonNull List<String> trustGrantedMessages) {
+    private void dispatchOnTrustChanged(boolean enabled, boolean newlyUnlocked, int userId,
+            int flags, @NonNull List<String> trustGrantedMessages) {
         if (DEBUG) {
-            Log.i(TAG, "onTrustChanged(" + enabled + ", " + userId + ", 0x"
+            Log.i(TAG, "onTrustChanged(" + enabled + ", " + newlyUnlocked + ", " + userId + ", 0x"
                     + Integer.toHexString(flags) + ")");
         }
         if (!enabled) flags = 0;
         for (int i = 0; i < mTrustListeners.size(); i++) {
             try {
-                mTrustListeners.get(i).onTrustChanged(enabled, userId, flags, trustGrantedMessages);
+                mTrustListeners.get(i).onTrustChanged(
+                        enabled, newlyUnlocked, userId, flags, trustGrantedMessages);
             } catch (DeadObjectException e) {
                 Slog.d(TAG, "Removing dead TrustListener.");
                 mTrustListeners.remove(i);
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 292ed95..4cc15c9 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -2772,6 +2772,13 @@
     @Override
     void getAnimationFrames(Rect outFrame, Rect outInsets, Rect outStableInsets,
             Rect outSurfaceInsets) {
+        // If this task has its adjacent task, it means they should animate together. Use display
+        // bounds for them could move same as full screen task.
+        if (getAdjacentTaskFragment() != null && getAdjacentTaskFragment().asTask() != null) {
+            super.getAnimationFrames(outFrame, outInsets, outStableInsets, outSurfaceInsets);
+            return;
+        }
+
         final WindowState windowState = getTopVisibleAppMainWindow();
         if (windowState != null) {
             windowState.getAnimationFrames(outFrame, outInsets, outStableInsets, outSurfaceInsets);
@@ -4920,6 +4927,11 @@
                     }
                     if (child.getVisibility(null /* starting */)
                             != TASK_FRAGMENT_VISIBILITY_VISIBLE) {
+                        if (child.topRunningActivity() == null) {
+                            // Skip the task if no running activity and continue resuming next task.
+                            continue;
+                        }
+                        // Otherwise, assuming everything behind this task should also be invisible.
                         break;
                     }
 
diff --git a/services/core/java/com/android/server/wm/TaskFragment.java b/services/core/java/com/android/server/wm/TaskFragment.java
index 78b008f..f6db3f7 100644
--- a/services/core/java/com/android/server/wm/TaskFragment.java
+++ b/services/core/java/com/android/server/wm/TaskFragment.java
@@ -227,11 +227,16 @@
     private TaskFragment mCompanionTaskFragment;
 
     /**
-     * Prevents duplicate calls to onTaskAppeared.
+     * Prevents duplicate calls to onTaskFragmentAppeared.
      */
     boolean mTaskFragmentAppearedSent;
 
     /**
+     * Prevents unnecessary callbacks after onTaskFragmentVanished.
+     */
+    boolean mTaskFragmentVanishedSent;
+
+    /**
      * The last running activity of the TaskFragment was finished due to clear task while launching
      * an activity in the Task.
      */
diff --git a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
index 6e4df79..90a0dff 100644
--- a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
@@ -553,6 +553,9 @@
 
     void onTaskFragmentAppeared(@NonNull ITaskFragmentOrganizer organizer,
             @NonNull TaskFragment taskFragment) {
+        if (taskFragment.mTaskFragmentVanishedSent) {
+            return;
+        }
         if (taskFragment.getTask() == null) {
             Slog.w(TAG, "onTaskFragmentAppeared failed because it is not attached tf="
                     + taskFragment);
@@ -574,6 +577,9 @@
 
     void onTaskFragmentInfoChanged(@NonNull ITaskFragmentOrganizer organizer,
             @NonNull TaskFragment taskFragment) {
+        if (taskFragment.mTaskFragmentVanishedSent) {
+            return;
+        }
         validateAndGetState(organizer);
         if (!taskFragment.mTaskFragmentAppearedSent) {
             // Skip if TaskFragment still not appeared.
@@ -586,10 +592,6 @@
                     .setTaskFragment(taskFragment)
                     .build();
         } else {
-            if (pendingEvent.mEventType == PendingTaskFragmentEvent.EVENT_VANISHED) {
-                // Skipped the info changed event if vanished event is pending.
-                return;
-            }
             // Remove and add for re-ordering.
             removePendingEvent(pendingEvent);
             // Reset the defer time when TaskFragment is changed, so that it can check again if
@@ -602,6 +604,10 @@
 
     void onTaskFragmentVanished(@NonNull ITaskFragmentOrganizer organizer,
             @NonNull TaskFragment taskFragment) {
+        if (taskFragment.mTaskFragmentVanishedSent) {
+            return;
+        }
+        taskFragment.mTaskFragmentVanishedSent = true;
         final TaskFragmentOrganizerState state = validateAndGetState(organizer);
         final List<PendingTaskFragmentEvent> pendingEvents = mPendingTaskFragmentEvents
                 .get(organizer.asBinder());
@@ -617,20 +623,18 @@
                 .setTaskFragment(taskFragment)
                 .build());
         state.removeTaskFragment(taskFragment);
+        // Make sure the vanished event will be dispatched if there are no other changes.
+        mAtmService.mWindowManager.mWindowPlacerLocked.requestTraversal();
     }
 
     void onTaskFragmentError(@NonNull ITaskFragmentOrganizer organizer,
             @Nullable IBinder errorCallbackToken, @Nullable TaskFragment taskFragment,
             int opType, @NonNull Throwable exception) {
-        validateAndGetState(organizer);
-        Slog.w(TAG, "onTaskFragmentError ", exception);
-        final PendingTaskFragmentEvent vanishedEvent = taskFragment != null
-                ? getPendingTaskFragmentEvent(taskFragment, PendingTaskFragmentEvent.EVENT_VANISHED)
-                : null;
-        if (vanishedEvent != null) {
-            // No need to notify if the TaskFragment has been removed.
+        if (taskFragment != null && taskFragment.mTaskFragmentVanishedSent) {
             return;
         }
+        validateAndGetState(organizer);
+        Slog.w(TAG, "onTaskFragmentError ", exception);
         addPendingEvent(new PendingTaskFragmentEvent.Builder(
                 PendingTaskFragmentEvent.EVENT_ERROR, organizer)
                 .setErrorCallbackToken(errorCallbackToken)
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 5c893de..2f81f93 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -18,6 +18,7 @@
 
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
+import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_BEHIND;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSET;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED;
@@ -3186,7 +3187,8 @@
         if (isOrganized()
                 // TODO(b/161711458): Clean-up when moved to shell.
                 && getWindowingMode() != WINDOWING_MODE_FULLSCREEN
-                && getWindowingMode() != WINDOWING_MODE_FREEFORM) {
+                && getWindowingMode() != WINDOWING_MODE_FREEFORM
+                && getWindowingMode() != WINDOWING_MODE_MULTI_WINDOW) {
             return null;
         }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java
index db65f49..140051d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java
@@ -247,6 +247,7 @@
         mController.onTaskFragmentVanished(mTaskFragment.getTaskFragmentOrganizer(), mTaskFragment);
         mController.dispatchPendingEvents();
 
+        assertTrue(mTaskFragment.mTaskFragmentVanishedSent);
         assertTaskFragmentVanishedTransaction();
     }
 
@@ -259,10 +260,12 @@
         mController.onTaskFragmentVanished(mTaskFragment.getTaskFragmentOrganizer(), mTaskFragment);
         mController.dispatchPendingEvents();
 
+        assertTrue(mTaskFragment.mTaskFragmentVanishedSent);
         assertTaskFragmentVanishedTransaction();
 
         // Not trigger onTaskFragmentInfoChanged.
         // Call onTaskFragmentAppeared before calling onTaskFragmentInfoChanged.
+        mTaskFragment.mTaskFragmentVanishedSent = false;
         mController.onTaskFragmentAppeared(mTaskFragment.getTaskFragmentOrganizer(), mTaskFragment);
         mController.dispatchPendingEvents();
         clearInvocations(mOrganizer);
diff --git a/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordAudioStreamCopier.java b/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordAudioStreamCopier.java
index b9d2ae6..244b3a0 100644
--- a/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordAudioStreamCopier.java
+++ b/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordAudioStreamCopier.java
@@ -19,13 +19,13 @@
 import static android.app.AppOpsManager.MODE_ALLOWED;
 import static android.service.voice.HotwordAudioStream.KEY_AUDIO_STREAM_COPY_BUFFER_LENGTH_BYTES;
 
-import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_CLOSE_ERROR_FROM_SYSTEM;
-import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_EMPTY_AUDIO_STREAM_LIST;
-import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_END;
-import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_ILLEGAL_COPY_BUFFER_SIZE;
-import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_INTERRUPTED_EXCEPTION;
-import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_NO_PERMISSION;
-import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_START;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__CLOSE_ERROR_FROM_SYSTEM;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__EMPTY_AUDIO_STREAM_LIST;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__ENDED;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__ILLEGAL_COPY_BUFFER_SIZE;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__INTERRUPTED_EXCEPTION;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__NO_PERMISSION;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__STARTED;
 import static com.android.server.voiceinteraction.HotwordDetectionConnection.DEBUG;
 
 import android.annotation.NonNull;
@@ -60,10 +60,10 @@
     private static final String OP_MESSAGE = "Streaming hotword audio to VoiceInteractionService";
     private static final String TASK_ID_PREFIX = "HotwordDetectedResult@";
     private static final String THREAD_NAME_PREFIX = "Copy-";
-    private static final int DEFAULT_COPY_BUFFER_LENGTH_BYTES = 2_560;
 
     // Corresponds to the OS pipe capacity in bytes
     private static final int MAX_COPY_BUFFER_LENGTH_BYTES = 65_536;
+    private static final int DEFAULT_COPY_BUFFER_LENGTH_BYTES = 32_768;
 
     private final AppOpsManager mAppOpsManager;
     private final int mDetectorType;
@@ -98,14 +98,17 @@
             throws IOException {
         List<HotwordAudioStream> audioStreams = result.getAudioStreams();
         if (audioStreams.isEmpty()) {
-            HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                    HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_EMPTY_AUDIO_STREAM_LIST,
-                    mVoiceInteractorUid);
+            HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                    HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__EMPTY_AUDIO_STREAM_LIST,
+                    mVoiceInteractorUid, /* streamSizeBytes= */ 0, /* bundleSizeBytes= */ 0,
+                    /* streamCount= */ 0);
             return result;
         }
 
+        final int audioStreamCount = audioStreams.size();
         List<HotwordAudioStream> newAudioStreams = new ArrayList<>(audioStreams.size());
         List<CopyTaskInfo> copyTaskInfos = new ArrayList<>(audioStreams.size());
+        int totalMetadataBundleSizeBytes = 0;
         for (HotwordAudioStream audioStream : audioStreams) {
             ParcelFileDescriptor[] clientPipe = ParcelFileDescriptor.createReliablePipe();
             ParcelFileDescriptor clientAudioSource = clientPipe[0];
@@ -117,12 +120,14 @@
 
             int copyBufferLength = DEFAULT_COPY_BUFFER_LENGTH_BYTES;
             PersistableBundle metadata = audioStream.getMetadata();
+            totalMetadataBundleSizeBytes += HotwordDetectedResult.getParcelableSize(metadata);
             if (metadata.containsKey(KEY_AUDIO_STREAM_COPY_BUFFER_LENGTH_BYTES)) {
                 copyBufferLength = metadata.getInt(KEY_AUDIO_STREAM_COPY_BUFFER_LENGTH_BYTES, -1);
                 if (copyBufferLength < 1 || copyBufferLength > MAX_COPY_BUFFER_LENGTH_BYTES) {
-                    HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                            HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_ILLEGAL_COPY_BUFFER_SIZE,
-                            mVoiceInteractorUid);
+                    HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                            HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__ILLEGAL_COPY_BUFFER_SIZE,
+                            mVoiceInteractorUid, /* streamSizeBytes= */ 0, /* bundleSizeBytes= */ 0,
+                            audioStreamCount);
                     Slog.w(TAG, "Attempted to set an invalid copy buffer length ("
                             + copyBufferLength + ") for: " + audioStream);
                     copyBufferLength = DEFAULT_COPY_BUFFER_LENGTH_BYTES;
@@ -139,7 +144,9 @@
         }
 
         String resultTaskId = TASK_ID_PREFIX + System.identityHashCode(result);
-        mExecutorService.execute(new HotwordDetectedResultCopyTask(resultTaskId, copyTaskInfos));
+        mExecutorService.execute(
+                new HotwordDetectedResultCopyTask(resultTaskId, copyTaskInfos,
+                        totalMetadataBundleSizeBytes));
 
         return result.buildUpon().setAudioStreams(newAudioStreams).build();
     }
@@ -159,11 +166,14 @@
     private class HotwordDetectedResultCopyTask implements Runnable {
         private final String mResultTaskId;
         private final List<CopyTaskInfo> mCopyTaskInfos;
+        private final int mTotalMetadataSizeBytes;
         private final ExecutorService mExecutorService = Executors.newCachedThreadPool();
 
-        HotwordDetectedResultCopyTask(String resultTaskId, List<CopyTaskInfo> copyTaskInfos) {
+        HotwordDetectedResultCopyTask(String resultTaskId, List<CopyTaskInfo> copyTaskInfos,
+                int totalMetadataSizeBytes) {
             mResultTaskId = resultTaskId;
             mCopyTaskInfos = copyTaskInfos;
+            mTotalMetadataSizeBytes = totalMetadataSizeBytes;
         }
 
         @Override
@@ -183,19 +193,38 @@
                     mVoiceInteractorUid, mVoiceInteractorPackageName,
                     mVoiceInteractorAttributionTag, OP_MESSAGE) == MODE_ALLOWED) {
                 try {
-                    HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                            HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_START,
-                            mVoiceInteractorUid);
+                    HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                            HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__STARTED,
+                            mVoiceInteractorUid, /* streamSizeBytes= */ 0, mTotalMetadataSizeBytes,
+                            size);
                     // TODO(b/244599891): Set timeout, close after inactivity
                     mExecutorService.invokeAll(tasks);
-                    HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                            HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_END,
-                            mVoiceInteractorUid);
+
+                    int totalStreamSizeBytes = 0;
+                    for (SingleAudioStreamCopyTask task : tasks) {
+                        totalStreamSizeBytes += task.mTotalCopiedBytes;
+                    }
+
+                    Slog.i(TAG, mResultTaskId + ": Task was completed. Total bytes streamed: "
+                            + totalStreamSizeBytes + ", total metadata bundle size bytes: "
+                            + mTotalMetadataSizeBytes);
+                    HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                            HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__ENDED,
+                            mVoiceInteractorUid, totalStreamSizeBytes, mTotalMetadataSizeBytes,
+                            size);
                 } catch (InterruptedException e) {
-                    HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                            HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_INTERRUPTED_EXCEPTION,
-                            mVoiceInteractorUid);
-                    Slog.e(TAG, mResultTaskId + ": Task was interrupted", e);
+                    int totalStreamSizeBytes = 0;
+                    for (SingleAudioStreamCopyTask task : tasks) {
+                        totalStreamSizeBytes += task.mTotalCopiedBytes;
+                    }
+
+                    HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                            HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__INTERRUPTED_EXCEPTION,
+                            mVoiceInteractorUid, totalStreamSizeBytes, mTotalMetadataSizeBytes,
+                            size);
+                    Slog.e(TAG, mResultTaskId + ": Task was interrupted. Total bytes streamed: "
+                            + totalStreamSizeBytes + ", total metadata bundle size bytes: "
+                            + mTotalMetadataSizeBytes);
                     bestEffortPropagateError(e.getMessage());
                 } finally {
                     mAppOpsManager.finishOp(AppOpsManager.OPSTR_RECORD_AUDIO_HOTWORD,
@@ -203,9 +232,10 @@
                             mVoiceInteractorAttributionTag);
                 }
             } else {
-                HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                        HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_NO_PERMISSION,
-                        mVoiceInteractorUid);
+                HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                        HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__NO_PERMISSION,
+                        mVoiceInteractorUid, /* streamSizeBytes= */ 0, /* bundleSizeBytes= */ 0,
+                        size);
                 bestEffortPropagateError(
                         "Failed to obtain RECORD_AUDIO_HOTWORD permission for voice interactor with"
                                 + " uid=" + mVoiceInteractorUid
@@ -220,9 +250,10 @@
                     copyTaskInfo.mSource.closeWithError(errorMessage);
                     copyTaskInfo.mSink.closeWithError(errorMessage);
                 }
-                HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                        HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_CLOSE_ERROR_FROM_SYSTEM,
-                        mVoiceInteractorUid);
+                HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                        HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__CLOSE_ERROR_FROM_SYSTEM,
+                        mVoiceInteractorUid, /* streamSizeBytes= */ 0, /* bundleSizeBytes= */ 0,
+                        mCopyTaskInfos.size());
             } catch (IOException e) {
                 Slog.e(TAG, mResultTaskId + ": Failed to propagate error", e);
             }
@@ -237,6 +268,8 @@
         private final int mDetectorType;
         private final int mUid;
 
+        private volatile int mTotalCopiedBytes = 0;
+
         SingleAudioStreamCopyTask(String streamTaskId, ParcelFileDescriptor audioSource,
                 ParcelFileDescriptor audioSink, int copyBufferLength, int detectorType, int uid) {
             mStreamTaskId = streamTaskId;
@@ -281,6 +314,7 @@
                                     Arrays.copyOfRange(buffer, 0, 20)));
                         }
                         fos.write(buffer, 0, bytesRead);
+                        mTotalCopiedBytes += bytesRead;
                     }
                     // TODO(b/244599891): Close PFDs after inactivity
                 }
@@ -288,8 +322,10 @@
                 mAudioSource.closeWithError(e.getMessage());
                 mAudioSink.closeWithError(e.getMessage());
                 Slog.e(TAG, mStreamTaskId + ": Failed to copy audio stream", e);
-                HotwordMetricsLogger.writeDetectorEvent(mDetectorType,
-                        HOTWORD_DETECTOR_EVENTS__EVENT__AUDIO_EGRESS_CLOSE_ERROR_FROM_SYSTEM, mUid);
+                HotwordMetricsLogger.writeAudioEgressEvent(mDetectorType,
+                        HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__EVENT__CLOSE_ERROR_FROM_SYSTEM,
+                        mUid, /* streamSizeBytes= */ 0, /* bundleSizeBytes= */ 0,
+                        /* streamCount= */ 0);
             } finally {
                 if (fis != null) {
                     fis.close();
diff --git a/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordMetricsLogger.java b/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordMetricsLogger.java
index 61c18be..c35d90f 100644
--- a/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordMetricsLogger.java
+++ b/services/voiceinteraction/java/com/android/server/voiceinteraction/HotwordMetricsLogger.java
@@ -16,6 +16,9 @@
 
 package com.android.server.voiceinteraction;
 
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__DETECTOR_TYPE__NORMAL_DETECTOR;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__DETECTOR_TYPE__TRUSTED_DETECTOR_DSP;
+import static com.android.internal.util.FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__DETECTOR_TYPE__TRUSTED_DETECTOR_SOFTWARE;
 import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTION_SERVICE_INIT_RESULT_REPORTED__DETECTOR_TYPE__NORMAL_DETECTOR;
 import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTION_SERVICE_INIT_RESULT_REPORTED__DETECTOR_TYPE__TRUSTED_DETECTOR_DSP;
 import static com.android.internal.util.FrameworkStatsLog.HOTWORD_DETECTION_SERVICE_INIT_RESULT_REPORTED__DETECTOR_TYPE__TRUSTED_DETECTOR_SOFTWARE;
@@ -47,6 +50,12 @@
             HOTWORD_DETECTION_SERVICE_INIT_RESULT_REPORTED__DETECTOR_TYPE__TRUSTED_DETECTOR_DSP;
     private static final int METRICS_INIT_NORMAL_DETECTOR =
             HOTWORD_DETECTION_SERVICE_INIT_RESULT_REPORTED__DETECTOR_TYPE__NORMAL_DETECTOR;
+    private static final int AUDIO_EGRESS_DSP_DETECTOR =
+            HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__DETECTOR_TYPE__TRUSTED_DETECTOR_DSP;
+    private static final int AUDIO_EGRESS_SOFTWARE_DETECTOR =
+            HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__DETECTOR_TYPE__TRUSTED_DETECTOR_SOFTWARE;
+    private static final int AUDIO_EGRESS_NORMAL_DETECTOR =
+            HOTWORD_AUDIO_EGRESS_EVENT_REPORTED__DETECTOR_TYPE__NORMAL_DETECTOR;
 
     private HotwordMetricsLogger() {
         // Class only contains static utility functions, and should not be instantiated
@@ -97,6 +106,16 @@
                 metricsDetectorType, event, uid);
     }
 
+    /**
+     * Logs information related to hotword audio egress events.
+     */
+    public static void writeAudioEgressEvent(int detectorType, int event, int uid,
+            int streamSizeBytes, int bundleSizeBytes, int streamCount) {
+        int metricsDetectorType = getAudioEgressDetectorType(detectorType);
+        FrameworkStatsLog.write(FrameworkStatsLog.HOTWORD_AUDIO_EGRESS_EVENT_REPORTED,
+                metricsDetectorType, event, uid, streamSizeBytes, bundleSizeBytes, streamCount);
+    }
+
     private static int getCreateMetricsDetectorType(int detectorType) {
         switch (detectorType) {
             case HotwordDetector.DETECTOR_TYPE_TRUSTED_HOTWORD_SOFTWARE:
@@ -151,4 +170,15 @@
                 return HOTWORD_DETECTOR_EVENTS__DETECTOR_TYPE__NORMAL_DETECTOR;
         }
     }
+
+    private static int getAudioEgressDetectorType(int detectorType) {
+        switch (detectorType) {
+            case HotwordDetector.DETECTOR_TYPE_TRUSTED_HOTWORD_SOFTWARE:
+                return AUDIO_EGRESS_SOFTWARE_DETECTOR;
+            case HotwordDetector.DETECTOR_TYPE_TRUSTED_HOTWORD_DSP:
+                return AUDIO_EGRESS_DSP_DETECTOR;
+            default:
+                return AUDIO_EGRESS_NORMAL_DETECTOR;
+        }
+    }
 }
diff --git a/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt b/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt
index 2031af2..1930a1c 100644
--- a/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt
+++ b/tests/TrustTests/src/android/trust/test/lib/LockStateTrackingRule.kt
@@ -63,6 +63,7 @@
     inner class Listener : TrustListener {
         override fun onTrustChanged(
             enabled: Boolean,
+            newlyUnlocked: Boolean,
             userId: Int,
             flags: Int,
             trustGrantedMessages: MutableList<String>