Pipe stream directly in the channel

Author: sandhose

Cargo.lock

@@ -18,6 +18,7 @@ serde.workspace = true
 thiserror.workspace = true
 thiserror-ext.workspace = true
 tokio.workspace = true
+tokio-util.workspace = true
 sqlx.workspace = true
 chrono.workspace = true
 compact_str.workspace = true

crates/syn2mas/Cargo.toml

@@ -20,11 +20,12 @@ use std::{pin::pin, time::Instant};
 
 use chrono::{DateTime, Utc};
 use compact_str::CompactString;
-use futures_util::StreamExt as _;
+use futures_util::{SinkExt, StreamExt as _, TryStreamExt as _};
 use mas_storage::Clock;
 use rand::{RngCore, SeedableRng};
 use thiserror::Error;
 use thiserror_ext::ContextInto;
+use tokio_util::sync::PollSender;
 use tracing::{info, Level, Span};
 use tracing_indicatif::span_ext::IndicatifSpanExt;
 use ulid::Ulid;
@@ -61,6 +62,8 @@ pub enum Error {
         source: ExtractLocalpartError,
         user: FullUserId,
     },
+    #[error("channel closed")]
+    ChannelClosed,
     #[error("user {user} was not found for migration but a row in {table} was found for them")]
     MissingUserFromDependentTable { table: String, user: FullUserId },
     #[error("missing a mapping for the auth provider with ID {synapse_id:?} (used by {user} and maybe other users)")]
@@ -121,6 +124,7 @@ pub async fn migrate(
         rng,
     )
     .await?;
+    let user_localparts_to_uuid = migrated_users.user_localparts_to_uuid;
 
     span.pb_set_message("migrating threepids");
     span.pb_inc(1);
@@ -133,7 +137,7 @@ pub async fn migrate(
             .expect("More than u64::MAX threepids — unable to handle this many!"),
         server_name,
         rng,
-        &migrated_users.user_localparts_to_uuid,
+        &user_localparts_to_uuid,
     )
     .await?;
 
@@ -148,13 +152,13 @@ pub async fn migrate(
             .expect("More than u64::MAX external IDs — unable to handle this many!"),
         server_name,
         rng,
-        &migrated_users.user_localparts_to_uuid,
+        &user_localparts_to_uuid,
         provider_id_mapping,
     )
     .await?;
 
     // `(MAS user_id, device_id)` mapped to `compat_session` ULID
