Use SystemUIDisplayComponent scope in PerDisplayRepository<CoroutineScope>

This allows classes outside the components to use the scope created inside the component.

As there are several places outside the display component using this scope, this is an intermediate step in the refactor.

Future classes requiring the coroutine scope are expected to be created inside the display subcomponent

Bug: 403481022
Test: DisplayScopeRepositoryInstanceProviderTest, DisplayComponentPerDisplayRepositoryTest
Flag: NONE - refactor without behaviour changes
Change-Id: I95cf815cbfdc3c334ff404434d7ff847a3342330

7 files changed, 126 insertions, 23 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/display/data/repository/DisplayComponentPerDisplayRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/display/data/repository/DisplayComponentPerDisplayRepositoryTest.kt
new file mode 100644
index 000000000000..c2ea6b994d11
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/display/data/repository/DisplayComponentPerDisplayRepositoryTest.kt
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.display.data.repository
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.useUnconfinedTestDispatcher
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.isActive
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class DisplayComponentPerDisplayRepositoryTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos().useUnconfinedTestDispatcher()
+    private val testScope = kosmos.testScope
+
+    private val underTest =
+        DisplayComponentInstanceProvider(kosmos.fakeSysuiDisplayComponentFactory)
+
+    @Test
+    fun createInstance_activeByDefault() =
+        testScope.runTest {
+            val scopeForDisplay = underTest.createInstance(displayId = 1)!!.displayCoroutineScope
+
+            assertThat(scopeForDisplay.isActive).isTrue()
+        }
+
+    @Test
+    fun destroyInstance_afterDisplayRemoved_scopeIsCancelled() =
+        testScope.runTest {
+            val component = underTest.createInstance(displayId = 1)
+            val scope = component!!.displayCoroutineScope
+
+            underTest.destroyInstance(component)
+
+            assertThat(scope.isActive).isFalse()
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryInstanceProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryInstanceProviderTest.kt
index 8afaea32273d..69ebee45a603 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryInstanceProviderTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryInstanceProviderTest.kt
@@ -16,11 +16,11 @@
 
 package com.android.systemui.display.data.repository
 
+import android.view.Display
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.kosmos.applicationCoroutineScope
-import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.kosmos.useUnconfinedTestDispatcher
 import com.android.systemui.testKosmos
@@ -36,28 +36,31 @@ class DisplayScopeRepositoryInstanceProviderTest : SysuiTestCase() {
 
     private val kosmos = testKosmos().useUnconfinedTestDispatcher()
     private val testScope = kosmos.testScope
+    private val displaySubcomponentRepository = kosmos.displaySubcomponentPerDisplayRepository
 
     private val underTest =
         DisplayScopeRepositoryInstanceProvider(
             kosmos.applicationCoroutineScope,
-            kosmos.testDispatcher,
+            displaySubcomponentRepository,
         )
 
     @Test
     fun createInstance_activeByDefault() =
         testScope.runTest {
-            val scopeForDisplay = underTest.createInstance(displayId = 1)
+            displaySubcomponentRepository.add(displayId = 1, kosmos.createFakeDisplaySubcomponent())
+            val scopeForDisplay = underTest.createInstance(displayId = 1)!!
 
             assertThat(scopeForDisplay.isActive).isTrue()
         }
 
     @Test
-    fun destroyInstance_afterDisplayRemoved_scopeIsCancelled() =
+    fun createInstance_forDefaultDisplay_returnsConstructorParam() =
         testScope.runTest {
-            val scopeForDisplay = underTest.createInstance(displayId = 1)
+            val scopeForDisplay = underTest.createInstance(displayId = Display.DEFAULT_DISPLAY)!!
 
-            underTest.destroyInstance(scopeForDisplay)
-
-            assertThat(scopeForDisplay.isActive).isFalse()
+            assertThat(scopeForDisplay).isEqualTo(kosmos.applicationCoroutineScope)
         }
+
+    // no test for destruction, as it's not handled by this class. The scope is meant to be
+    // destroyed by the PerDisplayRepository<SystemUIDisplaySubcomponent>
 }
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/PerDisplayRepositoriesModule.kt b/packages/SystemUI/src/com/android/systemui/dagger/PerDisplayRepositoriesModule.kt
index a8f21885907b..4bd275d50d23 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/PerDisplayRepositoriesModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/dagger/PerDisplayRepositoriesModule.kt
@@ -19,6 +19,7 @@ package com.android.systemui.dagger
 import com.android.app.displaylib.DefaultDisplayOnlyInstanceRepositoryImpl
 import com.android.app.displaylib.PerDisplayInstanceRepositoryImpl
 import com.android.app.displaylib.PerDisplayRepository
+import com.android.systemui.display.data.repository.DisplayComponentRepository
 import com.android.systemui.display.data.repository.PerDisplayCoroutineScopeRepositoryModule
 import com.android.systemui.model.SysUIStateInstanceProvider
 import com.android.systemui.model.SysUiState
@@ -27,7 +28,9 @@ import dagger.Module
 import dagger.Provides
 
 /** This module is meant to contain all the code to create the various [PerDisplayRepository<>]. */
-@Module(includes = [PerDisplayCoroutineScopeRepositoryModule::class])
+@Module(
+    includes = [PerDisplayCoroutineScopeRepositoryModule::class, DisplayComponentRepository::class]
+)
 class PerDisplayRepositoriesModule {
 
     @SysUISingleton
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index edee64e50b53..38c456f03c72 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -66,6 +66,7 @@ import com.android.systemui.demomode.dagger.DemoModeModule;
 import com.android.systemui.deviceentry.DeviceEntryModule;
 import com.android.systemui.display.DisplayModule;
 import com.android.app.displaylib.PerDisplayRepository;
+import com.android.systemui.display.dagger.SystemUIDisplaySubcomponent;
 import com.android.systemui.doze.dagger.DozeComponent;
 import com.android.systemui.dreams.dagger.DreamModule;
 import com.android.systemui.flags.FeatureFlags;
@@ -299,6 +300,7 @@ import javax.inject.Named;
                 NavigationBarComponent.class,
                 NotificationRowComponent.class,
                 WindowRootViewComponent.class,
+                SystemUIDisplaySubcomponent.class,
         })
 public abstract class SystemUIModule {
 
diff --git a/packages/SystemUI/src/com/android/systemui/display/dagger/SystemUIDisplaySubcomponent.kt b/packages/SystemUI/src/com/android/systemui/display/dagger/SystemUIDisplaySubcomponent.kt
index f878aa1ea87b..c83b99ca2d93 100644
--- a/packages/SystemUI/src/com/android/systemui/display/dagger/SystemUIDisplaySubcomponent.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/dagger/SystemUIDisplaySubcomponent.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.display.dagger
 
+import com.android.systemui.display.dagger.SystemUIDisplaySubcomponent.PerDisplaySingleton
 import dagger.BindsInstance
 import dagger.Subcomponent
 import javax.inject.Qualifier
@@ -31,10 +32,11 @@ import kotlinx.coroutines.CoroutineScope
  * cancelled in the background, so any teardown logic should be threadsafe. Cancelling on the main
  * thread is not feasible as it would cause jank.
  */
+@PerDisplaySingleton
 @Subcomponent(modules = [PerDisplayCommonModule::class])
 interface SystemUIDisplaySubcomponent {
 
-    @DisplayAware val displayCoroutineScope: CoroutineScope
+    @get:DisplayAware val displayCoroutineScope: CoroutineScope
 
     @Subcomponent.Factory
     interface Factory {
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayScopeRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayScopeRepository.kt
index ac59b16bfe3c..90dcd018e4fe 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayScopeRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayScopeRepository.kt
@@ -17,45 +17,48 @@
 package com.android.systemui.display.data.repository
 
 import android.view.Display
-import com.android.app.displaylib.PerDisplayInstanceProviderWithTeardown
+import com.android.app.displaylib.PerDisplayInstanceProvider
 import com.android.app.displaylib.PerDisplayInstanceRepositoryImpl
 import com.android.app.displaylib.PerDisplayRepository
-import com.android.systemui.coroutines.newTracingContext
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.display.dagger.SystemUIDisplaySubcomponent
+import com.android.systemui.display.dagger.SystemUIDisplaySubcomponent.PerDisplaySingleton
 import dagger.Module
 import dagger.Provides
 import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.cancel
 
 /**
  * Provides per display instances of [CoroutineScope].
  *
- * This is used to create a [PerDisplayRepository] of [CoroutineScope]
+ * This is used to create a [PerDisplayRepository] of [CoroutineScope].
+ *
+ * Note this scope is cancelled when the display is removed, see [DisplayComponentRepository]. This
+ * class is essentially only needed to reuse the application background scope for the default
+ * display, and get the display scope from the correct [SystemUIDisplaySubcomponent].
+ *
+ * We should eventually delete this, when all classes using display scoped instances are in the
+ * correct dagger scope ([PerDisplaySingleton])
  */
 @SysUISingleton
 class DisplayScopeRepositoryInstanceProvider
 @Inject
 constructor(
     @Background private val backgroundApplicationScope: CoroutineScope,
-    @Background private val backgroundDispatcher: CoroutineDispatcher,
-) : PerDisplayInstanceProviderWithTeardown<CoroutineScope> {
+    private val displayComponentRepository: PerDisplayRepository<SystemUIDisplaySubcomponent>,
+) : PerDisplayInstanceProvider<CoroutineScope> {
 
-    override fun createInstance(displayId: Int): CoroutineScope {
+    override fun createInstance(displayId: Int): CoroutineScope? {
         return if (displayId == Display.DEFAULT_DISPLAY) {
             // The default display is connected all the time, therefore we can optimise by reusing
             // the application scope, and don't need to create a new scope.
             backgroundApplicationScope
         } else {
-            CoroutineScope(backgroundDispatcher + newTracingContext("DisplayScope$displayId"))
+            // The scope is automatically cancelled from the component when the display is removed.
+            displayComponentRepository[displayId]?.displayCoroutineScope
         }
     }
-
-    override fun destroyInstance(instance: CoroutineScope) {
-        instance.cancel("DisplayContext has been cancelled.")
-    }
 }
 
 @Module
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayRepositoryKosmos.kt
index 048ea3c25388..5a91d01956cc 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/DisplayRepositoryKosmos.kt
@@ -16,7 +16,38 @@
 
 package com.android.systemui.display.data.repository
 
+import android.view.Display
+import com.android.systemui.display.dagger.SystemUIDisplaySubcomponent
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testScope
+import kotlinx.coroutines.CoroutineScope
 
 val Kosmos.displayRepository by Fixture { FakeDisplayRepository() }
+
+fun Kosmos.createFakeDisplaySubcomponent(
+    coroutineScope: CoroutineScope = testScope.backgroundScope
+): SystemUIDisplaySubcomponent {
+    return object : SystemUIDisplaySubcomponent {
+        override val displayCoroutineScope: CoroutineScope
+            get() = coroutineScope
+    }
+}
+
+val Kosmos.sysuiDefaultDisplaySubcomponent by Fixture {
+    createFakeDisplaySubcomponent(testScope.backgroundScope)
+}
+
+val Kosmos.fakeSysuiDisplayComponentFactory by Fixture {
+    object : SystemUIDisplaySubcomponent.Factory {
+        override fun create(displayId: Int): SystemUIDisplaySubcomponent {
+            return sysuiDefaultDisplaySubcomponent
+        }
+    }
+}
+
+val Kosmos.displaySubcomponentPerDisplayRepository by Fixture {
+    FakePerDisplayRepository<SystemUIDisplaySubcomponent>().apply {
+        add(Display.DEFAULT_DISPLAY, sysuiDefaultDisplaySubcomponent)
+    }
+}