Don't hold db conns when creating a device on the compat login API

Author: sandhose

crates/handlers/src/compat/login.rs

@@ -14,11 +14,12 @@ use mas_axum_utils::record_error;
 use mas_data_model::{CompatSession, CompatSsoLoginState, Device, SiteConfig, TokenType, User};
 use mas_matrix::HomeserverConnection;
 use mas_storage::{
-    BoxClock, BoxRepository, BoxRng, Clock, RepositoryAccess,
+    BoxClock, BoxRepository, BoxRepositoryFactory, BoxRng, Clock, RepositoryAccess,
     compat::{
         CompatAccessTokenRepository, CompatRefreshTokenRepository, CompatSessionRepository,
         CompatSsoLoginRepository,
     },
+    queue::{QueueJobRepositoryExt as _, SyncDevicesJob},
     user::{UserPasswordRepository, UserRepository},
 };
 use opentelemetry::{Key, KeyValue, metrics::Counter};
@@ -268,7 +269,7 @@ pub(crate) async fn post(
     mut rng: BoxRng,
     clock: BoxClock,
     State(password_manager): State<PasswordManager>,
-    mut repo: BoxRepository,
+    State(repository_factory): State<BoxRepositoryFactory>,
     activity_tracker: BoundActivityTracker,
     State(homeserver): State<Arc<dyn HomeserverConnection>>,
     State(site_config): State<SiteConfig>,
@@ -279,6 +280,7 @@ pub(crate) async fn post(
 ) -> Result<impl IntoResponse, RouteError> {
     let user_agent = user_agent.map(|ua| ua.as_str().to_owned());
     let login_type = input.credentials.login_type();
+    let mut repo = repository_factory.create().await?;
     let (mut session, user) = match (password_manager.is_enabled(), input.credentials) {
         (
             true,
@@ -301,15 +303,17 @@ pub(crate) async fn post(
                 }
             };
 
+            // Try getting the localpart out of the MXID
+            let username = homeserver.localpart(&user).unwrap_or(&user);
+
             user_password_login(
                 &mut rng,
                 &clock,
                 &password_manager,
                 &limiter,
                 requester,
                 &mut repo,
-                &homeserver,
-                user,
+                username,
                 password,
                 input.device_id, // TODO check for validity
                 input.initial_device_display_name,
@@ -322,7 +326,6 @@ pub(crate) async fn post(
                 &mut rng,
                 &clock,
                 &mut repo,
-                &homeserver,
                 &token,
                 input.device_id,
                 input.initial_device_display_name,
@@ -368,12 +371,53 @@ pub(crate) async fn post(
         None
     };
 
+    // Ideally, we'd keep the lock whilst we actually create the device, but we
+    // really want to stop holding the transaction while we talk to the
+    // homeserver.
+    //
+    // In practice, this is fine, because:
+    // - the session exists after we commited the transaction, so a sync job won't
+    //   try to delete it
+    // - we've acquired a lock on the user before creating the session, meaning
+    //   we've made sure that sync jobs finished before we create the new session
+    // - we're in the read-commited isolation level, which means the sync will see
+    //   what we've committed and won't try to delete the session once we release
+    //   the lock
     repo.save().await?;
 
     activity_tracker
         .record_compat_session(&clock, &session)
         .await;
 
+    // This session will have for sure the device on it, both methods create a
+    // device
+    let Some(device) = &session.device else {
+        unreachable!()
+    };
+
+    // Now we can create the device on the homeserver, without holding the
+    // transaction
+    if let Err(err) = homeserver
+        .create_device(&user_id, device.as_str(), session.human_name.as_deref())
+        .await
+    {
+        // Something went wrong, let's end this session and schedule a device sync
+        let mut repo = repository_factory.create().await?;
+        let session = repo.compat_session().finish(&clock, session).await?;
+
+        repo.queue_job()
+            .schedule_job(
+                &mut rng,
+                &clock,
+                SyncDevicesJob::new_for_id(session.user_id),
+            )
+            .await?;
+
+        repo.save().await?;
+
+        return Err(RouteError::ProvisionDeviceFailed(err));
+    }
+
     LOGIN_COUNTER.add(
         1,
         &[
@@ -395,7 +439,6 @@ async fn token_login(
     rng: &mut (dyn RngCore + Send),
     clock: &dyn Clock,
     repo: &mut BoxRepository,
-    homeserver: &dyn HomeserverConnection,
     token: &str,
     requested_device_id: Option<String>,
     initial_device_display_name: Option<String>,
@@ -461,7 +504,8 @@ async fn token_login(
         return Err(RouteError::InvalidLoginToken);
     }
 
-    // Lock the user sync to make sure we don't get into a race condition
+    // We're about to create a device, let's explicitly acquire a lock, so that
+    // any concurrent sync will read after we've committed
     repo.user()
         .acquire_lock_for_sync(&browser_session.user)
         .await?;
@@ -471,20 +515,14 @@ async fn token_login(
     } else {
         Device::generate(rng)
     };
-    let mxid = homeserver.mxid(&browser_session.user.username);
-    homeserver
-        .create_device(
-            &mxid,
-            device.as_str(),
-            initial_device_display_name.as_deref(),
-        )
-        .await
-        .map_err(RouteError::ProvisionDeviceFailed)?;
 
     repo.app_session()
         .finish_sessions_to_replace_device(clock, &browser_session.user, &device)
         .await?;
 
+    // We first create the session in the database, commit the transaction, then
+    // create it on the homeserver, scheduling a device sync job afterwards to
+    // make sure we don't end up in an inconsistent state.
     let compat_session = repo
         .compat_session()
         .add(
@@ -512,15 +550,11 @@ async fn user_password_login(
     limiter: &Limiter,
     requester: RequesterFingerprint,
     repo: &mut BoxRepository,
-    homeserver: &dyn HomeserverConnection,
-    username: String,
+    username: &str,
     password: String,
     requested_device_id: Option<String>,
     initial_device_display_name: Option<String>,
 ) -> Result<(CompatSession, User), RouteError> {
-    // Try getting the localpart out of the MXID
-    let username = homeserver.localpart(&username).unwrap_or(&username);
-
     // Find the user
     let user = repo
         .user()
@@ -566,25 +600,16 @@ async fn user_password_login(
             .await?;
     }
 
-    // Lock the user sync to make sure we don't get into a race condition
+    // We're about to create a device, let's explicitly acquire a lock, so that
+    // any concurrent sync will read after we've committed
     repo.user().acquire_lock_for_sync(&user).await?;
 
-    let mxid = homeserver.mxid(&user.username);
-
     // Now that the user credentials have been verified, start a new compat session
     let device = if let Some(requested_device_id) = requested_device_id {
         Device::from(requested_device_id)
     } else {
         Device::generate(&mut rng)
     };
-    homeserver
-        .create_device(
-            &mxid,
-            device.as_str(),
-            initial_device_display_name.as_deref(),
-        )
-        .await
-        .map_err(RouteError::ProvisionDeviceFailed)?;
 
     repo.app_session()
         .finish_sessions_to_replace_device(clock, &user, &device)

crates/handlers/src/lib.rs

@@ -42,7 +42,7 @@ use mas_keystore::{Encrypter, Keystore};
 use mas_matrix::HomeserverConnection;
 use mas_policy::Policy;
 use mas_router::{Route, UrlBuilder};
-use mas_storage::{BoxClock, BoxRepository, BoxRng};
+use mas_storage::{BoxClock, BoxRepository, BoxRepositoryFactory, BoxRng};
 use mas_templates::{ErrorContext, NotFoundContext, TemplateContext, Templates};
 use opentelemetry::metrics::Meter;
 use sqlx::PgPool;
@@ -265,6 +265,7 @@ where
     Arc<dyn HomeserverConnection>: FromRef<S>,
     PasswordManager: FromRef<S>,
     Limiter: FromRef<S>,
+    BoxRepositoryFactory: FromRef<S>,
     BoundActivityTracker: FromRequestParts<S>,
     RequesterFingerprint: FromRequestParts<S>,
     BoxRepository: FromRequestParts<S>,

crates/handlers/src/test_utils.rs

@@ -453,6 +453,12 @@ impl FromRef<TestState> for PgPool {
     }
 }
 
+impl FromRef<TestState> for BoxRepositoryFactory {
+    fn from_ref(input: &TestState) -> Self {
+        input.repository_factory.clone().boxed()
+    }
+}
+
 impl FromRef<TestState> for graphql::Schema {
     fn from_ref(input: &TestState) -> Self {
         input.graphql_schema.clone()