magnetar/src/service/emoji_cache.rs

use crate::web::ApiError;
use lru::LruCache;
use magnetar_calckey_model::emoji::{EmojiResolver, EmojiTag};
use magnetar_calckey_model::{ck, CalckeyDbError, CalckeyModel};
use std::collections::HashSet;
use std::sync::Arc;
use strum::EnumVariantNames;
use thiserror::Error;
use tokio::sync::Mutex;

#[derive(Debug, Error, EnumVariantNames)]
pub enum EmojiCacheError {
    #[error("Database error: {0}")]
    DbError(#[from] CalckeyDbError),
}

impl From<EmojiCacheError> for ApiError {
    fn from(err: EmojiCacheError) -> Self {
        let mut api_error: ApiError = match err {
            EmojiCacheError::DbError(err) => err.into(),
        };

        api_error.message = format!("Emoji cache error: {}", api_error.message);

        api_error
    }
}

#[derive(Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
struct EmojiLocator {
    name: String,
    host: Option<String>,
}

pub struct EmojiCacheService {
    cache: Mutex<LruCache<EmojiLocator, Arc<ck::emoji::Model>>>,
    db: EmojiResolver,
}

impl EmojiCacheService {
    pub(super) fn new(db: CalckeyModel) -> Self {
        const CACHE_SIZE: usize = 4096;
        Self {
            cache: Mutex::new(LruCache::new(CACHE_SIZE.try_into().unwrap())),
            db: EmojiResolver::new(db),
        }
    }

    pub async fn get(
        &self,
        name: &str,
        host: Option<&str>,
    ) -> Result<Option<Arc<ck::emoji::Model>>, EmojiCacheError> {
        let loc = EmojiLocator {
            name: name.to_string(),
            host: host.map(str::to_string),
        };
        let mut read = self.cache.lock().await;
        if let Some(emoji) = read.get(&loc) {
            return Ok(Some(emoji.clone()));
        }
        drop(read);

        let emoji = self.db.fetch_emoji(name, host).await?;

        if emoji.is_none() {
            return Ok(None);
        }

        let mut write = self.cache.lock().await;
        let emoji = Arc::new(emoji.unwrap());
        write.put(loc, emoji.clone());
        Ok(Some(emoji))
    }

    pub async fn get_many(
        &self,
        names: &[String],
        host: Option<&str>,
    ) -> Result<Vec<Arc<ck::emoji::Model>>, EmojiCacheError> {
        let locs = names
            .iter()
            .map(|n| EmojiLocator {
                name: n.clone(),
                host: host.map(str::to_string),
            })
            .collect::<HashSet<_>>();
        let mut to_resolve = Vec::new();
        let mut resolved = Vec::new();

        let mut read = self.cache.lock().await;
        for loc in locs {
            if let Some(emoji) = read.get(&loc) {
                resolved.push(emoji.clone());
            } else {
                to_resolve.push(loc.name);
            }
        }
        drop(read);

        let emoji = if to_resolve.is_empty() {
            Vec::new()
        } else {
            self.db
                .fetch_many_emojis(&to_resolve, host)
                .await?
                .into_iter()
                .map(Arc::new)
                .collect::<Vec<_>>()
        };
        resolved.extend(emoji.iter().cloned());

        let mut write = self.cache.lock().await;
        emoji.iter().for_each(|e| {
            write.put(
                EmojiLocator {
                    name: e.name.clone(),
                    host: e.host.clone(),
                },
                e.clone(),
            );
        });

        Ok(resolved)
    }

    pub async fn get_many_tagged(
        &self,
        tags: &[EmojiTag<'_>],
    ) -> Result<Vec<Arc<ck::emoji::Model>>, EmojiCacheError> {
        let locs = tags
            .iter()
            .map(|tag| EmojiLocator {
                name: tag.name.to_string(),
                host: tag.host.map(str::to_string),
            })
            .collect::<HashSet<_>>();
        let mut to_resolve = Vec::new();
        let mut resolved = Vec::new();

        let mut read = self.cache.lock().await;
        for loc in locs.iter() {
            if let Some(emoji) = read.get(loc) {
                resolved.push(emoji.clone());
            } else {
                to_resolve.push(EmojiTag {
                    name: &loc.name,
                    host: loc.host.as_deref(),
                });
            }
        }
        drop(read);

        let emoji = if to_resolve.is_empty() {
            Vec::new()
        } else {
            self.db
                .fetch_many_tagged_emojis(&to_resolve)
                .await?
                .into_iter()
                .map(Arc::new)
                .collect::<Vec<_>>()
        };
        resolved.extend(emoji.iter().cloned());

        let mut write = self.cache.lock().await;
        emoji.iter().for_each(|e| {
            write.put(
                EmojiLocator {
                    name: e.name.clone(),
                    host: e.host.clone(),
                },
                e.clone(),
            );
        });

        Ok(resolved)
    }
}