Throttle down the number of times we update the progress bars

Author: sandhose

Cargo.lock

@@ -14,6 +14,7 @@ anyhow.workspace = true
 bitflags.workspace = true
 camino.workspace = true
 figment.workspace = true
+pin-project-lite = "0.2.16"
 serde.workspace = true
 thiserror.workspace = true
 thiserror-ext.workspace = true

crates/syn2mas/Cargo.toml

@@ -7,6 +7,7 @@ mod mas_writer;
 mod synapse_reader;
 
 mod migration;
+mod progress_stream;
 
 pub use self::{
     mas_writer::{MasWriter, checks::mas_pre_migration_checks, locking::LockedMasDatabase},

crates/syn2mas/src/lib.rs

@@ -32,6 +32,7 @@ use crate::{
         MasNewEmailThreepid, MasNewUnsupportedThreepid, MasNewUpstreamOauthLink, MasNewUser,
         MasNewUserPassword, MasWriteBuffer, MasWriter,
     },
+    progress_stream::ProgressStreamExt,
     synapse_reader::{
         self, ExtractLocalpartError, FullUserId, SynapseAccessToken, SynapseDevice,
         SynapseExternalId, SynapseRefreshableTokenPair, SynapseThreepid, SynapseUser,
@@ -224,16 +225,11 @@ async fn migrate_users(
 ) -> Result<(MasWriter, MigrationState), Error> {
     let start = Instant::now();
 
-    let span = Span::current();
-    span.pb_set_length(count_hint as u64);
-
     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());
+    let mut users_stream = pin!(synapse.read_users().with_progress_bar(count_hint, 10_000));
 
     while let Some(user_res) = users_stream.next().await {
-        span.pb_inc(1);
-
         let user = user_res.into_synapse("reading user")?;
         let (mas_user, mas_password_opt) = transform_user(&user, &state.server_name, rng)?;
 
@@ -310,16 +306,15 @@ async fn migrate_threepids(
 ) -> Result<(MasWriter, MigrationState), Error> {
     let start = Instant::now();
 
-    let span = Span::current();
-    span.pb_set_length(count_hint as u64);
-
     let mut email_buffer = MasWriteBuffer::new(&mas, MasWriter::write_email_threepids);
     let mut unsupported_buffer = MasWriteBuffer::new(&mas, MasWriter::write_unsupported_threepids);
-    let mut users_stream = pin!(synapse.read_threepids());
+    let mut users_stream = pin!(
+        synapse
+            .read_threepids()
+            .with_progress_bar(count_hint, 10_000)
+    );
 
     while let Some(threepid_res) = users_stream.next().await {
-        span.pb_inc(1);
-
         let SynapseThreepid {
             user_id: synapse_user_id,
             medium,
@@ -405,15 +400,15 @@ async fn migrate_external_ids(
     state: MigrationState,
 ) -> Result<(MasWriter, MigrationState), Error> {
     let start = Instant::now();
-    let span = Span::current();
-    span.pb_set_length(count_hint as u64);
 
     let mut write_buffer = MasWriteBuffer::new(&mas, MasWriter::write_upstream_oauth_links);
-    let mut extids_stream = pin!(synapse.read_user_external_ids());
+    let mut extids_stream = pin!(
+        synapse
+            .read_user_external_ids()
+            .with_progress_bar(count_hint, 10_000)
+    );
 
     while let Some(extid_res) = extids_stream.next().await {
-        span.pb_inc(1);
-
         let SynapseExternalId {
             user_id: synapse_user_id,
             auth_provider,
@@ -494,15 +489,10 @@ async fn migrate_devices(
 ) -> Result<(MasWriter, MigrationState), Error> {
     let start = Instant::now();
 
-    let span = Span::current();
-    span.pb_set_length(count_hint as u64);
-
-    let mut devices_stream = pin!(synapse.read_devices());
+    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);
 
     while let Some(device_res) = devices_stream.next().await {
-        span.pb_inc(1);
-
         let SynapseDevice {
             user_id: synapse_user_id,
             device_id,
@@ -604,17 +594,16 @@ async fn migrate_unrefreshable_access_tokens(
 ) -> Result<(MasWriter, MigrationState), Error> {
     let start = Instant::now();
 
-    let span = Span::current();
-    span.pb_set_length(count_hint as u64);
-
-    let mut token_stream = pin!(synapse.read_unrefreshable_access_tokens());
+    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 {
-        span.pb_inc(1);
-
         let SynapseAccessToken {
             user_id: synapse_user_id,
             device_id,
@@ -729,18 +718,17 @@ async fn migrate_refreshable_token_pairs(
 ) -> Result<(MasWriter, MigrationState), Error> {
     let start = Instant::now();
 
-    let span = Span::current();
-    span.pb_set_length(count_hint as u64);
-
-    let mut token_stream = pin!(synapse.read_refreshable_token_pairs());
+    let mut token_stream = pin!(
+        synapse
+            .read_refreshable_token_pairs()
+            .with_progress_bar(count_hint, 10_000)
+    );
     let mut access_token_write_buffer =
         MasWriteBuffer::new(&mas, MasWriter::write_compat_access_tokens);
     let mut refresh_token_write_buffer =
         MasWriteBuffer::new(&mas, MasWriter::write_compat_refresh_tokens);
 
     while let Some(token_res) = token_stream.next().await {
-        span.pb_inc(1);
-
         let SynapseRefreshableTokenPair {
             user_id: synapse_user_id,
             device_id,

crates/syn2mas/src/migration.rs

@@ -0,0 +1,72 @@
+// Copyright 2025 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use futures_util::Stream;
+use tracing::Span;
+use tracing_indicatif::span_ext::IndicatifSpanExt as _;
+
+pin_project_lite::pin_project! {
+    pub struct ProgressStream<S> {
+        #[pin]
+        stream: S,
+        span: Span,
+        counter: usize,
+        batch_size: usize,
+    }
+}
+
+impl<S> ProgressStream<S> {
+    fn new(stream: S, span: Span, batch_size: usize) -> Self {
+        Self {
+            stream,
+            span,
+            counter: 0,
+            batch_size,
+        }
+    }
+}
+
+impl<S> Stream for ProgressStream<S>
+where
+    S: Stream,
+{
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        let item = this.stream.poll_next(cx);
+        if let Poll::Ready(Some(_)) = item {
+            *this.counter += 1;
+            if *this.counter % *this.batch_size == 0 {
+                this.span.pb_set_position(*this.counter as u64);
+            }
+        }
+        item
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.stream.size_hint()
+    }
+}
+
+/// Extension trait for [`Stream`] to add progress bar to the stream.
+pub trait ProgressStreamExt: Stream {
+    /// Add progress bar to the stream.
+    fn with_progress_bar(self, length: usize, batch_size: usize) -> ProgressStream<Self>
+    where
+        Self: Sized,
+    {
+        let span = Span::current();
+        span.pb_set_length(length as u64);
+        ProgressStream::new(self, span, batch_size)
+    }
+}
+
+impl<S> ProgressStreamExt for S where S: Stream {}