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

Author: sandhose

crates/cli/src/main.rs

@@ -15,6 +15,7 @@ use sentry_tracing::EventFilter;
 use tracing_indicatif::{
     IndicatifLayer,
     filter::{IndicatifFilter, hide_indicatif_span_fields},
+    style::ProgressStyle,
 };
 use tracing_subscriber::{
     EnvFilter, Layer, Registry, filter::LevelFilter, fmt::format::DefaultFields,
@@ -74,10 +75,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)
@@ -138,17 +152,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::{Level, Span, error, info, warn};
-use tracing_indicatif::{span_ext::IndicatifSpanExt, style::ProgressStyle};
+use tracing_indicatif::span_ext::IndicatifSpanExt;
 use uuid::Uuid;
 
 use self::{
@@ -550,14 +550,6 @@ impl MasWriter {
         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

@@ -21,7 +21,7 @@ use rand::RngCore;
 use thiserror::Error;
 use thiserror_ext::ContextInto;
 use tracing::{Level, Span, info};
-use tracing_indicatif::{span_ext::IndicatifSpanExt, style::ProgressStyle};
+use tracing_indicatif::span_ext::IndicatifSpanExt;
 use ulid::Ulid;
 use uuid::Uuid;
 
@@ -141,16 +141,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")?;
 
     let state = MigrationState {
@@ -233,7 +225,6 @@ async fn migrate_users(
     let start = Instant::now();
 
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(count_hint as u64);
 
     let mut user_buffer = MasWriteBuffer::new(&mas, MasWriter::write_users);
@@ -320,7 +311,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 as u64);
 
     let mut email_buffer = MasWriteBuffer::new(&mas, MasWriter::write_email_threepids);
@@ -416,7 +406,6 @@ async fn migrate_external_ids(
 ) -> Result<(MasWriter, MigrationState), Error> {
     let start = Instant::now();
     let span = Span::current();
-    span.pb_set_style(&ProgressStyle::default_bar());
     span.pb_set_length(count_hint as u64);
 
     let mut write_buffer = MasWriteBuffer::new(&mas, MasWriter::write_upstream_oauth_links);
@@ -506,7 +495,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 as u64);
 
     let mut devices_stream = pin!(synapse.read_devices());
@@ -617,7 +605,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 as u64);
 
     let mut token_stream = pin!(synapse.read_unrefreshable_access_tokens());
@@ -743,7 +730,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 as u64);
 
     let mut token_stream = pin!(synapse.read_refreshable_token_pairs());

crates/syn2mas/src/synapse_reader/mod.rs

@@ -340,8 +340,13 @@ 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.
@@ -369,11 +374,6 @@ impl<'conn> SynapseReader<'conn> {
         .try_into()
         .unwrap_or(usize::MAX);
 
-        // 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.
-
         let threepids = sqlx::query_scalar::<_, i64>(
             "
             SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'user_threepids'::regclass;