Optimize DisplayRepository to reduce lock contention

Despite all the calls to getDisplays() were going on in the background, they were often taking a lot of time, and ending up keeping a lock that was blocking other threads and processes.

This change reduces the binder calls to get displays to only one per addition.

Bug: 345472038
Bug: 344362341
Test: DisplayRepositoryTest
Flag: com.android.systemui.enable_efficient_display_repository
Change-Id: I77a463db42c5ae2c06d59f6bcef2f3f89fa05bef
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 85aa33a..c61f996 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -1076,6 +1076,16 @@
 }
 
 flag {
+   name: "enable_efficient_display_repository"
+   namespace: "systemui"
+   description: "Decide whether to use the new implementation of DisplayRepository that minimizes binder calls and background lock contention."
+   bug: "345472038"
+   metadata {
+     purpose: PURPOSE_BUGFIX
+   }
+}
+
+flag {
   name: "notification_media_manager_background_execution"
   namespace: "systemui"
   description: "Decide whether to execute binder calls in background thread"
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 4ac0c56..9d6c2bf 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
@@ -16,6 +16,7 @@
 
 package com.android.systemui.display.data.repository
 
+import android.annotation.SuppressLint
 import android.hardware.display.DisplayManager
 import android.hardware.display.DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED
 import android.hardware.display.DisplayManager.DisplayListener
@@ -27,6 +28,7 @@
 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
@@ -42,11 +44,12 @@
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.filterIsInstance
 import kotlinx.coroutines.flow.flowOn
 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
 
@@ -91,6 +94,7 @@
 }
 
 @SysUISingleton
+@SuppressLint("SharedFlowCreation")
 class DisplayRepositoryImpl
 @Inject
 constructor(
@@ -132,8 +136,50 @@
         allDisplayEvents.filterIsInstance<DisplayEvent.Changed>().map { event -> event.displayId }
 
     override val displayAdditionEvent: Flow<Display?> =
-        allDisplayEvents.filterIsInstance<DisplayEvent.Added>().map {
-            displayManager.getDisplay(it.displayId)
+        allDisplayEvents.filterIsInstance<DisplayEvent.Added>().map { getDisplay(it.displayId) }
+
+    private val oldEnabledDisplays: Flow<Set<Display>> =
+        allDisplayEvents
+            .map { getDisplays() }
+            .shareIn(bgApplicationScope, started = SharingStarted.WhileSubscribed(), replay = 1)
+
+    /** 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
+                        }
+                    }
+                    .shareIn(
+                        bgApplicationScope,
+                        started = SharingStarted.WhileSubscribed(),
+                        replay = 1
+                    )
+            } else {
+                oldEnabledDisplays.map { enabledDisplaysSet ->
+                    enabledDisplaysSet.map { it.displayId }.toSet()
+                }
+            }
+            .debugLog("enabledDisplayIds")
+
+    /**
+     * 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) }
+                .flowOn(backgroundCoroutineDispatcher)
+                .debugLog("enabledDisplayIds")
+        } else {
+            oldEnabledDisplays
         }
 
     /**
@@ -141,35 +187,26 @@
      *
      * Those are commonly the ones provided by [DisplayManager.getDisplays] by default.
      */
-    private val enabledDisplays =
-        allDisplayEvents
-            .map { getDisplays() }
-            .shareIn(bgApplicationScope, started = SharingStarted.WhileSubscribed(), replay = 1)
-
     override val displays: Flow<Set<Display>> = enabledDisplays
 
     private fun getDisplays(): Set<Display> =
         traceSection("$TAG#getDisplays()") { displayManager.displays?.toSet() ?: emptySet() }
 
-    /** Propagate to the listeners only enabled displays */
-    private val enabledDisplayIds: Flow<Set<Int>> =
-        enabledDisplays
-            .map { enabledDisplaysSet -> enabledDisplaysSet.map { it.displayId }.toSet() }
-            .debugLog("enabledDisplayIds")
-
     private val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
     private val ignoredDisplayIds: Flow<Set<Int>> = _ignoredDisplayIds.debugLog("ignoredDisplayIds")
 
     private fun getInitialConnectedDisplays(): Set<Int> =
-        displayManager
-            .getDisplays(DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED)
-            .map { it.displayId }
-            .toSet()
-            .also {
-                if (DEBUG) {
-                    Log.d(TAG, "getInitialConnectedDisplays: $it")
+        traceSection("$TAG#getInitialConnectedDisplays") {
+            displayManager
+                .getDisplays(DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED)
+                .map { it.displayId }
+                .toSet()
+                .also {
+                    if (DEBUG) {
+                        Log.d(TAG, "getInitialConnectedDisplays: $it")
+                    }
                 }
-            }
+        }
 
     /* keeps connected displays until they are disconnected. */
     private val connectedDisplayIds: StateFlow<Set<Int>> =
@@ -230,6 +267,9 @@
     private fun getDisplayType(displayId: Int): Int? =
         traceSection("$TAG#getDisplayType") { displayManager.getDisplay(displayId)?.type }
 
+    private fun getDisplay(displayId: Int): Display? =
+        traceSection("$TAG#getDisplay") { displayManager.getDisplay(displayId) }
+
     /**
      * Pending displays are the ones connected, but not enabled and not ignored.
      *
@@ -307,6 +347,30 @@
         }
     }
 
+    /**
+     * Maps a set of T to a set of V, minimizing the number of `createValue` calls taking into
+     * account the diff between each root flow emission.
+     *
+     * This is needed to minimize the number of [getDisplay] in this class. Note that if the
+     * [createValue] returns a null element, it will not be added in the output set.
+     */
+    private fun <T, V> Flow<Set<T>>.mapElementsLazily(createValue: (T) -> V?): Flow<Set<V>> {
+        var initialSet = emptySet<T>()
+        val currentMap = mutableMapOf<T, V>()
+        var resultSet = emptySet<V>()
+        return onEach { currentSet ->
+                val removed = initialSet - currentSet
+                val added = currentSet - initialSet
+                if (added.isNotEmpty() || removed.isNotEmpty()) {
+                    added.forEach { key: T -> createValue(key)?.let { currentMap[key] = it } }
+                    removed.forEach { key: T -> currentMap.remove(key) }
+                    resultSet = currentMap.values.toSet() // Creates a **copy** of values
+                }
+                initialSet = currentSet
+            }
+            .map { resultSet }
+    }
+
     private companion object {
         const val TAG = "DisplayRepository"
         val DEBUG = Log.isLoggable(TAG, Log.DEBUG) || Compile.IS_DEBUG