Fix race in DisplayRepository tests

while on real devices tests were passing fine, on cuttlefish there were some flakes.

Those were because sometimes the "initial value" of enabledDisplayId was being propagated before or after the initial value of the "scan". When the "scan" initial value was propagated first, only the default display was in the output for the first event. When the StateIn initial value was propagated first, everything worked as expected.

Now the same initial value is provided to the scan and the sharein for enabledDisplayIds.

Also, cleaning up the flag (now in next)

Bug: 345472038
Bug: 358145460
Test: DisplayRepositoryTest
Flag: com.android.systemui.enable_efficient_display_repository
Change-Id: Ie1e64271d01936ee064f780b1a1f390b769cca31

2 files changed, 56 insertions, 70 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
index 1f5878b4698f..a327e4a76c49 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
@@ -28,7 +28,6 @@ import android.util.Log
 import android.view.Display
 import com.android.app.tracing.FlowTracing.traceEach
 import com.android.app.tracing.traceSection
-import com.android.systemui.Flags
 import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
@@ -51,7 +50,6 @@ import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.scan
-import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 
 /** Provides a [Flow] of [Display] as returned by [DisplayManager]. */
@@ -102,7 +100,7 @@ constructor(
     private val displayManager: DisplayManager,
     @Background backgroundHandler: Handler,
     @Background bgApplicationScope: CoroutineScope,
-    @Background backgroundCoroutineDispatcher: CoroutineDispatcher
+    @Background backgroundCoroutineDispatcher: CoroutineDispatcher,
 ) : DisplayRepository {
     private val allDisplayEvents: Flow<DisplayEvent> =
         conflatedCallbackFlow {
@@ -139,70 +137,55 @@ constructor(
     override val displayAdditionEvent: Flow<Display?> =
         allDisplayEvents.filterIsInstance<DisplayEvent.Added>().map { getDisplay(it.displayId) }
 
-    // TODO: b/345472038 - Delete after the flag is ramped up.
-    private val oldEnabledDisplays: Flow<Set<Display>> =
-        allDisplayEvents
-            .map { getDisplays() }
-            .shareIn(bgApplicationScope, started = SharingStarted.WhileSubscribed(), replay = 1)
+    // This is necessary because there might be multiple displays, and we could
+    // have missed events for those added before this process or flow started.
+    // Note it causes a binder call from the main thread (it's traced).
+    private val initialDisplays: Set<Display> =
+        traceSection("$TAG#initialDisplays") { displayManager.displays?.toSet() ?: emptySet() }
+    private val initialDisplayIds = initialDisplays.map { display -> display.displayId }.toSet()
 
     /** Propagate to the listeners only enabled displays */
     private val enabledDisplayIds: Flow<Set<Int>> =
-        if (Flags.enableEfficientDisplayRepository()) {
-                allDisplayEvents
-                    .scan(initial = emptySet()) { previousIds: Set<Int>, event: DisplayEvent ->
-                        val id = event.displayId
-                        when (event) {
-                            is DisplayEvent.Removed -> previousIds - id
-                            is DisplayEvent.Added,
-                            is DisplayEvent.Changed -> previousIds + id
-                        }
-                    }
-                    .distinctUntilChanged()
-                    .stateIn(
-                        bgApplicationScope,
-                        SharingStarted.WhileSubscribed(),
-                        // This is necessary because there might be multiple displays, and we could
-                        // have missed events for those added before this process or flow started.
-                        // Note it causes a binder call from the main thread (it's traced).
-                        getDisplays().map { display -> display.displayId }.toSet(),
-                    )
-            } else {
-                oldEnabledDisplays.map { enabledDisplaysSet ->
-                    enabledDisplaysSet.map { it.displayId }.toSet()
+        allDisplayEvents
+            .scan(initial = initialDisplayIds) { previousIds: Set<Int>, event: DisplayEvent ->
+                val id = event.displayId
+                when (event) {
+                    is DisplayEvent.Removed -> previousIds - id
+                    is DisplayEvent.Added,
+                    is DisplayEvent.Changed -> previousIds + id
                 }
             }
+            .distinctUntilChanged()
+            .stateIn(bgApplicationScope, SharingStarted.WhileSubscribed(), initialDisplayIds)
             .debugLog("enabledDisplayIds")
 
     private val defaultDisplay by lazy {
         getDisplay(Display.DEFAULT_DISPLAY) ?: error("Unable to get default display.")
     }
+
     /**
      * Represents displays that went though the [DisplayListener.onDisplayAdded] callback.
      *
      * Those are commonly the ones provided by [DisplayManager.getDisplays] by default.
      */
     private val enabledDisplays: Flow<Set<Display>> =
-        if (Flags.enableEfficientDisplayRepository()) {
-            enabledDisplayIds
-                .mapElementsLazily { displayId -> getDisplay(displayId) }
-                .onEach {
-                    if (it.isEmpty()) Log.wtf(TAG, "No enabled displays. This should never happen.")
-                }
-                .flowOn(backgroundCoroutineDispatcher)
-                .debugLog("enabledDisplays")
-                .stateIn(
-                    bgApplicationScope,
-                    started = SharingStarted.WhileSubscribed(),
-                    // This triggers a single binder call on the UI thread per process. The
-                    // alternative would be to use sharedFlows, but they are prohibited due to
-                    // performance concerns.
-                    // Ultimately, this is a trade-off between a one-time UI thread binder call and
-                    // the constant overhead of sharedFlows.
-                    initialValue = getDisplays()
-                )
-        } else {
-            oldEnabledDisplays
-        }
+        enabledDisplayIds
+            .mapElementsLazily { displayId -> getDisplay(displayId) }
+            .onEach {
+                if (it.isEmpty()) Log.wtf(TAG, "No enabled displays. This should never happen.")
+            }
+            .flowOn(backgroundCoroutineDispatcher)
+            .debugLog("enabledDisplays")
+            .stateIn(
+                bgApplicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                // This triggers a single binder call on the UI thread per process. The
+                // alternative would be to use sharedFlows, but they are prohibited due to
+                // performance concerns.
+                // Ultimately, this is a trade-off between a one-time UI thread binder call and
+                // the constant overhead of sharedFlows.
+                initialValue = initialDisplays,
+            )
 
     /**
      * Represents displays that went though the [DisplayListener.onDisplayAdded] callback.
@@ -211,10 +194,7 @@ constructor(
      */
     override val displays: Flow<Set<Display>> = enabledDisplays
 
-    private fun getDisplays(): Set<Display> =
-        traceSection("$TAG#getDisplays()") { displayManager.displays?.toSet() ?: emptySet() }
-
-    private val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
+    val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
     private val ignoredDisplayIds: Flow<Set<Int>> = _ignoredDisplayIds.debugLog("ignoredDisplayIds")
 
     private fun getInitialConnectedDisplays(): Set<Int> =
@@ -271,7 +251,7 @@ constructor(
                 // the flow starts being collected. This is to ensure the call to get displays (an
                 // IPC) happens in the background instead of when this object
                 // is instantiated.
-                initialValue = emptySet()
+                initialValue = emptySet(),
             )
 
     private val connectedExternalDisplayIds: Flow<Set<Int>> =
@@ -308,7 +288,7 @@ constructor(
                         TAG,
                         "combining enabled=$enabledDisplaysIds, " +
                             "connectedExternalDisplayIds=$connectedExternalDisplayIds, " +
-                            "ignored=$ignoredDisplayIds"
+                            "ignored=$ignoredDisplayIds",
                     )
                 }
                 connectedExternalDisplayIds - enabledDisplaysIds - ignoredDisplayIds
@@ -382,7 +362,7 @@ constructor(
             val previousSet: Set<T>,
             // Caches T values from the previousSet that were already converted to V
             val valueMap: Map<T, V>,
-            val resultSet: Set<V>
+            val resultSet: Set<V>,
         )
 
         val emptyInitialState = State(emptySet<T>(), emptyMap(), emptySet<V>())
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
index 76539d776d3d..633efd8bfffd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
@@ -63,21 +63,27 @@ class DisplayRepositoryTest : SysuiTestCase() {
     private val defaultDisplay =
         display(type = TYPE_INTERNAL, id = DEFAULT_DISPLAY, state = Display.STATE_ON)
 
-    private lateinit var displayRepository: DisplayRepositoryImpl
+    // This is Lazy as displays could be set before the instance is created, and we want to verify
+    // that the initial state (soon after construction) contains the expected ones set in every
+    // test.
+    private val displayRepository: DisplayRepositoryImpl by lazy {
+        DisplayRepositoryImpl(
+                displayManager,
+                testHandler,
+                TestScope(UnconfinedTestDispatcher()),
+                UnconfinedTestDispatcher(),
+            )
+            .also {
+                verify(displayManager, never()).registerDisplayListener(any(), any())
+                // It needs to be called, just once, for the initial value.
+                verify(displayManager).getDisplays()
+            }
+    }
 
     @Before
     fun setup() {
         setDisplays(listOf(defaultDisplay))
         setAllDisplaysIncludingDisabled(DEFAULT_DISPLAY)
-        displayRepository =
-            DisplayRepositoryImpl(
-                displayManager,
-                testHandler,
-                TestScope(UnconfinedTestDispatcher()),
-                UnconfinedTestDispatcher()
-            )
-        verify(displayManager, never()).registerDisplayListener(any(), any())
-        verify(displayManager, never()).getDisplays(any())
     }
 
     @Test
@@ -502,7 +508,7 @@ class DisplayRepositoryTest : SysuiTestCase() {
             .registerDisplayListener(
                 connectedDisplayListener.capture(),
                 eq(testHandler),
-                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED)
+                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED),
             )
         return flowValue
     }
@@ -522,7 +528,7 @@ class DisplayRepositoryTest : SysuiTestCase() {
                     DisplayManager.EVENT_FLAG_DISPLAY_ADDED or
                         DisplayManager.EVENT_FLAG_DISPLAY_CHANGED or
                         DisplayManager.EVENT_FLAG_DISPLAY_REMOVED
-                )
+                ),
             )
     }