-    let mut devices_to_compat_sessions: HashMap<(Uuid, CompactString), Uuid> =
+    let devices_to_compat_sessions: HashMap<(Uuid, CompactString), Uuid> =
         HashMap::with_capacity_and_hasher(
             usize::try_from(counts.devices)
                 .expect("More than usize::MAX devices — unable to handle this many!")
@@ -166,20 +170,21 @@ pub async fn migrate(
 
     span.pb_set_message("migrating access tokens");
     span.pb_inc(1);
-    migrate_unrefreshable_access_tokens(
-        &mut synapse,
-        &mut mas,
-        counts
-            .access_tokens
-            .try_into()
-            .expect("More than u64::MAX access tokens — unable to handle this many!"),
-        server_name,
-        clock,
-        rng,
-        &migrated_users.user_localparts_to_uuid,
-        &mut devices_to_compat_sessions,
-    )
-    .await?;
+    let (mut mas, user_localparts_to_uuid, mut devices_to_compat_sessions) =
+        migrate_unrefreshable_access_tokens(
+            &mut synapse,
+            mas,
+            counts
+                .access_tokens
+                .try_into()
+                .expect("More than u64::MAX access tokens — unable to handle this many!"),
+            server_name,
+            clock,
+            rng,
+            user_localparts_to_uuid,
+            devices_to_compat_sessions,
+        )
+        .await?;
 
     span.pb_set_message("migrating refresh tokens");
     span.pb_inc(1);
@@ -193,7 +198,7 @@ pub async fn migrate(
         server_name,
         clock,
         rng,
-        &migrated_users.user_localparts_to_uuid,
+        &user_localparts_to_uuid,
         &mut devices_to_compat_sessions,
     )
     .await?;
@@ -209,7 +214,7 @@ pub async fn migrate(
             .expect("More than u64::MAX devices — unable to handle this many!"),
         server_name,
         rng,
-        &migrated_users.user_localparts_to_uuid,
+        &user_localparts_to_uuid,
         &mut devices_to_compat_sessions,
         &migrated_users.synapse_admins,
     )
@@ -244,9 +249,6 @@ async fn migrate_users(
 ) -> Result<(UsersMigrated, MasWriter), Error> {
     let start = Instant::now();
 
-    let mut users_stream = pin!(synapse
-        .read_users()
-        .with_progress_bar(user_count_hint as u64, 10_000));
     let (tx, mut rx) = tokio::sync::mpsc::channel(1024 * 1024);
 
     let mut rng = rand_chacha::ChaCha8Rng::from_rng(rng).expect("failed to seed rng");
@@ -299,10 +301,12 @@ async fn migrate_users(
         Ok((synapse_admins, user_localparts_to_uuid, mas))
     });
 
-    while let Some(user_res) = users_stream.next().await {
-        let user = user_res.into_synapse("reading user")?;
-        tx.send(user).await.expect("failed to send in channel");
-    }
+    synapse
+        .read_users()
+        .with_progress_bar(user_count_hint as u64, 10_000)
+        .map_err(|e| e.into_synapse("reading users"))
+        .forward(PollSender::new(tx).sink_map_err(|_| Error::ChannelClosed))
+        .await?;
 
     let (synapse_admins, user_localparts_to_uuid, mas) = task.await.expect("task panicked")?;
 
@@ -608,122 +612,144 @@ async fn migrate_devices(
 /// Migrates unrefreshable access tokens (those without an associated refresh
 /// token). Some of these may be deviceless.
 #[tracing::instrument(skip_all, fields(indicatif.pb_show), level = Level::INFO)]
-#[allow(clippy::too_many_arguments)]
+#[allow(clippy::too_many_arguments, clippy::type_complexity)]
 async fn migrate_unrefreshable_access_tokens(
     synapse: &mut SynapseReader<'_>,
-    mas: &mut MasWriter,
+    mut mas: MasWriter,
     count_hint: u64,
     server_name: &str,
     clock: &dyn Clock,
     rng: &mut impl RngCore,
-    user_localparts_to_uuid: &HashMap<CompactString, Uuid>,
-    devices: &mut HashMap<(Uuid, CompactString), Uuid>,
-) -> Result<(), Error> {
+    user_localparts_to_uuid: HashMap<CompactString, Uuid>,
+    mut devices: HashMap<(Uuid, CompactString), Uuid>,
+) -> Result<
+    (
+        MasWriter,
+        HashMap<CompactString, Uuid>,
+        HashMap<(Uuid, CompactString), Uuid>,
+    ),
+    Error,
+> {
     let start = Instant::now();
 
-    let mut token_stream = pin!(synapse
-        .read_unrefreshable_access_tokens()
-        .with_progress_bar(count_hint, 10_000));
-    let mut write_buffer = MasWriteBuffer::new(mas, MasWriter::write_compat_access_tokens);
-    let mut deviceless_session_write_buffer =
-        MasWriteBuffer::new(mas, MasWriter::write_compat_sessions);
-
-    while let Some(token_res) = token_stream.next().await {
-        let SynapseAccessToken {
-            user_id: synapse_user_id,
-            device_id,
-            token,
-            valid_until_ms,
-            last_validated,
-        } = token_res.into_synapse("reading Synapse access token")?;
+    let (tx, mut rx) = tokio::sync::mpsc::channel(1024 * 1024);
 
-        let username = synapse_user_id
-            .extract_localpart(server_name)
-            .into_extract_localpart(synapse_user_id.clone())?
-            .to_owned();
-        let Some(user_id) = 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: "access_tokens".to_owned(),
-                user: synapse_user_id,
-            });
-        };
+    let now = clock.now();
+    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_access_tokens);
+        let mut deviceless_session_write_buffer =
+            MasWriteBuffer::new(&mas, MasWriter::write_compat_sessions);
+
+        while let Some(token) = rx.recv().await {
+            let SynapseAccessToken {
+                user_id: synapse_user_id,
+                device_id,
+                token,
+                valid_until_ms,
+                last_validated,
+            } = token;
+
+            let username = synapse_user_id
+                .extract_localpart(&server_name)
+                .into_extract_localpart(synapse_user_id.clone())?
+                .to_owned();
+            let Some(user_id) = 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: "access_tokens".to_owned(),
+                    user: synapse_user_id,
+                });
+            };
+
+            // It's not always accurate, but last_validated is *often* the creation time of
+            // the device If we don't have one, then use the current time as a
+            // fallback.
+            let created_at = last_validated.map_or_else(|| now, DateTime::from);
+
+            let session_id = if let Some(device_id) = device_id {
+                // Use the existing device_id if this is the second token for a device
+                *devices
+                    .entry((user_id, CompactString::new(&device_id)))
+                    .or_insert_with(|| {
+                        Uuid::from(Ulid::from_datetime_with_source(created_at.into(), &mut rng))
+                    })
+            } else {
+                // If this is a deviceless access token, create a deviceless compat session
+                // for it (since otherwise we won't create one whilst migrating devices)
+                let deviceless_session_id =
+                    Uuid::from(Ulid::from_datetime_with_source(created_at.into(), &mut rng));
+
+                deviceless_session_write_buffer
+                    .write(
+                        &mut mas,
+                        MasNewCompatSession {
+                            session_id: deviceless_session_id,
+                            user_id,
+                            device_id: None,
+                            human_name: None,
+                            created_at,
+                            is_synapse_admin: false,
+                            last_active_at: None,
+                            last_active_ip: None,
+                            user_agent: None,
+                        },
+                    )
+                    .await
+                    .into_mas("failed to write deviceless compat sessions")?;
 
-        // It's not always accurate, but last_validated is *often* the creation time of
-        // the device If we don't have one, then use the current time as a
-        // fallback.
-        let created_at = last_validated.map_or_else(|| clock.now(), DateTime::from);
+                deviceless_session_id
+            };
 
