Read and write users in different tasks

This also makes the connection in MasWriter owned

Author: sandhose

Cargo.lock

@@ -149,7 +149,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/cli/src/commands/syn2mas.rs

@@ -27,6 +27,7 @@ futures-util = "0.3.30"
 rustc-hash = "2.1.1"
 
 rand.workspace = true
+rand_chacha = "0.3.1"
 uuid = "1.13.1"
 ulid = { workspace = true, features = ["uuid"] }
 

crates/syn2mas/Cargo.toml

@@ -48,9 +48,7 @@ pub enum Error {
 /// - If we can't check whether syn2mas is already in progress on this database
 ///   or not.
 #[tracing::instrument(name = "syn2mas.mas_pre_migration_checks", 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/checks.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/locking.rs

@@ -237,8 +237,8 @@ impl FinishCheckerHandle {
     }
 }
 
-pub struct MasWriter<'c> {
-    conn: LockedMasDatabase<'c>,
+pub struct MasWriter {
+    conn: LockedMasDatabase,
     writer_pool: WriterConnectionPool,
 
     indices_to_restore: Vec<IndexDescription>,
@@ -378,7 +378,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
@@ -389,7 +389,7 @@ impl<'conn> MasWriter<'conn> {
     #[allow(clippy::missing_panics_doc)] // not real
     #[tracing::instrument(name = "syn2mas.mas_writer.new", 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,
@@ -544,7 +544,7 @@ impl<'conn> MasWriter<'conn> {
 
     #[tracing::instrument(skip_all, fields(indicatif.pb_show))]
     async fn restore_indices(
-        conn: &mut LockedMasDatabase<'_>,
+        conn: &mut LockedMasDatabase,
         indices_to_restore: &[IndexDescription],
         constraints_to_restore: &[ConstraintDescription],
     ) -> Result<(), Error> {
@@ -577,7 +577,7 @@ impl<'conn> MasWriter<'conn> {
     ///
     /// - 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> {
         self.write_buffer_finish_checker.check_all_finished()?;
 
         // Commit all writer transactions to the database.
@@ -614,12 +614,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.
@@ -1087,33 +1088,33 @@ const WRITE_BUFFER_BATCH_SIZE: usize = 4096;
 
 /// A function that can accept and flush buffers from a `MasWriteBuffer`.
 /// Intended uses are the methods on `MasWriter` such as `write_users`.
-type WriteBufferFlusher<'conn, T> =
-    for<'a> fn(&'a mut MasWriter<'conn>, Vec<T>) -> BoxFuture<'a, Result<(), Error>>;
+type WriteBufferFlusher<T> =
+    for<'a> fn(&'a mut MasWriter, Vec<T>) -> BoxFuture<'a, Result<(), Error>>;
 
 /// A buffer for writing rows to the MAS database.
 /// Generic over the type of rows.
-pub struct MasWriteBuffer<'conn, T> {
+pub struct MasWriteBuffer<T> {
     rows: Vec<T>,
-    flusher: WriteBufferFlusher<'conn, T>,
+    flusher: WriteBufferFlusher<T>,
     finish_checker_handle: FinishCheckerHandle,
 }
 
-impl<'conn, T> MasWriteBuffer<'conn, T> {
-    pub fn new(writer: &MasWriter, flusher: WriteBufferFlusher<'conn, T>) -> Self {
+impl<T> MasWriteBuffer<T> {
+    pub fn new(writer: &MasWriter, flusher: WriteBufferFlusher<T>) -> Self {
         MasWriteBuffer {
             rows: Vec::with_capacity(WRITE_BUFFER_BATCH_SIZE),
             flusher,
             finish_checker_handle: writer.write_buffer_finish_checker.handle(),
         }
     }
 
-    pub async fn finish(mut self, writer: &mut MasWriter<'conn>) -> Result<(), Error> {
+    pub async fn finish(mut self, writer: &mut MasWriter) -> Result<(), Error> {
         self.flush(writer).await?;
         self.finish_checker_handle.declare_finished();
         Ok(())
     }
 
-    pub async fn flush(&mut self, writer: &mut MasWriter<'conn>) -> Result<(), Error> {
+    pub async fn flush(&mut self, writer: &mut MasWriter) -> Result<(), Error> {
         if self.rows.is_empty() {
             return Ok(());
         }
@@ -1123,7 +1124,7 @@ impl<'conn, T> MasWriteBuffer<'conn, T> {
         Ok(())
     }
 
-    pub async fn write(&mut self, writer: &mut MasWriter<'conn>, row: T) -> Result<(), Error> {
+    pub async fn write(&mut self, writer: &mut MasWriter, row: T) -> Result<(), Error> {
         self.rows.push(row);
         if self.rows.len() >= WRITE_BUFFER_BATCH_SIZE {
             self.flush(writer).await?;
@@ -1235,10 +1236,7 @@ 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, main_conn: PgConnection) -> MasWriter {
         let mut writer_conns = Vec::new();
         for _ in 0..2 {
             writer_conns.push(
@@ -1260,8 +1258,8 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1275,7 +1273,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);
     }
@@ -1285,8 +1283,8 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1309,16 +1307,16 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1342,7 +1340,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);
     }
@@ -1351,8 +1349,8 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1376,7 +1374,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);
     }
@@ -1386,8 +1384,8 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1412,16 +1410,16 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1450,16 +1448,16 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1499,7 +1497,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);
     }
@@ -1508,8 +1506,8 @@ 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 conn = pool.acquire().await.unwrap().detach();
+        let mut writer = make_mas_writer(&pool, conn).await;
 
         writer
             .write_users(vec![MasNewUser {
@@ -1560,7 +1558,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);
     }