(fix): connections (again)

src/rho/client_connection.rs

@@ -75,6 +75,10 @@ impl ClientConnection {
         self.rho_connection.read().await.clone()
     }
 
+    pub async fn set_rho_connection(&self, rho_connection: Arc<RhoConnection>) {
+        *self.rho_connection.write().await = Some(rho_connection);
+    }
+
     /// Send a CommunicationValue to the client
     pub async fn send_message(self: Arc<Self>, cv: &CommunicationValue) {
         if !*self.is_open.read().await {

src/rho/connection.rs

@@ -234,6 +234,12 @@ impl GeneralConnection {
         let iota = IotaConnection::from_general(self.clone(), id).await;
         let rho = Arc::new(RhoConnection::new(iota.clone(), Vec::new()).await);
         iota.set_rho_connection(rho.clone()).await;
+        if let Some(previous) = self.state.rho.get_by_iota(id as i64).await {
+            for client in previous.drain_client_connections() {
+                client.set_rho_connection(rho.clone()).await;
+                rho.add_client_connection(client).await;
+            }
+        }
         self.state.rho.add(rho).await;
         self.load_iota_users(&iota, id).await;
         iota.start();

src/rho/iota_connection.rs

@@ -66,7 +66,14 @@ impl IotaConnection {
                             iota_for_closure.handle_message(cv).await;
                         });
                     }
-                    Err(_) => {
+                    Err(error) => {
+                        log_err!(
+                            self_clone.iota_id as i64,
+                            PrintType::Iota,
+                            "Iota receive loop ended: {}; transport close reason: {:?}",
+                            error,
+                            self_clone.receiver.close_reason()
+                        );
                         break;
                     }
                 }

src/rho/rho_connection.rs

@@ -63,6 +63,17 @@ impl RhoConnection {
             .collect()
     }
 
+    pub fn drain_client_connections(&self) -> Vec<Arc<ClientConnection>> {
+        let keys: Vec<_> = self
+            .client_connections
+            .iter()
+            .map(|entry| *entry.key())
+            .collect();
+        keys.into_iter()
+            .filter_map(|key| self.client_connections.remove(&key).map(|(_, value)| value))
+            .collect()
+    }
+
     /// Get client connections for a specific user
     pub async fn get_client_connections_for_user(
         &self,
@@ -198,25 +209,26 @@ impl RhoConnection {
 
     /// Close the Iota connection and all associated client connections
     pub async fn close_iota_connection(&self) {
+        let iota_id = self.get_iota_id().await as i64;
+        if self
+            .iota_connection
+            .state
+            .rho
+            .remove_if_current(iota_id, self)
+            .await
+            .is_none()
+        {
+            return;
+        }
+
         // Close all client connections
         let connections = self.get_client_connections().await;
         for connection in connections {
             connection.close().await;
         }
 
-        // Remove from manager
-        self.iota_connection
-            .state
-            .rho
-            .remove(self.get_iota_id().await as i64)
-            .await;
-
         // Notify OmegaConnection
-        self.iota_connection
-            .state
-            .omega
-            .close_iota(self.get_iota_id().await as i64)
-            .await;
+        self.iota_connection.state.omega.close_iota(iota_id).await;
     }
 
     /// Send message from Iota to specific client

src/rho/rho_manager.rs

@@ -89,8 +89,17 @@ impl RhoManager {
         }
     }
 
-    pub async fn remove(&self, iota_id: i64) -> Option<Arc<RhoConnection>> {
-        let rho = self.connections.remove(&iota_id).map(|(_, rho)| rho);
+    pub async fn remove_if_current(
+        &self,
+        iota_id: i64,
+        expected: &RhoConnection,
+    ) -> Option<Arc<RhoConnection>> {
+        let rho = self
+            .connections
+            .remove_if(&iota_id, |_, current| {
+                std::ptr::eq(current.as_ref(), expected)
+            })
+            .map(|(_, rho)| rho);
         if let Some(rho) = rho.as_ref() {
             self.users.retain(|_, mapped| !Arc::ptr_eq(mapped, rho));
         }

src/rho/server.rs

@@ -124,7 +124,7 @@ pub async fn start(state: Arc<AppState>) -> Result<(), Box<dyn std::error::Error
                 Duration::from_millis(30_000),
             )
             .with_keep_alive(Some(Duration::from_secs(6)))
-            .with_max_idle_timeout(Some(Duration::from_secs(30)))
+            .with_max_idle_timeout(None)
             .with_receiver_queue_capacity(1000)
             .with_max_concurrent_stream_tasks(10)
             .with_persistent_stream_retries(5, Duration::from_secs(5)),