228 lines
6.5 KiB
Rust
228 lines
6.5 KiB
Rust
use std::collections::HashSet;
|
|
use std::future::Future;
|
|
use std::time::Duration;
|
|
use cave::{
|
|
feed::{Feed, Post},
|
|
store::Store,
|
|
};
|
|
use futures::{StreamExt, future};
|
|
use crate::trend_setter::UpdateSet;
|
|
|
|
pub struct RobotsTxt {
|
|
robot: Option<texting_robots::Robot>,
|
|
}
|
|
|
|
impl RobotsTxt {
|
|
pub async fn fetch(
|
|
client: &reqwest::Client,
|
|
host: &str,
|
|
) -> Self {
|
|
let url = format!("https://{}/robots.txt", host);
|
|
let robot = async {
|
|
let body = client.get(url)
|
|
.send().await
|
|
.ok()?
|
|
.text().await
|
|
.ok()?;
|
|
texting_robots::Robot::new(
|
|
env!("CARGO_PKG_NAME"),
|
|
body.as_bytes(),
|
|
).ok()
|
|
}.await;
|
|
RobotsTxt { robot }
|
|
}
|
|
|
|
pub fn allowed(&self, url: &str) -> bool {
|
|
if let Some(robot) = &self.robot {
|
|
robot.allowed(url)
|
|
} else {
|
|
true
|
|
}
|
|
}
|
|
|
|
pub fn delay(&self) -> Option<Duration> {
|
|
self.robot.as_ref()
|
|
.and_then(|robot| robot.delay)
|
|
.map(|delay| Duration::from_secs(delay as u64))
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum Message {
|
|
WorkerDone,
|
|
Fetched {
|
|
host: String,
|
|
new_post_ratio: Option<f64>,
|
|
mean_interval: Option<Duration>,
|
|
},
|
|
IntroduceHosts { hosts: Vec<String> },
|
|
}
|
|
|
|
pub async fn run(
|
|
message_tx: tokio::sync::mpsc::UnboundedSender<Message>,
|
|
store: Store,
|
|
trend_setter_tx: crate::trend_setter::Tx,
|
|
client: reqwest::Client,
|
|
host: String,
|
|
) {
|
|
// Fetch /robots.txt
|
|
let robots_txt = RobotsTxt::fetch(&client, &host).await;
|
|
// Fetch posts and open stream
|
|
let ((new_post_ratio, mut mean_interval), stream_result) = future::join(
|
|
fetch_timeline(
|
|
message_tx.clone(), store.clone(),
|
|
trend_setter_tx.clone(),
|
|
&client, &robots_txt, &host
|
|
),
|
|
open_stream(
|
|
message_tx.clone(), store.clone(),
|
|
trend_setter_tx.clone(),
|
|
&client, &robots_txt, host.clone()
|
|
),
|
|
).await;
|
|
|
|
// Next worker can start
|
|
message_tx.send(Message::WorkerDone).unwrap();
|
|
|
|
// Process stream
|
|
if let Ok(stream) = stream_result {
|
|
log::info!("Processing stream for {}", host);
|
|
stream.await;
|
|
log::warn!("Ended stream for {}", host);
|
|
}
|
|
|
|
// Ready for reenqueue
|
|
if let Some(mean_interval) = &mut mean_interval {
|
|
if let Some(delay) = robots_txt.delay() {
|
|
*mean_interval = (*mean_interval).max(delay);
|
|
}
|
|
}
|
|
message_tx.send(Message::Fetched {
|
|
host: host.to_string(),
|
|
new_post_ratio,
|
|
mean_interval,
|
|
}).unwrap();
|
|
}
|
|
|
|
async fn fetch_timeline(
|
|
message_tx: tokio::sync::mpsc::UnboundedSender<Message>,
|
|
mut store: Store,
|
|
trend_setter_tx: crate::trend_setter::Tx,
|
|
client: &reqwest::Client,
|
|
robots_txt: &RobotsTxt,
|
|
host: &str,
|
|
) -> (Option<f64>, Option<Duration>) {
|
|
let url = format!("https://{}/api/v1/timelines/public?limit=40", host);
|
|
if ! robots_txt.allowed(&url) {
|
|
log::warn!("Timeline of {} forbidden by robots.txt", host);
|
|
return (None, None);
|
|
}
|
|
|
|
match Feed::fetch(&client, &url).await {
|
|
Ok(feed) => {
|
|
let mean_interval = feed.mean_post_interval();
|
|
|
|
let (new_post_ratio, introduce_hosts) = process_posts(&mut store, &trend_setter_tx, &host, feed.posts.into_iter()).await;
|
|
let _ = message_tx.send(Message::IntroduceHosts {
|
|
hosts: introduce_hosts.into_iter().collect(),
|
|
});
|
|
|
|
// successfully fetched, save for future run
|
|
store.save_host(&host).await.unwrap();
|
|
|
|
(new_post_ratio, mean_interval)
|
|
}
|
|
Err(e) => {
|
|
log::error!("Failed fetching {}: {}", host, e);
|
|
(None, None)
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn process_posts(
|
|
store: &mut Store,
|
|
trend_setter_tx: &crate::trend_setter::Tx,
|
|
host: &str,
|
|
posts: impl Iterator<Item = Post>,
|
|
) -> (Option<f64>, HashSet<String>) {
|
|
// introduce new hosts, validate posts
|
|
let mut introduce_hosts = HashSet::new();
|
|
let mut new_posts = 0;
|
|
let mut posts_len = 0;
|
|
for post in posts {
|
|
posts_len += 1;
|
|
|
|
// introduce instances from mentions
|
|
for mention in &post.mentions {
|
|
if let Some(user_host) = mention.user_host() {
|
|
introduce_hosts.insert(user_host);
|
|
}
|
|
}
|
|
|
|
// check if it's an actual post, not a repost
|
|
if let Some(author_host) = post.account.host() {
|
|
// send away to redis
|
|
let update_set = UpdateSet::from(&post);
|
|
if store.save_post(post).await == Ok(true) {
|
|
new_posts += 1;
|
|
|
|
if ! update_set.is_empty() {
|
|
trend_setter_tx.send(update_set).await.unwrap();
|
|
}
|
|
}
|
|
|
|
// introduce instances from authors
|
|
introduce_hosts.insert(author_host);
|
|
} else {
|
|
log::warn!("drop repost ({:?} on {})", post.account.host(), host);
|
|
}
|
|
}
|
|
log::info!("{}: {}/{} new posts", host, new_posts, posts_len);
|
|
let new_post_ratio = if posts_len > 0 {
|
|
Some((new_posts as f64) / (posts_len as f64))
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let introduce_hosts = introduce_hosts.into_iter()
|
|
.map(|host| host.to_owned())
|
|
.collect();
|
|
|
|
(new_post_ratio, introduce_hosts)
|
|
}
|
|
|
|
async fn open_stream(
|
|
message_tx: tokio::sync::mpsc::UnboundedSender<Message>,
|
|
store: Store,
|
|
trend_setter_tx: crate::trend_setter::Tx,
|
|
client: &reqwest::Client,
|
|
robots_txt: &RobotsTxt,
|
|
host: String,
|
|
) -> Result<impl Future<Output = ()>, ()> {
|
|
let url = format!("https://{}/api/v1/streaming/public", host);
|
|
if ! robots_txt.allowed(&url) {
|
|
log::warn!("Streaming of {} forbidden by robots.txt", host);
|
|
return Err(());
|
|
}
|
|
|
|
let stream = Feed::stream(client, &url).await
|
|
.map_err(|e| {
|
|
log::error!("Stream error for {}: {}", host, e);
|
|
})?;
|
|
|
|
Ok(stream.for_each(move |post| {
|
|
let message_tx = message_tx.clone();
|
|
let mut store = store.clone();
|
|
let trend_setter_tx = trend_setter_tx.clone();
|
|
let host = host.clone();
|
|
async move {
|
|
log::info!("Stream {} new post: {}", host, &post.uri);
|
|
let (_, introduce_hosts) =
|
|
process_posts(&mut store, &trend_setter_tx, &host, [post].into_iter()).await;
|
|
let _ = message_tx.send(Message::IntroduceHosts {
|
|
hosts: introduce_hosts.into_iter().collect(),
|
|
});
|
|
}
|
|
}))
|
|
}
|