Rotate quick restore QR code and web socket.

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

@@ -169,8 +169,12 @@ private fun RestoreViaQrScreen(
           ) {
             AnimatedContent(
               targetState = state.qrState,
+              contentKey = { it::class },
               contentAlignment = Alignment.Center,
-              label = "qr-code-progress"
+              label = "qr-code-progress",
+              modifier = Modifier
+                .fillMaxWidth()
+                .fillMaxHeight()
             ) { qrState ->
               when (qrState) {
                 is RestoreViaQrViewModel.QrState.Loaded -> {
@@ -184,7 +188,9 @@ private fun RestoreViaQrScreen(
                 }
 
                 RestoreViaQrViewModel.QrState.Loading -> {
-                  CircularProgressIndicator(modifier = Modifier.size(48.dp))
+                  Box(contentAlignment = Alignment.Center) {
+                    CircularProgressIndicator(modifier = Modifier.size(48.dp))
+                  }
                 }
 
                 is RestoreViaQrViewModel.QrState.Scanned,

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

@@ -6,10 +6,15 @@
 package org.thoughtcrime.securesms.registrationv3.ui.restore
 
 import androidx.lifecycle.ViewModel
-import kotlinx.coroutines.CoroutineExceptionHandler
+import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.isActive
+import kotlinx.coroutines.launch
 import org.signal.core.util.logging.Log
 import org.signal.registration.proto.RegistrationProvisionMessage
 import org.thoughtcrime.securesms.backup.v2.MessageBackupTier
@@ -31,15 +36,30 @@ class RestoreViaQrViewModel : ViewModel() {
 
   val state: StateFlow<RestoreViaQrState> = store
 
-  private var socketHandle: Closeable
+  private var socketHandles: MutableList<Closeable> = mutableListOf()
+  private var startNewSocketJob: Job? = null
 
   init {
-    socketHandle = start()
+    restart()
   }
 
   fun restart() {
-    socketHandle.close()
-    socketHandle = start()
+    SignalStore.registration.restoreMethodToken = null
+    shutdown()
+
+    startNewSocket()
+
+    startNewSocketJob = viewModelScope.launch(Dispatchers.IO) {
+      var count = 0
+      while (count < 5 && isActive) {
+        delay(ProvisioningSocket.LIFESPAN / 2)
+        if (isActive) {
+          startNewSocket()
+          count++
+          Log.d(TAG, "Started next websocket count: $count")
+        }
+      }
+    }
   }
 
   fun handleRegistrationFailure() {
@@ -61,20 +81,66 @@ class RestoreViaQrViewModel : ViewModel() {
   }
 
   override fun onCleared() {
-    socketHandle.close()
+    shutdown()
+  }
+
+  private fun startNewSocket() {
+    synchronized(socketHandles) {
+      socketHandles += start()
+
+      if (socketHandles.size > 2) {
+        socketHandles.removeAt(0).close()
+      }
+    }
+  }
+
+  private fun shutdown() {
+    startNewSocketJob?.cancel()
+    synchronized(socketHandles) {
+      socketHandles.forEach { it.close() }
+      socketHandles.clear()
+    }
   }
 
   private fun start(): Closeable {
-    SignalStore.registration.restoreMethodToken = null
-    store.update { it.copy(qrState = QrState.Loading) }
+    store.update {
+      if (it.qrState !is QrState.Loaded) {
+        it.copy(qrState = QrState.Loading)
+      } else {
+        it
+      }
+    }
 
     return ProvisioningSocket.start(
       identityKeyPair = IdentityKeyUtil.generateIdentityKeyPair(),
       configuration = AppDependencies.signalServiceNetworkAccess.getConfiguration(),
-      handler = CoroutineExceptionHandler { _, _ -> store.update { it.copy(qrState = QrState.Failed) } }
+      handler = { id, t ->
+        store.update {
+          if (it.currentSocketId == null || it.currentSocketId == id) {
+            Log.w(TAG, "Current socket [$id] has failed, stopping automatic connects", t)
+            shutdown()
+            it.copy(currentSocketId = null, qrState = QrState.Failed)
+          } else {
+            Log.i(TAG, "Old socket [$id] failed, ignoring")
+            it
+          }
+        }
+      }
     ) { socket ->
       val url = socket.getProvisioningUrl()
-      store.update { it.copy(qrState = QrState.Loaded(qrData = QrCodeData.forData(data = url, supportIconOverlay = false))) }
+      store.update {
+        Log.d(TAG, "Updating QR code with data from [${socket.id}]")
+
+        it.copy(
+          currentSocketId = socket.id,
+          qrState = QrState.Loaded(
+            qrData = QrCodeData.forData(
+              data = url,
+              supportIconOverlay = false
+            )
+          )
+        )
+      }
 
       val result = socket.getRegistrationProvisioningMessage()
 
@@ -94,8 +160,15 @@ class RestoreViaQrViewModel : ViewModel() {
           SignalStore.backup.usedBackupMediaSpace = result.message.backupSizeBytes
         }
         store.update { it.copy(isRegistering = true, provisioningMessage = result.message, qrState = QrState.Scanned) }
+        shutdown()
       } else {
-        store.update { it.copy(showProvisioningError = true, qrState = QrState.Scanned) }
+        store.update {
+          if (it.currentSocketId == socket.id) {
+            it.copy(showProvisioningError = true, qrState = QrState.Scanned)
+          } else {
+            it
+          }
+        }
       }
     }
   }
@@ -105,7 +178,8 @@ class RestoreViaQrViewModel : ViewModel() {
     val qrState: QrState = QrState.Loading,
     val provisioningMessage: RegistrationProvisionMessage? = null,
     val showProvisioningError: Boolean = false,
-    val showRegistrationError: Boolean = false
+    val showRegistrationError: Boolean = false,
+    val currentSocketId: Int? = null
   )
 
   sealed interface QrState {