Serialize creation of media/thumbnails

This way, no two tasks will put effort into creating the same exact
media at the same time.

Cargo.lock

@@ -840,6 +840,7 @@ dependencies = [
  "base64 0.22.1",
  "bytes",
  "clap",
+ "either",
  "futures-util",
  "hmac",
  "html-escape",

Cargo.toml

@@ -5,7 +5,6 @@ explicit_outlives_requirements = "warn"
 macro_use_extern_crate = "warn"
 missing_abi = "warn"
 noop_method_call = "warn"
-pointer_structural_match = "warn"
 single_use_lifetimes = "warn"
 unreachable_pub = "warn"
 unsafe_op_in_unsafe_fn = "warn"
@@ -17,7 +16,7 @@ unused_qualifications = "warn"
 
 [workspace.lints.clippy]
 # Groups. Keep alphabetically sorted
-pedantic = "warn"
+pedantic = { level = "warn", priority = -1 }
 
 # Lints. Keep alphabetically sorted
 as_conversions = "warn"
@@ -80,7 +79,7 @@ version = "0.1.0"
 edition = "2021"
 
 # See also `rust-toolchain.toml`
-rust-version = "1.78.0"
+rust-version = "1.81.0"
 
 [lints]
 workspace = true
@@ -95,6 +94,7 @@ axum-server = { version = "0.6.0", features = ["tls-rustls"] }
 base64 = "0.22.1"
 bytes = "1.6.0"
 clap = { version = "4.5.4", default-features = false, features = ["std", "derive", "help", "usage", "error-context", "string", "wrap_help"] }
+either = "1.12"
 futures-util = { version = "0.3.30", default-features = false }
 hmac = "0.12.1"
 html-escape = "0.2.13"

flake.nix

@@ -11,7 +11,7 @@
 
     rust-manifest = {
       # Keep version in sync with rust-toolchain.toml
-      url = "https://static.rust-lang.org/dist/channel-rust-1.78.0.toml";
+      url = "https://static.rust-lang.org/dist/channel-rust-1.81.0.toml";
       flake = false;
     };
   };

rust-toolchain.toml

@@ -9,7 +9,7 @@
 # If you're having trouble making the relevant changes, bug a maintainer.
 
 [toolchain]
