For performance, switch to a row count estimate for users and devices

Author: reivilibre

crates/syn2mas/src/migration.rs

@@ -167,10 +167,11 @@ pub async fn migrate(
     // `(MAS user_id, device_id)` mapped to `compat_session` ULID
     let mut devices_to_compat_sessions: HashMap<(Uuid, CompactString), Uuid> =
         HashMap::with_capacity(
-            counts
-                .devices
-                .try_into()
-                .expect("More than usize::MAX devices — unable to handle this many!"),
+            usize::try_from(counts.devices)
+                .expect("More than usize::MAX devices — unable to handle this many!")
+            // Oversize the capacity, because the count is only an estimate and
+            // we would like to avoid a reallocation
+            * 9 / 8,
         );
 
     span.pb_set_message("migrating access tokens");
@@ -258,8 +259,9 @@ async fn migrate_users(
     let mut user_buffer = MasWriteBuffer::new(mas, MasWriter::write_users);
     let mut password_buffer = MasWriteBuffer::new(mas, MasWriter::write_passwords);
     let mut users_stream = pin!(synapse.read_users());
-    // TODO is 1:1 capacity enough for a hashmap?
-    let mut user_localparts_to_uuid = HashMap::with_capacity(user_count_hint);
+    // Oversize the capacity, because the count is only an estimate and
+    // we would like to avoid a reallocation
+    let mut user_localparts_to_uuid = HashMap::with_capacity(user_count_hint * 9 / 8);
     let mut synapse_admins = HashSet::new();
 
     while let Some(user_res) = users_stream.next().await {

crates/syn2mas/src/synapse_reader/mod.rs

@@ -338,25 +338,28 @@ impl<'conn> SynapseReader<'conn> {
     ///
     /// - An underlying database error
     pub async fn count_rows(&mut self) -> Result<SynapseRowCounts, Error> {
+        // We don't get to filter out application service users by using this estimate,
+        // which is a shame, but on a large database this is way faster.
+        // On matrix.org, counting users and devices properly takes around 1m10s,
+        // which is unnecessary extra downtime during the migration, just to
+        // show a more accurate progress bar and size a hash map accurately.
         let users: i64 = sqlx::query_scalar(
             "
-            SELECT COUNT(1) FROM users
-            WHERE appservice_id IS NULL
+            SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'users'::regclass;
             ",
         )
         .fetch_one(&mut *self.txn)
         .await
-        .into_database("counting Synapse users")?;
+        .into_database("estimating count of users")?;
 
         let devices = sqlx::query_scalar(
             "
-            SELECT COUNT(1) FROM devices
-            WHERE NOT hidden
+            SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'devices'::regclass;
             ",
         )
         .fetch_one(&mut *self.txn)
         .await
-        .into_database("counting Synapse devices")?;
+        .into_database("estimating count of devices")?;
 
         // For other rows, we don't particularly care about the number except for
         // progress bars, so retrieve a fast estimate from the statistics system