Move restore messages out of durable job.

app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt

@@ -12,6 +12,8 @@ import android.os.StatFs
 import androidx.annotation.Discouraged
 import androidx.annotation.WorkerThread
 import androidx.core.app.NotificationCompat
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.isActive
 import kotlinx.coroutines.withContext
 import okio.ByteString
 import okio.ByteString.Companion.toByteString
@@ -48,6 +50,7 @@ import org.thoughtcrime.securesms.attachments.Cdn
 import org.thoughtcrime.securesms.attachments.DatabaseAttachment
 import org.thoughtcrime.securesms.backup.ArchiveUploadProgress
 import org.thoughtcrime.securesms.backup.DeletionState
+import org.thoughtcrime.securesms.backup.RestoreState
 import org.thoughtcrime.securesms.backup.v2.BackupRepository.copyAttachmentToArchive
 import org.thoughtcrime.securesms.backup.v2.importer.ChatItemArchiveImporter
 import org.thoughtcrime.securesms.backup.v2.processor.AccountDataArchiveProcessor
@@ -84,8 +87,10 @@ import org.thoughtcrime.securesms.database.model.InAppPaymentSubscriberRecord
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.groups.GroupId
 import org.thoughtcrime.securesms.jobmanager.Job
+import org.thoughtcrime.securesms.jobmanager.impl.DataRestoreConstraint
 import org.thoughtcrime.securesms.jobs.AvatarGroupsV2DownloadJob
 import org.thoughtcrime.securesms.jobs.BackupDeleteJob
+import org.thoughtcrime.securesms.jobs.BackupRestoreMediaJob
 import org.thoughtcrime.securesms.jobs.CheckRestoreMediaLeftJob
 import org.thoughtcrime.securesms.jobs.CreateReleaseChannelJob
 import org.thoughtcrime.securesms.jobs.RequestGroupV2InfoJob
@@ -99,8 +104,10 @@ import org.thoughtcrime.securesms.keyvalue.isDecisionPending
 import org.thoughtcrime.securesms.net.SignalNetwork
 import org.thoughtcrime.securesms.notifications.NotificationChannels
 import org.thoughtcrime.securesms.notifications.NotificationIds
+import org.thoughtcrime.securesms.providers.BlobProvider
 import org.thoughtcrime.securesms.recipients.Recipient
 import org.thoughtcrime.securesms.recipients.RecipientId
+import org.thoughtcrime.securesms.service.BackupProgressService
 import org.thoughtcrime.securesms.storage.StorageSyncHelper
 import org.thoughtcrime.securesms.util.RemoteConfig
 import org.thoughtcrime.securesms.util.ServiceUtil
@@ -121,6 +128,7 @@ import org.whispersystems.signalservice.api.backup.MediaName
 import org.whispersystems.signalservice.api.backup.MediaRootBackupKey
 import org.whispersystems.signalservice.api.backup.MessageBackupKey
 import org.whispersystems.signalservice.api.crypto.AttachmentCipherStreamUtil
+import org.whispersystems.signalservice.api.messages.AttachmentTransferProgress
 import org.whispersystems.signalservice.api.messages.SignalServiceAttachment.ProgressListener
 import org.whispersystems.signalservice.api.push.ServiceId.ACI
 import org.whispersystems.signalservice.api.push.ServiceId.PNI
@@ -1723,6 +1731,75 @@ object BackupRepository {
     )
   }
 
