Remove MutableStateFlow for transition current info

When multiple threads are attempting to begin transitions,
the MutableStateFlow can get stuck on a background thread
while emitting the update the consumers. This leaves a
small chance for other transitions to jump the queue, ones
that are scheduled onto a more favorable thread.

Avoid MutableStateFlow completely. There's no need, just
use a volatile to store state.

Fixes: 358533338
Test: atest KeyguardTransitionRepositoryTest
Flag: com.android.systemui.transition_race_condition
Change-Id: I1daf9d5074c966445a1ed54acc10fb91193d6c85
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 2278789..5cde322 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -1546,6 +1546,16 @@
 }
 
 flag {
+  name: "transition_race_condition"
+  namespace: "systemui"
+  description: "Thread-safe keyguard transitions"
+  bug: "358533338"
+  metadata {
+       purpose: PURPOSE_BUGFIX
+  }
+}
+
+flag {
    name: "media_projection_request_attribution_fix"
    namespace: "systemui"
    description: "Ensure MediaProjection consent requests are properly attributed"
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt
index 7b8c19c..ec55401 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt
@@ -73,7 +73,7 @@
     private var progressJob: Job? = null
 
     private val currentToState: KeyguardState
-        get() = internalTransitionInteractor.currentTransitionInfoInternal.value.to
+        get() = internalTransitionInteractor.currentTransitionInfoInternal().to
 
     /**
      * The next keyguard state to trigger when exiting [CommunalScenes.Communal]. This is only used
@@ -197,7 +197,7 @@
         val newTransition =
             TransitionInfo(
                 ownerName = this::class.java.simpleName,
-                from = internalTransitionInteractor.currentTransitionInfoInternal.value.to,
+                from = internalTransitionInteractor.currentTransitionInfoInternal().to,
                 to = state,
                 animator = null,
                 modeOnCanceled = TransitionModeOnCanceled.REVERSE,
@@ -273,7 +273,7 @@
     }
 
     private suspend fun startTransitionToGlanceableHub() {
-        val currentState = internalTransitionInteractor.currentTransitionInfoInternal.value.to
+        val currentState = internalTransitionInteractor.currentTransitionInfoInternal().to
         val newTransition =
             TransitionInfo(
                 ownerName = this::class.java.simpleName,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt
index b7d0d45..3a5614f 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardTransitionRepository.kt
@@ -25,6 +25,7 @@
 import android.util.Log
 import com.android.app.animation.Interpolators
 import com.android.app.tracing.coroutines.withContextTraced as withContext
+import com.android.systemui.Flags.transitionRaceCondition
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.keyguard.shared.model.KeyguardState
@@ -77,6 +78,8 @@
 
     /** The [TransitionInfo] of the most recent call to [startTransition]. */
     val currentTransitionInfoInternal: StateFlow<TransitionInfo>
+    /** The [TransitionInfo] of the most recent call to [startTransition]. */
+    val currentTransitionInfo: TransitionInfo
 
     /**
      * Interactors that require information about changes between [KeyguardState]s will call this to
@@ -132,7 +135,7 @@
     private var lastStep: TransitionStep = TransitionStep()
     private var lastAnimator: ValueAnimator? = null
 
-    private val _currentTransitionMutex = Mutex()
+    private val withContextMutex = Mutex()
     private val _currentTransitionInfo: MutableStateFlow<TransitionInfo> =
         MutableStateFlow(
             TransitionInfo(
@@ -144,6 +147,16 @@
         )
     override var currentTransitionInfoInternal = _currentTransitionInfo.asStateFlow()
 
+    @Volatile
+    override var currentTransitionInfo: TransitionInfo =
+        TransitionInfo(
+            ownerName = "",
+            from = KeyguardState.OFF,
+            to = KeyguardState.OFF,
+            animator = null,
+        )
+        private set
+
     /*
      * When manual control of the transition is requested, a unique [UUID] is used as the handle
      * to permit calls to [updateTransition]
@@ -163,13 +176,17 @@
     }
 
     override suspend fun startTransition(info: TransitionInfo): UUID? {
-        _currentTransitionInfo.value = info
+        if (transitionRaceCondition()) {
+            currentTransitionInfo = info
+        } else {
+            _currentTransitionInfo.value = info
+        }
         Log.d(TAG, "(Internal) Setting current transition info: $info")
 
         // There is no fairness guarantee with 'withContext', which means that transitions could
         // be processed out of order. Use a Mutex to guarantee ordering. [updateTransition]
         // requires the same lock
-        _currentTransitionMutex.lock()
+        withContextMutex.lock()
         // Only used in a test environment
         if (forceDelayForRaceConditionTest) {
             delay(50L)
@@ -177,7 +194,7 @@
 
         // Animators must be started on the main thread.
         return withContext("$TAG#startTransition", mainDispatcher) {
-            _currentTransitionMutex.unlock()
+            withContextMutex.unlock()
             if (lastStep.from == info.from && lastStep.to == info.to) {
                 Log.i(TAG, "Duplicate call to start the transition, rejecting: $info")
                 return@withContext null
@@ -265,9 +282,9 @@
         // There is no fairness guarantee with 'withContext', which means that transitions could
         // be processed out of order. Use a Mutex to guarantee ordering. [startTransition]
         // requires the same lock
-        _currentTransitionMutex.lock()
+        withContextMutex.lock()
         withContext("$TAG#updateTransition", mainDispatcher) {
-            _currentTransitionMutex.unlock()
+            withContextMutex.unlock()
 
             updateTransitionInternal(transitionId, value, state)
         }
@@ -302,13 +319,23 @@
         // Tests runs on testDispatcher, which is not the main thread, causing the animator thread
         // check to fail
         if (testSetup) {
-            _currentTransitionInfo.value =
-                TransitionInfo(
-                    ownerName = ownerName,
-                    from = KeyguardState.OFF,
-                    to = to,
-                    animator = null,
-                )
+            if (transitionRaceCondition()) {
+                currentTransitionInfo =
+                    TransitionInfo(
+                        ownerName = ownerName,
+                        from = KeyguardState.OFF,
+                        to = to,
+                        animator = null,
+                    )
+            } else {
+                _currentTransitionInfo.value =
+                    TransitionInfo(
+                        ownerName = ownerName,
+                        from = KeyguardState.OFF,
+                        to = to,
+                        animator = null,
+                    )
+            }
             emitTransition(
                 TransitionStep(
                     KeyguardState.OFF,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
index a7dde34e..8b75545 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
@@ -40,6 +40,7 @@
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.data.repository.ShadeRepository
 import com.android.systemui.util.kotlin.Utils.Companion.sample as sampleCombine
+import com.android.systemui.util.kotlin.sample
 import java.util.UUID
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.milliseconds
@@ -132,11 +133,10 @@
         scope.launch("$TAG#listenForLockscreenToDreaming") {
             keyguardInteractor.isAbleToDream
                 .filterRelevantKeyguardState()
-                .sampleCombine(
-                    internalTransitionInteractor.currentTransitionInfoInternal,
-                    transitionInteractor.isFinishedIn(KeyguardState.LOCKSCREEN),
-                )
-                .collect { (isAbleToDream, transitionInfo, isOnLockscreen) ->
+                .sample(transitionInteractor.isFinishedIn(KeyguardState.LOCKSCREEN), ::Pair)
+                .collect { (isAbleToDream, isOnLockscreen) ->
+                    val transitionInfo =
+                        internalTransitionInteractor.currentTransitionInfoInternal()
                     val isTransitionInterruptible =
                         transitionInfo.to == KeyguardState.LOCKSCREEN &&
                             !invalidFromStates.contains(transitionInfo.from)
@@ -179,7 +179,6 @@
             shadeRepository.legacyShadeExpansion
                 .sampleCombine(
                     transitionInteractor.startedKeyguardTransitionStep,
-                    internalTransitionInteractor.currentTransitionInfoInternal,
                     keyguardInteractor.statusBarState,
                     keyguardInteractor.isKeyguardDismissible,
                     keyguardInteractor.isKeyguardOccluded,
@@ -188,11 +187,12 @@
                     (
                         shadeExpansion,
                         startedStep,
-                        currentTransitionInfo,
                         statusBarState,
                         isKeyguardUnlocked,
                         isKeyguardOccluded) ->
                     val id = transitionId
+                    val currentTransitionInfo =
+                        internalTransitionInteractor.currentTransitionInfoInternal()
                     if (id != null) {
                         if (startedStep.to == KeyguardState.PRIMARY_BOUNCER) {
                             // An existing `id` means a transition is started, and calls to
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InternalKeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InternalKeyguardTransitionInteractor.kt
index 2cc6afa..0507834 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InternalKeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/InternalKeyguardTransitionInteractor.kt
@@ -17,13 +17,13 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.annotation.FloatRange
+import com.android.systemui.Flags.transitionRaceCondition
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.TransitionInfo
 import com.android.systemui.keyguard.shared.model.TransitionState
 import java.util.UUID
 import javax.inject.Inject
-import kotlinx.coroutines.flow.StateFlow
 
 /**
  * This interactor provides direct access to [KeyguardTransitionRepository] internals and exposes
@@ -32,9 +32,7 @@
 @SysUISingleton
 class InternalKeyguardTransitionInteractor
 @Inject
-constructor(
-    private val repository: KeyguardTransitionRepository,
-) {
+constructor(private val repository: KeyguardTransitionRepository) {
 
     /**
      * The [TransitionInfo] of the most recent call to
@@ -58,14 +56,19 @@
      * *will* be emitted, and therefore that it can safely request an AOD -> LOCKSCREEN transition
      * which will subsequently cancel GONE -> AOD.
      */
-    internal val currentTransitionInfoInternal: StateFlow<TransitionInfo> =
-        repository.currentTransitionInfoInternal
+    internal fun currentTransitionInfoInternal(): TransitionInfo {
+        return if (transitionRaceCondition()) {
+            repository.currentTransitionInfo
+        } else {
+            repository.currentTransitionInfoInternal.value
+        }
+    }
 
     suspend fun startTransition(info: TransitionInfo) = repository.startTransition(info)
 
     suspend fun updateTransition(
         transitionId: UUID,
         @FloatRange(from = 0.0, to = 1.0) value: Float,
-        state: TransitionState
+        state: TransitionState,
     ) = repository.updateTransition(transitionId, value, state)
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissTransitionInteractor.kt
index c19bbbc..4793d95 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissTransitionInteractor.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.util.Log
+import com.android.systemui.Flags.transitionRaceCondition
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
@@ -51,7 +52,13 @@
     fun startDismissKeyguardTransition(reason: String = "") {
         if (SceneContainerFlag.isEnabled) return
         Log.d(TAG, "#startDismissKeyguardTransition(reason=$reason)")
-        when (val startedState = repository.currentTransitionInfoInternal.value.to) {
+        val startedState =
+            if (transitionRaceCondition()) {
+                repository.currentTransitionInfo.to
+            } else {
+                repository.currentTransitionInfoInternal.value.to
+            }
+        when (startedState) {
             LOCKSCREEN -> fromLockscreenTransitionInteractor.dismissKeyguard()
             PRIMARY_BOUNCER -> fromPrimaryBouncerTransitionInteractor.dismissPrimaryBouncer()
             ALTERNATE_BOUNCER -> fromAlternateBouncerTransitionInteractor.dismissAlternateBouncer()
@@ -61,7 +68,7 @@
             KeyguardState.GONE ->
                 Log.i(
                     TAG,
-                    "Already transitioning to GONE; ignoring startDismissKeyguardTransition."
+                    "Already transitioning to GONE; ignoring startDismissKeyguardTransition.",
                 )
             else -> Log.e(TAG, "We don't know how to dismiss keyguard from state $startedState.")
         }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
index 5f08aa3..631e44a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
@@ -22,7 +22,7 @@
 import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
-import com.android.systemui.util.kotlin.Utils.Companion.sample as sampleCombine
+import com.android.systemui.util.kotlin.sample
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
@@ -74,11 +74,9 @@
             .onEach { SceneContainerFlag.assertInLegacyMode() }
             // Whenever the keyguard is disabled...
             .filter { enabled -> !enabled }
-            .sampleCombine(
-                internalTransitionInteractor.currentTransitionInfoInternal,
-                biometricSettingsRepository.isCurrentUserInLockdown,
-            )
-            .map { (_, transitionInfo, inLockdown) ->
+            .sample(biometricSettingsRepository.isCurrentUserInLockdown, ::Pair)
+            .map { (_, inLockdown) ->
+                val transitionInfo = internalTransitionInteractor.currentTransitionInfoInternal()
                 // ...we hide the keyguard, if it's showing and we're not in lockdown. In that case,
                 // we want to remember that and re-show it when keyguard is enabled again.
                 transitionInfo.to != KeyguardState.GONE && !inLockdown
@@ -93,11 +91,10 @@
             if (!SceneContainerFlag.isEnabled) {
                 repository.isKeyguardEnabled
                     .filter { enabled -> !enabled }
-                    .sampleCombine(
-                        biometricSettingsRepository.isCurrentUserInLockdown,
-                        internalTransitionInteractor.currentTransitionInfoInternal,
-                    )
-                    .collect { (_, inLockdown, currentTransitionInfo) ->
+                    .sample(biometricSettingsRepository.isCurrentUserInLockdown, ::Pair)
+                    .collect { (_, inLockdown) ->
+                        val currentTransitionInfo =
+                            internalTransitionInteractor.currentTransitionInfoInternal()
                         if (currentTransitionInfo.to != KeyguardState.GONE && !inLockdown) {
                             keyguardDismissTransitionInteractor.startDismissKeyguardTransition(
                                 "keyguard disabled"
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardOcclusionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardOcclusionInteractor.kt
index 7f1e881..278a98f 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardOcclusionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardOcclusionInteractor.kt
@@ -80,7 +80,7 @@
         // *_BOUNCER -> LOCKSCREEN.
         return powerInteractor.detailedWakefulness.value.powerButtonLaunchGestureTriggered &&
             KeyguardState.deviceIsAsleepInState(
-                internalTransitionInteractor.currentTransitionInfoInternal.value.to
+                internalTransitionInteractor.currentTransitionInfoInternal().to
             )
     }
 
@@ -100,13 +100,13 @@
                             scene = Scenes.Gone,
                             stateWithoutSceneContainer = KeyguardState.GONE,
                         ),
-                        ::Pair
+                        ::Pair,
                     )
                     .map { (wakefulness, isOnGone) ->
                         wakefulness.powerButtonLaunchGestureTriggered && !isOnGone
                     },
                 // Emit false once that activity goes away.
-                isShowWhenLockedActivityOnTop.filter { !it }.map { false }
+                isShowWhenLockedActivityOnTop.filter { !it }.map { false },
             )
             .stateIn(applicationScope, SharingStarted.Eagerly, false)
 
@@ -134,7 +134,7 @@
      */
     fun setWmNotifiedShowWhenLockedActivityOnTop(
         showWhenLockedActivityOnTop: Boolean,
-        taskInfo: RunningTaskInfo? = null
+        taskInfo: RunningTaskInfo? = null,
     ) {
         repository.setShowWhenLockedActivityInfo(showWhenLockedActivityOnTop, taskInfo)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionBootInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionBootInteractor.kt
index cddeaaf..b986d56 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionBootInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionBootInteractor.kt
@@ -61,7 +61,7 @@
 
     fun start() {
         scope.launch {
-            if (internalTransitionInteractor.currentTransitionInfoInternal.value.from != OFF) {
+            if (internalTransitionInteractor.currentTransitionInfoInternal().from != OFF) {
                 Log.e(
                     "KeyguardTransitionInteractor",
                     "showLockscreenOnBoot emitted, but we've already " +
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
index 249982d..abd7f90 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
@@ -71,14 +71,14 @@
         ownerReason: String = "",
     ): UUID? {
         toState.checkValidState()
-        if (fromState != internalTransitionInteractor.currentTransitionInfoInternal.value.to) {
+        if (fromState != internalTransitionInteractor.currentTransitionInfoInternal().to) {
             Log.e(
                 name,
                 "Ignoring startTransition: This interactor asked to transition from " +
                     "$fromState -> $toState, but we last transitioned to " +
-                    "${internalTransitionInteractor.currentTransitionInfoInternal.value.to}, not" +
+                    "${internalTransitionInteractor.currentTransitionInfoInternal().to}, not" +
                     " $fromState. This should never happen - check currentTransitionInfoInternal" +
-                    " or use filterRelevantKeyguardState before starting transitions."
+                    " or use filterRelevantKeyguardState before starting transitions.",
             )
             return null
         }
@@ -149,7 +149,7 @@
             if (keyguardInteractor.isKeyguardDismissible.value) {
                 startTransitionTo(
                     KeyguardState.GONE,
-                    ownerReason = "Power button gesture while keyguard is dismissible"
+                    ownerReason = "Power button gesture while keyguard is dismissible",
                 )
 
                 return true
@@ -159,7 +159,7 @@
                 // should transition to GONE.
                 startTransitionTo(
                     KeyguardState.GONE,
-                    ownerReason = "Power button gesture on dismissable keyguard"
+                    ownerReason = "Power button gesture on dismissable keyguard",
                 )
 
                 return true
@@ -190,16 +190,13 @@
                 startTransitionTo(
                     toState = keyguardInteractor.asleepKeyguardState.value,
                     modeOnCanceled = modeOnCanceled,
-                    ownerReason = "Sleep transition triggered"
+                    ownerReason = "Sleep transition triggered",
                 )
             }
     }
 
     /** This signal may come in before the occlusion signal, and can provide a custom transition */
-    fun listenForTransitionToCamera(
-        scope: CoroutineScope,
-        keyguardInteractor: KeyguardInteractor,
-    ) {
+    fun listenForTransitionToCamera(scope: CoroutineScope, keyguardInteractor: KeyguardInteractor) {
         if (!KeyguardWmStateRefactor.isEnabled) {
             scope.launch {
                 keyguardInteractor.onCameraLaunchDetected.filterRelevantKeyguardState().collect {
@@ -223,7 +220,7 @@
      * [startedKeyguardState] as it does not wait for the emission of the first STARTED step.
      */
     fun inOrTransitioningToRelevantKeyguardState(): Boolean {
-        return internalTransitionInteractor.currentTransitionInfoInternal.value.to == fromState
+        return internalTransitionInteractor.currentTransitionInfoInternal().to == fromState
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
index a09cd7c..a1f6067 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
@@ -19,6 +19,7 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.systemui.Flags.transitionRaceCondition
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
@@ -65,7 +66,7 @@
         combine(
             transitionInteractor.isFinishedIn(
                 scene = Scenes.Gone,
-                stateWithoutSceneContainer = KeyguardState.GONE
+                stateWithoutSceneContainer = KeyguardState.GONE,
             ),
             wakeToGoneInteractor.canWakeDirectlyToGone,
         ) { isOnGone, canWakeDirectlyToGone ->
@@ -197,11 +198,11 @@
             combine(
                     transitionInteractor.isInTransition(
                         edge = Edge.create(to = Scenes.Gone),
-                        edgeWithoutSceneContainer = Edge.create(to = KeyguardState.GONE)
+                        edgeWithoutSceneContainer = Edge.create(to = KeyguardState.GONE),
                     ),
                     transitionInteractor.isFinishedIn(
                         scene = Scenes.Gone,
-                        stateWithoutSceneContainer = KeyguardState.GONE
+                        stateWithoutSceneContainer = KeyguardState.GONE,
                     ),
                     surfaceBehindInteractor.isAnimatingSurface,
                     notificationLaunchAnimationInteractor.isLaunchAnimationRunning,
@@ -231,7 +232,7 @@
             combine(
                     transitionInteractor.currentKeyguardState,
                     wakeToGoneInteractor.canWakeDirectlyToGone,
-                    ::Pair
+                    ::Pair,
                 )
                 .sample(transitionInteractor.startedStepWithPrecedingStep, ::toTriple)
                 .map { (currentState, canWakeDirectlyToGone, startedWithPrev) ->
@@ -242,7 +243,12 @@
                             startedFromStep.transitionState == TransitionState.CANCELED &&
                             startedFromStep.from == KeyguardState.GONE
 
-                    val transitionInfo = transitionRepository.currentTransitionInfoInternal.value
+                    val transitionInfo =
+                        if (transitionRaceCondition()) {
+                            transitionRepository.currentTransitionInfo
+                        } else {
+                            transitionRepository.currentTransitionInfoInternal.value
+                        }
                     val wakingDirectlyToGone =
                         deviceIsAsleepInState(transitionInfo.from) &&
                             transitionInfo.to == KeyguardState.GONE
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt
index 5524b20..aa44b6d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt
@@ -106,7 +106,7 @@
 
     private suspend fun handleIdle(
         prevTransition: ObservableTransitionState,
-        idle: ObservableTransitionState.Idle
+        idle: ObservableTransitionState.Idle,
     ) {
         if (currentTransitionId == null) return
         if (prevTransition !is ObservableTransitionState.Transition) return
@@ -133,10 +133,10 @@
         val newTransition =
             TransitionInfo(
                 ownerName = this::class.java.simpleName,
-                from = internalTransitionInteractor.currentTransitionInfoInternal.value.to,
+                from = internalTransitionInteractor.currentTransitionInfoInternal().to,
                 to = state,
                 animator = null,
-                modeOnCanceled = TransitionModeOnCanceled.REVERSE
+                modeOnCanceled = TransitionModeOnCanceled.REVERSE,
             )
         currentTransitionId = internalTransitionInteractor.startTransition(newTransition)
         internalTransitionInteractor.updateTransition(currentTransitionId!!, 1f, FINISHED)
@@ -152,8 +152,7 @@
     private suspend fun handleTransition(transition: ObservableTransitionState.Transition) {
         if (transition.fromContent == Scenes.Lockscreen) {
             if (currentTransitionId != null) {
-                val currentToState =
-                    internalTransitionInteractor.currentTransitionInfoInternal.value.to
+                val currentToState = internalTransitionInteractor.currentTransitionInfoInternal().to
                 if (currentToState == UNDEFINED) {
                     transitionKtfTo(transitionInteractor.startedKeyguardTransitionStep.value.from)
                 }
@@ -197,21 +196,21 @@
                 from = UNDEFINED,
                 to = repository.nextLockscreenTargetState.value,
                 animator = null,
-                modeOnCanceled = TransitionModeOnCanceled.RESET
+                modeOnCanceled = TransitionModeOnCanceled.RESET,
             )
         repository.nextLockscreenTargetState.value = DEFAULT_STATE
         startTransition(newTransition)
     }
 
     private suspend fun startTransitionFromLockscreen() {
-        val currentState = internalTransitionInteractor.currentTransitionInfoInternal.value.to
+        val currentState = internalTransitionInteractor.currentTransitionInfoInternal().to
         val newTransition =
             TransitionInfo(
                 ownerName = this::class.java.simpleName,
                 from = currentState,
                 to = UNDEFINED,
                 animator = null,
-                modeOnCanceled = TransitionModeOnCanceled.RESET
+                modeOnCanceled = TransitionModeOnCanceled.RESET,
             )
         startTransition(newTransition)
     }
@@ -228,7 +227,7 @@
         internalTransitionInteractor.updateTransition(
             currentTransitionId!!,
             progress.coerceIn(0f, 1f),
-            RUNNING
+            RUNNING,
         )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DozingToOccludedTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DozingToOccludedTransitionViewModel.kt
index 12bcc7e..b15cacf 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DozingToOccludedTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DozingToOccludedTransitionViewModel.kt
@@ -35,9 +35,7 @@
 @SysUISingleton
 class DozingToOccludedTransitionViewModel
 @Inject
-constructor(
-    animationFlow: KeyguardTransitionAnimationFlow,
-) : DeviceEntryIconTransition {
+constructor(animationFlow: KeyguardTransitionAnimationFlow) : DeviceEntryIconTransition {
     private val transitionAnimation =
         animationFlow.setup(
             duration = FromAodTransitionInteractor.TO_OCCLUDED_DURATION,
@@ -56,11 +54,7 @@
         var currentAlpha = 0f
         return transitionAnimation.sharedFlow(
             duration = 250.milliseconds,
-            startTime = if (lightRevealMigration()) {
-                100.milliseconds // Wait for the light reveal to "hit" the LS elements.
-            } else {
-                0.milliseconds
-            },
+            startTime = 0.milliseconds,
             onStart = {
                 if (lightRevealMigration()) {
                     currentAlpha = viewState.alpha()
@@ -69,7 +63,6 @@
                 }
             },
             onStep = { MathUtils.lerp(currentAlpha, 0f, it) },
-            onCancel = { 0f },
         )
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
index 70b4f79..4976cc2 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
@@ -18,6 +18,7 @@
 package com.android.systemui.keyguard.data.repository
 
 import android.annotation.FloatRange
+import com.android.systemui.Flags.transitionRaceCondition
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionInfo
@@ -88,6 +89,13 @@
             )
         )
     override var currentTransitionInfoInternal = _currentTransitionInfo.asStateFlow()
+    override var currentTransitionInfo =
+        TransitionInfo(
+            ownerName = "",
+            from = KeyguardState.OFF,
+            to = KeyguardState.LOCKSCREEN,
+            animator = null,
+        )
 
     init {
         // Seed with a FINISHED transition in OFF, same as the real repository.
@@ -261,8 +269,13 @@
         validateStep: Boolean = true,
     ) {
         if (step.transitionState == TransitionState.STARTED) {
-            _currentTransitionInfo.value =
-                TransitionInfo(from = step.from, to = step.to, animator = null, ownerName = "")
+            if (transitionRaceCondition()) {
+                currentTransitionInfo =
+                    TransitionInfo(from = step.from, to = step.to, animator = null, ownerName = "")
+            } else {
+                _currentTransitionInfo.value =
+                    TransitionInfo(from = step.from, to = step.to, animator = null, ownerName = "")
+            }
         }
 
         _transitions.replayCache.last().let { lastStep ->
@@ -308,7 +321,11 @@
     }
 
     override suspend fun startTransition(info: TransitionInfo): UUID? {
-        _currentTransitionInfo.value = info
+        if (transitionRaceCondition()) {
+            currentTransitionInfo = info
+        } else {
+            _currentTransitionInfo.value = info
+        }
 
         if (sendTransitionStepsOnStartTransition) {
             sendTransitionSteps(from = info.from, to = info.to, testScope = testScope)