fix: close websocket handler when web socket is dropped (#320)

* When `QsWebSocket` was dropped, the handler was still running. Now, it
  is also cancelled.
* `UserCubit` considered web socket disconnect events as an error, and
  reconnected again. This is not needed since internally `QsWebSocket`
  automatically reconnects.

Additionally, replace the broadcast channel by mpsc one, since we have
only a single subscriber. Also remove unused public APIs.

Author: boxdot

Cargo.lock

@@ -53,6 +53,9 @@ tracing-subscriber = { version = "0.3", features = [
     "parking_lot",
 ] }
 
+tokio-util = "0.7.13"
+
+
 [patch.crates-io]
 #opaque-ke = { git = "https://github.com/facebook/opaque-ke", branch = "dependabot/cargo/voprf-eq-0.5.0" }
 

Cargo.toml

@@ -14,6 +14,7 @@ reqwest = { workspace = true }
 thiserror = "1"
 phnxtypes = { path = "../types" }
 tokio = { version = "1.18.2", features = ["macros"] }
+tokio-util = { workspace = true }
 tokio-tungstenite = { version = "0.23", features = ["rustls-tls-webpki-roots"] }
 futures-util = "0.3.21"
 http = "1"
@@ -23,6 +24,7 @@ mls-assist = { workspace = true }
 privacypass = { workspace = true }
 tls_codec = { workspace = true }
 url = "2"
+uuid = { version = "1", features = ["v4"] }
 
 [dev-dependencies]
 tokio = { version = "1.18.2", features = ["macros"] }

apiclient/Cargo.toml

@@ -20,6 +20,7 @@ use phnxtypes::{
     messages::{client_ds::QsWsMessage, client_qs::QsOpenWsParams},
 };
 use tls_codec::Serialize;
+use tokio_util::sync::CancellationToken;
 use tracing::{error, info};
 use uuid::Uuid;
 
@@ -49,8 +50,9 @@ async fn ws_lifecycle() {
     let client = ApiClient::with_default_http_client(address).expect("Failed to initialize client");
 
     // Spawn the websocket connection task
+    let cancel = CancellationToken::new();
     let mut ws = client
-        .spawn_websocket(queue_id, timeout, retry_interval)
+        .spawn_websocket(queue_id, timeout, retry_interval, cancel)
         .await
         .expect("Failed to execute request");
 

apiclient/src/qs_api/tests.rs

@@ -18,16 +18,17 @@ use thiserror::*;
 use tls_codec::DeserializeBytes;
 use tokio::{
     net::TcpStream,
-    sync::broadcast::{self, Receiver, Sender},
-    task::JoinHandle,
+    sync::mpsc,
     time::{sleep, Instant},
 };
 use tokio_tungstenite::{
     connect_async,
     tungstenite::{client::IntoClientRequest, protocol::Message},
     MaybeTlsStream, WebSocketStream,
 };
+use tokio_util::sync::{CancellationToken, DropGuard};
 use tracing::{error, info};
+use uuid::Uuid;
 
 use crate::{ApiClient, Protocol};
 
@@ -53,9 +54,12 @@ impl ConnectionStatus {
         Self { connected: false }
     }
 
-    fn set_connected(&mut self, tx: &Sender<WsEvent>) -> Result<(), ConnectionStatusError> {
+    async fn set_connected(
+        &mut self,
+        tx: &mpsc::Sender<WsEvent>,
+    ) -> Result<(), ConnectionStatusError> {
         if !self.connected {
-            if let Err(error) = tx.send(WsEvent::ConnectedEvent) {
+            if let Err(error) = tx.send(WsEvent::ConnectedEvent).await {
                 error!(%error, "Error sending to channel");
                 self.connected = false;
                 return Err(ConnectionStatusError::ChannelClosed);
@@ -65,9 +69,12 @@ impl ConnectionStatus {
         Ok(())
     }
 
-    fn set_disconnected(&mut self, tx: &Sender<WsEvent>) -> Result<(), ConnectionStatusError> {
+    async fn set_disconnected(
+        &mut self,
+        tx: &mpsc::Sender<WsEvent>,
+    ) -> Result<(), ConnectionStatusError> {
         if self.connected {
-            if let Err(error) = tx.send(WsEvent::DisconnectedEvent) {
+            if let Err(error) = tx.send(WsEvent::DisconnectedEvent).await {
                 error!(%error, "Error sending to channel");
                 return Err(ConnectionStatusError::ChannelClosed);
             }
@@ -79,48 +86,30 @@ impl ConnectionStatus {
 
 /// A websocket connection to the QS server. See the
 /// [`ApiClient::spawn_websocket`] method for more information.
+///
+/// When dropped, the websocket connection will be closed.
 pub struct QsWebSocket {
-    rx: Receiver<WsEvent>,
-    tx: Sender<WsEvent>,
-    handle: JoinHandle<()>,
+    rx: mpsc::Receiver<WsEvent>,
+    _cancel: DropGuard,
 }
 
 impl QsWebSocket {
     /// Returns the next [`WsEvent`] event. This will block until an event is
     /// sent or the connection is closed (in which case a final `None` is
     /// returned).
     pub async fn next(&mut self) -> Option<WsEvent> {
-        match self.rx.recv().await {
-            Ok(message) => Some(message),
-            Err(error) => {
-                error!(%error, "Error receiving from channel");
-                None
-            }
-        }
-    }
-
-    /// Subscribe to the event stream
-    pub fn subscribe(&self) -> Receiver<WsEvent> {
-        self.tx.subscribe()
-    }
-
-    /// Join the websocket connection task. This will block until the task has
-    /// completed.
-    pub async fn join(self) -> Result<(), tokio::task::JoinError> {
-        self.handle.await
-    }
-
-    /// Abort the websocket connection task. This will close the websocket connection.
-    pub fn abort(&mut self) {
-        self.handle.abort();
+        self.rx.recv().await
     }
 
     /// Internal helper function to handle an established websocket connection
+    ///
+    /// Returns `true` if the connection should be re-established, otherwise `false`.
     async fn handle_connection(
         ws_stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
-        tx: &Sender<WsEvent>,
+        tx: &mpsc::Sender<WsEvent>,
         timeout: u64,
-    ) {
+        cancel: &CancellationToken,
+    ) -> bool {
         let mut last_ping = Instant::now();
 
         // Watchdog to monitor the connection.
@@ -131,26 +120,31 @@ impl QsWebSocket {
 
         // Initialize the connection status
         let mut connection_status = ConnectionStatus::new();
-        if connection_status.set_connected(tx).is_err() {
+        if connection_status.set_connected(tx).await.is_err() {
             // Close the stream if all subscribers of the watch have been dropped
             let _ = ws_stream.close().await;
-            return;
+            return false;
         }
 
         // Loop while the connection is open
         loop {
             tokio::select! {
+                // Check is the handler is cancelled
+                _ = cancel.cancelled() => {
+                    info!("QS WebSocket connection cancelled");
+                    break false;
+                },
                 // Check if the connection is still alive
                 _ = interval.tick() => {
                     let now = Instant::now();
                     // Check if we have reached the timeout
                     if now.duration_since(last_ping) > Duration::from_secs(timeout) {
                         // Change the status to Disconnected and send an event
                         let _ = ws_stream.close().await;
-                        if connection_status.set_disconnected(tx).is_err() {
+                        if connection_status.set_disconnected(tx).await.is_err() {
                             // Close the stream if all subscribers of the watch have been dropped
                             info!("Closing the connection because all subscribers are dropped");
-                            return;
+                            return false;
                         }
                     }
                 },
@@ -163,53 +157,53 @@ impl QsWebSocket {
                                 // Reset the last ping time
                                 last_ping = Instant::now();
                                 // Change the status to Connected and send an event
-                                if connection_status.set_connected(tx).is_err() {
+                                if connection_status.set_connected(tx).await.is_err() {
                                     // Close the stream if all subscribers of the watch have been dropped
                                     info!("Closing the connection because all subscribers are dropped");
                                     let _ = ws_stream.close().await;
-                                    return;
+                                    return false;
                                 }
                                 // Try to deserialize the message
                                 if let Ok(QsWsMessage::QueueUpdate) =
                                     QsWsMessage::tls_deserialize_exact_bytes(&data)
                                 {
                                     // We received a new message notification from the QS
                                     // Send the event to the channel
-                                    if tx.send(WsEvent::MessageEvent(QsWsMessage::QueueUpdate)).is_err() {
+                                    if tx.send(WsEvent::MessageEvent(QsWsMessage::QueueUpdate)).await.is_err() {
                                         info!("Closing the connection because all subscribers are dropped");
                                         // Close the stream if all subscribers of the watch have been dropped
                                         let _ = ws_stream.close().await;
-                                        return;
+                                        return false;
                                     }
                                 }
                             },
                             // We received a ping
                             Message::Ping(_) => {
                                 // We update the last ping time
                                 last_ping = Instant::now();
-                                if connection_status.set_connected(tx).is_err() {
+                                if connection_status.set_connected(tx).await.is_err() {
                                     // Close the stream if all subscribers of the watch have been dropped
                                     info!("Closing the connection because all subscribers are dropped");
                                     let _ = ws_stream.close().await;
-                                    return;
+                                    return false;
                                 }
                             }
                             Message::Close(_) => {
                                 // Change the status to Disconnected and send an
                                 // event
-                                let _ = connection_status.set_disconnected(tx);
+                                let _ = connection_status.set_disconnected(tx).await;
                                 // We close the websocket
                                 let _ = ws_stream.close().await;
-                                return;
+                                return true;
                             }
                             _ => {
                             }
                         }
                     } else {
                         // It seems the connection is closed, send disconnect
                         // event
-                        let _ = connection_status.set_disconnected(tx);
-                        break;
+                        let _ = connection_status.set_disconnected(tx).await;
+                        break true;
                     }
                 },
             }
@@ -255,13 +249,14 @@ impl ApiClient {
     /// [`WsEvent::ConnectedEvent].
     ///
     /// The connection will be closed if all subscribers of the [`QsWebSocket`]
-    /// have been dropped, or when it is manually closed with using the
-    /// [`QsWebSocket::abort()`] function.
+    /// have been dropped, or when it is manually closed by cancelling the token
+    /// `cancel`.
     ///
     /// # Arguments
     ///  -  `queue_id` - The ID of the queue monitor.
     ///  - `timeout` - The timeout for the connection in seconds.
     ///  - `retry_interval` - The interval between connection attempts in seconds.
+    ///  - `cancel` - The cancellation token to stop the socket.
     ///
     /// # Returns
     /// A new [`QsWebSocket`] that represents the websocket connection.
@@ -270,6 +265,7 @@ impl ApiClient {
         queue_id: QsClientId,
         timeout: u64,
         retry_interval: u64,
+        cancel: CancellationToken,
     ) -> Result<QsWebSocket, SpawnWsError> {
         // Set the request parameter
         let qs_ws_open_params = QsOpenWsParams { queue_id };
@@ -289,19 +285,19 @@ impl ApiClient {
             })?;
 
         // We create a channel to send events to
-        let (tx, rx) = broadcast::channel(100);
+        let (tx, rx) = mpsc::channel(100);
 
-        // We clone the sender, so that we can subscribe to more receivers
-        let tx_clone = tx.clone();
+        let connection_id = Uuid::new_v4();
+        info!(%connection_id, "Spawning the websocket connection...");
 
-        info!("Spawning the websocket connection...");
+        let cancel_guard = cancel.clone().drop_guard();
 
         // Spawn the connection task
-        let handle = tokio::spawn(async move {
+        tokio::spawn(async move {
             // Connection loop
             #[cfg(test)]
             let mut counter = 0;
-            loop {
+            while !cancel.is_cancelled() {
                 // We build the request and set a custom header
                 let req = match address.clone().into_client_request() {
                     Ok(mut req) => {
@@ -319,13 +315,15 @@ impl ApiClient {
                 match connect_async(req).await {
                     // The connection was established
                     Ok((ws_stream, _)) => {
-                        info!("Connected to QS WebSocket");
+                        info!(%connection_id, "Connected to QS WebSocket");
                         // Hand over the connection to the handler
-                        QsWebSocket::handle_connection(ws_stream, &tx, timeout).await;
+                        if !QsWebSocket::handle_connection(ws_stream, &tx, timeout, &cancel).await {
+                            break;
+                        }
                     }
                     // The connection was not established, wait and try again
-                    Err(e) => {
-                        error!("Error connecting to QS WebSocket: {}", e);
+                    Err(error) => {
+                        error!(%error, "Error connecting to QS WebSocket");
                         #[cfg(test)]
                         {
                             counter += 1;
@@ -336,17 +334,20 @@ impl ApiClient {
                     }
                 }
                 info!(
+                    %connection_id,
                     retry_in_sec = retry_interval,
+                    is_cancelled = cancel.is_cancelled(),
                     "The websocket was closed, will reconnect...",
                 );
                 sleep(time::Duration::from_secs(retry_interval)).await;
             }
+
+            info!(%connection_id, "QS WebSocket closed");
         });
 
         Ok(QsWebSocket {
             rx,
-            tx: tx_clone,
-            handle,
+            _cancel: cancel_guard,
         })
     }
 }

apiclient/src/qs_api/ws.rs

@@ -37,6 +37,6 @@ flutter_rust_bridge = { version = "=2.7.0", features = ["chrono", "uuid"] }
 notify-rust = "4"
 chrono = { workspace = true }
 jni = "0.21"
-tokio-util = "0.7.13"
+tokio-util = { workspace = true }
 tokio-stream = "0.1.17"
 blake3 = "1.5.5"