Improve lifecycle of RegistrationRepository.

demo/registration/src/main/java/org/signal/registration/sample/MainActivity.kt

@@ -29,9 +29,7 @@ import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Surface
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.platform.LocalContext
 import androidx.lifecycle.compose.LifecycleResumeEffect
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import androidx.lifecycle.viewmodel.compose.viewModel
@@ -46,10 +44,8 @@ import androidx.navigation3.ui.NavDisplay
 import com.google.accompanist.permissions.ExperimentalPermissionsApi
 import kotlinx.serialization.Serializable
 import org.signal.core.ui.compose.theme.SignalTheme
-import org.signal.core.util.billing.OneTimePurchaseApi
 import org.signal.registration.RegistrationDependencies
 import org.signal.registration.RegistrationNavHost
-import org.signal.registration.RegistrationRepository
 import org.signal.registration.sample.debug.NetworkDebugOverlay
 import org.signal.registration.sample.screens.RegistrationCompleteScreen
 import org.signal.registration.sample.screens.main.MainScreen
@@ -130,20 +126,6 @@ private fun SampleNavHost(
   backStack: NavBackStack<NavKey>,
   modifier: Modifier = Modifier
 ) {
-  val context = LocalContext.current
-
-  val registrationRepository = remember {
-    RegistrationRepository(
-      context = context.applicationContext,
-      networkController = registrationDependencies.networkController,
-      storageController = registrationDependencies.storageController,
-      isLinkAndSyncAvailable = registrationDependencies.isLinkAndSyncAvailable,
-      isPhoneNumberlessRegistrationAvailable = registrationDependencies.isPhoneNumberlessRegistrationAvailable,
-      // The demo app is never published to the Play Store, so there is nothing to buy from.
-      signalLoginPurchaseApi = OneTimePurchaseApi.Empty
-    )
-  }
-
   val entryProvider: (NavKey) -> NavEntry<NavKey> = entryProvider {
     entry<SampleRoute.Main> {
       val viewModel: MainScreenViewModel = viewModel(
@@ -170,7 +152,6 @@ private fun SampleNavHost(
 
     entry<SampleRoute.Registration> {
       RegistrationNavHost(
-        registrationRepository,
         modifier = Modifier.fillMaxSize(),
         onRegistrationComplete = {
           backStack.add(SampleRoute.RegistrationComplete)

feature/registration/src/main/java/org/signal/registration/RegistrationActivity.kt

@@ -13,7 +13,6 @@ import androidx.compose.material3.Surface
 import androidx.compose.ui.Modifier
 import androidx.core.content.IntentCompat
 import com.google.accompanist.permissions.ExperimentalPermissionsApi
-import org.signal.billing.BillingFactory
 import org.signal.core.ui.compose.theme.SignalTheme
 
 /**
@@ -54,22 +53,6 @@ class RegistrationActivity : ComponentActivity() {
     }
   }
 
-  private val repositoryLazy = lazy {
-    RegistrationRepository(
-      context = this.application,
-      networkController = RegistrationDependencies.get().networkController,
-      storageController = RegistrationDependencies.get().storageController,
-      isLinkAndSyncAvailable = RegistrationDependencies.get().isLinkAndSyncAvailable,
-      isPhoneNumberlessRegistrationAvailable = RegistrationDependencies.get().isPhoneNumberlessRegistrationAvailable,
-      isGooglePlayBillingAvailable = RegistrationDependencies.get().isGooglePlayBillingAvailable,
-      signalLoginPurchaseApi = BillingFactory.createOneTimePurchaseApi(
-        context = this.application,
-        isAvailable = RegistrationDependencies.get().isGooglePlayBillingAvailable
-      )
-    )
-  }
-  private val repository: RegistrationRepository by repositoryLazy
-
   @OptIn(ExperimentalPermissionsApi::class)
   override fun onCreate(savedInstanceState: Bundle?) {
     enableEdgeToEdge()
@@ -82,7 +65,6 @@ class RegistrationActivity : ComponentActivity() {
       SignalTheme(incognitoKeyboardEnabled = false) {
         Surface(modifier = Modifier.fillMaxSize()) {
           RegistrationNavHost(
-            registrationRepository = repository,
             startDestination = startDestination,
             startFresh = startFresh,
             modifier = Modifier
@@ -99,13 +81,6 @@ class RegistrationActivity : ComponentActivity() {
     }
   }
 
-  override fun onDestroy() {
-    super.onDestroy()
-    if (repositoryLazy.isInitialized()) {
-      repository.close()
-    }
-  }
-
   /**
    * Activity result contract for launching the registration flow.
    *

feature/registration/src/main/java/org/signal/registration/RegistrationNavigation.kt

@@ -423,8 +423,8 @@ private fun openUrl(context: Context, url: String) {
 /**
  * Sets up the navigation graph for the registration flow using Navigation 3.
  *
- * @param registrationRepository The repository for registration data.
- * @param registrationViewModel Optional ViewModel for testing. If null, creates one internally.
+ * @param registrationViewModel Optional ViewModel for testing. If null, creates one internally, along with the
+ *   [RegistrationRepository] it owns.
  * @param startFresh When true, any persisted registration data is not restored and the user starts the flow fresh from
  *   the beginning.
  * @param permissionsState Optional permissions state for testing. If null, creates one internally.
@@ -435,7 +435,6 @@ private fun openUrl(context: Context, url: String) {
 @OptIn(ExperimentalPermissionsApi::class)
 @Composable
 fun RegistrationNavHost(
-  registrationRepository: RegistrationRepository,
   registrationViewModel: RegistrationViewModel? = null,
   startFresh: Boolean = false,
   permissionsState: MultiplePermissionsState? = null,
@@ -443,8 +442,9 @@ fun RegistrationNavHost(
   modifier: Modifier = Modifier,
   onRegistrationComplete: () -> Unit = {}
 ) {
+  val context = LocalContext.current
   val viewModel: RegistrationViewModel = registrationViewModel ?: viewModel {
-    RegistrationViewModel(registrationRepository, createSavedStateHandle(), startDestination, startFresh)
+    RegistrationViewModel(RegistrationRepository.create(context), createSavedStateHandle(), startDestination, startFresh)
   }
 
   val registrationState by viewModel.state.collectAsStateWithLifecycle()
@@ -465,7 +465,7 @@ fun RegistrationNavHost(
 
   val entryProvider = entryProvider {
     navigationEntries(
-      registrationRepository = registrationRepository,
+      registrationRepository = viewModel.repository,
       registrationViewModel = viewModel,
       permissionsState = permissionsState,
       onRegistrationComplete = onRegistrationComplete

feature/registration/src/main/java/org/signal/registration/RegistrationRepository.kt

@@ -32,6 +32,7 @@ import kotlinx.serialization.json.Json
 import okio.ByteString
 import okio.ByteString.Companion.toByteString
 import org.signal.archive.LocalBackupRestoreProgress
+import org.signal.billing.BillingFactory
 import org.signal.core.models.AccountEntropyPool
 import org.signal.core.models.MasterKey
 import org.signal.core.models.ServiceId.ACI
@@ -132,6 +133,25 @@ class RegistrationRepository(
     private val TAG = Log.tag(RegistrationRepository::class)
     private val json = Json { ignoreUnknownKeys = true }
 
+    /** Builds a repository from the module's injected [RegistrationDependencies]. */
+    fun create(context: Context): RegistrationRepository {
+      val dependencies = RegistrationDependencies.get()
+      val application = context.applicationContext
+
+      return RegistrationRepository(
+        context = application,
+        networkController = dependencies.networkController,
+        storageController = dependencies.storageController,
+        isLinkAndSyncAvailable = dependencies.isLinkAndSyncAvailable,
+        isPhoneNumberlessRegistrationAvailable = dependencies.isPhoneNumberlessRegistrationAvailable,
+        isGooglePlayBillingAvailable = dependencies.isGooglePlayBillingAvailable,
+        signalLoginPurchaseApi = BillingFactory.createOneTimePurchaseApi(
+          context = application,
+          isAvailable = dependencies.isGooglePlayBillingAvailable
+        )
+      )
+    }
+
     /**
      * The purchase option to buy within the service-provided Signal Login product. The service names the product but
      * not the option within it, so this half stays a client constant.

feature/registration/src/main/java/org/signal/registration/RegistrationViewModel.kt

@@ -27,7 +27,7 @@ import org.signal.registration.screens.restoreselection.RegisteredState
  * Manages state and logic for registration screens.
  */
 class RegistrationViewModel(
-  private val repository: RegistrationRepository,
+  val repository: RegistrationRepository,
   private val savedStateHandle: SavedStateHandle,
   startDestination: RegistrationRoute? = null,
   private val startFresh: Boolean = false
@@ -86,6 +86,10 @@ class RegistrationViewModel(
     }
   }
 
+  override fun onCleared() {
+    repository.close()
+  }
+
   override suspend fun processEvent(event: RegistrationFlowEvent) {
     _state.value = applyEvent(_state.value, event)
     persistFlowState(event)

feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt

@@ -2045,7 +2045,6 @@ class RegistrationEndToEndTest {
       SignalTheme {
         ActivityResultInterceptor(folderPickerResult) {
           RegistrationNavHost(
-            registrationRepository = repository,
             registrationViewModel = viewModel,
             permissionsState = createMockPermissionsState(),
             onRegistrationComplete = onRegistrationComplete

feature/registration/src/test/java/org/signal/registration/RegistrationNavigationTest.kt

@@ -72,7 +72,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme(incognitoKeyboardEnabled = false) {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -91,7 +90,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -116,7 +114,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -139,7 +136,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -166,7 +162,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -198,7 +193,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -225,7 +219,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -246,7 +239,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -269,7 +261,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -291,7 +282,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -314,7 +304,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -341,7 +330,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -369,7 +357,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )
@@ -395,7 +382,6 @@ class RegistrationNavigationTest {
     composeTestRule.setContent {
       SignalTheme {
         RegistrationNavHost(
-          registrationRepository = mockRepository,
           registrationViewModel = viewModel,
           permissionsState = permissionsState
         )

feature/registration/src/test/java/org/signal/registration/RegistrationViewModelTest.kt

@@ -6,6 +6,10 @@
 package org.signal.registration
 
 import androidx.lifecycle.SavedStateHandle
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.ViewModelStore
+import androidx.lifecycle.viewmodel.initializer
+import androidx.lifecycle.viewmodel.viewModelFactory
 import assertk.assertThat
 import assertk.assertions.isEqualTo
 import assertk.assertions.isNotNull
@@ -14,6 +18,7 @@ import assertk.assertions.isTrue
 import io.mockk.coEvery
 import io.mockk.coVerify
 import io.mockk.mockk
+import io.mockk.verify
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
@@ -912,6 +917,26 @@ class RegistrationViewModelTest {
     advanceUntilIdle()
   }
 
+  // ==================== Repository Ownership Tests ====================
+
+  @Test
+  fun `repository is closed when the view model is cleared`() = runTest(testDispatcher) {
+    val store = ViewModelStore()
+    val provider = ViewModelProvider.create(
+      store,
+      viewModelFactory {
+        initializer { RegistrationViewModel(mockRepository, SavedStateHandle()) }
+      }
+    )
+
+    provider[RegistrationViewModel::class]
+    advanceUntilIdle()
+    verify(exactly = 0) { mockRepository.close() }
+
+    store.clear()
+    verify(exactly = 1) { mockRepository.close() }
+  }
+
   // ==================== Helpers ====================
 
   private fun createSessionMetadata(id: String = "test-session"): SessionMetadata {