Make the MAS connection owned in the locked database struct

Author: sandhose

crates/cli/src/commands/syn2mas.rs

@@ -142,7 +142,7 @@ impl Options {
             .await?;
         }
 
-        let Either::Left(mut mas_connection) = LockedMasDatabase::try_new(&mut mas_connection)
+        let Either::Left(mut mas_connection) = LockedMasDatabase::try_new(mas_connection)
             .await
             .context("failed to issue query to lock database")?
         else {

crates/syn2mas/src/mas_writer/checks.rs

@@ -47,9 +47,7 @@ pub enum Error {
 /// - If we can't check whether syn2mas is already in progress on this database
 ///   or not.
 #[tracing::instrument(skip_all)]
-pub async fn mas_pre_migration_checks<'a>(
-    mas_connection: &mut LockedMasDatabase<'a>,
-) -> Result<(), Error> {
+pub async fn mas_pre_migration_checks(mas_connection: &mut LockedMasDatabase) -> Result<(), Error> {
     if is_syn2mas_in_progress(mas_connection.as_mut())
         .await
         .map_err(Error::UnableToCheckInProgress)?

crates/syn2mas/src/mas_writer/locking.rs

@@ -15,11 +15,11 @@ static SYN2MAS_ADVISORY_LOCK: LazyLock<PgAdvisoryLock> =
 
 /// A wrapper around a Postgres connection which holds a session-wide advisory
 /// lock preventing concurrent access by other syn2mas instances.
-pub struct LockedMasDatabase<'conn> {
-    inner: PgAdvisoryLockGuard<'static, &'conn mut PgConnection>,
+pub struct LockedMasDatabase {
+    inner: PgAdvisoryLockGuard<'static, PgConnection>,
 }
 
-impl<'conn> LockedMasDatabase<'conn> {
+impl LockedMasDatabase {
     /// Attempts to lock the MAS database against concurrent access by other
     /// syn2mas instances.
     ///
@@ -31,8 +31,8 @@ impl<'conn> LockedMasDatabase<'conn> {
     ///
     /// Errors are returned for underlying database errors.
     pub async fn try_new(
-        mas_connection: &'conn mut PgConnection,
-    ) -> Result<Either<Self, &'conn mut PgConnection>, sqlx::Error> {
+        mas_connection: PgConnection,
+    ) -> Result<Either<Self, PgConnection>, sqlx::Error> {
         SYN2MAS_ADVISORY_LOCK
             .try_acquire(mas_connection)
             .await
@@ -48,12 +48,12 @@ impl<'conn> LockedMasDatabase<'conn> {
     /// # Errors
     ///
     /// Errors are returned for underlying database errors.
-    pub async fn unlock(self) -> Result<&'conn mut PgConnection, sqlx::Error> {
+    pub async fn unlock(self) -> Result<PgConnection, sqlx::Error> {
         self.inner.release_now().await
     }
 }
 
-impl AsMut<PgConnection> for LockedMasDatabase<'_> {
+impl AsMut<PgConnection> for LockedMasDatabase {
     fn as_mut(&mut self) -> &mut PgConnection {
         self.inner.as_mut()
     }

crates/syn2mas/src/mas_writer/mod.rs

@@ -186,7 +186,9 @@ impl WriterConnectionPool {
 }
 
 pub struct MasWriter<'c> {
-    conn: LockedMasDatabase<'c>,
+    conn: LockedMasDatabase,
+    // Temporary phantom data, so that we don't remove the lifetime parameter yet
+    phantom: std::marker::PhantomData<&'c ()>,
     writer_pool: WriterConnectionPool,
 
     indices_to_restore: Vec<IndexDescription>,
@@ -324,7 +326,7 @@ pub async fn is_syn2mas_in_progress(conn: &mut PgConnection) -> Result<bool, Err
     }
 }
 
-impl<'conn> MasWriter<'conn> {
+impl MasWriter<'_> {
     /// Creates a new MAS writer.
     ///
     /// # Errors
@@ -335,7 +337,7 @@ impl<'conn> MasWriter<'conn> {
     #[allow(clippy::missing_panics_doc)] // not real
     #[tracing::instrument(skip_all)]
     pub async fn new(
-        mut conn: LockedMasDatabase<'conn>,
+        mut conn: LockedMasDatabase,
         mut writer_connections: Vec<PgConnection>,
     ) -> Result<Self, Error> {
         // Given that we don't have any concurrent transactions here,
@@ -446,6 +448,7 @@ impl<'conn> MasWriter<'conn> {
 
         Ok(Self {
             conn,
+            phantom: std::marker::PhantomData,
             writer_pool: WriterConnectionPool::new(writer_connections),
             indices_to_restore,
             constraints_to_restore,
@@ -488,7 +491,7 @@ impl<'conn> MasWriter<'conn> {
     }
 
     async fn restore_indices(
-        conn: &mut LockedMasDatabase<'_>,
+        conn: &mut LockedMasDatabase,
         indices_to_restore: &[IndexDescription],
         constraints_to_restore: &[ConstraintDescription],
     ) -> Result<(), Error> {
@@ -507,14 +510,15 @@ impl<'conn> MasWriter<'conn> {
     }
 
     /// Finish writing to the MAS database, flushing and committing all changes.
+    /// It returns the unlocked underlying connection.
     ///
     /// # Errors
     ///
     /// Errors are returned in the following conditions:
     ///
     /// - If the database connection experiences an error.
     #[tracing::instrument(skip_all)]
-    pub async fn finish(mut self) -> Result<(), Error> {
+    pub async fn finish(mut self) -> Result<PgConnection, Error> {
         // Commit all writer transactions to the database.
         self.writer_pool
             .finish()
@@ -549,12 +553,13 @@ impl<'conn> MasWriter<'conn> {
             .await
             .into_database("ending MAS transaction")?;
 
-        self.conn
+        let conn = self
+            .conn
             .unlock()
             .await
             .into_database("could not unlock MAS database")?;
 
-        Ok(())
+        Ok(conn)
     }
 
     /// Write a batch of users to the database.
@@ -1180,10 +1185,8 @@ mod test {
     /// Runs some code with a `MasWriter`.
     ///
     /// The callback is responsible for `finish`ing the `MasWriter`.
-    async fn make_mas_writer<'conn>(
-        pool: &PgPool,
-        main_conn: &'conn mut PgConnection,
-    ) -> MasWriter<'conn> {
+    async fn make_mas_writer(pool: &PgPool) -> MasWriter<'static> {
+        let main_conn = pool.acquire().await.unwrap().detach();
         let mut writer_conns = Vec::new();
         for _ in 0..2 {
             writer_conns.push(
@@ -1205,8 +1208,7 @@ mod test {
     /// Tests writing a single user, without a password.
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
     async fn test_write_user(pool: PgPool) {
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1220,7 +1222,7 @@ mod test {
             .await
             .expect("failed to write user");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }
@@ -1230,8 +1232,7 @@ mod test {
     async fn test_write_user_with_password(pool: PgPool) {
         const USER_ID: Uuid = Uuid::from_u128(1u128);
 
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1254,16 +1255,15 @@ mod test {
             .await
             .expect("failed to write password");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }
 
     /// Tests writing a single user, with an e-mail address associated.
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
     async fn test_write_user_with_email(pool: PgPool) {
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1287,7 +1287,7 @@ mod test {
             .await
             .expect("failed to write e-mail");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }
@@ -1296,8 +1296,7 @@ mod test {
     /// associated.
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
     async fn test_write_user_with_unsupported_threepid(pool: PgPool) {
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1321,7 +1320,7 @@ mod test {
             .await
             .expect("failed to write phone number (unsupported threepid)");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }
@@ -1331,8 +1330,7 @@ mod test {
     /// real migration, this is done by running a provider sync first.
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR", fixtures("upstream_provider"))]
     async fn test_write_user_with_upstream_provider_link(pool: PgPool) {
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1357,16 +1355,15 @@ mod test {
             .await
             .expect("failed to write link");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }
 
     /// Tests writing a single user, with a device (compat session).
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
     async fn test_write_user_with_device(pool: PgPool) {
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1395,16 +1392,15 @@ mod test {
             .await
             .expect("failed to write compat session");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }
 
     /// Tests writing a single user, with a device and an access token.
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
     async fn test_write_user_with_access_token(pool: PgPool) {
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1444,7 +1440,7 @@ mod test {
             .await
             .expect("failed to write access token");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }
@@ -1453,8 +1449,7 @@ mod test {
     /// refresh token.
     #[sqlx::test(migrator = "mas_storage_pg::MIGRATOR")]
     async fn test_write_user_with_refresh_token(pool: PgPool) {
-        let mut conn = pool.acquire().await.unwrap();
-        let mut writer = make_mas_writer(&pool, &mut conn).await;
+        let mut writer = make_mas_writer(&pool).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1505,7 +1500,7 @@ mod test {
             .await
             .expect("failed to write refresh token");
 
-        writer.finish().await.expect("failed to finish MasWriter");
+        let mut conn = writer.finish().await.expect("failed to finish MasWriter");
 
         assert_db_snapshot!(&mut conn);
     }