Merge "[Unfold transition] Fully manage the vignette on the background thread" into tm-qpr-dev

12 files changed, 123 insertions, 128 deletions

diff --git a/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt b/packages/SystemUI/src/com/android/keyguard/mediator/ScreenOnCoordinator.kt
index 4b7e9a57ef6a..98ac2c0bd026 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 @@ class ScreenOnCoordinator @Inject constructor(
         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 @@ class ScreenOnCoordinator @Inject constructor(
         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/keyguard/KeyguardLifecyclesDispatcher.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardLifecyclesDispatcher.java
index 822b1cfdf877..757afb616fd1 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 @@ package com.android.systemui.keyguard;
 
 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 @@ public class KeyguardLifecyclesDispatcher {
     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 c332a0d66294..f4a1227a467c 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.IKeyguardDrawnCallback;
 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 @@ public class KeyguardService extends Service {
 
     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 @@ public class KeyguardService extends Service {
     @Inject
     public KeyguardService(KeyguardViewMediator keyguardViewMediator,
                            KeyguardLifecyclesDispatcher keyguardLifecyclesDispatcher,
+                           ScreenOnCoordinator screenOnCoordinator,
                            ShellTransitions shellTransitions) {
         super();
         mKeyguardViewMediator = keyguardViewMediator;
         mKeyguardLifecyclesDispatcher = keyguardLifecyclesDispatcher;
+        mScreenOnCoordinator = screenOnCoordinator;
         mShellTransitions = shellTransitions;
     }
 
@@ -583,6 +587,31 @@ public class KeyguardService extends Service {
             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 @@ public class KeyguardService extends Service {
             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/LifecycleScreenStatusProvider.kt b/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt
index 0a55294dfe8a..0b04fb43c2d7 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 @@ class LifecycleScreenStatusProvider @Inject constructor(screenLifecycle: ScreenL
         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 b3481219a85d..8535eda93f96 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 @@ package com.android.systemui.keyguard;
 
 import android.os.Trace;
 
-import androidx.annotation.NonNull;
-
 import com.android.systemui.Dumpable;
 import com.android.systemui.dump.DumpManager;
 
@@ -50,14 +48,9 @@ public class ScreenLifecycle extends Lifecycle<ScreenLifecycle.Observer> impleme
         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 class ScreenLifecycle extends Lifecycle<ScreenLifecycle.Observer> impleme
     }
 
     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/recents/OverviewProxyService.java b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
index 00d129ae70ca..4d005bebd99e 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 @@ public class OverviewProxyService implements CallbackController<OverviewProxyLis
          * 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/statusbar/phone/CentralSurfacesImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
index 005cd1bff90c..93da0b806a78 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 @@ public class CentralSurfacesImpl implements CoreStartable, CentralSurfaces {
 
     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 6216acd6081e..101bd4483cb3 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 java.util.function.Consumer
 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 @@ import kotlinx.coroutines.launch
 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 @@ constructor(
     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 @@ constructor(
      * @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 @@ constructor(
         }
     }
 
-    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 b2ec27c8ce43..9cca95028729 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.SurfaceControlViewHost
 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.UnfoldLightRevealOverlayAnimation.AddOverlayR
 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 @@ constructor(
     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 @@ constructor(
     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 @@ constructor(
      * @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 @@ constructor(
             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 @@ constructor(
 
                 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 @@ constructor(
     }
 
     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 @@ constructor(
 
     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 @@ constructor(
     }
 
     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 @@ constructor(
         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/util/concurrency/PendingTasksContainer.kt b/packages/SystemUI/src/com/android/systemui/util/concurrency/PendingTasksContainer.kt
index 6cd384f17803..ceebcb77fde2 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 @@ import java.util.concurrent.atomic.AtomicReference
  */
 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/mediator/ScreenOnCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/mediator/ScreenOnCoordinatorTest.kt
index 5734c3de70e0..34e78eb8c2eb 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 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
     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 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
                 .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 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
     fun testUnfoldTransitionDisabledDrawnTasksReady_onScreenTurningOn_callsDrawnCallback() {
         // Recreate with empty unfoldComponent
         screenOnCoordinator = ScreenOnCoordinator(
-            screenLifecycle,
             Optional.empty(),
-            FakeExecution()
         )
         screenOnCoordinator.onScreenTurningOn(runnable)
 
@@ -105,11 +96,11 @@ class ScreenOnCoordinatorTest : SysuiTestCase() {
 
     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 f46d58d679b5..70a0415d2e35 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 @@ public class ScreenLifecycleTest extends SysuiTestCase {
     @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 @@ public class ScreenLifecycleTest extends SysuiTestCase {
 
     @Test
     public void screenTurningOff() throws Exception {
-        mScreen.dispatchScreenTurningOn(null);
+        mScreen.dispatchScreenTurningOn();
         mScreen.dispatchScreenTurnedOn();
         mScreen.dispatchScreenTurningOff();
 
@@ -85,7 +85,7 @@ public class ScreenLifecycleTest extends SysuiTestCase {
 
     @Test
     public void screenTurnedOff() throws Exception {
-        mScreen.dispatchScreenTurningOn(null);
+        mScreen.dispatchScreenTurningOn();
         mScreen.dispatchScreenTurnedOn();
         mScreen.dispatchScreenTurningOff();
         mScreen.dispatchScreenTurnedOff();