Improve progress bar configuration (and stop it clobbering the logs)

Author: reivilibre

crates/cli/src/main.rs

@@ -14,6 +14,7 @@ use mas_config::{ConfigurationSection, TelemetryConfig};
 use sentry_tracing::EventFilter;
 use tracing_indicatif::{
     filter::{hide_indicatif_span_fields, IndicatifFilter},
+    style::ProgressStyle,
     IndicatifLayer,
 };
 use tracing_subscriber::{
@@ -72,10 +73,23 @@ async fn try_main() -> anyhow::Result<ExitCode> {
         // Display the error if it is something other than the .env file not existing
         .or_else(|e| if e.not_found() { Ok(None) } else { Err(e) });
 
+    // Set up progress bars, used in syn2mas for example
+    let progress_layer = IndicatifLayer::new()
+        .with_span_field_formatter(hide_indicatif_span_fields(DefaultFields::new()))
+        // TODO this progress style would likely benefit from being tweaked,
+        // but we want to use a custom one to show both a bar and the message at the same time.
+        .with_progress_style(
+            ProgressStyle::with_template(
+                "[{elapsed_precise}] {bar:40.cyan/blue} {pos:>10}/{len:10} {msg}",
+            )
+            .unwrap(),
+        );
+
     // Setup logging
     // This writes logs to stderr
-    let output = std::io::stderr();
-    let with_ansi = output.is_terminal();
+    let with_ansi = std::io::stderr().is_terminal(); // TODO not sure this is the best way
+    let output = progress_layer.get_stderr_writer();
+
     let (log_writer, _guard) = tracing_appender::non_blocking(output);
     let fmt_layer = tracing_subscriber::fmt::layer()
         .with_writer(log_writer)
@@ -136,17 +150,12 @@ async fn try_main() -> anyhow::Result<ExitCode> {
             .with_filter(LevelFilter::INFO)
     });
 
-    // Set up progress bars, used in syn2mas for example
-    let progress_layer = IndicatifLayer::new()
-        .with_span_field_formatter(hide_indicatif_span_fields(DefaultFields::new()))
-        .with_filter(IndicatifFilter::new(false));
-
     let subscriber = Registry::default()
         .with(sentry_layer)
         .with(telemetry_layer)
         .with(filter_layer)
         .with(fmt_layer)
-        .with(progress_layer);
+        .with(progress_layer.with_filter(IndicatifFilter::new(false)));
     subscriber
         .try_init()
         .context("could not initialize logging")?;

crates/syn2mas/src/mas_writer/mod.rs

@@ -23,7 +23,7 @@ use thiserror::Error;
 use thiserror_ext::{Construct, ContextInto};
 use tokio::sync::mpsc::{self, Receiver, Sender};
 use tracing::{error, info, warn, Level, Span};
-use tracing_indicatif::{span_ext::IndicatifSpanExt, style::ProgressStyle};
+use tracing_indicatif::span_ext::IndicatifSpanExt;
 use uuid::Uuid;
 
 use self::{
@@ -546,14 +546,6 @@ impl<'conn> MasWriter<'conn> {
         constraints_to_restore: &[ConstraintDescription],
     ) -> Result<(), Error> {
         let span = Span::current();
-        // TODO this style is a quick workaround for showing the message, but
-        // might not be optimal for other purposes
-        span.pb_set_style(
-            &ProgressStyle::with_template(
-                "[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}",
-            )
-            .unwrap(),
-        );
         span.pb_set_length((indices_to_restore.len() + constraints_to_restore.len()) as u64);
 
         // First restore all indices. The order is not important as far as I know.

crates/syn2mas/src/migration.rs

@@ -30,7 +30,7 @@ use rand::RngCore;
 use thiserror::Error;
 use thiserror_ext::ContextInto;
 use tracing::{info, Level, Span};
-use tracing_indicatif::{span_ext::IndicatifSpanExt, style::ProgressStyle};
+use tracing_indicatif::span_ext::IndicatifSpanExt;
 use ulid::Ulid;
 use uuid::Uuid;
 
@@ -107,16 +107,8 @@ pub async fn migrate(
     provider_id_mapping: &HashMap<String, Uuid>,
 ) -> Result<(), Error> {
     let span = Span::current();
-    // TODO this style is inconsistent with the child spans; it's just used because
-    // the default style doesn't seem to include the message?
-    span.pb_set_style(
-        &ProgressStyle::with_template(
-            "[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}",
-        )
-        .unwrap(),
-    );
     span.pb_set_message("counting work");
-    span.pb_set_length(7);
+    span.pb_set_length(8);
     let counts = synapse.count_rows().await.into_synapse("counting rows")?;
 
     span.pb_set_message("migrating user rows");
@@ -255,7 +247,6 @@ async fn migrate_users(
     let start = Instant::now();
 
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(user_count_hint as u64);
 
     let mut user_buffer = MasWriteBuffer::new(mas, MasWriter::write_users);
@@ -323,7 +314,6 @@ async fn migrate_threepids(
     let start = Instant::now();
 
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(count_hint);
 
     let mut email_buffer = MasWriteBuffer::new(mas, MasWriter::write_email_threepids);
@@ -426,7 +416,6 @@ async fn migrate_external_ids(
     let start = Instant::now();
 
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(count_hint);
 
     let mut write_buffer = MasWriteBuffer::new(mas, MasWriter::write_upstream_oauth_links);
@@ -519,7 +508,6 @@ async fn migrate_devices(
     let start = Instant::now();
 
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(count_hint);
 
     let mut devices_stream = pin!(synapse.read_devices());
@@ -627,7 +615,6 @@ async fn migrate_unrefreshable_access_tokens(
     let start = Instant::now();
 
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(count_hint);
 
     let mut token_stream = pin!(synapse.read_unrefreshable_access_tokens());
@@ -752,7 +739,6 @@ async fn migrate_refreshable_token_pairs(
     let start = Instant::now();
 
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(count_hint);
 
     let mut token_stream = pin!(synapse.read_refreshable_token_pairs());

crates/syn2mas/src/synapse_reader/mod.rs

@@ -338,33 +338,35 @@ impl<'conn> SynapseReader<'conn> {
     ///
     /// - An underlying database error
     pub async fn count_rows(&mut self) -> Result<SynapseRowCounts, Error> {
+        // For speed, retrieve a fast estimate from the statistics system
+        // of Postgres instead. https://wiki.postgresql.org/wiki/Count_estimate
+        // If the estimates are not up to date, the result might be `-1`, so clamp to 0.
+
         // 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.
+        // which is a shame, but on a large database this is way faster than counting
+        // exactly.
         // 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(
+        let users: i64 = sqlx::query_scalar::<_, i64>(
             "
             SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'users'::regclass;
             ",
         )
         .fetch_one(&mut *self.txn)
         .await
-        .into_database("estimating count of users")?;
+        .into_database("estimating count of users")?
+        .max(0);
 
-        let devices = sqlx::query_scalar(
+        let devices = sqlx::query_scalar::<_, i64>(
             "
             SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'devices'::regclass;
             ",
         )
         .fetch_one(&mut *self.txn)
         .await
-        .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
-        // of Postgres instead. https://wiki.postgresql.org/wiki/Count_estimate
-        // If the estimates are not up to date, the result might be `-1`, so clamp to 0.
+        .into_database("estimating count of devices")?
+        .max(0);
 
         let threepids = sqlx::query_scalar::<_, i64>(
             "