buzzrelay/src/stream.rs

68 lines
1.8 KiB
Rust
Raw Normal View History

2022-12-18 20:31:50 +00:00
use std::time::Duration;
use futures::{Stream, StreamExt};
use eventsource_stream::Eventsource;
use tokio::{
sync::mpsc::{channel, Receiver},
time::sleep,
};
#[derive(Debug)]
pub enum StreamError {
Http(reqwest::Error),
HttpStatus(reqwest::StatusCode),
InvalidContentType,
}
2023-01-11 17:45:45 +00:00
async fn run(url: &str) -> Result<impl Stream<Item = String>, StreamError> {
2022-12-18 20:31:50 +00:00
let client = reqwest::Client::new();
let res = client.get(url)
.timeout(Duration::MAX)
.send()
.await
.map_err(StreamError::Http)?;
if res.status() != 200 {
return Err(StreamError::HttpStatus(res.status()));
}
let ct = res.headers().get("content-type")
.and_then(|c| c.to_str().ok());
if ct.map_or(true, |ct| ct != "text/event-stream") {
return Err(StreamError::InvalidContentType);
}
let src = res.bytes_stream()
.eventsource()
.filter_map(|result| async {
let result = result.ok()?;
if result.event == "update" {
Some(result)
} else {
None
}
})
2022-12-19 20:20:13 +00:00
.map(|event| event.data);
2022-12-18 20:31:50 +00:00
Ok(src)
}
2023-01-11 17:45:45 +00:00
pub fn spawn(hosts: impl Iterator<Item = impl Into<String>>) -> Receiver<String> {
2022-12-18 20:31:50 +00:00
let (tx, rx) = channel(1024);
2023-01-11 17:45:45 +00:00
for host in hosts {
let host = host.into();
let tx = tx.clone();
tokio::spawn(async move {
loop {
match run(&host).await {
Ok(stream) =>
stream.for_each(|post| async {
tx.send(post).await.unwrap();
}).await,
Err(e) =>
tracing::error!("stream: {:?}", e),
}
2022-12-18 20:31:50 +00:00
2023-01-11 17:45:45 +00:00
sleep(Duration::from_secs(1)).await;
}
});
}
2022-12-18 20:31:50 +00:00
rx
}