+  suspend fun restoreRemoteBackup(): RemoteRestoreResult {
+    val context = AppDependencies.application
+    SignalStore.backup.restoreState = RestoreState.PENDING
+
+    try {
+      DataRestoreConstraint.isRestoringData = true
+      return withContext(Dispatchers.IO) {
+        return@withContext BackupProgressService.start(context, context.getString(R.string.BackupProgressService_title)).use {
+          restoreRemoteBackup(controller = it, cancellationSignal = { !isActive })
+        }
+      }
+    } finally {
+      DataRestoreConstraint.isRestoringData = false
+    }
+  }
+
+  private fun restoreRemoteBackup(controller: BackupProgressService.Controller, cancellationSignal: () -> Boolean): RemoteRestoreResult {
+    SignalStore.backup.restoreState = RestoreState.RESTORING_DB
+
+    val progressListener = object : ProgressListener {
+      override fun onAttachmentProgress(progress: AttachmentTransferProgress) {
+        controller.update(
+          title = AppDependencies.application.getString(R.string.BackupProgressService_title_downloading),
+          progress = progress.value,
+          indeterminate = false
+        )
+        EventBus.getDefault().post(RestoreV2Event(RestoreV2Event.Type.PROGRESS_DOWNLOAD, progress.transmitted, progress.total))
+      }
+
+      override fun shouldCancel() = cancellationSignal()
+    }
+
+    Log.i(TAG, "[remoteRestore] Downloading backup")
+    val tempBackupFile = BlobProvider.getInstance().forNonAutoEncryptingSingleSessionOnDisk(AppDependencies.application)
+    when (val result = downloadBackupFile(tempBackupFile, progressListener)) {
+      is NetworkResult.Success -> Log.i(TAG, "[remoteRestore] Download successful")
+      else -> {
+        Log.w(TAG, "[remoteRestore] Failed to download backup file", result.getCause())
+        return RemoteRestoreResult.NetworkError
+      }
+    }
+
+    if (cancellationSignal()) {
+      return RemoteRestoreResult.Canceled
+    }
+
+    controller.update(
+      title = AppDependencies.application.getString(R.string.BackupProgressService_title),
+      progress = 0f,
+      indeterminate = true
+    )
+
+    val self = Recipient.self()
+    val selfData = SelfData(self.aci.get(), self.pni.get(), self.e164.get(), ProfileKey(self.profileKey))
+    Log.i(TAG, "[remoteRestore] Importing backup")
+    val result = import(length = tempBackupFile.length(), inputStreamFactory = tempBackupFile::inputStream, selfData = selfData, backupKey = SignalStore.backup.messageBackupKey, cancellationSignal = cancellationSignal)
+    if (result == ImportResult.Failure) {
+      Log.w(TAG, "[remoteRestore] Failed to import backup")
+      return RemoteRestoreResult.Failure
+    }
+
+    SignalStore.backup.restoreState = RestoreState.RESTORING_MEDIA
+
+    AppDependencies.jobManager.add(BackupRestoreMediaJob())
+
+    Log.i(TAG, "[remoteRestore] Restore successful")
+    return RemoteRestoreResult.Success
+  }
+
   interface ExportProgressListener {
     fun onAccount()
     fun onRecipient()
@@ -1794,6 +1871,13 @@ sealed class ImportResult {
   data object Failure : ImportResult()
 }
 
+sealed interface RemoteRestoreResult {
+  data object Success : RemoteRestoreResult
+  data object NetworkError : RemoteRestoreResult
+  data object Canceled : RemoteRestoreResult
+  data object Failure : RemoteRestoreResult
+}
+
 /**
  * Iterator that reads values from the given cursor. Expects that REMOTE_DIGEST is present and non-null, and ARCHIVE_CDN is present.
  *

app/src/main/java/org/thoughtcrime/securesms/components/settings/app/internal/backup/InternalBackupPlaygroundViewModel.kt

@@ -39,6 +39,7 @@ import org.thoughtcrime.securesms.backup.v2.ArchiveValidator
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
 import org.thoughtcrime.securesms.backup.v2.DebugBackupMetadata
 import org.thoughtcrime.securesms.backup.v2.MessageBackupTier
+import org.thoughtcrime.securesms.backup.v2.RemoteRestoreResult
 import org.thoughtcrime.securesms.backup.v2.local.ArchiveFileSystem
 import org.thoughtcrime.securesms.backup.v2.local.ArchiveResult
 import org.thoughtcrime.securesms.backup.v2.local.LocalArchiver
@@ -50,8 +51,6 @@ import org.thoughtcrime.securesms.database.AttachmentTable.DebugAttachmentStats
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.jobs.BackupMessagesJob
-import org.thoughtcrime.securesms.jobs.BackupRestoreJob
-import org.thoughtcrime.securesms.jobs.BackupRestoreMediaJob
 import org.thoughtcrime.securesms.jobs.RestoreLocalAttachmentJob
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.providers.BlobProvider
@@ -69,7 +68,6 @@ import javax.crypto.Cipher
 import javax.crypto.CipherInputStream
 import javax.crypto.spec.IvParameterSpec
 import javax.crypto.spec.SecretKeySpec
-import kotlin.time.Duration.Companion.seconds
 
 class InternalBackupPlaygroundViewModel : ViewModel() {
 
@@ -342,18 +340,16 @@ class InternalBackupPlaygroundViewModel : ViewModel() {
   private fun restoreFromRemote() {
     _state.value = _state.value.copy(statusMessage = "Importing from remote...")
 
-    disposables += Single.fromCallable {
-      AppDependencies
-        .jobManager
-        .startChain(BackupRestoreJob())
-        .then(BackupRestoreMediaJob())
-        .enqueueAndBlockUntilCompletion(120.seconds.inWholeMilliseconds)
-    }
-      .subscribeOn(Schedulers.io())
-      .observeOn(AndroidSchedulers.mainThread())
-      .subscribeBy {
-        _state.value = _state.value.copy(statusMessage = "Import complete!")
+    viewModelScope.launch {
+      when (val result = BackupRepository.restoreRemoteBackup()) {
+        RemoteRestoreResult.Success -> _state.value = _state.value.copy(statusMessage = "Import complete!")
+        RemoteRestoreResult.Canceled,
+        RemoteRestoreResult.Failure,
+        RemoteRestoreResult.NetworkError -> {
+          _state.value = _state.value.copy(statusMessage = "Import failed! $result")
+        }
       }
+    }
   }
 
   fun loadStats() {

app/src/main/java/org/thoughtcrime/securesms/jobs/BackupRestoreJob.kt

@@ -1,122 +0,0 @@
-/*
- * Copyright 2024 Signal Messenger, LLC
- * SPDX-License-Identifier: AGPL-3.0-only
- */
-
-package org.thoughtcrime.securesms.jobs
-
-import org.greenrobot.eventbus.EventBus
-import org.signal.core.util.logging.Log
-import org.signal.libsignal.zkgroup.profiles.ProfileKey
-import org.thoughtcrime.securesms.R
-import org.thoughtcrime.securesms.backup.RestoreState
-import org.thoughtcrime.securesms.backup.v2.BackupRepository
-import org.thoughtcrime.securesms.backup.v2.ImportResult
-import org.thoughtcrime.securesms.backup.v2.RestoreV2Event
-import org.thoughtcrime.securesms.dependencies.AppDependencies
-import org.thoughtcrime.securesms.jobmanager.Job
-import org.thoughtcrime.securesms.jobmanager.impl.NetworkConstraint
-import org.thoughtcrime.securesms.keyvalue.SignalStore
-import org.thoughtcrime.securesms.net.NotPushRegisteredException
-import org.thoughtcrime.securesms.providers.BlobProvider
-import org.thoughtcrime.securesms.recipients.Recipient
-import org.thoughtcrime.securesms.service.BackupProgressService
-import org.whispersystems.signalservice.api.NetworkResult
-import org.whispersystems.signalservice.api.messages.AttachmentTransferProgress
-import org.whispersystems.signalservice.api.messages.SignalServiceAttachment.ProgressListener
-import java.io.IOException
-
-/**
- * Job that is responsible for restoring a backup from the server
- */
-class BackupRestoreJob private constructor(parameters: Parameters) : BaseJob(parameters) {
-
-  companion object {
-    private val TAG = Log.tag(BackupRestoreJob::class.java)
-
-    const val KEY = "BackupRestoreJob"
-  }
-
-  constructor() : this(
-    Parameters.Builder()
-      .addConstraint(NetworkConstraint.KEY)
-      .setMaxAttempts(Parameters.UNLIMITED)
-      .setMaxInstancesForFactory(1)
-      .setQueue("BackupRestoreJob")
-      .build()
-  )
-
-  override fun serialize(): ByteArray? = null
-
-  override fun getFactoryKey(): String = KEY
-
-  override fun onFailure() = Unit
-
-  override fun onAdded() {
-    SignalStore.backup.restoreState = RestoreState.PENDING
-  }
-
-  override fun onRun() {
-    if (!SignalStore.account.isRegistered) {
-      Log.e(TAG, "Not registered, cannot restore!")
-      throw NotPushRegisteredException()
-    }
-
-    BackupProgressService.start(context, context.getString(R.string.BackupProgressService_title)).use {
-      restore(it)
-    }
-  }
-
-  private fun restore(controller: BackupProgressService.Controller) {
-    SignalStore.backup.restoreState = RestoreState.RESTORING_DB
-
-    val progressListener = object : ProgressListener {
-      override fun onAttachmentProgress(progress: AttachmentTransferProgress) {
-        controller.update(
-          title = context.getString(R.string.BackupProgressService_title_downloading),
-          progress = progress.value,
-          indeterminate = false
-        )
-        EventBus.getDefault().post(RestoreV2Event(RestoreV2Event.Type.PROGRESS_DOWNLOAD, progress.transmitted, progress.total))
-      }
-
-      override fun shouldCancel() = isCanceled
-    }
-
-    val tempBackupFile = BlobProvider.getInstance().forNonAutoEncryptingSingleSessionOnDisk(AppDependencies.application)
-    when (val result = BackupRepository.downloadBackupFile(tempBackupFile, progressListener)) {
-      is NetworkResult.Success -> Log.i(TAG, "Download successful")
-      else -> {
-        Log.w(TAG, "Failed to download backup file", result.getCause())
-        throw IOException(result.getCause())
-      }
-    }
-
-    if (isCanceled) {
-      return
-    }
-
-    controller.update(
-      title = context.getString(R.string.BackupProgressService_title),
-      progress = 0f,
-      indeterminate = true
-    )
-
-    val self = Recipient.self()
-    val selfData = BackupRepository.SelfData(self.aci.get(), self.pni.get(), self.e164.get(), ProfileKey(self.profileKey))
-    val result = BackupRepository.import(length = tempBackupFile.length(), inputStreamFactory = tempBackupFile::inputStream, selfData = selfData, backupKey = SignalStore.backup.messageBackupKey, cancellationSignal = { isCanceled })
-    if (result == ImportResult.Failure) {
-      throw IOException("Failed to import backup")
-    }
-
-    SignalStore.backup.restoreState = RestoreState.RESTORING_MEDIA
-  }
-
-  override fun onShouldRetry(e: Exception): Boolean = false
-
-  class Factory : Job.Factory<BackupRestoreJob> {
-    override fun create(parameters: Parameters, serializedData: ByteArray?): BackupRestoreJob {
-      return BackupRestoreJob(parameters)
-    }
-  }
-}

app/src/main/java/org/thoughtcrime/securesms/jobs/JobManagerFactories.java

@@ -9,7 +9,6 @@ import org.thoughtcrime.securesms.jobmanager.Constraint;
 import org.thoughtcrime.securesms.jobmanager.ConstraintObserver;
 import org.thoughtcrime.securesms.jobmanager.Job;
 import org.thoughtcrime.securesms.jobmanager.JobMigration;
-import org.thoughtcrime.securesms.jobmanager.impl.NoRemoteArchiveGarbageCollectionPendingConstraint;
 import org.thoughtcrime.securesms.jobmanager.impl.AutoDownloadEmojiConstraint;
 import org.thoughtcrime.securesms.jobmanager.impl.BatteryNotLowConstraint;
 import org.thoughtcrime.securesms.jobmanager.impl.CellServiceConstraintObserver;
@@ -24,6 +23,7 @@ import org.thoughtcrime.securesms.jobmanager.impl.DecryptionsDrainedConstraintOb
 import org.thoughtcrime.securesms.jobmanager.impl.NetworkConstraint;
 import org.thoughtcrime.securesms.jobmanager.impl.NetworkConstraintObserver;
 import org.thoughtcrime.securesms.jobmanager.impl.NetworkOrCellServiceConstraint;
+import org.thoughtcrime.securesms.jobmanager.impl.NoRemoteArchiveGarbageCollectionPendingConstraint;
 import org.thoughtcrime.securesms.jobmanager.impl.NotInCallConstraint;
 import org.thoughtcrime.securesms.jobmanager.impl.NotInCallConstraintObserver;
 import org.thoughtcrime.securesms.jobmanager.impl.RestoreAttachmentConstraint;
@@ -135,7 +135,6 @@ public final class JobManagerFactories {
       put(BackfillDigestsForDataFileJob.KEY,           new BackfillDigestsForDataFileJob.Factory());
       put(BackupDeleteJob.KEY,                         new BackupDeleteJob.Factory());
       put(BackupMessagesJob.KEY,                       new BackupMessagesJob.Factory());
-      put(BackupRestoreJob.KEY,                        new BackupRestoreJob.Factory());
       put(BackupRestoreMediaJob.KEY,                   new BackupRestoreMediaJob.Factory());
       put(BackupSubscriptionCheckJob.KEY,              new BackupSubscriptionCheckJob.Factory());
       put(BuildExpirationConfirmationJob.KEY,          new BuildExpirationConfirmationJob.Factory());
@@ -394,6 +393,7 @@ public final class JobManagerFactories {
       put("AttachmentMarkUploadedJob",                   new FailingJob.Factory());
       put("BackupMediaSnapshotSyncJob",                  new FailingJob.Factory());
       put("PnpInitializeDevicesJob",                     new FailingJob.Factory());
+      put("BackupRestoreJob",                            new FailingJob.Factory());
     }};
   }
 