-channel = "1.78.0"
+channel = "1.81.0"
 components = [
     # For rust-analyzer
     "rust-src",

src/api/client_server/account.rs

@@ -316,8 +316,7 @@ pub(crate) async fn register_route(
 /// - Requires UIAA to verify user password
 /// - Changes the password of the sender user
 /// - The password hash is calculated using argon2 with 32 character salt, the
-///   plain password is
-/// not saved
+///   plain password is not saved
 ///
 /// If `logout_devices` is true it does the following for each device except the
 /// sender device:

src/api/client_server/context.rs

@@ -16,8 +16,7 @@ use crate::{services, Ar, Error, Ra, Result};
 /// Allows loading room history around an event.
 ///
 /// - Only works if the user is joined (TODO: always allow, but only show events
-///   if the user was
-/// joined, depending on `history_visibility`)
+///   if the user was joined, depending on `history_visibility`)
 #[allow(clippy::too_many_lines)]
 pub(crate) async fn get_context_route(
     body: Ar<get_context::v3::Request>,

src/api/client_server/media.rs

@@ -1,6 +1,7 @@
 use std::time::Duration;
 
 use axum::response::IntoResponse;
+use either::Either;
 use http::{
     header::{CONTENT_DISPOSITION, CONTENT_SECURITY_POLICY, CONTENT_TYPE},
     HeaderName, HeaderValue, Method,
@@ -20,10 +21,8 @@ use ruma::{
 use tracing::{debug, error, info, warn};
 
 use crate::{
-    service::media::FileMeta,
-    services,
-    utils::{self, MxcData},
-    Ar, Error, Ra, Result,
+    service::media::{FileMeta, MediaCreationToken, MxcData, ThumbnailResult},
+    services, utils, Ar, Error, Ra, Result,
 };
 
 const MXC_LENGTH: usize = 32;
@@ -156,10 +155,15 @@ pub(crate) async fn create_content_route(
     let media_id = utils::random_string(MXC_LENGTH);
     let mxc = MxcData::new(services().globals.server_name(), &media_id)?;
 
+    // this is a fresh MXC, there ought to not be anyone else creating it at the
+    // same time
+    let Either::Right(token) = services().media.get(mxc).await? else {
+        return Err(Error::bad_database("Media for fresh mxc already exists"));
+    };
     services()
         .media
         .create(
-            mxc.to_string(),
+            token,
             body.filename
                 .clone()
                 .map(|filename| ContentDisposition {
@@ -257,7 +261,7 @@ async fn get_redirected_content(
 
 #[tracing::instrument(skip_all)]
 async fn get_remote_content_via_federation_api(
-    mxc: &MxcData<'_>,
+    mxc: MxcData<'_>,
 ) -> Result<RemoteResponse, Error> {
     let authenticated_media_fed::get_content::v1::Response {
         metadata,
@@ -293,7 +297,7 @@ async fn get_remote_content_via_federation_api(
 #[allow(deprecated)] // unauthenticated media
 #[tracing::instrument(skip_all)]
 async fn get_remote_content_via_legacy_api(
-    mxc: &MxcData<'_>,
+    mxc: MxcData<'_>,
 ) -> Result<RemoteResponse, Error> {
     let content_response = services()
         .sending
@@ -321,9 +325,9 @@ async fn get_remote_content_via_legacy_api(
 
 #[tracing::instrument]
 pub(crate) async fn get_remote_content(
-    mxc: &MxcData<'_>,
+    token: MediaCreationToken,
 ) -> Result<RemoteResponse, Error> {
-    let fed_result = get_remote_content_via_federation_api(mxc).await;
+    let fed_result = get_remote_content_via_federation_api(token.mxc()).await;
 
     let response = match fed_result {
         Ok(response) => {
@@ -340,7 +344,7 @@ pub(crate) async fn get_remote_content(
                  back to deprecated API"
             );
 
-            get_remote_content_via_legacy_api(mxc).await?
+            get_remote_content_via_legacy_api(token.mxc()).await?
         }
         Err(e) => {
             return Err(e);
@@ -350,7 +354,7 @@ pub(crate) async fn get_remote_content(
     services()
         .media
         .create(
-            mxc.to_string(),
+            token,
             response.content.content_disposition.as_ref(),
             response.content.content_type.clone(),
             &response.content.file,
@@ -449,37 +453,45 @@ async fn get_content_route_ruma(
 ) -> Result<authenticated_media_client::get_content::v1::Response> {
     let mxc = MxcData::new(&body.server_name, &body.media_id)?;
 
-    if let Some((
-        FileMeta {
-            content_type,
-            ..
-        },
-        file,
-    )) = services().media.get(mxc.to_string()).await?
-    {
-        Ok(authenticated_media_client::get_content::v1::Response {
+    let token = match services().media.get(mxc).await? {
+        Either::Left((
+            FileMeta {
+                content_type,
+                ..
+            },
             file,
-            content_disposition: Some(content_disposition_for(
-                content_type.as_deref(),
-                None,
-            )),
-            content_type,
-        })
-    } else if &*body.server_name != services().globals.server_name()
-        && allow_remote == AllowRemote::Yes
+        )) => {
+            return Ok(authenticated_media_client::get_content::v1::Response {
+                file,
+                content_disposition: Some(content_disposition_for(
+                    content_type.as_deref(),
+                    None,
+                )),
+                content_type,
+            });
+        }
+
+        Either::Right(token) => token,
+    };
+
+    if &*body.server_name == services().globals.server_name()
+        || allow_remote != AllowRemote::Yes
     {
-        let remote_response = get_remote_content(&mxc).await?;
-        Ok(authenticated_media_client::get_content::v1::Response {
-            file: remote_response.content.file,
-            content_disposition: Some(content_disposition_for(
-                remote_response.content.content_type.as_deref(),
-                None,
-            )),
-            content_type: remote_response.content.content_type,
-        })
-    } else {
-        Err(Error::BadRequest(ErrorKind::NotYetUploaded, "Media not found."))
+        return Err(Error::BadRequest(
+            ErrorKind::NotYetUploaded,
+            "Media not found.",
+        ));
     }
+
+    let remote_response = get_remote_content(token).await?;
+    Ok(authenticated_media_client::get_content::v1::Response {
+        file: remote_response.content.file,
+        content_disposition: Some(content_disposition_for(
+            remote_response.content.content_type.as_deref(),
+            None,
+        )),
+        content_type: remote_response.content.content_type,
+    })
 }
 
 /// # `GET /_matrix/media/r0/download/{serverName}/{mediaId}/{fileName}`
@@ -573,30 +585,36 @@ async fn get_content_as_filename_route_ruma(
     body: Ar<authenticated_media_client::get_content_as_filename::v1::Request>,
     allow_remote: AllowRemote,
 ) -> Result<authenticated_media_client::get_content_as_filename::v1::Response> {
+    use authenticated_media_client::get_content_as_filename::v1::Response;
+
     let mxc = MxcData::new(&body.server_name, &body.media_id)?;
 
-    if let Some((
-        FileMeta {
-            content_type,
-            ..
-        },
-        file,
-    )) = services().media.get(mxc.to_string()).await?
-    {
-        Ok(authenticated_media_client::get_content_as_filename::v1::Response {
+    let token = match services().media.get(mxc).await? {
+        Either::Left((
+            FileMeta {
+                content_type,
+                ..
+            },
             file,
-            content_disposition: Some(content_disposition_for(
-                content_type.as_deref(),
-                Some(body.filename.clone()),
-            )),
-            content_type,
-        })
-    } else if &*body.server_name != services().globals.server_name()
+        )) => {
+            return Ok(Response {
+                file,
+                content_disposition: Some(content_disposition_for(
+                    content_type.as_deref(),
+                    Some(body.filename.clone()),
+                )),
+                content_type,
+            });
+        }
+        Either::Right(token) => token,
+    };
+
+    if &*body.server_name != services().globals.server_name()
         && allow_remote == AllowRemote::Yes
     {
-        let remote_response = get_remote_content(&mxc).await?;
+        let remote_response = get_remote_content(token).await?;
 
-        Ok(authenticated_media_client::get_content_as_filename::v1::Response {
+        Ok(Response {
             content_disposition: Some(content_disposition_for(
                 remote_response.content.content_type.as_deref(),
                 Some(body.filename.clone()),
@@ -836,81 +854,99 @@ async fn get_content_thumbnail_route_ruma(
         }
     };
 
-    if let Some((
+    let lookup_result =
+        services().media.get_thumbnail(mxc, width, height).await?;
+
+    if let ThumbnailResult::Data(
         FileMeta {
             content_type,
             ..
         },
         file,
-    )) =
-        services().media.get_thumbnail(mxc.to_string(), width, height).await?
+    ) = lookup_result
     {
         return Ok(make_response(file, content_type));
     }
 
-    if &*body.server_name != services().globals.server_name()
-        && allow_remote == AllowRemote::Yes
+    if &*body.server_name == services().globals.server_name()
+        || allow_remote != AllowRemote::Yes
     {
-        let get_thumbnail_response = get_remote_thumbnail(
-            &body.server_name,
-            authenticated_media_fed::get_content_thumbnail::v1::Request {
-                height: body.height,
-                width: body.width,
-                method: body.method.clone(),
-                media_id: body.media_id.clone(),
-                timeout_ms: Duration::from_secs(20),
-                // we don't support animated thumbnails, so don't try requesting
-                // one - we're allowed to ignore the client's request for an
-                // animated thumbnail
-                animated: Some(false),
-            },
-        )
-        .await;
-
-        match get_thumbnail_response {
-            Ok(resp) => {
-                services()
-                    .media
-                    .upload_thumbnail(
-                        mxc.to_string(),
-                        None,
-                        resp.content.content_type.clone(),
-                        width,
-                        height,
-                        &resp.content.file,
-                    )
-                    .await?;
-
-                return Ok(make_response(
-                    resp.content.file,
-                    resp.content.content_type,
-                ));
-            }
-            Err(error) => warn!(
-                %error,
-                "Failed to fetch thumbnail via federation, trying to fetch \
-                 original media and create thumbnail ourselves"
-            ),
-        }
-
-        get_remote_content(&mxc).await?;
-
-        if let Some((
-            FileMeta {
-                content_type,
-                ..
-            },
-            file,
-        )) = services()
-            .media
-            .get_thumbnail(mxc.to_string(), width, height)
-            .await?
-        {
-            return Ok(make_response(file, content_type));
-        }
-
-        error!("Source media doesn't exist even after fetching it from remote");
+        return Err(Error::BadRequest(
+            ErrorKind::NotYetUploaded,
+            "Media not found.",
+        ));
     }
 
+    let source_token = match lookup_result {
+        ThumbnailResult::Data(..) => unreachable!(),
+        ThumbnailResult::NeedSource(token) => {
+            // we need to fetch the whole media
+            token
+        }
+        ThumbnailResult::NeedThumbnail(thumbnail_token, source_token) => {
+            // try to fetch thumbnail from remote, falling back to fetching the
+            // whole media anyway
+
+            let get_thumbnail_response = get_remote_thumbnail(
+                &body.server_name,
+                authenticated_media_fed::get_content_thumbnail::v1::Request {
+                    height: thumbnail_token.height().into(),
+                    width: thumbnail_token.width().into(),
+                    method: body.method.clone(),
+                    media_id: body.media_id.clone(),
+                    timeout_ms: Duration::from_secs(20),
+                    // we don't support animated thumbnails, so don't try
+                    // requesting one - we're allowed to
+                    // ignore the client's request for an
+                    // animated thumbnail
+                    animated: Some(false),
+                },
+            )
+            .await;
+
+            match get_thumbnail_response {
+                Ok(resp) => {
+                    services()
+                        .media
+                        .create_thumbnail(
+                            thumbnail_token,
+                            None,
+                            resp.content.content_type.clone(),
+                            &resp.content.file,
+                        )
+                        .await?;
+
+                    return Ok(make_response(
+                        resp.content.file,
+                        resp.content.content_type,
+                    ));
+                }
+                Err(error) => {
+                    warn!(
+                        %error,
+                        "Failed to fetch thumbnail via federation, trying to fetch \
+                         original media and create thumbnail ourselves"
+                    );
+                }
+            }
+
+            source_token
+        }
+    };
+
+    get_remote_content(source_token).await?;
+
+    if let ThumbnailResult::Data(
+        FileMeta {
+            content_type,
+            ..
+        },
+        file,
+    ) = services().media.get_thumbnail(mxc, width, height).await?
+    {
+        return Ok(make_response(file, content_type));
+    }
+
+    error!("Source media doesn't exist even after fetching it from remote");
     Err(Error::BadRequest(ErrorKind::NotYetUploaded, "Media not found."))
 }

src/api/client_server/message.rs

@@ -117,8 +117,7 @@ pub(crate) async fn send_message_event_route(
 /// Allows paginating through room history.
 ///
 /// - Only works if the user is joined (TODO: always allow, but only show events
-///   where the user was
-/// joined, depending on `history_visibility`)
+///   where the user was joined, depending on `history_visibility`)
 #[allow(clippy::too_many_lines)]
 pub(crate) async fn get_message_events_route(
     body: Ar<get_message_events::v3::Request>,

src/api/client_server/sync.rs

@@ -36,8 +36,7 @@ use crate::{
 /// Synchronize the client's state with the latest state on the server.
 ///
 /// - This endpoint takes a `since` parameter which should be the `next_batch`
-///   value from a
-/// previous request for incremental syncs.
+///   value from a previous request for incremental syncs.
 ///
 /// Calling this endpoint without a `since` parameter returns:
 /// - Some of the most recent events of each timeline
@@ -50,11 +49,9 @@ use crate::{
 /// - Some of the most recent events of each timeline that happened after
 ///   `since`
 /// - If user joined the room after `since`: All state events (unless lazy
-///   loading is activated) and
-/// all device list updates in that room
+///   loading is activated) and all device list updates in that room
 /// - If the user was already in the room: A list of all events that are in the
-///   state now, but were
-/// not in the state at `since`
+///   state now, but were not in the state at `since`
 /// - If the state we send contains a member event: Joined and invited member
 ///   counts, heroes
 /// - Device list updates that happened after `since`

src/api/client_server/user_directory.rs

@@ -13,8 +13,7 @@ use crate::{services, Ar, Ra, Result};
 /// Searches all known users for a match.
 ///
 /// - Hides any local users that aren't in any public rooms (i.e. those that
-///   have the join rule set to public)
-/// and don't share a room with the sender
+///   have the join rule set to public) and don't share a room with the sender
 pub(crate) async fn search_users_route(
     body: Ar<search_users::v3::Request>,
 ) -> Result<Ra<search_users::v3::Response>> {

src/api/server_server.rs

@@ -10,6 +10,7 @@ use std::{
 use axum::{response::IntoResponse, Json};
 use axum_extra::headers::{Authorization, HeaderMapExt};
 use base64::Engine as _;
+use either::Either;
 use get_profile_information::v1::ProfileField;
 use ruma::{
     api::{
@@ -70,9 +71,12 @@ use super::appservice_server;
 use crate::{
     api::client_server::{self, claim_keys_helper, get_keys_helper},
     observability::{FoundIn, Lookup, METRICS},
-    service::pdu::{gen_event_id_canonical_json, PduBuilder},
+    service::{
+        media::{MxcData, ThumbnailResult},
+        pdu::{gen_event_id_canonical_json, PduBuilder},
+    },
     services,
-    utils::{self, dbg_truncate_str, MxcData},
+    utils::{self, dbg_truncate_str},
     Ar, Error, PduEvent, Ra, Result,
 };
 
@@ -532,7 +536,7 @@ async fn request_well_known(destination: &str) -> Option<String> {
     let response = services()
         .globals
         .default_client()
-        .get(&format!("https://{destination}/.well-known/matrix/server"))
+        .get(format!("https://{destination}/.well-known/matrix/server"))
         .send()
         .await;
     debug!("Got well known response");
@@ -565,8 +569,7 @@ pub(crate) async fn get_server_version_route(
 /// Gets the public signing keys of this server.
 ///
 /// - Matrix does not support invalidating public keys, so the key returned by
-///   this will be valid
-/// forever.
+///   this will be valid forever.
 // Response type for this endpoint is Json because we need to calculate a
 // signature for the response
 pub(crate) async fn get_server_keys_route() -> Result<impl IntoResponse> {
@@ -617,8 +620,7 @@ pub(crate) async fn get_server_keys_route() -> Result<impl IntoResponse> {
 /// Gets the public signing keys of this server.
 ///
 /// - Matrix does not support invalidating public keys, so the key returned by
-///   this will be valid
-/// forever.
+///   this will be valid forever.
 pub(crate) async fn get_server_keys_deprecated_route() -> impl IntoResponse {
     get_server_keys_route().await
 }
@@ -2046,13 +2048,13 @@ pub(crate) async fn media_download_route(
     body: Ar<authenticated_media::get_content::v1::Request>,
 ) -> Result<Ra<authenticated_media::get_content::v1::Response>> {
     let mxc = MxcData::new(services().globals.server_name(), &body.media_id)?;
-    let Some((
+    let Either::Left((
         crate::service::media::FileMeta {
             content_disposition,
             content_type,
         },
         file,
-    )) = services().media.get(mxc.to_string()).await?
+    )) = services().media.get(mxc).await?
     else {
         return Err(Error::BadRequest(
             ErrorKind::NotYetUploaded,
@@ -2093,13 +2095,13 @@ pub(crate) async fn media_thumbnail_route(
         Error::BadRequest(ErrorKind::InvalidParam, "Height is invalid.")
     })?;
 
-    let Some((
+    let ThumbnailResult::Data(
         crate::service::media::FileMeta {
             content_type,
             ..
         },
         file,
-    )) = services().media.get_thumbnail(mxc.to_string(), width, height).await?
+    ) = services().media.get_thumbnail(mxc, width, height).await?
     else {
         return Err(Error::BadRequest(
             ErrorKind::NotYetUploaded,

src/database/key_value/media.rs

@@ -4,7 +4,7 @@ use crate::{
     database::KeyValueDatabase,
     service::{
         self,
-        media::{FileMeta, MediaFileKey},
+        media::{FileMeta, MediaFileKey, MxcData},
     },
     utils, Error, Result,
 };
@@ -12,12 +12,12 @@ use crate::{
 impl service::media::Data for KeyValueDatabase {
     fn create_file_metadata(
         &self,
-        mxc: String,
+        mxc: MxcData<'_>,
         width: u32,
         height: u32,
         meta: &FileMeta,
     ) -> Result<MediaFileKey> {
-        let mut key = mxc.as_bytes().to_vec();
+        let mut key = mxc.to_string().as_bytes().to_vec();
         key.push(0xFF);
         key.extend_from_slice(&width.to_be_bytes());
         key.extend_from_slice(&height.to_be_bytes());
@@ -45,11 +45,11 @@ impl service::media::Data for KeyValueDatabase {
 
     fn search_file_metadata(
         &self,
-        mxc: String,
+        mxc: MxcData<'_>,
         width: u32,
         height: u32,
     ) -> Result<(FileMeta, MediaFileKey)> {
-        let mut prefix = mxc.as_bytes().to_vec();
+        let mut prefix = mxc.to_string().as_bytes().to_vec();
         prefix.push(0xFF);
         prefix.extend_from_slice(&width.to_be_bytes());
         prefix.extend_from_slice(&height.to_be_bytes());

src/database/key_value/rooms/edus/read_receipt.rs

@@ -1,5 +1,3 @@
-use std::mem;
-
 use ruma::{
     events::receipt::ReceiptEvent, serde::Raw, CanonicalJsonObject,
     OwnedUserId, RoomId, UserId,
@@ -83,7 +81,7 @@ impl service::rooms::edus::read_receipt::Data for KeyValueDatabase {
                 .take_while(move |(k, _)| k.starts_with(&prefix2))
                 .map(move |(k, v)| {
                     let count = utils::u64_from_bytes(
-                        &k[prefix.len()..prefix.len() + mem::size_of::<u64>()],
+                        &k[prefix.len()..prefix.len() + size_of::<u64>()],
                     )
                     .map_err(|_| {
                         Error::bad_database(
@@ -92,7 +90,7 @@ impl service::rooms::edus::read_receipt::Data for KeyValueDatabase {
                     })?;
                     let user_id = UserId::parse(
                         utils::string_from_bytes(
-                            &k[prefix.len() + mem::size_of::<u64>() + 1..],
+                            &k[prefix.len() + size_of::<u64>() + 1..],
                         )
                         .map_err(|_| {
                             Error::bad_database(

src/database/key_value/rooms/pdu_metadata.rs

@@ -1,4 +1,4 @@
-use std::{mem, sync::Arc};
+use std::sync::Arc;
 
 use ruma::{EventId, RoomId, UserId};
 
@@ -47,12 +47,13 @@ impl service::rooms::pdu_metadata::Data for KeyValueDatabase {
                 .iter_from(&current, true)
                 .take_while(move |(k, _)| k.starts_with(&prefix))
                 .map(move |(tofrom, _data)| {
-                    let from = utils::u64_from_bytes(
-                        &tofrom[(mem::size_of::<u64>())..],
-                    )
-                    .map_err(|_| {
-                        Error::bad_database("Invalid count in tofrom_relation.")
-                    })?;
+                    let from =
+                        utils::u64_from_bytes(&tofrom[(size_of::<u64>())..])
+                            .map_err(|_| {
+                                Error::bad_database(
+                                    "Invalid count in tofrom_relation.",
+                                )
+                            })?;
 
                     let mut pduid = shortroomid.get().to_be_bytes().to_vec();
                     pduid.extend_from_slice(&from.to_be_bytes());

src/database/key_value/rooms/threads.rs

@@ -1,5 +1,3 @@
-use std::mem;
-
 use ruma::{
     api::client::threads::get_threads::v1::IncludeThreads, OwnedUserId, RoomId,
     UserId,
@@ -36,14 +34,13 @@ impl service::rooms::threads::Data for KeyValueDatabase {
                 .iter_from(&current, true)
                 .take_while(move |(k, _)| k.starts_with(&prefix))
                 .map(move |(pduid, _users)| {
-                    let count = utils::u64_from_bytes(
-                        &pduid[(mem::size_of::<u64>())..],
-                    )
-                    .map_err(|_| {
-                        Error::bad_database(
-                            "Invalid pduid in threadid_userids.",
-                        )
-                    })?;
+                    let count =
+                        utils::u64_from_bytes(&pduid[(size_of::<u64>())..])
+                            .map_err(|_| {
+                                Error::bad_database(
+                                    "Invalid pduid in threadid_userids.",
+                                )
+                            })?;
 
                     let pduid = PduId::new(pduid);
 

src/service.rs

@@ -145,9 +145,7 @@ impl Services {
             account_data: db,
             admin: admin::Service::build(),
             key_backups: db,
-            media: media::Service {
-                db,
-            },
+            media: media::Service::new(db),
             sending: sending::Service::build(db, &config),
 
             globals: globals::Service::load(db, config, reload_handles)?,

src/service/media.rs

@@ -1,14 +1,22 @@
-use std::io::Cursor;
+use std::{fmt, io::Cursor};
 
+use either::Either;
 use image::imageops::FilterType;
-use ruma::http_headers::ContentDisposition;
+use ruma::{
+    api::client::error::ErrorKind, http_headers::ContentDisposition, MxcUri,
+    MxcUriError, OwnedMxcUri, OwnedServerName,
+};
 use tokio::{
     fs::File,
     io::{AsyncReadExt, AsyncWriteExt},
 };
-use tracing::{debug, warn};
+use tracing::{debug, instrument, warn};
 
-use crate::{services, Result};
+use crate::{
+    services,
+    utils::on_demand_hashmap::{KeyToken, TokenSet},
+    Error, Result,
+};
 
 mod data;
 
@@ -37,16 +45,209 @@ impl MediaFileKey {
     }
 }
 
+/// Data that makes up an `mxc://` URL.
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct MxcData<'a> {
+    pub(crate) server_name: &'a ruma::ServerName,
+    pub(crate) media_id: &'a str,
+}
+
+impl<'a> MxcData<'a> {
+    pub(crate) fn new(
+        server_name: &'a ruma::ServerName,
+        media_id: &'a str,
+    ) -> Result<Self> {
+        if !media_id.bytes().all(|b| {
+            matches!(b,
+                b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'-' | b'_'
+            )
+        }) {
+            return Err(Error::BadRequest(
+                ErrorKind::InvalidParam,
+                "Invalid MXC media id",
+            ));
+        }
+
+        Ok(Self {
+            server_name,
+            media_id,
+        })
+    }
+}
+
+impl fmt::Display for MxcData<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "mxc://{}/{}", self.server_name, self.media_id)
+    }
+}
+
+impl From<MxcData<'_>> for OwnedMxcUri {
+    fn from(value: MxcData<'_>) -> Self {
+        value.to_string().into()
+    }
+}
+
+impl<'a> TryFrom<&'a MxcUri> for MxcData<'a> {
+    type Error = MxcUriError;
+
+    fn try_from(value: &'a MxcUri) -> Result<Self, Self::Error> {
+        Ok(Self::new(value.server_name()?, value.media_id()?)
+            .expect("validated MxcUri should always be valid MxcData"))
+    }
+}
+
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
+struct ThumbnailDimensions {
+    width: u32,
+    height: u32,
+    should_crop: bool,
+}
+
+impl ThumbnailDimensions {
+    /// Returns None when the server should send the original file.
+    fn new(width: u32, height: u32) -> Option<Self> {
+        match (width, height) {
+            (0..=32, 0..=32) => Some(Self {
+                width: 32,
+                height: 32,
+                should_crop: true,
+            }),
+            (0..=96, 0..=96) => Some(Self {
+                width: 96,
+                height: 96,
+                should_crop: true,
+            }),
+            (0..=320, 0..=240) => Some(Self {
+                width: 320,
+                height: 240,
+                should_crop: false,
+            }),
+            (0..=640, 0..=480) => Some(Self {
+                width: 640,
+                height: 480,
+                should_crop: false,
+            }),
+            (0..=800, 0..=600) => Some(Self {
+                width: 800,
+                height: 600,
+                should_crop: false,
+            }),
+            _ => None,
+        }
+    }
+}
+
+#[derive(Clone, Eq, Hash, PartialEq)]
+struct MediaCreationData {
+    server_name: OwnedServerName,
+    media_id: String,
+    thumbnail_dimensions: Option<ThumbnailDimensions>,
+}
+
+impl MediaCreationData {
+    fn mxc(&self) -> MxcData<'_> {
+        MxcData {
+            server_name: &self.server_name,
+            media_id: &self.media_id,
+        }
+    }
+}
+
+impl fmt::Debug for MediaCreationData {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?}", self.mxc())?;
+        if let Some(ThumbnailDimensions {
+            width,
+            height,
+            ..
+        }) = self.thumbnail_dimensions
+        {
+            write!(f, " @ {width}x{height}")?;
+        }
+
+        Ok(())
+    }
+}
+
+/// Used to serialize media creation.
+///
+/// There can only be up to one holder of a `MediaCreationToken` for a given mxc
+/// at any one time.
+///
+/// Returned by [`Service::get()`] when the file is not available locally.
+#[derive(Debug)]
+pub(crate) struct MediaCreationToken(KeyToken<MediaCreationData>);
+
+impl MediaCreationToken {
+    pub(crate) fn mxc(&self) -> MxcData<'_> {
+        self.0.mxc()
+    }
+}
+
+/// Used to serialize thumbnail creation.
+///
+/// There can only be up to one holder of a `ThumbnailCreationToken` for a given
+/// mxc and resolution at any one time.
+///
+/// Returned by [`Service::get_thumbnail()`] when the thumbnail is not available
+/// locally.
+#[derive(Debug)]
+pub(crate) struct ThumbnailCreationToken(KeyToken<MediaCreationData>);
+
+impl ThumbnailCreationToken {
+    pub(crate) fn width(&self) -> u32 {
+        self.0
+            .thumbnail_dimensions
+            .as_ref()
+            .expect("thumbnail creation token should have thumbnail dimensions")
+            .width
+    }
+
+    pub(crate) fn height(&self) -> u32 {
+        self.0
+            .thumbnail_dimensions
+            .as_ref()
+            .expect("thumbnail creation token should have thumbnail dimensions")
+            .height
+    }
+}
+
+pub(crate) enum ThumbnailResult {
+    /// Thumbnail available, corresponding metadata and file content returned.
+    Data(FileMeta, Vec<u8>),
+    /// Source file for thumbnail is required.
+    NeedSource(MediaCreationToken),
+    /// Thumbnail data is required (or can be generated from source file, once
+    /// present)
+    NeedThumbnail(ThumbnailCreationToken, MediaCreationToken),
+}
+
 pub(crate) struct Service {
     pub(crate) db: &'static dyn Data,
+
+    creation_tokens: TokenSet<MediaCreationData>,
 }
 
 impl Service {
-    /// Uploads a file.
-    #[tracing::instrument(skip(self, file))]
-    pub(crate) async fn create(
+    pub(crate) fn new(db: &'static dyn Data) -> Service {
+        Self {
+            db,
+            creation_tokens: TokenSet::new("media_creation".to_owned()),
+        }
+    }
+
+    #[instrument(skip(self))]
+    async fn lock_creation(
         &self,
-        mxc: String,
+        creation_data: MediaCreationData,
+    ) -> KeyToken<MediaCreationData> {
+        self.creation_tokens.lock_key(creation_data).await
+    }
+
+    #[tracing::instrument(skip(self, file))]
+    async fn create_inner(
+        &self,
+        creation_data: &MediaCreationData,
         content_disposition: Option<&ContentDisposition>,
         content_type: Option<String>,
         file: &[u8],
@@ -56,8 +257,23 @@ impl Service {
                 .map(ContentDisposition::to_string),
             content_type,
         };
-        // Width, Height = 0 if it's not a thumbnail
-        let key = self.db.create_file_metadata(mxc, 0, 0, &meta)?;
+
+        let key = if let Some(ThumbnailDimensions {
+            width,
+            height,
+            ..
+        }) = creation_data.thumbnail_dimensions
+        {
+            self.db.create_file_metadata(
+                creation_data.mxc(),
+                width,
+                height,
+                &meta,
+            )?
+        } else {
+            // Width, Height = 0 if it's not a thumbnail
+            self.db.create_file_metadata(creation_data.mxc(), 0, 0, &meta)?
+        };
 
         let path = services().globals.get_media_file(&key);
         let mut f = File::create(path).await?;
@@ -65,65 +281,102 @@ impl Service {
         Ok(meta)
     }
 
+    /// Uploads or replaces a file.
+    pub(crate) async fn create(
+        &self,
+        MediaCreationToken(creation_data): MediaCreationToken,
+        content_disposition: Option<&ContentDisposition>,
+        content_type: Option<String>,
+        file: &[u8],
+    ) -> Result<FileMeta> {
+        self.create_inner(
+            &creation_data,
+            content_disposition,
+            content_type,
+            file,
+        )
+        .await
+    }
+
     /// Uploads or replaces a file thumbnail.
-    #[tracing::instrument(skip(self, file))]
-    pub(crate) async fn upload_thumbnail(
+    pub(crate) async fn create_thumbnail(
         &self,
-        mxc: String,
-        content_disposition: Option<String>,
+        ThumbnailCreationToken(creation_data): ThumbnailCreationToken,
+        content_disposition: Option<&ContentDisposition>,
         content_type: Option<String>,
-        width: u32,
-        height: u32,
         file: &[u8],
     ) -> Result<FileMeta> {
-        let meta = FileMeta {
+        self.create_inner(
+            &creation_data,
             content_disposition,
             content_type,
-        };
-        let key = self.db.create_file_metadata(mxc, width, height, &meta)?;
-
-        let path = services().globals.get_media_file(&key);
-        let mut f = File::create(path).await?;
-        f.write_all(file).await?;
-
-        Ok(meta)
+            file,
+        )
+        .await
     }
 
-    /// Downloads a file.
     #[tracing::instrument(skip(self))]
+    // internal function, and not actually that complex
+    #[allow(clippy::type_complexity)]
+    async fn get_inner(
+        &self,
+        mxc: MxcData<'_>,
+        thumbnail_dimensions: Option<ThumbnailDimensions>,
+    ) -> Result<Either<(FileMeta, Vec<u8>), KeyToken<MediaCreationData>>> {
+        let try_get = || async {
+            if let Ok((meta, key)) = self.db.search_file_metadata(mxc, 0, 0) {
+                let path = services().globals.get_media_file(&key);
+                let mut file_data = Vec::new();
+                let Ok(mut file) = File::open(path).await else {
+                    return Ok::<_, Error>(None);
+                };
+
+                file.read_to_end(&mut file_data).await?;
+
+                Ok(Some((meta, file_data)))
+            } else {
+                Ok(None)
+            }
+        };
+
+        if let Some(ret) = try_get().await? {
+            return Ok(Either::Left(ret));
+        }
+
+        let token = self
+            .lock_creation(MediaCreationData {
+                server_name: mxc.server_name.to_owned(),
+                media_id: mxc.media_id.to_owned(),
+                thumbnail_dimensions,
+            })
+            .await;
+
+        // check again in case it has been created in the meantime
+        if let Some(ret) = try_get().await? {
+            Ok(Either::Left(ret))
+        } else {
+            Ok(Either::Right(token))
+        }
+    }
+
+    /// Loads a file from local storage.
     pub(crate) async fn get(
         &self,
-        mxc: String,
-    ) -> Result<Option<(FileMeta, Vec<u8>)>> {
-        if let Ok((meta, key)) = self.db.search_file_metadata(mxc, 0, 0) {
-            let path = services().globals.get_media_file(&key);
-            let mut file_data = Vec::new();
-            let Ok(mut file) = File::open(path).await else {
-                return Ok(None);
-            };
-
-            file.read_to_end(&mut file_data).await?;
-
-            Ok(Some((meta, file_data)))
-        } else {
-            Ok(None)
-        }
+        mxc: MxcData<'_>,
+    ) -> Result<Either<(FileMeta, Vec<u8>), MediaCreationToken>> {
+        Ok(self.get_inner(mxc, None).await?.map_right(MediaCreationToken))
     }
 
-    /// Returns width, height of the thumbnail and whether it should be cropped.
-    /// Returns None when the server should send the original file.
-    fn thumbnail_properties(
-        width: u32,
-        height: u32,
-    ) -> Option<(u32, u32, bool)> {
-        match (width, height) {
-            (0..=32, 0..=32) => Some((32, 32, true)),
-            (0..=96, 0..=96) => Some((96, 96, true)),
-            (0..=320, 0..=240) => Some((320, 240, false)),
-            (0..=640, 0..=480) => Some((640, 480, false)),
-            (0..=800, 0..=600) => Some((800, 600, false)),
-            _ => None,
-        }
+    /// Loads a thumbnail from local storage.
+    async fn get_thumbnail_local(
+        &self,
+        mxc: MxcData<'_>,
+        dimensions: ThumbnailDimensions,
+    ) -> Result<Either<(FileMeta, Vec<u8>), ThumbnailCreationToken>> {
+        Ok(self
+            .get_inner(mxc, Some(dimensions))
+            .await?
+            .map_right(ThumbnailCreationToken))
     }
 
     /// Generates a thumbnail from the given image file contents. Returns
@@ -171,23 +424,23 @@ impl Service {
                         / u64::from(original_height)
                 };
                 if use_width {
-                    if intermediate <= u64::from(::std::u32::MAX) {
-                        (width, intermediate.try_into().unwrap_or(u32::MAX))
+                    if let Ok(intermediate) = u32::try_from(intermediate) {
+                        (width, intermediate)
                     } else {
                         (
-                            (u64::from(width) * u64::from(::std::u32::MAX)
+                            (u64::from(width) * u64::from(u32::MAX)
                                 / intermediate)
                                 .try_into()
                                 .unwrap_or(u32::MAX),
-                            ::std::u32::MAX,
+                            u32::MAX,
                         )
                     }
-                } else if intermediate <= u64::from(::std::u32::MAX) {
-                    (intermediate.try_into().unwrap_or(u32::MAX), height)
+                } else if let Ok(intermediate) = u32::try_from(intermediate) {
+                    (intermediate, height)
                 } else {
                     (
-                        ::std::u32::MAX,
-                        (u64::from(height) * u64::from(::std::u32::MAX)
+                        u32::MAX,
+                        (u64::from(height) * u64::from(u32::MAX)
                             / intermediate)
                             .try_into()
                             .unwrap_or(u32::MAX),
@@ -208,7 +461,7 @@ impl Service {
         Ok(Some(thumbnail_bytes))
     }
 
-    /// Downloads a file's thumbnail.
+    /// Loads or generates a file's thumbnail.
     ///
     /// Here's an example on how it works:
     ///
@@ -224,34 +477,39 @@ impl Service {
     #[tracing::instrument(skip(self))]
     pub(crate) async fn get_thumbnail(
         &self,
-        mxc: String,
+        mxc: MxcData<'_>,
         width: u32,
         height: u32,
-    ) -> Result<Option<(FileMeta, Vec<u8>)>> {
-        // 0, 0 because that's the original file
-        let (width, height, crop) =
-            Self::thumbnail_properties(width, height).unwrap_or((0, 0, false));
-
-        if let Ok((meta, key)) =
-            self.db.search_file_metadata(mxc.clone(), width, height)
-        {
-            debug!("Using saved thumbnail");
-            let path = services().globals.get_media_file(&key);
-            let mut file = Vec::new();
-            File::open(path).await?.read_to_end(&mut file).await?;
-
-            return Ok(Some((meta, file.clone())));
-        }
-
-        let Ok((meta, key)) = self.db.search_file_metadata(mxc.clone(), 0, 0)
-        else {
-            debug!("Original image not found, can't generate thumbnail");
-            return Ok(None);
+    ) -> Result<ThumbnailResult> {
+        let Some(dimensions) = ThumbnailDimensions::new(width, height) else {
+            // image should be used as-is
+            return match self.get(mxc).await? {
+                Either::Left((meta, file)) => {
+                    Ok(ThumbnailResult::Data(meta, file))
+                }
+                Either::Right(token) => Ok(ThumbnailResult::NeedSource(token)),
+            };
         };
 
-        let path = services().globals.get_media_file(&key);
-        let mut file = Vec::new();
-        File::open(path).await?.read_to_end(&mut file).await?;
+        let thumbnail_token =
+            match self.get_thumbnail_local(mxc, dimensions).await? {
+                Either::Left((meta, file)) => {
+                    debug!("Using saved thumbnail");
+                    return Ok(ThumbnailResult::Data(meta, file));
+                }
+                Either::Right(token) => token,
+            };
+
+        let (meta, file) = match self.get(mxc).await? {
+            Either::Left(ret) => ret,
+            Either::Right(media_token) => {
+                debug!("Original image not found, can't generate thumbnail");
+                return Ok(ThumbnailResult::NeedThumbnail(
+                    thumbnail_token,
+                    media_token,
+                ));
+            }
+        };
 
         debug!("Generating thumbnail");
         let thumbnail_result = {
@@ -260,7 +518,12 @@ impl Service {
 
             tokio::task::spawn_blocking(move || {
                 outer_span.in_scope(|| {
-                    Self::generate_thumbnail(&file, width, height, crop)
+                    Self::generate_thumbnail(
+                        &file,
+                        width,
+                        height,
+                        dimensions.should_crop,
+                    )
                 })
             })
             .await
@@ -269,20 +532,20 @@ impl Service {
 
         let Some(thumbnail_bytes) = thumbnail_result? else {
             debug!("Returning source image as-is");
-            return Ok(Some((meta, file)));
+            return Ok(ThumbnailResult::Data(meta, file));
         };
 
         debug!("Saving created thumbnail");
 
-        // Save thumbnail in database so we don't have to generate it
-        // again next time
-        let thumbnail_key =
-            self.db.create_file_metadata(mxc, width, height, &meta)?;
+        let meta = self
+            .create_thumbnail(
+                thumbnail_token,
+                None,
+                meta.content_type,
+                &thumbnail_bytes,
+            )
+            .await?;
 
-        let path = services().globals.get_media_file(&thumbnail_key);
-        let mut f = File::create(path).await?;
-        f.write_all(&thumbnail_bytes).await?;
-
-        Ok(Some((meta, thumbnail_bytes.clone())))
+        Ok(ThumbnailResult::Data(meta, thumbnail_bytes.clone()))
     }
 }

src/service/media/data.rs

@@ -1,10 +1,10 @@
-use super::{FileMeta, MediaFileKey};
+use super::{FileMeta, MediaFileKey, MxcData};
 use crate::Result;
 
 pub(crate) trait Data: Send + Sync {
     fn create_file_metadata(
         &self,
-        mxc: String,
+        mxc: MxcData<'_>,
         width: u32,
         height: u32,
         meta: &FileMeta,
@@ -12,7 +12,7 @@ pub(crate) trait Data: Send + Sync {
 
     fn search_file_metadata(
         &self,
-        mxc: String,
+        mxc: MxcData<'_>,
         width: u32,
         height: u32,
     ) -> Result<(FileMeta, MediaFileKey)>;

src/utils.rs

@@ -13,11 +13,10 @@ use cmp::Ordering;
 use rand::{prelude::*, rngs::OsRng};
 use ring::digest;
 use ruma::{
-    api::client::error::ErrorKind, canonical_json::try_from_json_map,
-    CanonicalJsonError, CanonicalJsonObject, MxcUri, MxcUriError, OwnedMxcUri,
+    canonical_json::try_from_json_map, CanonicalJsonError, CanonicalJsonObject,
 };
 
-use crate::{Error, Result};
+use crate::Result;
 
 // Hopefully we have a better chat protocol in 530 years
 #[allow(clippy::as_conversions, clippy::cast_possible_truncation)]
@@ -255,57 +254,6 @@ pub(crate) fn dbg_truncate_str(s: &str, mut max_len: usize) -> Cow<'_, str> {
     }
 }
 
-/// Data that makes up an `mxc://` URL.
-#[derive(Debug, Clone)]
-pub(crate) struct MxcData<'a> {
-    pub(crate) server_name: &'a ruma::ServerName,
-    pub(crate) media_id: &'a str,
-}
-
-impl<'a> MxcData<'a> {
-    pub(crate) fn new(
-        server_name: &'a ruma::ServerName,
-        media_id: &'a str,
-    ) -> Result<Self> {
-        if !media_id.bytes().all(|b| {
-            matches!(b,
-                b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b'-' | b'_'
-            )
-        }) {
-            return Err(Error::BadRequest(
-                ErrorKind::InvalidParam,
-                "Invalid MXC media id",
-            ));
-        }
-
-        Ok(Self {
-            server_name,
-            media_id,
-        })
-    }
-}
-
-impl fmt::Display for MxcData<'_> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "mxc://{}/{}", self.server_name, self.media_id)
-    }
-}
-
-impl From<MxcData<'_>> for OwnedMxcUri {
-    fn from(value: MxcData<'_>) -> Self {
-        value.to_string().into()
-    }
-}
-
-impl<'a> TryFrom<&'a MxcUri> for MxcData<'a> {
-    type Error = MxcUriError;
-
-    fn try_from(value: &'a MxcUri) -> Result<Self, Self::Error> {
-        Ok(Self::new(value.server_name()?, value.media_id()?)
-            .expect("validated MxcUri should always be valid MxcData"))
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use crate::utils::dbg_truncate_str;