ref: extract serve_file_variant for static file handlers

nexus-common/src/models/user/mod.rs

@@ -14,8 +14,8 @@ pub use influencers::Influencers;
 pub use relationship::Relationship;
 pub use search::{UserSearch, USER_NAME_KEY_PARTS};
 pub use stream::{
-    UserIdStream, UserStream, UserStreamInput, UserStreamSource, USER_INFLUENCERS_KEY_PARTS,
-    USER_MOSTFOLLOWED_KEY_PARTS,
+    UserIdStream, UserStream, UserStreamInput, UserStreamSource, CACHE_USER_RECOMMENDED_KEY_PARTS,
+    USER_INFLUENCERS_KEY_PARTS, USER_MOSTFOLLOWED_KEY_PARTS,
 };
 pub use tags::ProfileTag;
 pub use tags::UserTags;

nexus-webapi/Cargo.toml

@@ -35,6 +35,7 @@ utoipa-swagger-ui = { version = "9.0.2", features = ["axum"] }
 [dev-dependencies]
 anyhow = { workspace = true }
 criterion = { version = "0.8", features = ["async_tokio"] }
+http-body-util = "0.1.3"
 httpc-test = "0.1.10"
 pubky-testnet = { workspace = true }
 tempfile = { workspace = true }
@@ -72,6 +73,11 @@ harness = false
 [[bench]]
 name = "search"
 harness = false
+
+[[bench]]
+name = "avatar"
+harness = false
+
 [[bench]]
 name = "events"
 harness = false

nexus-webapi/benches/avatar.rs

