Fix race condition with manual transitions

Transitions are triggered on background scope. Calls to
startTransition were transfered to the main dispatcher, but,
manual transitions that used updateTransition were not. This
could lead to a very rare race condition where a RUNNING step
could be sent after the corresponding CANCEL.

Test: atest KeyguardTransitionRepositoryTest
Fixes: 353698157
Flag: com.android.systemui.migrate_clocks_to_blueprint
Change-Id: I515358f0a9230d660fba4b0703eaad38477edd46

8 files changed, 108 insertions, 46 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
index a8eccc5add1a..03647b9a3956 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
@@ -136,30 +136,6 @@ class KeyguardTransitionInteractorTest : SysuiTestCase() {
         }
 
     @Test
-    fun finishedKeyguardTransitionStepTests() =
-        testScope.runTest {
-            val finishedSteps by collectValues(underTest.finishedKeyguardTransitionStep)
-            val steps = mutableListOf<TransitionStep>()
-
-            steps.add(TransitionStep(LOCKSCREEN, AOD, 0f, STARTED))
-            steps.add(TransitionStep(LOCKSCREEN, AOD, 0.9f, RUNNING))
-            steps.add(TransitionStep(LOCKSCREEN, AOD, 1f, FINISHED))
-            steps.add(TransitionStep(AOD, LOCKSCREEN, 0f, STARTED))
-            steps.add(TransitionStep(AOD, LOCKSCREEN, 0.5f, RUNNING))
-            steps.add(TransitionStep(AOD, LOCKSCREEN, 1f, FINISHED))
-            steps.add(TransitionStep(AOD, GONE, 1f, STARTED))
-
-            steps.forEach {
-                repository.sendTransitionStep(it)
-                runCurrent()
-            }
-
-            // Ignore the default state.
-            assertThat(finishedSteps.subList(1, finishedSteps.size))
-                .isEqualTo(listOf(steps[2], steps[5]))
-        }
-
-    @Test
     fun startedKeyguardTransitionStepTests() =
         testScope.runTest {
             val startedSteps by collectValues(underTest.startedKeyguardTransitionStep)
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 6a20610da3a6..c780aac5aaca 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
@@ -175,7 +175,7 @@ constructor(
         }
     }
 
