Merge "Make ClockRegistry and its provider singleton" into udc-dev
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 89f0233..3c85afb 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -99,11 +99,6 @@
     private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
         null
     private var notificationsSnapshotRestorer: NotificationsSnapshotRestorer? = null
-    /**
-     * Mapping from LifeCycleOwner's hashcode to ClockRegistry as we need to keep different
-     * ClockRegistries per LifeCycle to ensure proper cleanup
-     */
-    private var clockRegistries: MutableMap<Int, ClockRegistry> = HashMap()
     private var clockPickerInteractor: ClockPickerInteractor? = null
     private var clockSectionViewModel: ClockSectionViewModel? = null
     private var clockCarouselViewModelFactory: ClockCarouselViewModel.Factory? = null
@@ -122,6 +117,7 @@
     private var gridInteractor: GridInteractor? = null
     private var gridSnapshotRestorer: GridSnapshotRestorer? = null
     private var gridScreenViewModelFactory: GridScreenViewModel.Factory? = null
+    private var clockRegistryProvider: ClockRegistryProvider? = null
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         return customizationSections
@@ -344,26 +340,15 @@
     }
 
     override fun getClockRegistry(context: Context, lifecycleOwner: LifecycleOwner): ClockRegistry {
-        return clockRegistries[lifecycleOwner.hashCode()]
-            ?: ClockRegistryProvider(
-                    context = context,
-                    coroutineScope = getApplicationCoroutineScope(),
-                    mainDispatcher = Dispatchers.Main,
-                    backgroundDispatcher = Dispatchers.IO,
-                )
-                .get()
-                .also {
-                    clockRegistries[lifecycleOwner.hashCode()] = it
-                    lifecycleOwner.lifecycle.addObserver(
-                        object : DefaultLifecycleObserver {
-                            override fun onDestroy(owner: LifecycleOwner) {
-                                super.onDestroy(owner)
-                                clockRegistries[lifecycleOwner.hashCode()]?.unregisterListeners()
-                                clockRegistries.remove(lifecycleOwner.hashCode())
-                            }
-                        }
+        return (clockRegistryProvider
+                ?: ClockRegistryProvider(
+                        context = context,
+                        coroutineScope = getApplicationCoroutineScope(),
+                        mainDispatcher = Dispatchers.Main,
+                        backgroundDispatcher = Dispatchers.IO,
                     )
-                }
+                    .also { clockRegistryProvider = it })
+            .getForOwner(lifecycleOwner)
     }
 
     override fun getClockPickerInteractor(
diff --git a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
index e359117..52c3c4e 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
@@ -19,6 +19,8 @@
 import android.content.ComponentName
 import android.content.Context
 import android.view.LayoutInflater
+import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.LifecycleOwner
 import com.android.systemui.plugins.Plugin
 import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
@@ -43,6 +45,7 @@
     private val mainDispatcher: CoroutineDispatcher,
     private val backgroundDispatcher: CoroutineDispatcher,
 ) {
+    private val lifecycleOwners = mutableSetOf<Int>()
     private val pluginManager: PluginManager by lazy { createPluginManager(context) }
     private val clockRegistry: ClockRegistry by lazy {
         ClockRegistry(
@@ -60,10 +63,32 @@
             .apply { registerListeners() }
     }
 
-    fun get(): ClockRegistry {
+    fun getForOwner(lifecycleOwner: LifecycleOwner): ClockRegistry {
+        registerLifecycleOwner(lifecycleOwner)
         return clockRegistry
     }
 
+    private fun registerLifecycleOwner(lifecycleOwner: LifecycleOwner) {
+        lifecycleOwners.add(lifecycleOwner.hashCode())
+
+        lifecycleOwner.lifecycle.addObserver(
+            object : DefaultLifecycleObserver {
+                override fun onDestroy(owner: LifecycleOwner) {
+                    super.onDestroy(owner)
+                    unregisterLifecycleOwner(owner)
+                }
+            }
+        )
+    }
+
+    private fun unregisterLifecycleOwner(lifecycleOwner: LifecycleOwner) {
+        lifecycleOwners.remove(lifecycleOwner.hashCode())
+
+        if (lifecycleOwners.isEmpty()) {
+            clockRegistry.unregisterListeners()
+        }
+    }
+
     private fun createPluginManager(context: Context): PluginManager {
         val privilegedPlugins = listOf<String>()
         val isDebugDevice = true