diff --git a/src/ingest/firehose.rs b/src/ingest/firehose.rs --- a/src/ingest/firehose.rs +++ b/src/ingest/firehose.rs @@ -101,11 +101,12 @@ if self.state.pds_meta.load().is_banned(host) { break Ok(()); } - self.enabled.wait_enabled("firehose").await; + // sleep stream backoff out if we have any tokio::time::sleep(backoff).await; + // get cursor let start_cursor = self .state .firehose_cursors @@ -114,7 +115,6 @@ (val > 0).then_some(val) }) .flatten(); - match start_cursor { Some(c) => info!(cursor = %c, "resuming from cursor"), None => info!("no cursor found, live tailing"), @@ -155,8 +155,7 @@ let res = loop { tokio::select! { msg = stream.next() => { - let Some(bytes_res) = msg else { break Err(FirehoseError::EmptyFrame); }; - let bytes = match bytes_res { + let bytes = match msg { Ok(b) => b, Err(e) => break Err(e), }; diff --git a/src/ingest/stream.rs b/src/ingest/stream.rs --- a/src/ingest/stream.rs +++ b/src/ingest/stream.rs @@ -2,7 +2,7 @@ use axum::http::Uri; use bytes::Bytes; -use futures::StreamExt; +use futures::{SinkExt, StreamExt}; use jacquard_common::error::DecodeError; use jacquard_common::{ CowStr, @@ -15,7 +15,7 @@ use serde::{Deserialize, Serialize}; use smol_str::format_smolstr; use thiserror::Error; -use tokio_websockets::{ClientBuilder, WebSocketStream}; +use tokio_websockets::{ClientBuilder, Message as WsMsg, WebSocketStream}; use tracing::trace; use url::Url; @@ -42,6 +42,8 @@ UnknownType(String), #[error("cbor decode error: {0}")] Cbor(String), + #[error("stream closed")] + StreamClosed, } impl From> for FirehoseError { @@ -76,19 +78,27 @@ } /// gets the next message bytes from the firehose - pub async fn next(&mut self) -> Option> { + /// none means the stream is closed + pub async fn next(&mut self) -> Result { loop { - match self.ws.next().await? { - Err(e) => return Some(Err(e.into())), - Ok(msg) if msg.is_binary() => { + let res = self + .ws + .next() + .await + .map(|m| m.map_err(Into::into)) + .unwrap_or_else(|| Err(FirehoseError::StreamClosed))?; + match res { + msg if msg.is_binary() => { let bytes: Bytes = msg.into_payload().into(); if bytes.is_empty() { - return Some(Err(FirehoseError::EmptyFrame)); + return Err(FirehoseError::EmptyFrame); } - return Some(Ok(bytes)); + return Ok(bytes); } - Ok(msg) if msg.is_close() => return None, - Ok(x) => { + msg if msg.is_ping() => self.ws.send(WsMsg::pong(msg.into_payload())).await?, + // if ws closed treat it as an error, since why would a host close the stream?? + msg if msg.is_close() => return Err(FirehoseError::StreamClosed), + x => { trace!(msg = ?x, "relay sent unexpected message"); continue; }