-    private fun finishCurrentTransition() {
+    private suspend fun finishCurrentTransition() {
         internalTransitionInteractor.updateTransition(
             currentTransitionId!!,
             1f,
@@ -285,7 +285,7 @@ constructor(
         currentTransitionId = internalTransitionInteractor.startTransition(transitionInfo)
     }
 
-    private fun updateProgress(progress: Float) {
+    private suspend fun updateProgress(progress: Float) {
         if (currentTransitionId == null) return
         internalTransitionInteractor.updateTransition(
             currentTransitionId!!,
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 aaeeb3911797..de60c1117c19 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
@@ -105,7 +105,7 @@ interface KeyguardTransitionRepository {
      * When the transition is over, TransitionState.FINISHED must be passed into the [state]
      * parameter.
      */
-    fun updateTransition(
+    suspend fun updateTransition(
         transitionId: UUID,
         @FloatRange(from = 0.0, to = 1.0) value: Float,
         state: TransitionState
@@ -173,9 +173,9 @@ constructor(
         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.
+        // be processed out of order. Use a Mutex to guarantee ordering. [updateTransition]
+        // requires the same lock
         _currentTransitionMutex.lock()
-
         // Only used in a test environment
         if (forceDelayForRaceConditionTest) {
             delay(50L)
@@ -184,7 +184,6 @@ constructor(
         // Animators must be started on the main thread.
         return withContext("$TAG#startTransition", mainDispatcher) {
             _currentTransitionMutex.unlock()
-
             if (lastStep.from == info.from && lastStep.to == info.to) {
                 Log.i(TAG, "Duplicate call to start the transition, rejecting: $info")
                 return@withContext null
@@ -206,7 +205,7 @@ constructor(
 
             // Cancel any existing manual transitions
             updateTransitionId?.let { uuid ->
-                updateTransition(uuid, lastStep.value, TransitionState.CANCELED)
+                updateTransitionInternal(uuid, lastStep.value, TransitionState.CANCELED)
             }
 
             info.animator?.let { animator ->
@@ -264,7 +263,23 @@ constructor(
         }
     }
 
-    override fun updateTransition(
+    override suspend fun updateTransition(
+        transitionId: UUID,
+        @FloatRange(from = 0.0, to = 1.0) value: Float,
+        state: TransitionState
+    ) {
+        // 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()
+        withContext("$TAG#updateTransition", mainDispatcher) {
+            _currentTransitionMutex.unlock()
+
+            updateTransitionInternal(transitionId, value, state)
+        }
+    }
+
+    private suspend fun updateTransitionInternal(
         transitionId: UUID,
         @FloatRange(from = 0.0, to = 1.0) value: Float,
         state: TransitionState
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 a51421c10309..2cc6afa2f407 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
@@ -63,7 +63,7 @@ constructor(
 
     suspend fun startTransition(info: TransitionInfo) = repository.startTransition(info)
 
-    fun updateTransition(
+    suspend fun updateTransition(
         transitionId: UUID,
         @FloatRange(from = 0.0, to = 1.0) value: Float,
         state: TransitionState
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
index 797d4667c56d..afbe3579315d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
@@ -47,6 +47,7 @@ import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
@@ -55,6 +56,7 @@ import kotlinx.coroutines.flow.mapLatest
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.transform
 import kotlinx.coroutines.launch
 
 /** Encapsulates business-logic related to the keyguard transitions. */
@@ -150,6 +152,12 @@ constructor(
                         startedStep.to != prevStep.from
                 ) {
                     getTransitionValueFlow(prevStep.from).emit(0f)
+                } else if (prevStep.transitionState == TransitionState.RUNNING) {
+                    Log.e(
+                        TAG,
+                        "STARTED step ($startedStep) was preceded by a RUNNING step " +
+                            "($prevStep), which should never happen. Things could go badly here."
+                    )
                 }
             }
         }
@@ -252,15 +260,12 @@ constructor(
     val startedKeyguardTransitionStep: Flow<TransitionStep> =
         repository.transitions.filter { step -> step.transitionState == TransitionState.STARTED }
 
-    /** The last [TransitionStep] with a [TransitionState] of FINISHED */
-    val finishedKeyguardTransitionStep: Flow<TransitionStep> =
-        repository.transitions.filter { step -> step.transitionState == TransitionState.FINISHED }
-
     /** The destination state of the last [TransitionState.STARTED] transition. */
     @SuppressLint("SharedFlowCreation")
     val startedKeyguardState: SharedFlow<KeyguardState> =
         startedKeyguardTransitionStep
             .map { step -> step.to }
+            .buffer(2, BufferOverflow.DROP_OLDEST)
             .shareIn(scope, SharingStarted.Eagerly, replay = 1)
 
     /** The from state of the last [TransitionState.STARTED] transition. */
@@ -269,6 +274,7 @@ constructor(
     val startedKeyguardFromState: SharedFlow<KeyguardState> =
         startedKeyguardTransitionStep
             .map { step -> step.from }
+            .buffer(2, BufferOverflow.DROP_OLDEST)
             .shareIn(scope, SharingStarted.Eagerly, replay = 1)
 
     /** Which keyguard state to use when the device goes to sleep. */
@@ -310,8 +316,13 @@ constructor(
      */
     @SuppressLint("SharedFlowCreation")
     val finishedKeyguardState: SharedFlow<KeyguardState> =
-        finishedKeyguardTransitionStep
-            .map { step -> step.to }
+        repository.transitions
+            .transform { step ->
+                if (step.transitionState == TransitionState.FINISHED) {
+                    emit(step.to)
+                }
+            }
+            .buffer(2, BufferOverflow.DROP_OLDEST)
             .shareIn(scope, SharingStarted.Eagerly, replay = 1)
 
     /**
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 324811443e9d..b8500952d90a 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
@@ -124,7 +124,7 @@ constructor(
         }
     }
 
-    private fun finishCurrentTransition() {
+    private suspend fun finishCurrentTransition() {
         internalTransitionInteractor.updateTransition(currentTransitionId!!, 1f, FINISHED)
         resetTransitionData()
     }
@@ -223,7 +223,7 @@ constructor(
         currentTransitionId = internalTransitionInteractor.startTransition(transitionInfo)
     }
 
-    private fun updateProgress(progress: Float) {
+    private suspend fun updateProgress(progress: Float) {
         if (currentTransitionId == null) return
         internalTransitionInteractor.updateTransition(
             currentTransitionId!!,
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 48a5df91d47c..2af4d872f6a0 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
@@ -45,6 +45,7 @@ import java.math.BigDecimal
 import java.math.RoundingMode
 import java.util.UUID
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.dropWhile
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
@@ -332,10 +333,11 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
         }
 
     @Test
-    fun attemptTomanuallyUpdateTransitionWithInvalidUUIDthrowsException() {
-        underTest.updateTransition(UUID.randomUUID(), 0f, TransitionState.RUNNING)
-        assertThat(wtfHandler.failed).isTrue()
-    }
+    fun attemptTomanuallyUpdateTransitionWithInvalidUUIDthrowsException() =
+        testScope.runTest {
+            underTest.updateTransition(UUID.randomUUID(), 0f, TransitionState.RUNNING)
+            assertThat(wtfHandler.failed).isTrue()
+        }
 
     @Test
     fun attemptToManuallyUpdateTransitionAfterFINISHEDstateThrowsException() =
@@ -414,6 +416,64 @@ class KeyguardTransitionRepositoryTest : SysuiTestCase() {
                 )
         }
 
+    @Test
+    fun simulateRaceConditionIsProcessedInOrderUsingUpdateTransition() =
+        testScope.runTest {
+            val ktr = KeyguardTransitionRepositoryImpl(kosmos.testDispatcher)
+            val steps by collectValues(ktr.transitions.dropWhile { step -> step.from == OFF })
+
+            // Begin a manual transition
+            val info1 = TransitionInfo(OWNER_NAME, AOD, LOCKSCREEN, animator = null)
+            launch {
+                ktr.forceDelayForRaceConditionTest = false
+                val uuid = ktr.startTransition(info1)
+
+                // Pause here to allow another transition to start
+                delay(20)
+
+                // Attempt to send an update, which should fail
+                ktr.updateTransition(uuid!!, 0.5f, TransitionState.RUNNING)
+            }
+
+            // Now start another transition, which should acquire the preempt the first
+            val info2 = TransitionInfo(OWNER_NAME, LOCKSCREEN, OCCLUDED, animator = null)
+            launch {
+                delay(10)
+                ktr.forceDelayForRaceConditionTest = true
+                ktr.startTransition(info2)
+            }
+
+            runCurrent()
+
+            // Manual transition has started
+            assertThat(steps[0])
+                .isEqualTo(
+                    TransitionStep(info1.from, info1.to, 0f, TransitionState.STARTED, OWNER_NAME)
+                )
+
+            // The second transition has requested to start, and grabbed the mutex. But it is
+            // delayed
+            advanceTimeBy(15L)
+
+            // Advancing another 10ms should now trigger the first transition to request an update,
+            // which should not happen as the second transition has the mutex
+            advanceTimeBy(10L)
+
+            // Finally, advance past the delay in the second transition so it can run
+            advanceTimeBy(50L)
+
+            assertThat(steps[1])
+                .isEqualTo(
+                    TransitionStep(info1.from, info1.to, 0f, TransitionState.CANCELED, OWNER_NAME)
+                )
+            assertThat(steps[2])
+                .isEqualTo(
+                    TransitionStep(info2.from, info2.to, 0f, TransitionState.STARTED, OWNER_NAME)
+                )
+
+            assertThat(steps.size).isEqualTo(3)
+        }
+
     private fun listWithStep(
         step: BigDecimal,
         start: BigDecimal = BigDecimal.ZERO,
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 b5ea619cd57f..616f2b688746 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
@@ -283,7 +283,7 @@ class FakeKeyguardTransitionRepository(
         )
     }
 
-    override fun updateTransition(
+    override suspend fun updateTransition(
         transitionId: UUID,
         @FloatRange(from = 0.0, to = 1.0) value: Float,
         state: TransitionState