caveman/hunter/src/worker.rs

use std::collections::HashSet;
use std::future::Future;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use cave::{
    feed::{Feed, Post},
    store::Store,
};
use futures::{StreamExt, future};
use crate::trend_setter::UpdateSet;

#[derive(Clone)]
pub struct RobotsTxt {
    robot: Arc<RwLock<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: Arc::new(RwLock::new(robot)),
        }
    }

    pub fn allowed(&self, url: &str) -> bool {
        if let Some(robot) = self.robot.read().unwrap().as_ref() {
            robot.allowed(url)
        } else {
            true
        }
    }

    pub fn delay(&self) -> Option<Duration> {
        self.robot.read().unwrap().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;
    let robots_delay = robots_txt.delay();
    // 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.clone(), &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(robots_delay) = robots_delay {
            *mean_interval = (*mean_interval).max(robots_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);
    }
    // free as early as possible
    drop(robots_txt);

    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(());
    }
    // free as early as possible
    drop(robots_txt);

    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(),
            });
        }
    }))
}