websocket: automatically reconnect if not heard from for a while

kordophone/src/api/http_client.rs

@@ -12,7 +12,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use tokio::net::TcpStream;
 
-use futures_util::{StreamExt, TryStreamExt};
+use futures_util::{SinkExt, StreamExt, TryStreamExt};
 use futures_util::stream::{SplitStream, SplitSink, Stream};
 use futures_util::stream::BoxStream;
 
@@ -124,24 +124,21 @@ impl<B> AuthSetting for hyper::http::Request<B> {
     }
 }
 
-type WebsocketSink = SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message>;
-type WebsocketStream = SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>;
-
 pub struct WebsocketEventSocket {
-    _sink: WebsocketSink,
-    stream: WebsocketStream,
+    socket: WebSocketStream<MaybeTlsStream<TcpStream>>,
 }
 
 impl WebsocketEventSocket {
     pub fn new(socket: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
-        let (sink, stream) = socket.split();
-        Self { _sink: sink, stream }
+        Self { socket }
     }
 }
 
 impl WebsocketEventSocket {
     fn raw_update_stream(self) -> impl Stream<Item = Result<Vec<UpdateItem>, Error>> {
-        self.stream
+        let (_, stream) = self.socket.split();
+
+        stream
             .map_err(Error::from)
             .try_filter_map(|msg| async move {
                 match msg {
@@ -150,6 +147,15 @@ impl WebsocketEventSocket {
                             .map(Some)
                             .map_err(Error::from)
                     }
+                    tungstenite::Message::Ping(_) => {
+                        // Borrowing issue here with the sink, need to handle pings at the client level (whomever
+                        // is consuming these updateitems, should be a union type of updateitem | ping).
+                        Ok(None)
+                    }
+                    tungstenite::Message::Close(_) => {
+                        // Connection was closed cleanly
+                        Err(Error::ClientError("WebSocket connection closed".into()))
+                    }
                     _ => Ok(None)
                 }
             })
@@ -246,11 +252,16 @@ impl<K: AuthenticationStore + Send + Sync> APIInterface for HTTPAPIClient<K> {
         Ok(message)
     }
 
-    async fn open_event_socket(&mut self) -> Result<WebsocketEventSocket, Self::Error> {
+    async fn open_event_socket(&mut self, update_seq: Option<u64>) -> Result<WebsocketEventSocket, Self::Error> {
         use tungstenite::handshake::client::Request as TungsteniteRequest;
         use tungstenite::handshake::client::generate_key;
 
-        let uri = self.uri_for_endpoint("updates", Some(self.websocket_scheme()));
+        let endpoint = match update_seq {
+            Some(seq) => format!("updates?seq={}", seq),
+            None => "updates".to_string(),
+        };
+
+        let uri = self.uri_for_endpoint(&endpoint, Some(self.websocket_scheme()));
 
         log::debug!("Connecting to websocket: {:?}", uri);
 
@@ -501,7 +512,7 @@ mod test {
         let mut client = local_mock_client();
 
         // We just want to see if the connection is established, we won't wait for any events
-        let _ = client.open_event_socket().await.unwrap();
+        let _ = client.open_event_socket(None).await.unwrap();
         assert!(true);
     }
 }

kordophone/src/api/mod.rs

@@ -46,5 +46,5 @@ pub trait APIInterface {
     async fn authenticate(&mut self, credentials: Credentials) -> Result<JwtToken, Self::Error>;
 
     // (WS) /updates
-    async fn open_event_socket(&mut self) -> Result<impl EventSocket, Self::Error>;
+    async fn open_event_socket(&mut self, update_seq: Option<u64>) -> Result<impl EventSocket, Self::Error>;
 }