Also migrate devices through a stream

Author: sandhose

crates/syn2mas/src/migration.rs

@@ -204,7 +204,7 @@ pub async fn migrate(
 
     span.pb_set_message("migrating devices");
     span.pb_inc(1);
-    let mas = migrate_devices(
+    let (mas, _devices) = migrate_devices(
         &mut synapse,
         mas,
         counts
@@ -214,7 +214,7 @@ pub async fn migrate(
         server_name,
         rng,
         migrated_users,
-        &mut devices_to_compat_sessions,
+        devices_to_compat_sessions,
     )
     .await?;
 
@@ -524,100 +524,116 @@ async fn migrate_devices(
     server_name: &str,
     rng: &mut impl RngCore,
     user_migrated: UsersMigrated,
-    devices: &mut HashMap<(Uuid, CompactString), Uuid>,
-) -> Result<MasWriter, Error> {
+    mut devices: HashMap<(Uuid, CompactString), Uuid>,
+) -> Result<(MasWriter, HashMap<(Uuid, CompactString), Uuid>), Error> {
     let start = Instant::now();
 
-    let mut devices_stream = pin!(synapse.read_devices().with_progress_bar(count_hint, 10_000));
-    let mut write_buffer = MasWriteBuffer::new(&mas, MasWriter::write_compat_sessions);
+    let (tx, mut rx) = tokio::sync::mpsc::channel(1024 * 1024);
 
-    while let Some(device_res) = devices_stream.next().await {
-        let SynapseDevice {
-            user_id: synapse_user_id,
-            device_id,
-            display_name,
-            last_seen,
-            ip,
-            user_agent,
-        } = device_res.into_synapse("reading Synapse device")?;
+    let server_name = server_name.to_owned();
+    let mut rng = rand_chacha::ChaChaRng::from_rng(rng).expect("failed to seed rng");
+    let task = tokio::spawn(async move {
+        let mut write_buffer = MasWriteBuffer::new(&mas, MasWriter::write_compat_sessions);
 
-        let username = synapse_user_id
-            .extract_localpart(server_name)
-            .into_extract_localpart(synapse_user_id.clone())?
-            .to_owned();
-        let Some(user_id) = user_migrated
-            .user_localparts_to_uuid
-            .get(username.as_str())
-            .copied()
-        else {
-            if true || is_likely_appservice(&username) {
-                // HACK can we do anything better
-                continue;
-            }
-            return Err(Error::MissingUserFromDependentTable {
-                table: "devices".to_owned(),
-                user: synapse_user_id,
-            });
-        };
+        while let Some(device) = rx.recv().await {
+            let SynapseDevice {
+                user_id: synapse_user_id,
+                device_id,
+                display_name,
+                last_seen,
+                ip,
+                user_agent,
+            } = device;
 
-        let session_id = *devices
-            .entry((user_id, CompactString::new(&device_id)))
-            .or_insert_with(||
+            let username = synapse_user_id
+                .extract_localpart(&server_name)
+                .into_extract_localpart(synapse_user_id.clone())?
+                .to_owned();
+            let Some(user_id) = user_migrated
+                .user_localparts_to_uuid
+                .get(username.as_str())
+                .copied()
+            else {
+                if true || is_likely_appservice(&username) {
+                    // HACK can we do anything better
+                    continue;
+                }
+                return Err(Error::MissingUserFromDependentTable {
+                    table: "devices".to_owned(),
+                    user: synapse_user_id,
+                });
+            };
+
+            let session_id = *devices
+                .entry((user_id, CompactString::new(&device_id)))
+                .or_insert_with(||
                 // We don't have a creation time for this device (as it has no access token),
                 // so use now as a least-evil fallback.
-                Ulid::with_source(rng).into());
-        let created_at = Ulid::from(session_id).datetime().into();
-
-        // As we're using a real IP type in the MAS database, it is possible
-        // that we encounter invalid IP addresses in the Synapse database.
-        // In that case, we should ignore them, but still log a warning.
-        // One special case: Synapse will record '-' as IP in some cases, we don't want
-        // to log about those
-        let last_active_ip = ip.filter(|ip| ip != "-").and_then(|ip| {
-            ip.parse()
-                .map_err(|e| {
-                    tracing::warn!(
-                        error = &e as &dyn std::error::Error,
-                        mxid = %synapse_user_id,
-                        %device_id,
-                        %ip,
-                        "Failed to parse device IP, ignoring"
-                    );
-                })
-                .ok()
-        });
+                Ulid::with_source(&mut rng).into());
+            let created_at = Ulid::from(session_id).datetime().into();
+
+            // As we're using a real IP type in the MAS database, it is possible
+            // that we encounter invalid IP addresses in the Synapse database.
+            // In that case, we should ignore them, but still log a warning.
+            // One special case: Synapse will record '-' as IP in some cases, we don't want
+            // to log about those
+            let last_active_ip = ip.filter(|ip| ip != "-").and_then(|ip| {
+                ip.parse()
+                    .map_err(|e| {
+                        tracing::warn!(
+                            error = &e as &dyn std::error::Error,
+                            mxid = %synapse_user_id,
+                            %device_id,
+                            %ip,
+                            "Failed to parse device IP, ignoring"
+                        );
+                    })
+                    .ok()
+            });
+
+            // TODO skip access tokens for deactivated users
+            write_buffer
+                .write(
+                    &mut mas,
+                    MasNewCompatSession {
+                        session_id,
+                        user_id,
+                        device_id: Some(device_id),
+                        human_name: display_name,
+                        created_at,
+                        is_synapse_admin: user_migrated.synapse_admins.contains(&user_id),
+                        last_active_at: last_seen.map(DateTime::from),
+                        last_active_ip,
+                        user_agent,
+                    },
+                )
+                .await
+                .into_mas("writing compat sessions")?;
+        }
 
-        // TODO skip access tokens for deactivated users
         write_buffer
-            .write(
-                &mut mas,
-                MasNewCompatSession {
-                    session_id,
-                    user_id,
-                    device_id: Some(device_id),
-                    human_name: display_name,
-                    created_at,
-                    is_synapse_admin: user_migrated.synapse_admins.contains(&user_id),
-                    last_active_at: last_seen.map(DateTime::from),
-                    last_active_ip,
-                    user_agent,
-                },
-            )
+            .finish(&mut mas)
             .await
             .into_mas("writing compat sessions")?;
-    }
 
-    write_buffer
-        .finish(&mut mas)
-        .await
-        .into_mas("writing compat sessions")?;
+        Ok((mas, devices))
+    });
+
+    synapse
+        .read_devices()
+        .with_progress_bar(count_hint, 10_000)
+        .map_err(|e| e.into_synapse("reading devices"))
+        .forward(PollSender::new(tx).sink_map_err(|_| Error::ChannelClosed))
+        .await?;
+
+    let (mas, devices) = task.await.expect("task panicked")?;
 
     info!(
         "devices migrated in {:.1}s",
         Instant::now().duration_since(start).as_secs_f64()
     );
 
-    Ok(mas)
+    Ok((mas, devices))
 }
 
 /// Migrates unrefreshable access tokens (those without an associated refresh