Handle case where no new messages are available

ProjectLighthouseCAU · Jul 17, 2024 · 3cd3c2b · 3cd3c2b
1 parent 1ad4f39
commit 3cd3c2b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/lighthouse-client/src/error.rs b/lighthouse-client/src/error.rs
@@ -17,6 +17,8 @@ pub enum Error {
     Value(#[from] ValueError),
     #[error("Server error: {} {} (warnings: {:?})", code, message.clone().unwrap_or_else(|| "(no message)".to_string()), warnings)]
     Server { code: i32, message: Option<String>, warnings: Vec<String> },
+    #[error("No next message available")]
+    NoNextMessage,
     #[error("Custom error")]
     Custom(String),
 }

diff --git a/lighthouse-client/src/lighthouse.rs b/lighthouse-client/src/lighthouse.rs
@@ -85,6 +85,10 @@ impl<S> Lighthouse<S>
                         warn!("Got message without request id from server: {:?}", msg);
                     }
                 },
+                Err(Error::NoNextMessage) => {
+                    info!("No next message available, closing receive loop");
+                    break
+                },
                 Err(e) => error!("Bad message: {:?}", e),
             }
         }
@@ -104,7 +108,7 @@ impl<S> Lighthouse<S>
     #[tracing::instrument(skip(ws_stream))]
     async fn receive_raw_from(ws_stream: &mut SplitStream<S>) -> Result<Vec<u8>> {
         loop {
-            let message = ws_stream.next().await.ok_or_else(|| Error::custom("Got no message"))??;
+            let message = ws_stream.next().await.ok_or_else(|| Error::NoNextMessage)??;
             match message {
                 Message::Binary(bytes) => break Ok(bytes),
                 // We ignore pings for now