app/src/main/java/org/thoughtcrime/securesms/registrationv3/ui/restore/RemoteRestoreActivity.kt

@@ -311,7 +311,8 @@ private fun BackupAvailableContent(
     when (state.importState) {
       RemoteRestoreViewModel.ImportState.None -> Unit
       RemoteRestoreViewModel.ImportState.InProgress -> RestoreProgressDialog(state.restoreProgress)
-      is RemoteRestoreViewModel.ImportState.Restored -> Unit
+      RemoteRestoreViewModel.ImportState.Restored -> Unit
+      RemoteRestoreViewModel.ImportState.NetworkFailure -> RestoreNetworkFailedDialog(onDismiss = onImportErrorDialogDismiss)
       RemoteRestoreViewModel.ImportState.Failed -> {
         if (SignalStore.backup.hasInvalidBackupVersion) {
           InvalidBackupVersionDialog(onUpdateSignal = onUpdateSignal, onDismiss = onImportErrorDialogDismiss)
@@ -490,6 +491,19 @@ fun RestoreFailedDialog(
   )
 }
 
+@Composable
+fun RestoreNetworkFailedDialog(
+  onDismiss: () -> Unit = {}
+) {
+  Dialogs.SimpleAlertDialog(
+    title = stringResource(R.string.RemoteRestoreActivity__couldnt_transfer),
+    body = stringResource(R.string.RegistrationActivity_error_connecting_to_service),
+    confirm = stringResource(android.R.string.ok),
+    onConfirm = onDismiss,
+    onDismiss = onDismiss
+  )
+}
+
 @Composable
 fun TierRestoreFailedDialog(
   loadAttempts: Int = 0,

app/src/main/java/org/thoughtcrime/securesms/registrationv3/ui/restore/RemoteRestoreViewModel.kt

@@ -8,11 +8,9 @@ package org.thoughtcrime.securesms.registrationv3.ui.restore
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
@@ -21,12 +19,9 @@ import org.signal.core.util.bytes
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
 import org.thoughtcrime.securesms.backup.v2.MessageBackupTier
+import org.thoughtcrime.securesms.backup.v2.RemoteRestoreResult
 import org.thoughtcrime.securesms.backup.v2.RestoreV2Event
 import org.thoughtcrime.securesms.database.model.databaseprotos.RestoreDecisionState
-import org.thoughtcrime.securesms.dependencies.AppDependencies
-import org.thoughtcrime.securesms.jobmanager.JobTracker
-import org.thoughtcrime.securesms.jobs.BackupRestoreJob
-import org.thoughtcrime.securesms.jobs.BackupRestoreMediaJob
 import org.thoughtcrime.securesms.keyvalue.Completed
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.keyvalue.Skipped
@@ -97,44 +92,25 @@ class RemoteRestoreViewModel(isOnlyRestoreOption: Boolean) : ViewModel() {
       withContext(Dispatchers.IO) {
         QuickRegistrationRepository.setRestoreMethodForOldDevice(RestoreMethod.REMOTE_BACKUP)
 
-        val jobStateFlow = callbackFlow {
-          val listener = JobTracker.JobListener { _, jobState ->
-            trySend(jobState)
+        when (val result = BackupRepository.restoreRemoteBackup()) {
+          RemoteRestoreResult.Success -> {
+            Log.i(TAG, "Restore successful")
+            SignalStore.registration.restoreDecisionState = RestoreDecisionState.Completed
+
+            StorageServiceRestore.restore()
+
+            store.update { it.copy(importState = ImportState.Restored) }
           }
 
-          AppDependencies
-            .jobManager
-            .startChain(BackupRestoreJob())
-            .then(BackupRestoreMediaJob())
-            .enqueue(listener)
-
-          awaitClose {
-            AppDependencies.jobManager.removeListener(listener)
+          RemoteRestoreResult.NetworkError -> {
+            Log.w(TAG, "Restore failed to download")
+            store.update { it.copy(importState = ImportState.NetworkFailure) }
           }
-        }
 
-        jobStateFlow.collect { state ->
-          when (state) {
-            JobTracker.JobState.SUCCESS -> {
-              Log.i(TAG, "Restore successful")
-              SignalStore.registration.restoreDecisionState = RestoreDecisionState.Completed
-
-              StorageServiceRestore.restore()
-
-              store.update { it.copy(importState = ImportState.Restored) }
-            }
-
-            JobTracker.JobState.PENDING,
-            JobTracker.JobState.RUNNING -> {
-              Log.i(TAG, "Restore job states updated: $state")
-            }
-
-            JobTracker.JobState.FAILURE,
-            JobTracker.JobState.IGNORED -> {
-              Log.w(TAG, "Restore failed with $state")
-
-              store.update { it.copy(importState = ImportState.Failed) }
-            }
+          RemoteRestoreResult.Canceled,
+          RemoteRestoreResult.Failure -> {
+            Log.w(TAG, "Restore failed with $result")
+            store.update { it.copy(importState = ImportState.Failed) }
           }
         }
       }
@@ -195,6 +171,7 @@ class RemoteRestoreViewModel(isOnlyRestoreOption: Boolean) : ViewModel() {
     data object None : ImportState
     data object InProgress : ImportState
     data object Restored : ImportState
+    data object NetworkFailure : ImportState
     data object Failed : ImportState
   }
 }