-        let session_id = if let Some(device_id) = device_id {
-            // Use the existing device_id if this is the second token for a device
-            *devices
-                .entry((user_id, CompactString::new(&device_id)))
-                .or_insert_with(|| {
-                    Uuid::from(Ulid::from_datetime_with_source(created_at.into(), rng))
-                })
-        } else {
-            // If this is a deviceless access token, create a deviceless compat session
-            // for it (since otherwise we won't create one whilst migrating devices)
-            let deviceless_session_id =
-                Uuid::from(Ulid::from_datetime_with_source(created_at.into(), rng));
+            let token_id = Uuid::from(Ulid::from_datetime_with_source(created_at.into(), &mut rng));
 
-            deviceless_session_write_buffer
+            // TODO skip access tokens for deactivated users
+            write_buffer
                 .write(
-                    mas,
-                    MasNewCompatSession {
-                        session_id: deviceless_session_id,
-                        user_id,
-                        device_id: None,
-                        human_name: None,
+                    &mut mas,
+                    MasNewCompatAccessToken {
+                        token_id,
+                        session_id,
+                        access_token: token,
                         created_at,
-                        is_synapse_admin: false,
-                        last_active_at: None,
-                        last_active_ip: None,
-                        user_agent: None,
+                        expires_at: valid_until_ms.map(DateTime::from),
                     },
                 )
                 .await
-                .into_mas("failed to write deviceless compat sessions")?;
-
-            deviceless_session_id
-        };
-
-        let token_id = Uuid::from(Ulid::from_datetime_with_source(created_at.into(), rng));
+                .into_mas("writing compat access tokens")?;
+        }
 
-        // TODO skip access tokens for deactivated users
         write_buffer
-            .write(
-                mas,
-                MasNewCompatAccessToken {
-                    token_id,
-                    session_id,
-                    access_token: token,
-                    created_at,
-                    expires_at: valid_until_ms.map(DateTime::from),
-                },
-            )
+            .finish(&mut mas)
             .await
             .into_mas("writing compat access tokens")?;
-    }
+        deviceless_session_write_buffer
+            .finish(&mut mas)
+            .await
+            .into_mas("writing deviceless compat sessions")?;
 
-    write_buffer
-        .finish(mas)
-        .await
-        .into_mas("writing compat access tokens")?;
-    deviceless_session_write_buffer
-        .finish(mas)
-        .await
-        .into_mas("writing deviceless compat sessions")?;
+        Ok((mas, user_localparts_to_uuid, devices))
+    });
+
+    synapse
+        .read_unrefreshable_access_tokens()
+        .with_progress_bar(count_hint, 10_000)
+        .map_err(|e| e.into_synapse("reading tokens"))
+        .forward(PollSender::new(tx).sink_map_err(|_| Error::ChannelClosed))
+        .await?;
+
+    let (mas, user_localparts_to_uuid, devices) = task.await.expect("task crashed")?;
 
     info!(
         "non-refreshable access tokens migrated in {:.1}s",
         Instant::now().duration_since(start).as_secs_f64()
     );
 
-    Ok(())
+    Ok((mas, user_localparts_to_uuid, devices))
 }
 
 /// Migrates (access token, refresh token) pairs.