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)")]
@@ -154,7 +157,7 @@ pub async fn migrate(
     .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,9 +169,9 @@ pub async fn migrate(
 
     span.pb_set_message("migrating access tokens");
     span.pb_inc(1);
-    migrate_unrefreshable_access_tokens(
+    let (mas, devices_to_compat_sessions) = migrate_unrefreshable_access_tokens(
         &mut synapse,
-        &mut mas,
+        mas,
         counts
             .access_tokens
             .try_into()
@@ -177,7 +180,7 @@ pub async fn migrate(
         clock,
         rng,
         &migrated_users.user_localparts_to_uuid,
-        &mut devices_to_compat_sessions,
+        devices_to_compat_sessions,
     )
     .await?;
 
@@ -244,9 +247,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 +299,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")?;
 
@@ -611,22 +613,22 @@ async fn migrate_devices(
 #[allow(clippy::too_many_arguments)]
 async fn migrate_unrefreshable_access_tokens(
     synapse: &mut SynapseReader<'_>,
-    mas: &mut MasWriter,
+    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> {
+    devices: HashMap<(Uuid, CompactString), Uuid>,
+) -> Result<(MasWriter, 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 write_buffer = MasWriteBuffer::new(&mas, MasWriter::write_compat_access_tokens);
     let mut deviceless_session_write_buffer =
-        MasWriteBuffer::new(mas, MasWriter::write_compat_sessions);
+        MasWriteBuffer::new(&mas, MasWriter::write_compat_sessions);
 
     while let Some(token_res) = token_stream.next().await {
         let SynapseAccessToken {
@@ -672,7 +674,7 @@ async fn migrate_unrefreshable_access_tokens(
 
             deviceless_session_write_buffer
                 .write(
-                    mas,
+                    &mut mas,
                     MasNewCompatSession {
                         session_id: deviceless_session_id,
                         user_id,
@@ -696,7 +698,7 @@ async fn migrate_unrefreshable_access_tokens(
         // TODO skip access tokens for deactivated users
         write_buffer
             .write(
-                mas,
+                &mut mas,
                 MasNewCompatAccessToken {
                     token_id,
                     session_id,
@@ -710,11 +712,11 @@ async fn migrate_unrefreshable_access_tokens(
     }
 
     write_buffer
-        .finish(mas)
+        .finish(&mut mas)
         .await
         .into_mas("writing compat access tokens")?;
     deviceless_session_write_buffer
-        .finish(mas)
+        .finish(&mut mas)
         .await
         .into_mas("writing deviceless compat sessions")?;
 
@@ -723,7 +725,7 @@ async fn migrate_unrefreshable_access_tokens(
         Instant::now().duration_since(start).as_secs_f64()
     );
 
-    Ok(())
+    Ok((mas, devices_to_compat_sessions))
 }
 
 /// Migrates (access token, refresh token) pairs.