@@ -0,0 +1,190 @@
+use std::{path::PathBuf, sync::Arc, time::Duration};
+
+use axum::{
+    body::{Body, Bytes},
+    extract::State,
+    http::{Request, StatusCode},
+    response::Response,
+};
+use criterion::{criterion_group, criterion_main, Criterion};
+use http_body_util::BodyExt;
+use nexus_common::models::{
+    file::{FileDetails, FileUrls},
+    traits::Collection,
+    user::UserDetails,
+};
+use nexus_webapi::{
+    models::PubkyId,
+    routes::{r#static::user_avatar_handler, AppState, Path},
+};
+use tempfile::TempDir;
+use tokio::{fs, runtime::Runtime};
+use tower_http::services::fs::ServeFileSystemResponseBody;
+
+mod setup;
+
+use setup::run_setup;
+
+const AVATAR_BLOB_NAME: &str = "avatar.png";
+const BLOB_PATH: &str = "tests/user/blobs";
+const USER_PUBKY: &str = "4snwyct86m383rsduhw5xgcxpw7c63j3pq8x4ycqikxgik8y64ro";
+const FILE_ID: &str = "003286NSMY490";
+
+type AvatarResponse = Response<ServeFileSystemResponseBody>;
+
+struct AvatarBenchSetup {
+    _temp_dir: TempDir,
+    app_state: AppState,
+    user_id: PubkyId,
+}
+
+impl AvatarBenchSetup {
+    async fn new() -> Self {
+        let user_id = PubkyId::try_from(USER_PUBKY).unwrap();
+        let temp_dir = TempDir::new().unwrap();
+        let files_path = temp_dir.path().to_path_buf();
+        let image_dir = files_path.join(USER_PUBKY).join(FILE_ID);
+        fs::create_dir_all(&image_dir).await.unwrap();
+
+        let source_path = PathBuf::from(BLOB_PATH).join(AVATAR_BLOB_NAME);
+        let source_size = fs::copy(source_path, image_dir.join("main")).await.unwrap();
+
+        Self::seed_avatar_records(&user_id, source_size).await;
+
+        let setup = Self {
+            _temp_dir: temp_dir,
+            app_state: AppState {
+                files_path: Arc::new(files_path),
+            },
+            user_id,
+        };
+
+        let response = setup.call_avatar_handler().await.unwrap();
+        std::hint::black_box(consume_avatar_response(response).await);
+        setup
+    }
+
+    async fn seed_avatar_records(user_id: &PubkyId, source_size: u64) {
+        let avatar_uri = format!("pubky://{USER_PUBKY}/pub/pubky.app/files/{FILE_ID}");
+
+        let user = UserDetails {
+            name: "Avatar Bench User".to_string(),
+            bio: None,
+            id: user_id.clone(),
+            links: None,
+            status: None,
+            image: Some(avatar_uri.clone()),
+            indexed_at: 1_724_134_095_000,
+        };
+
+        UserDetails::put_to_index(&[USER_PUBKY], vec![Some(user)])
+            .await
+            .unwrap();
+
+        let file = FileDetails {
+            id: FILE_ID.to_string(),
+            uri: avatar_uri,
+            owner_id: USER_PUBKY.to_string(),
+            indexed_at: 1_724_134_095_000,
+            created_at: 1_784_134_095_000,
+            src: format!("pubky://{USER_PUBKY}/pub/pubky.app/blobs/{FILE_ID}"),
+            name: AVATAR_BLOB_NAME.to_string(),
+            size: source_size as i64,
+            content_type: "image/png".to_string(),
+            urls: FileUrls {
+                main: format!("{USER_PUBKY}/{FILE_ID}"),
+                feed: None,
+                small: None,
+            },
+            metadata: None,
+        };
+
+        FileDetails::put_to_index(&[&[USER_PUBKY, FILE_ID]], vec![Some(file)])
+            .await
+            .unwrap();
+    }
+
+    fn small_variant_path(&self) -> PathBuf {
+        self.app_state
+            .files_path
+            .join(USER_PUBKY)
+            .join(FILE_ID)
+            .join("small")
+    }
+
+    async fn clear_small_variant(&self) {
+        if let Err(err) = fs::remove_file(self.small_variant_path()).await {
+            if err.kind() != std::io::ErrorKind::NotFound {
+                panic!("failed to remove small avatar variant: {err}");
+            }
+        }
+    }
+
+    async fn call_avatar_handler(&self) -> nexus_webapi::Result<AvatarResponse> {
+        user_avatar_handler(
+            Path(self.user_id.clone()),
+            State(self.app_state.clone()),
+            Request::new(Body::empty()),
+        )
+        .await
+    }
+
+    async fn request_avatar(&self) {
+        let response = self.call_avatar_handler().await.unwrap();
+        std::hint::black_box(consume_avatar_response(response).await);
+    }
+}
+
+async fn consume_avatar_response(response: AvatarResponse) -> Bytes {
+    let status = response.status();
+    let bytes = response.into_body().collect().await.unwrap().to_bytes();
+
+    assert_eq!(status, StatusCode::OK);
+    assert!(
+        !bytes.is_empty(),
+        "avatar response body should not be empty"
+    );
+
+    bytes
+}
+
+fn bench_avatar_handler(c: &mut Criterion) {
+    println!("******************************************************************************");
+    println!("Benchmarking avatar handler chain without an HTTP server.");
+    println!("******************************************************************************");
+
+    run_setup();
+
+    let rt = Runtime::new().unwrap();
+    // Reuse one TempDir because PubkyServeDir stores the first files_path
+    // in a process-global OnceLock.
+    let setup = rt.block_on(AvatarBenchSetup::new());
+
+    c.bench_function("avatar_handler_warm", |b| {
+        b.to_async(&rt).iter(|| async {
+            setup.request_avatar().await;
+        });
+    });
+
+    c.bench_function("avatar_handler_cold", |b| {
+        b.to_async(&rt).iter(|| async {
+            setup.clear_small_variant().await;
+            setup.request_avatar().await;
+        });
+    });
+}
+
+fn configure_criterion() -> Criterion {
+    Criterion::default()
+        .measurement_time(Duration::new(5, 0))
+        .sample_size(100)
+        .warm_up_time(Duration::new(1, 0))
+}
+
+criterion_group! {
+    name = avatar;
+    config = configure_criterion();
+    targets = bench_avatar_handler
+}
+
+criterion_main!(avatar);

nexus-webapi/src/routes/static/avatar.rs

@@ -1,15 +1,14 @@
 use std::path::PathBuf;
 
 use super::endpoints::USER_AVATAR_ROUTE;
+use super::serve_dir::serve_file_variant;
 use crate::models::PubkyId;
-use crate::routes::r#static::PubkyServeDir;
 use crate::routes::AppState;
 use crate::routes::Path;
 use crate::{Error, Result};
 use axum::extract::{Request, State};
 use axum::response::Response;
 use nexus_common::media::FileVariant;
-use nexus_common::models::file::Blob;
 use nexus_common::models::{file::FileDetails, traits::Collection, user::UserDetails};
 use tower_http::services::fs::ServeFileSystemResponseBody;
 use tracing::{debug, error};
@@ -38,73 +37,40 @@ pub async fn user_avatar_handler(
 
     let file_path: &PathBuf = &app_state.files_path;
 
-    // 1. Get user details
     let details = match UserDetails::get_by_id(&user_id).await? {
         None => return Err(Error::user_not_found(user_id)),
         Some(d) => d,
     };
 
-    // 2. Check if user has image. If not, 404
     let Some(image_uri) = details.image else {
         return Err(Error::FileNotFound {});
     };
 
-    // 3. Parse user_id + file_id from the "pubky://owner_id/file_id" style URI
     let (owner_id, file_id) =
         FileDetails::file_key_from_uri(&image_uri).ok_or(Error::InternalServerError {
             source: format!("Invalid file URI: {image_uri}").into(),
         })?;
 
-    // 4. Look up FileDetails in Redis/Neo4j using get_by_ids
     let file_list = FileDetails::get_by_ids(&[&[&owner_id, &file_id]]).await?;
 
-    // We expect only one result in file_list, a Vec<Option<FileDetails>>
     let Some(file_details) = file_list.into_iter().flatten().next() else {
         return Err(Error::FileNotFound {});
     };
 
-    // 5. ensure small variant is created
-    let small_variant_content_type =
-        Blob::get_by_id(&file_details, &FileVariant::Small, file_path.clone())
-            .await
-            .inspect_err(|_| {
-                error!(
-            "Error while processing small variant for user: {user_id} avatar with file: {file_id}"
-        )
-            })?;
-
-    // serve the file using ServeDir
-    // Create a new request with a modified path to serve the file using ServeDir
-    // 6. Build the url using small variant
-    let file_uri_path = format!(
-        "/{}/{}/{}", // /{owner_id}/{file_id}/{variant}
-        user_id,
-        file_details.id,
-        FileVariant::Small,
-    );
-
-    // 7. Serve the file. Then remove/replace any default Cache-Control header.
-    let mut response = PubkyServeDir::try_call(
+    serve_file_variant(
         request,
-        file_uri_path,
-        small_variant_content_type,
+        &file_details,
+        &owner_id,
+        &FileVariant::Small,
         file_path.clone(),
+        false,
     )
-    .await?;
-
-    // Remove any default "cache-control" header
-    response.headers_mut().remove("cache-control");
-
-    // Insert a new Cache-Control header (e.g., 1 hour)
-    let cache_control_header = "public, max-age=3600"
-        .parse()
-        .inspect_err(|err| error!("Failed to parse Cache-Control header value: {}", err))?;
-
-    response
-        .headers_mut()
-        .insert("cache-control", cache_control_header);
-
-    Ok(response)
+    .await
+    .inspect_err(|_| {
+        error!(
+            "Error while processing small variant for user: {user_id} avatar with file: {file_id}"
+        )
+    })
 }
 
 #[derive(OpenApi)]