Cached emoji resolution

2023-10-27 21:55:08 +02:00 · 2023-10-27 21:55:08 +02:00 · 5572695515
parent 18d526cf8c
commit 5572695515
9 changed files with 309 additions and 15 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -263,6 +263,15 @@ dependencies = [
 "rustc-demangle",
 ]

+[[package]]
+name = "backtrace-ext"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "537beee3be4a18fb023b570f80e3ae28003db9167a751266b259926e25539d50"
+dependencies = [
+ "backtrace",
+]
+
 [[package]]
 name = "base64"
 version = "0.13.1"
@ -1333,6 +1342,12 @@ dependencies = [
 "windows-sys",
 ]

+[[package]]
+name = "is_ci"
+version = "1.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "616cde7c720bb2bb5824a224687d8f77bfd38922027f01d825cd7453be5099fb"
+
 [[package]]
 name = "itertools"
 version = "0.10.5"
@ -1342,6 +1357,15 @@ dependencies = [
 "either",
 ]

+[[package]]
+name = "itertools"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.9"
@ -1411,6 +1435,15 @@ version = "0.4.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b06a4cde4c0f271a446782e3eff8de789548ce57dbc8eca9292c27f4a42004b4"

+[[package]]
+name = "lru"
+version = "0.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1efa59af2ddfad1854ae27d75009d538d0998b4b2fd47083e743ac1a10e46c60"
+dependencies = [
+ "hashbrown 0.14.0",
+]
+
 [[package]]
 name = "magnetar"
 version = "0.2.1-alpha"
@ -1422,6 +1455,8 @@ dependencies = [
 "dotenvy",
 "headers",
 "hyper",
+ "itertools 0.11.0",
+ "lru",
 "magnetar_calckey_model",
 "magnetar_common",
 "magnetar_core",
@ -1599,8 +1634,17 @@ version = "5.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "59bb584eaeeab6bd0226ccf3509a69d7936d148cf3d036ad350abe35e8c6856e"
 dependencies = [
+ "backtrace",
+ "backtrace-ext",
+ "is-terminal",
 "miette-derive",
 "once_cell",
+ "owo-colors",
+ "supports-color",
+ "supports-hyperlinks",
+ "supports-unicode",
+ "terminal_size",
+ "textwrap",
 "thiserror",
 "unicode-width",
 ]
@ -1812,6 +1856,12 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"

+[[package]]
+name = "owo-colors"
+version = "3.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c1b04fb49957986fdce4d6ee7a65027d55d4b6d2265e5848bbb507b58ccfdb6f"
+
 [[package]]
 name = "parking_lot"
 version = "0.12.1"
@ -2702,6 +2752,12 @@ version = "1.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "62bb4feee49fdd9f707ef802e22365a35de4b7b299de4763d44bfea899442ff9"

+[[package]]
+name = "smawk"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f67ad224767faa3c7d8b6d91985b78e70a1324408abcb1cfcc2be4c06bc06043"
+
 [[package]]
 name = "socket2"
 version = "0.4.9"
@ -2743,7 +2799,7 @@ version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0c12bc9199d1db8234678b7051747c07f517cdcf019262d1847b94ec8b1aee3e"
 dependencies = [
- "itertools",
+ "itertools 0.10.5",
 "nom",
 "unicode_categories",
 ]
@ -3015,6 +3071,34 @@ version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "81cdd64d312baedb58e21336b31bc043b77e01cc99033ce76ef539f78e965ebc"

+[[package]]
+name = "supports-color"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4950e7174bffabe99455511c39707310e7e9b440364a2fcb1cc21521be57b354"
+dependencies = [
+ "is-terminal",
+ "is_ci",
+]
+
+[[package]]
+name = "supports-hyperlinks"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f84231692eb0d4d41e4cdd0cabfdd2e6cd9e255e65f80c9aa7c98dd502b4233d"
+dependencies = [
+ "is-terminal",
+]
+
+[[package]]
+name = "supports-unicode"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b6c2cb240ab5dd21ed4906895ee23fe5a48acdbd15a3ce388e7b62a9b66baf7"
+dependencies = [
+ "is-terminal",
+]
+
 [[package]]
 name = "syn"
 version = "1.0.109"
@ -3093,6 +3177,27 @@ dependencies = [
 "winapi-util",
 ]

+[[package]]
+name = "terminal_size"
+version = "0.1.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "633c1a546cee861a1a6d0dc69ebeca693bf4296661ba7852b9d21d159e0506df"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
+[[package]]
+name = "textwrap"
+version = "0.15.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b7b3e525a49ec206798b40326a44121291b530c963cfb01018f63e135bac543d"
+dependencies = [
+ "smawk",
+ "unicode-linebreak",
+ "unicode-width",
+]
+
 [[package]]
 name = "thiserror"
 version = "1.0.44"
@ -3514,6 +3619,16 @@ version = "1.0.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "301abaae475aa91687eb82514b328ab47a211a533026cb25fc3e519b86adfc3c"

+[[package]]
+name = "unicode-linebreak"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c5faade31a542b8b35855fff6e8def199853b2da8da256da52f52f1316ee3137"
+dependencies = [
+ "hashbrown 0.12.3",
+ "regex",
+]
+
 [[package]]
 name = "unicode-normalization"
 version = "0.1.22"
--- a/Cargo.toml
+++ b/Cargo.toml
@ -37,8 +37,10 @@ futures-util = "0.3"
 headers = "0.3"
 http = "0.2"
 hyper = "0.14"
+itertools = "0.11"
 js-sys = "0.3"
 log = "0.4"
+lru = "0.12"
 miette = "5.9"
 nom = "7"
 nom_locate = "4"
@ -78,6 +80,8 @@ magnetar_calckey_model = { path = "./ext_calckey_model" }
 magnetar_sdk = { path = "./magnetar_sdk" }

 cached = { workspace = true }
+lru = { workspace = true }
+
 chrono = { workspace = true }
 dotenvy = { workspace = true }

@ -93,9 +97,11 @@ tracing = { workspace = true }

 cfg-if = { workspace = true }

+itertools = { workspace = true }
+
 strum = { workspace = true, features = ["derive"] }
 thiserror = { workspace = true }
-miette = { workspace = true }
+miette = { workspace = true, features = ["fancy"] }

 percent-encoding = { workspace = true }

--- a/ext_calckey_model/src/lib.rs
+++ b/ext_calckey_model/src/lib.rs
@ -125,6 +125,22 @@ impl CalckeyModel {
            .await?)
    }

