Initialize enabledDisplays with the correct displays

It could have happened that the list of displays was not up to date when DisplayRepository was being initialized.
This makes a binder call to get them from the ui thread once DisplayRepository class is instantiated to ensure correctness immediately (and not eventual, as before)

It would be possible to also have correctness immediately avoiding the binder call using SharedFlows, but we banned them due to performance concerns. The ui thread binder call here happens once per process start, so it shouldn't be that problematic.

Fixes: 355364683
Test: DisplayRepositoryTest
Flag: com.android.systemui.enable_efficient_display_repository
Change-Id: If379d754fa378eaaeddb29fc5f94f5536f0b5c01

2 files changed, 23 insertions, 3 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 69ddb62cc05a..40e2f174cfb7 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
@@ -160,7 +160,10 @@ constructor(
                     .stateIn(
                         bgApplicationScope,
                         SharingStarted.WhileSubscribed(),
-                        emptySet(),
+                        // 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 ->
@@ -186,8 +189,12 @@ constructor(
                 .stateIn(
                     bgApplicationScope,
                     started = SharingStarted.WhileSubscribed(),
-                    initialValue = setOf(defaultDisplay)
-                )
+                    // 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
         }
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 20cb1e129f49..0ac04b61f13d 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
@@ -28,6 +28,7 @@ import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.FlowValue
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.coroutines.collectValues
 import com.android.systemui.util.mockito.kotlinArgumentCaptor
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
@@ -456,8 +457,20 @@ class DisplayRepositoryTest : SysuiTestCase() {
             assertThat(value?.ids()).containsExactly(DEFAULT_DISPLAY)
         }
 
+    @Test
+    fun displayFlow_emitsCorrectDisplaysAtFirst() =
+        testScope.runTest {
+            setDisplays(0, 1, 2)
+
+            val values: List<Set<Display>> by collectValues(displayRepository.displays)
+
+            assertThat(values.toIdSets()).containsExactly(setOf(0, 1, 2))
+        }
+
     private fun Iterable<Display>.ids(): List<Int> = map { it.displayId }
 
+    private fun Iterable<Set<Display>>.toIdSets(): List<Set<Int>> = map { it.ids().toSet() }
+
     // Wrapper to capture the displayListener.
     private fun TestScope.latestDisplayFlowValue(): FlowValue<Set<Display>?> {
         val flowValue = collectLastValue(displayRepository.displays)