Generic caching and basic user fetching in the backend

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)

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))
+    }
+}

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

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

src/model/mod.rs

@@ -1,3 +1,4 @@
+use crate::model::processing::PackResult;
 use crate::service::MagnetarService;
 use magnetar_calckey_model::ck;
 use std::sync::Arc;
@@ -29,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()
     }

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())))
+    }
+}

src/model/processing/emoji.rs

@@ -1,5 +1,6 @@
 use crate::model::processing::PackResult;
 use crate::model::{PackType, PackingContext};
+use itertools::Itertools;
 use magnetar_calckey_model::ck;
 use magnetar_sdk::types::emoji::{EmojiBase, PackEmojiBase};
 use magnetar_sdk::types::Id;
@@ -26,4 +27,13 @@ impl EmojiModel {
 
         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<_>>()
+    }
 }

src/model/processing/mod.rs

@@ -1,13 +1,16 @@
 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),
@@ -15,6 +18,10 @@ pub enum PackError {
     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>;
@@ -22,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 {

src/model/processing/user.rs

@@ -2,39 +2,37 @@ 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_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;
 
 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(ctx.service.db.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 shortcodes = get_mm_token_emoji(&username_mm)
-            .into_iter()
-            .sorted()
-            .dedup()
-            .take(ctx.limits.max_emojis)
-            .collect::<Vec<_>>();
-        let emojis = EmojiModel
+        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, &shortcodes, user.host.as_deref())
             .await?;
         let emoji_context = EmojiContext(emojis);
@@ -44,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,
             },
         );

src/service/emoji_cache.rs

@@ -5,7 +5,6 @@ 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)]

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))
+    }
+}

src/service/mod.rs

@@ -1,9 +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;
 
@@ -14,6 +16,7 @@ pub struct MagnetarService {
     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 {
@@ -42,6 +45,8 @@ impl MagnetarService {
             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,
@@ -50,6 +55,7 @@ impl MagnetarService {
             auth_cache,
             instance_meta_cache,
             emoji_cache,
+            drive_file_cache,
         })
     }
 }

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