Unregister listeners from ClockRegistry when destroyed

We now need to keep track of ClockRegistry per lifecycleOwner or we
could risk having a ClockRegistry with already unregistered listeners
used elsewhere (eg, if the injector is called from 2 different Activities)

Fixes: 283124260
Test: manually verified going through all screens with clocks and checking no leaks
Change-Id: I1390646a1dfbf745a0c7b90dcd1c6a988ed4421d
diff --git a/src/com/android/customization/module/CustomizationInjector.kt b/src/com/android/customization/module/CustomizationInjector.kt
index dbcff27..b943925 100644
--- a/src/com/android/customization/module/CustomizationInjector.kt
+++ b/src/com/android/customization/module/CustomizationInjector.kt
@@ -18,6 +18,7 @@
 import android.content.Context
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.FragmentActivity
+import androidx.lifecycle.LifecycleOwner
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
 import com.android.customization.model.theme.ThemeManager
@@ -47,11 +48,17 @@
         context: Context,
     ): KeyguardQuickAffordancePickerInteractor
 
-    fun getClockRegistry(context: Context): ClockRegistry
+    fun getClockRegistry(context: Context, lifecycleOwner: LifecycleOwner): ClockRegistry
 
-    fun getClockPickerInteractor(context: Context): ClockPickerInteractor
+    fun getClockPickerInteractor(
+        context: Context,
+        lifecycleOwner: LifecycleOwner
+    ): ClockPickerInteractor
 
-    fun getClockSectionViewModel(context: Context): ClockSectionViewModel
+    fun getClockSectionViewModel(
+        context: Context,
+        lifecycleOwner: LifecycleOwner
+    ): ClockSectionViewModel
 
     fun getColorPickerInteractor(
         context: Context,
@@ -73,5 +80,6 @@
         context: Context,
         wallpaperColorsViewModel: WallpaperColorsViewModel,
         clockViewFactory: ClockViewFactory,
+        lifecycleOwner: LifecycleOwner,
     ): ClockSettingsViewModel.Factory
 }
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 5d9880d..66814c5 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -27,7 +27,6 @@
 import androidx.lifecycle.DefaultLifecycleObserver
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.ViewModelProvider
-import androidx.lifecycle.get
 import com.android.customization.model.color.ColorCustomizationManager
 import com.android.customization.model.color.ColorOptionsProvider
 import com.android.customization.model.grid.GridOptionsManager
@@ -99,7 +98,11 @@
     private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
         null
     private var notificationsSnapshotRestorer: NotificationsSnapshotRestorer? = null
-    private var clockRegistry: ClockRegistry? = 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
@@ -132,7 +135,7 @@
                     ),
                     getFlags(),
                     getClockCarouselViewModelFactory(
-                        getClockPickerInteractor(activity.applicationContext),
+                        getClockPickerInteractor(activity.applicationContext, activity),
                     ),
                     getClockViewFactory(activity),
                     getDarkModeSnapshotRestorer(activity),
@@ -314,8 +317,8 @@
                 .also { notificationsSnapshotRestorer = it }
     }
 
-    override fun getClockRegistry(context: Context): ClockRegistry {
-        return clockRegistry
+    override fun getClockRegistry(context: Context, lifecycleOwner: LifecycleOwner): ClockRegistry {
+        return clockRegistries[lifecycleOwner.hashCode()]
             ?: ClockRegistryProvider(
                     context = context,
                     coroutineScope = getApplicationCoroutineScope(),
@@ -323,28 +326,42 @@
                     backgroundDispatcher = Dispatchers.IO,
                 )
                 .get()
-                .also { clockRegistry = it }
+                .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())
+                            }
+                        }
+                    )
+                }
     }
 
     override fun getClockPickerInteractor(
         context: Context,
+        lifecycleOwner: LifecycleOwner,
     ): ClockPickerInteractor {
         return clockPickerInteractor
             ?: ClockPickerInteractor(
                     ClockPickerRepositoryImpl(
                         secureSettingsRepository = getSecureSettingsRepository(context),
-                        registry = getClockRegistry(context),
+                        registry = getClockRegistry(context, lifecycleOwner),
                         scope = getApplicationCoroutineScope(),
                     ),
                 )
                 .also { clockPickerInteractor = it }
     }
 
-    override fun getClockSectionViewModel(context: Context): ClockSectionViewModel {
+    override fun getClockSectionViewModel(
+        context: Context,
+        lifecycleOwner: LifecycleOwner
+    ): ClockSectionViewModel {
         return clockSectionViewModel
-            ?: ClockSectionViewModel(context, getClockPickerInteractor(context)).also {
-                clockSectionViewModel = it
-            }
+            ?: ClockSectionViewModel(context, getClockPickerInteractor(context, lifecycleOwner))
+                .also { clockSectionViewModel = it }
     }
 
     override fun getClockCarouselViewModelFactory(
@@ -363,9 +380,7 @@
                     activity.applicationContext,
                     ScreenSizeCalculator.getInstance()
                         .getScreenSize(activity.windowManager.defaultDisplay),
-                    getClockRegistry(
-                        activity.applicationContext,
-                    ),
+                    getClockRegistry(activity.applicationContext, activity),
                 )
                 .also {
                     clockViewFactories[activityHashCode] = it
@@ -484,11 +499,12 @@
         context: Context,
         wallpaperColorsViewModel: WallpaperColorsViewModel,
         clockViewFactory: ClockViewFactory,
+        lifecycleOwner: LifecycleOwner,
     ): ClockSettingsViewModel.Factory {
         return clockSettingsViewModelFactory
             ?: ClockSettingsViewModel.Factory(
                     context,
-                    getClockPickerInteractor(context),
+                    getClockPickerInteractor(context, lifecycleOwner),
                     getColorPickerInteractor(
                         context,
                         wallpaperColorsViewModel,
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
index 7e53ac4..6203e24 100644
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
@@ -36,7 +36,8 @@
         setUpToolbar(view)
         clockRegistry =
             (InjectorProvider.getInjector() as ThemePickerInjector).getClockRegistry(
-                requireContext()
+                requireContext(),
+                this
             )
         val listInUse = clockRegistry.getClocks().filter { "NOT_IN_USE" !in it.clockId }
 
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
index 7edaecf..50840cf 100644
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
@@ -123,6 +123,7 @@
                         context,
                         injector.getWallpaperColorsViewModel(),
                         injector.getClockViewFactory(activity),
+                        activity,
                     ),
                 )
                 .get(),