caveman/cave/src/store.rs

use redis::{Value, RedisError};
use crate::{feed::Post, trend_tag::TrendTag};

const POST_EXPIRE: usize = 86400;
const TAG_EXPIRE: i64 = 30 * 24;

pub const TREND_POOL_SIZE: usize = 30;

pub type Error = RedisError;

#[derive(Clone)]
pub struct Store {
    man: redis::aio::ConnectionManager,
}

impl Store {
    pub async fn new(redis_url: &str) -> Self {
        crate::systemd::status("Starting redis client");
        let client = redis::Client::open(redis_url)
            .expect("redis::Client");
        let man = redis::aio::ConnectionManager::new(client).await
            .expect("redis::aio::ConnectionManager");
        Self { man }
    }

    pub async fn save_post(&mut self, post: Post) -> Result<bool, RedisError> {
        let post_key = format!("p:{}", post.uri);
        let check = redis::pipe()
            .getset(&post_key, "1")
            .expire(post_key, POST_EXPIRE)
            .ignore()
            .query_async::<_, Value>(&mut self.man)
            .await?;
        if check != Value::Bulk(vec![Value::Nil]) {
            // post is not new
            return Ok(false);
        }
        log::info!("New post ({}{} tags): {}",
                   if post.account.bot { "bot, " } else { "" },
                   post.tags.len(), post.uri);

        self.save_post_tags(post).await;

        // post was new
        Ok(true)
    }

    async fn save_post_tags(&mut self, post: Post) {
        if post.account.bot || post.tags.is_empty() {
            // irrelevant
            return;
        }

        let host = match post.uri_host() {
            Some(host) => host,
            None => {
                log::warn!("no uri_host");
                return;
            },
        };

        let timestamp = match post.timestamp() {
            Some(timestamp) => timestamp,
            None => {
                log::warn!("no timestamp");
                return;
            }
        };
        let hour = timestamp.naive_utc().timestamp() / 3600;
        let until = chrono::offset::Utc::now().naive_utc().timestamp() / 3600;
        if hour > until {
            log::warn!("future post from {}", timestamp);
            return;
        }
        let from = until - TAG_EXPIRE;
        if hour < from {
            log::warn!("ancient post from {}", timestamp);
            return;
        }

        // clip "en-us" to "en"
        let language = post.lang();

        let mut cmd = redis::pipe();
        if let Some(language) = &language {
            cmd.hincr("r", language, 1)
                .ignore();
        }
        let mut store_tags = |spellings, tag_key| {
            // by hour
            cmd.hincr(
                &tag_key,
                format!("t:{}", hour),
                1
            ).ignore();
            // by spelling
            for spelling in spellings {
                cmd.hincr(
                    &tag_key,
                    format!("s:{}", spelling),
                    1
                ).ignore();
            }
            // by instance
            cmd.hincr(
                tag_key,
                format!("h:{}", host),
                1
            ).ignore();
        };
        for (name, spellings) in post.tags_set() {
            log::debug!("tag {}", name);
            // global
            store_tags(spellings.clone(), format!("g:{}", name));
            // by language
            if let Some(language) = &language {
                store_tags(spellings, format!("l:{}:{}", language, name));
            }
        }

        match cmd.query_async(&mut self.man).await {
            Ok(()) => {}
            Err(e) => {
                log::error!("redis error: {:?}", e);
            }
        }
    }

    pub async fn save_host(&mut self, host: &str) -> Result<(), RedisError> {
        redis::Cmd::set(format!("h:{}", host), "1")
            .query_async::<_, ()>(&mut self.man)
            .await
    }

    pub async fn get_hosts(&mut self) -> Result<Vec<String>, RedisError> {
        let mut results = vec![];
        self.scan("h:*", |key| results.push(key[2..].to_string())).await?;
        Ok(results)
    }

    pub async fn get_languages(&mut self) -> Result<Vec<(String, usize)>, RedisError> {
        let values: Vec<String> = redis::Cmd::hgetall("r")
            .query_async(&mut self.man)
            .await?;
        let mut results = Vec::with_capacity(values.len() / 2);
        let mut key = None;
        for value in values {
            if let Some(key) = key.take() {
                if let Ok(count) = str::parse(&value) {
                    results.push((key, count));
                }
            } else {
                key = Some(value);
            }
        }
        Ok(results)
    }

    async fn scan<F: FnMut(String)>(
        &mut self,
        pattern: &str,
        mut f: F,
    ) -> Result<(), RedisError> {
        let mut cmd = redis::cmd("SCAN");
        cmd.cursor_arg(0)
            .arg("MATCH").arg(pattern)
            .arg("COUNT").arg(100);
        let mut iter = cmd.iter_async::<String>(&mut self.man)
            .await?;
        while let Some(key) = iter.next_item().await {
            f(key);
        }
        Ok(())
    }

    pub async fn get_trend_tags(
        &mut self,
        language: &Option<String>,
        names_in: impl Iterator<Item = String>,
    ) -> Result<Vec<TrendTag>, RedisError> {
        let mut cmd = redis::pipe();
        let prefix = match language {
            Some(language) => format!("l:{}:", language),
            None => "g:".to_string(),
        };
        let names: Vec<String> = names_in.map(|name| {
            cmd.hgetall(format!("{}{}", prefix, name));
            name
        }).collect();
        let hashes = cmd.query_async::<_, Vec<Vec<String>>>(&mut self.man).await?;
        let results = names.into_iter()
            .zip(hashes)
            .map(|(name, hash_values)| TrendTag::from_hash(name, hash_values))
            .collect();
        Ok(results)
    }

    pub async fn get_trend_pools(
        &mut self,
        language: &Option<String>,
        periods: &[u64],
    ) -> Result<Vec<(u64, Vec<String>)>, RedisError> {
        let mut cmd = redis::pipe();
        for period in periods {
            cmd.smembers(pool_key(language, *period));
        }
        let sets: Vec<Vec<String>> = cmd.query_async(&mut self.man)
            .await?;
        let results = periods.iter().cloned()
            .zip(sets.into_iter())
            .collect();
        Ok(results)
    }

    pub async fn update_trend_pools(
        &mut self,
        remove: impl Iterator<Item = (&Option<String>, u64, Vec<&str>)>,
        add: impl Iterator<Item = (&Option<String>, u64, Vec<&str>)>,
    ) -> Result<(), RedisError> {
        let mut cmd = redis::pipe();
        for (language, period, tags) in remove {
            if ! tags.is_empty() {
                let pool_key = pool_key(language, period);
                cmd.srem(&pool_key, tags)
                    .ignore()
                    .expire(pool_key, period as usize * 3600)
                    .ignore();
            }
        }
        for (language, period, tags) in add {
            if ! tags.is_empty() {
                let pool_key = pool_key(language, period);
                cmd.sadd(&pool_key, tags)
                    .ignore()
                    .expire(pool_key, period as usize * 3600)
                    .ignore();
            }
        }
        cmd.query_async(&mut self.man)
            .await?;
        Ok(())
    }
}

fn pool_key(language: &Option<String>, period: u64) -> String {
    match language {
        Some(language) =>
            format!("q:{}:{}", period, language),
        None =>
            format!("q:{}", period),
    }
}