Generic caching and basic user fetching in the backend

Cached emoji resolution
2023-10-28 00:33:09 +02:00 · 2023-10-27 21:55:08 +02:00
16 changed files with 561 additions and 40 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/magnetar_mmm_parser/src/lib.rs
+++ b/magnetar_mmm_parser/src/lib.rs
@ -233,16 +233,14 @@ impl Token {
        }
    }

-    pub fn walk_map_collect<T>(&self, func: impl Fn(&Token) -> Option<T>, out: &mut Vec<T>) {
+    pub fn walk_map_collect<T>(&self, func: &impl Fn(&Token) -> Option<T>, out: &mut Vec<T>) {
        if let Some(v) = func(self) {
            out.push(v)
        }

        match self {
            Token::Sequence(items) => {
-                items
-                    .iter()
-                    .for_each(|tok| tok.walk_map_collect(&func, out));
+                items.iter().for_each(|tok| tok.walk_map_collect(func, out));
            }
            Token::Quote(inner)
            | Token::Small(inner)
--- a/magnetar_sdk/src/types/user.rs
+++ b/magnetar_sdk/src/types/user.rs
@ -171,6 +171,12 @@ pack!(
        & Option<UserAuthOverviewExt> as auth
 );

+impl From<PackUserBase> for PackUserMaybeAll {
+    fn from(value: PackUserBase) -> Self {
+        Self::pack_from((value.id, value.user, None, None, None, None, None))
+    }
+}
+
 pack!(
    PackUserSelfMaybeAll,
    Required<Id> as id
@ -180,3 +186,9 @@ pack!(
        & Option<UserDetailExt> as detail
        & Option<UserSecretsExt> as secrets
 );
+
+impl From<PackUserBase> for PackUserSelfMaybeAll {
+    fn from(value: PackUserBase) -> Self {
+        Self::pack_from((value.id, value.user, None, None, None, None))
+    }
+}
--- a/src/api_v1/user.rs
+++ b/src/api_v1/user.rs
@ -1,8 +1,11 @@
+use crate::model::processing::user::UserModel;
+use crate::model::PackingContext;
 use crate::service::MagnetarService;
 use crate::web::auth::{AuthenticatedUser, MaybeUser};
-use crate::web::ApiError;
+use crate::web::{ApiError, ObjectNotFound};
 use axum::extract::{Path, Query, State};
 use axum::Json;
+use magnetar_calckey_model::ck;
 use magnetar_sdk::endpoints::user::{GetUserById, GetUserSelf, UserByIdReq, UserSelfReq};
 use magnetar_sdk::endpoints::{Req, Res};
 use std::sync::Arc;
@ -14,14 +17,18 @@ pub async fn handle_user_info_self(
        profile: _,
        secrets: _,
    }): Query<Req<GetUserSelf>>,
-    State(_service): State<Arc<MagnetarService>>,
-    AuthenticatedUser(_user): AuthenticatedUser,
+    State(service): State<Arc<MagnetarService>>,
+    AuthenticatedUser(user): AuthenticatedUser,
 ) -> Result<Json<Res<GetUserSelf>>, ApiError> {
-    todo!()
+    // TODO: Extended properties!
+
+    let ctx = PackingContext::new(service, Some(user.clone())).await?;
+    let user = UserModel.base_from_existing(&ctx, user.as_ref()).await?;
+    Ok(Json(user.into()))
 }

 pub async fn handle_user_info(
-    Path(_id): Path<String>,
+    Path(id): Path<String>,
    Query(UserByIdReq {
        detail: _,
        pins: _,
@ -29,8 +36,18 @@ pub async fn handle_user_info(
        relation: _,
        auth: _,
    }): Query<Req<GetUserById>>,
-    State(_service): State<Arc<MagnetarService>>,
-    MaybeUser(_user): MaybeUser,
+    State(service): State<Arc<MagnetarService>>,
+    MaybeUser(self_user): MaybeUser,
 ) -> Result<Json<Res<GetUserById>>, ApiError> {
-    todo!()
+    // TODO: Extended properties!
+
+    let ctx = PackingContext::new(service.clone(), self_user).await?;
+    let user_model = service
+        .db
+        .get_user_by_id(&id)
+        .await?
+        .ok_or_else(|| ObjectNotFound(id))?;
+
+    let user = UserModel.base_from_existing(&ctx, &user_model).await?;
+    Ok(Json(user.into()))
 }
--- a/src/model/data/user.rs
+++ b/src/model/data/user.rs
@ -19,7 +19,7 @@ impl PackType<&[PackEmojiBase]> for EmojiContext {
 pub struct UserBaseSource<'a> {
    pub user: &'a ck::user::Model,
    pub username_mm: Option<&'a MmXml>,
-    pub avatar: &'a Option<ck::drive_file::Model>,
+    pub avatar: Option<&'a ck::drive_file::Model>,
    pub emoji_context: &'a EmojiContext,
 }

@ -49,8 +49,8 @@ impl PackType<UserBaseSource<'_>> for UserBase {
                SpeechTransform::None
            },
            created_at: user.created_at.into(),
-            avatar_url: avatar.as_ref().map(|v| v.url.clone()),
-            avatar_blurhash: avatar.as_ref().and_then(|v| v.blurhash.clone()),
+            avatar_url: avatar.map(|v| v.url.clone()),
+            avatar_blurhash: avatar.and_then(|v| v.blurhash.clone()),
            avatar_color: None,
            avatar_decoration: if user.is_cat {
                AvatarDecoration::CatEars
--- a/src/model/mod.rs
+++ b/src/model/mod.rs
@ -1,13 +1,28 @@
+use crate::model::processing::PackResult;
+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 {
@ -15,6 +30,18 @@ pub trait PackType<I>: 'static {
 }

 impl PackingContext {
+    pub async fn new(
+        service: Arc<MagnetarService>,
+        self_user: Option<Arc<ck::user::Model>>,
+    ) -> PackResult<PackingContext> {
+        Ok(Self {
+            instance_meta: service.instance_meta_cache.get().await?,
+            self_user,
+            service,
+            limits: Default::default(),
+        })
+    }
+
    fn self_user(&self) -> Option<&ck::user::Model> {
        self.self_user.as_deref()
    }
--- a/src/model/processing/drive.rs
+++ b/src/model/processing/drive.rs
@ -0,0 +1,33 @@
+use crate::model::processing::PackResult;
+use crate::model::{PackType, PackingContext};
+use magnetar_calckey_model::ck;
+use magnetar_sdk::types::drive::{DriveFileBase, PackDriveFileBase};
+use magnetar_sdk::types::Id;
+use magnetar_sdk::{Packed, Required};
+
+pub struct DriveModel;
+
+impl DriveModel {
+    pub fn pack_existing(
+        &self,
+        ctx: &PackingContext,
+        file: &ck::drive_file::Model,
+    ) -> PackDriveFileBase {
+        PackDriveFileBase::pack_from((
+            Required(Id::from(&file.id)),
+            Required(DriveFileBase::extract(ctx, &file)),
+        ))
+    }
+
+    pub async fn get_cached_base(
+        &self,
+        ctx: &PackingContext,
+        id: &str,
+    ) -> PackResult<Option<PackDriveFileBase>> {
+        let Some(file) = ctx.service.drive_file_cache.get(id).await? else {
+            return Ok(None);
+        };
+
+        Ok(Some(self.pack_existing(ctx, file.as_ref())))
+    }
+}
--- a/src/model/processing/emoji.rs
+++ b/src/model/processing/emoji.rs
@ -1,17 +1,14 @@
 use crate::model::processing::PackResult;
 use crate::model::{PackType, PackingContext};
-use magnetar_calckey_model::{ck, CalckeyModel};
+use itertools::Itertools;
+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,9 +22,18 @@ 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)
    }
+
+    pub fn deduplicate_emoji(&self, ctx: &PackingContext, emoji_list: Vec<String>) -> Vec<String> {
+        emoji_list
+            .into_iter()
+            .sorted()
+            .dedup()
+            .take(ctx.limits.max_emojis)
+            .collect::<Vec<_>>()
+    }
 }
--- a/src/model/processing/mod.rs
+++ b/src/model/processing/mod.rs
@ -1,17 +1,27 @@
+use crate::service::emoji_cache::EmojiCacheError;
+use crate::service::generic_id_cache::GenericIdCacheError;
+use crate::service::instance_meta_cache::InstanceMetaCacheError;
 use magnetar_calckey_model::sea_orm::DbErr;
 use magnetar_calckey_model::CalckeyDbError;
 use magnetar_sdk::mmm::Token;
 use thiserror::Error;

+pub mod drive;
 pub mod emoji;
 pub mod user;

-#[derive(Debug, Error)]
+#[derive(Debug, Error, strum::IntoStaticStr)]
 pub enum PackError {
    #[error("Database error: {0}")]
    DbError(#[from] DbErr),
    #[error("Calckey database wrapper error: {0}")]
    CalckeyDbError(#[from] CalckeyDbError),
+    #[error("Emoji cache error: {0}")]
+    EmojiCacheError(#[from] EmojiCacheError),
+    #[error("Instance cache error: {0}")]
+    InstanceMetaCacheError(#[from] InstanceMetaCacheError),
+    #[error("Generic cache error: {0}")]
+    GenericCacheError(#[from] GenericIdCacheError),
 }

 pub type PackResult<T> = Result<T, PackError>;
@ -19,7 +29,7 @@ pub type PackResult<T> = Result<T, PackError>;
 fn get_mm_token_emoji(token: &Token) -> Vec<String> {
    let mut v = Vec::new();
    token.walk_map_collect(
-        |t| {
+        &|t| {
            if let Token::ShortcodeEmoji(e) = t {
                Some(e.to_owned())
            } else {
--- a/src/model/processing/user.rs
+++ b/src/model/processing/user.rs
@ -2,34 +2,38 @@ 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 magnetar_calckey_model::ck;
 use magnetar_calckey_model::sea_orm::EntityTrait;
-use magnetar_calckey_model::{ck, CalckeyModel};
+use magnetar_sdk::mmm::Token;
 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};
+use std::sync::Arc;

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

 impl UserModel {
+    pub fn tokenize_username(&self, user: &ck::user::Model) -> Token {
+        mmm::Context::default().parse_ui(user.name.as_deref().unwrap_or(&user.username))
+    }
+
    pub async fn base_from_existing(
        &self,
        ctx: &PackingContext,
        user: &ck::user::Model,
    ) -> 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())
-                .await?
-        } else {
-            None
+        let avatar = match &user.avatar_id {
+            Some(av_id) => ctx.service.drive_file_cache.get(av_id).await?,
+            None => None,
        };

-        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 username_mm = self.tokenize_username(&user);
+
+        let emoji_model = EmojiModel;
+        let shortcodes = emoji_model.deduplicate_emoji(ctx, get_mm_token_emoji(&username_mm));
        let emojis = emoji_model
-            .fetch_many_emojis(ctx, &get_mm_token_emoji(&username_mm), user.host.as_deref())
+            .fetch_many_emojis(ctx, &shortcodes, user.host.as_deref())
            .await?;
        let emoji_context = EmojiContext(emojis);

@ -38,7 +42,7 @@ impl UserModel {
            UserBaseSource {
                user,
                username_mm: mmm::to_xml_string(&username_mm).map(MmXml).as_ref().ok(),
-                avatar: &avatar,
+                avatar: avatar.as_deref(),
                emoji_context: &emoji_context,
            },
        );
--- a/src/service/emoji_cache.rs
+++ b/src/service/emoji_cache.rs
@ -0,0 +1,122 @@
+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::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/generic_id_cache.rs
+++ b/src/service/generic_id_cache.rs
@ -0,0 +1,96 @@
+use crate::web::ApiError;
+use lru::LruCache;
+use magnetar_calckey_model::sea_orm::{EntityTrait, PrimaryKeyTrait};
+use magnetar_calckey_model::{CalckeyDbError, CalckeyModel};
+use std::marker::PhantomData;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use strum::EnumVariantNames;
+use thiserror::Error;
+use tokio::sync::Mutex;
+
+#[derive(Debug, Error, EnumVariantNames)]
+pub enum GenericIdCacheError {
+    #[error("Database error: {0}")]
+    DbError(#[from] CalckeyDbError),
+}
+
+impl From<GenericIdCacheError> for ApiError {
+    fn from(err: GenericIdCacheError) -> Self {
+        let mut api_error: ApiError = match err {
+            GenericIdCacheError::DbError(err) => err.into(),
+        };
+
+        api_error.message = format!("Generic ID cache error: {}", api_error.message);
+
+        api_error
+    }
+}
+
+#[derive(Debug)]
+struct CacheEntry<E: EntityTrait> {
+    created: Instant,
+    data: Arc<E::Model>,
+}
+
+impl<E: EntityTrait> CacheEntry<E> {
+    fn new(data: Arc<E::Model>) -> Self {
+        Self {
+            created: Instant::now(),
+            data,
+        }
+    }
+}
+
+pub struct GenericIdCacheService<E: EntityTrait> {
+    cache: Mutex<LruCache<String, CacheEntry<E>>>,
+    lifetime_max: Duration,
+    db: CalckeyModel,
+    _entity_type: PhantomData<E>,
+}
+
+impl<E: EntityTrait> GenericIdCacheService<E> {
+    pub(super) fn new(db: CalckeyModel, cache_size: usize, entry_lifetime: Duration) -> Self {
+        const CACHE_SIZE: usize = 4096;
+
+        Self {
+            cache: Mutex::new(LruCache::new(
+                cache_size
+                    .try_into()
+                    .unwrap_or(CACHE_SIZE.try_into().unwrap()),
+            )),
+            lifetime_max: entry_lifetime,
+            db,
+            _entity_type: PhantomData,
+        }
+    }
+
+    pub async fn get<'a>(&self, id: &'a str) -> Result<Option<Arc<E::Model>>, GenericIdCacheError>
+    where
+        <<E as EntityTrait>::PrimaryKey as PrimaryKeyTrait>::ValueType: From<&'a str>,
+    {
+        let mut read = self.cache.lock().await;
+        if let Some(item) = read.peek(id) {
+            if item.created + self.lifetime_max >= Instant::now() {
+                let data = item.data.clone();
+                read.promote(id);
+                return Ok(Some(data));
+            }
+        }
+        drop(read);
+
+        let val = E::find_by_id(id)
+            .one(self.db.inner())
+            .await
+            .map_err(CalckeyDbError::from)?;
+
+        if val.is_none() {
+            return Ok(None);
+        }
+
+        let mut write = self.cache.lock().await;
+        let data = Arc::new(val.unwrap());
+        write.put(id.to_string(), CacheEntry::new(data.clone()));
+        Ok(Some(data))
+    }
+}
--- a/src/service/mod.rs
+++ b/src/service/mod.rs
@ -1,7 +1,11 @@
-use magnetar_calckey_model::{CalckeyCache, CalckeyModel};
+use magnetar_calckey_model::{ck, CalckeyCache, CalckeyModel};
 use magnetar_common::config::MagnetarConfig;
+use std::fmt::{Debug, Formatter};
+use std::time::Duration;
 use thiserror::Error;

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

@ -11,6 +15,18 @@ 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,
+    pub drive_file_cache: generic_id_cache::GenericIdCacheService<ck::drive_file::Entity>,
+}
+
+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 +44,9 @@ 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());
+        let drive_file_cache =
+            generic_id_cache::GenericIdCacheService::new(db.clone(), 128, Duration::from_secs(10));

        Ok(Self {
            db,
@ -35,6 +54,8 @@ impl MagnetarService {
            config,
            auth_cache,
            instance_meta_cache,
+            emoji_cache,
+            drive_file_cache,
        })
    }
 }
--- a/src/web/mod.rs
+++ b/src/web/mod.rs
@ -1,3 +1,4 @@
+use crate::model::processing::PackError;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
 use axum::Json;
@ -81,3 +82,40 @@ impl From<CalckeyCacheError> for ApiError {
        }
    }
 }
+
+impl From<PackError> for ApiError {
+    fn from(err: PackError) -> Self {
+        Self {
+            status: StatusCode::INTERNAL_SERVER_ERROR,
+            code: err.error_code(),
+            message: if cfg!(debug_assertions) {
+                format!("Data transformation error: {}", err)
+            } else {
+                "Data transformation error".to_string()
+            },
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct ObjectNotFound(pub String);
+
+impl From<&ObjectNotFound> for &str {
+    fn from(_: &ObjectNotFound) -> Self {
+        "ObjectNotFound"
+    }
+}
+
+impl From<ObjectNotFound> for ApiError {
+    fn from(err: ObjectNotFound) -> Self {
+        Self {
+            status: StatusCode::NOT_FOUND,
+            code: err.error_code(),
+            message: if cfg!(debug_assertions) {
+                format!("Object not found: {}", err.0)
+            } else {
+                "Object not found".to_string()
+            },
+        }
+    }
+}
Author	SHA1	Message	Date
Natty	f0e56deca9	Generic caching and basic user fetching in the backend	2023-10-28 00:33:09 +02:00
Natty	5572695515	Cached emoji resolution	2023-10-27 21:55:08 +02:00