+    pub async fn fetch_emoji(
+        &self,
+        shortcode: &str,
+        host: Option<&str>,
+    ) -> Result<Option<emoji::Model>, CalckeyDbError> {
+        let host_filter = if let Some(host) = host {
+            emoji::Column::Host.eq(host)
+        } else {
+            emoji::Column::Host.is_null()
+        };
+
+        let name_filter = emoji::Column::Name.eq(shortcode);
+        let filter = host_filter.and(name_filter);
+
+        Ok(emoji::Entity::find().filter(filter).one(&self.0).await?)
+    }
    pub async fn fetch_many_emojis(
        &self,
        shortcodes: &[String],
--- a/src/model/mod.rs
+++ b/src/model/mod.rs
@ -1,13 +1,27 @@
+use crate::service::MagnetarService;
 use magnetar_calckey_model::ck;
 use std::sync::Arc;

 pub mod data;
 pub mod processing;

+#[derive(Clone, Debug)]
+pub struct ProcessingLimits {
+    max_emojis: usize,
+}
+
+impl Default for ProcessingLimits {
+    fn default() -> Self {
+        ProcessingLimits { max_emojis: 500 }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct PackingContext {
    instance_meta: Arc<ck::meta::Model>,
    self_user: Option<Arc<ck::user::Model>>,
+    service: Arc<MagnetarService>,
+    limits: ProcessingLimits,
 }

 pub trait PackType<I>: 'static {
--- a/src/model/processing/emoji.rs
+++ b/src/model/processing/emoji.rs
@ -1,17 +1,13 @@
 use crate::model::processing::PackResult;
 use crate::model::{PackType, PackingContext};
-use magnetar_calckey_model::{ck, CalckeyModel};
+use magnetar_calckey_model::ck;
 use magnetar_sdk::types::emoji::{EmojiBase, PackEmojiBase};
 use magnetar_sdk::types::Id;
 use magnetar_sdk::{Packed, Required};

-pub struct EmojiModel(CalckeyModel);
+pub struct EmojiModel;

 impl EmojiModel {
-    pub fn new(model: CalckeyModel) -> Self {
-        EmojiModel(model)
-    }
-
    pub fn pack_existing(&self, ctx: &PackingContext, emoji: &ck::emoji::Model) -> PackEmojiBase {
        PackEmojiBase::pack_from((
            Required(Id::from(&emoji.id)),
@ -25,7 +21,7 @@ impl EmojiModel {
        shortcodes: &[String],
        host: Option<&str>,
    ) -> PackResult<Vec<PackEmojiBase>> {
-        let emojis = self.0.fetch_many_emojis(shortcodes, host).await?;
+        let emojis = ctx.service.emoji_cache.get_many(shortcodes, host).await?;
        let packed_emojis = emojis.iter().map(|e| self.pack_existing(ctx, &e)).collect();

        Ok(packed_emojis)
--- a/src/model/processing/mod.rs
+++ b/src/model/processing/mod.rs
@ -1,3 +1,4 @@
+use crate::service::emoji_cache::EmojiCacheError;
 use magnetar_calckey_model::sea_orm::DbErr;
 use magnetar_calckey_model::CalckeyDbError;
 use magnetar_sdk::mmm::Token;
@ -12,6 +13,8 @@ pub enum PackError {
    DbError(#[from] DbErr),
    #[error("Calckey database wrapper error: {0}")]
    CalckeyDbError(#[from] CalckeyDbError),
+    #[error("Emoji cache error: {0}")]
+    EmojiCacheError(#[from] EmojiCacheError),
 }

 pub type PackResult<T> = Result<T, PackError>;
--- a/src/model/processing/user.rs
+++ b/src/model/processing/user.rs
@ -2,14 +2,15 @@ use crate::model::data::user::UserBaseSource;
 use crate::model::processing::emoji::EmojiModel;
 use crate::model::processing::{get_mm_token_emoji, PackResult};
 use crate::model::{PackType, PackingContext};
+use itertools::Itertools;
+use magnetar_calckey_model::ck;
 use magnetar_calckey_model::sea_orm::EntityTrait;
-use magnetar_calckey_model::{ck, CalckeyModel};
 use magnetar_sdk::types::emoji::EmojiContext;
 use magnetar_sdk::types::user::{PackUserBase, UserBase};
 use magnetar_sdk::types::{Id, MmXml};
 use magnetar_sdk::{mmm, Packed, Required};

-pub struct UserModel(CalckeyModel);
+pub struct UserModel;

 impl UserModel {
    pub async fn base_from_existing(
@ -19,7 +20,7 @@ impl UserModel {
    ) -> PackResult<PackUserBase> {
        let avatar = if let Some(avatar_id) = user.avatar_id.as_ref() {
            ck::drive_file::Entity::find_by_id(avatar_id)
-                .one(self.0.inner())
+                .one(ctx.service.db.inner())
                .await?
        } else {
            None
@ -27,9 +28,14 @@ impl UserModel {

        let username_mm =
            mmm::Context::default().parse_ui(user.name.as_deref().unwrap_or(&user.username));
-        let emoji_model = EmojiModel::new(self.0.clone());
-        let emojis = emoji_model
-            .fetch_many_emojis(ctx, &get_mm_token_emoji(&username_mm), user.host.as_deref())
+        let shortcodes = get_mm_token_emoji(&username_mm)
+            .into_iter()
+            .sorted()
+            .dedup()
+            .take(ctx.limits.max_emojis)
+            .collect::<Vec<_>>();
+        let emojis = EmojiModel
+            .fetch_many_emojis(ctx, &shortcodes, user.host.as_deref())
            .await?;
        let emoji_context = EmojiContext(emojis);

--- a/src/service/emoji_cache.rs
+++ b/src/service/emoji_cache.rs
@ -0,0 +1,123 @@
+use crate::web::ApiError;
+use lru::LruCache;
+use magnetar_calckey_model::{ck, CalckeyDbError, CalckeyModel};
+use std::collections::HashSet;
+use std::sync::Arc;
+use strum::EnumVariantNames;
+use thiserror::Error;
+use tokio::io::AsyncWriteExt;
+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: CalckeyModel,
+}
+
+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,
+        }
+    }
+
+    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
+            .into_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 = 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)
+    }
+}
--- a/src/service/mod.rs
+++ b/src/service/mod.rs
@ -1,7 +1,9 @@
 use magnetar_calckey_model::{CalckeyCache, CalckeyModel};
 use magnetar_common::config::MagnetarConfig;
+use std::fmt::{Debug, Formatter};
 use thiserror::Error;

+pub mod emoji_cache;
 pub mod instance_meta_cache;
 pub mod user_cache;

@ -11,6 +13,17 @@ pub struct MagnetarService {
    pub config: &'static MagnetarConfig,
    pub auth_cache: user_cache::UserCacheService,
    pub instance_meta_cache: instance_meta_cache::InstanceMetaCacheService,
+    pub emoji_cache: emoji_cache::EmojiCacheService,
+}
+
+impl Debug for MagnetarService {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MagnetarService")
+            .field("db", &self.db)
+            .field("cache", &self.cache)
+            .field("config", &self.config)
+            .finish_non_exhaustive()
+    }
 }

 #[derive(Debug, Error)]
@ -28,6 +41,7 @@ impl MagnetarService {
        let auth_cache =
            user_cache::UserCacheService::new(config, db.clone(), cache.clone()).await?;
        let instance_meta_cache = instance_meta_cache::InstanceMetaCacheService::new(db.clone());
+        let emoji_cache = emoji_cache::EmojiCacheService::new(db.clone());

        Ok(Self {
            db,
@ -35,6 +49,7 @@ impl MagnetarService {
            config,
            auth_cache,
            instance_meta_cache,
+            emoji_cache,
        })
    }
 }