#[error("future cursor: {message:?}")]
FutureCursor { message: Option<String> },
#[error("connect timed out")]
ConnectTimeout,
#[error("host went silent: no frame or pong for {0:?}")]
Silent(Duration),
}
impl From<serde_ipld_dagcbor::DecodeError<Infallible>> for FirehoseError {
fn from(e: serde_ipld_dagcbor::DecodeError<Infallible>) -> Self {
} else {
Duration::from_secs(30)
};
const PING_INTERVAL: Duration = if cfg!(test) {
Duration::from_millis(100)
} else {
Duration::from_secs(45)
};
/// two ping intervals with nothing heard means the connection is dead, because every host pongs
const SILENCE_DEADLINE: Duration = PING_INTERVAL.saturating_mul(2);
pub struct FirehoseStream {
ws: WebSocketStream<tokio_websockets::MaybeTlsStream<tokio::net::TcpStream>>,
ping_timer: tokio::time::Interval,
last_heard: tokio::time::Instant,
}
impl FirehoseStream {
pub async fn connect(
let (ws, _) = tokio::time::timeout(CONNECT_TIMEOUT, connect)
.await
.map_err(|_| FirehoseError::ConnectTimeout)??;
let mut ping_timer = tokio::time::interval(Duration::from_secs(45));
let mut ping_timer = tokio::time::interval(PING_INTERVAL);
ping_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
ping_timer.tick().await; // consume the first tick
Ok(Self { ws, ping_timer })
Ok(Self {
ws,
ping_timer,
last_heard: tokio::time::Instant::now(),
})
}
/// gets the next message bytes from the firehose
/// none means the stream is closed
pub async fn next(&mut self) -> Result<Bytes, FirehoseError> {
loop {
let res = tokio::select! {
_ = self.ping_timer.tick() => {
let silent = self.last_heard.elapsed();
if silent >= SILENCE_DEADLINE {
return Err(FirehoseError::Silent(silent));
}
self.ping().await?;
continue;
}
res = self.ws.next() => {
res.map(|m| m.map_err(Into::into))
.unwrap_or(Err(FirehoseError::TcpDropped))?
}
};
self.last_heard = tokio::time::Instant::now();
match res {
msg if msg.is_binary() => {
let bytes: Bytes = msg.into_payload().into();
if bytes.is_empty() {
191 unmodified lines
Expand all
let res = FirehoseStream::connect(url, None, false).await;
assert!(matches!(res, Err(FirehoseError::ConnectTimeout)));
held.abort();
}
#[tokio::test]
async fn a_host_that_stops_answering_pings_is_dropped() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = Url::parse(&format!("ws://{}/", listener.local_addr().unwrap())).unwrap();
// tokio-websockets only pongs while reading, and this server never reads
let held = tokio::spawn(async move {
let (tcp, _) = listener.accept().await.unwrap();
let _ws = tokio_websockets::ServerBuilder::new()
.accept(tcp)
.await
.unwrap();
std::future::pending::<()>().await
});
let mut stream = FirehoseStream::connect(url, None, false).await.unwrap();
let res = tokio::time::timeout(SILENCE_DEADLINE * 5, stream.next()).await;
assert!(matches!(res, Ok(Err(FirehoseError::Silent(_)))));
held.abort();
}
#[tokio::test]
async fn a_host_that_answers_pings_stays_connected() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = Url::parse(&format!("ws://{}/", listener.local_addr().unwrap())).unwrap();
let held = tokio::spawn(async move {
let (tcp, _) = listener.accept().await.unwrap();
let (_, mut ws) = tokio_websockets::ServerBuilder::new()
.accept(tcp)
.await
.unwrap();
while ws.next().await.is_some() {}
});
let mut stream = FirehoseStream::connect(url, None, false).await.unwrap();
let res = tokio::time::timeout(SILENCE_DEADLINE * 5, stream.next()).await;
assert!(
res.is_err(),
"a quiet host that pongs isn't dropped: {res:?}"
);
held.abort();
}
}