From ebad8a77aa089516cd0ed5e08e47c578097259b7 Mon Sep 17 00:00:00 2001 From: Olivier 'reivilibre Date: Wed, 12 Mar 2025 18:30:09 +0000 Subject: [PATCH] syn2mas: Add progress reporting to log and to opentelemetry metrics --- crates/cli/src/commands/syn2mas.rs | 85 ++++++++++++++++++++++++++++-- crates/syn2mas/src/lib.rs | 2 + crates/syn2mas/src/migration.rs | 85 +++++++++++++++++++++++++++--- crates/syn2mas/src/progress.rs | 54 +++++++++++++++++++ 4 files changed, 214 insertions(+), 12 deletions(-) create mode 100644 crates/syn2mas/src/progress.rs diff --git a/crates/cli/src/commands/syn2mas.rs b/crates/cli/src/commands/syn2mas.rs index 0d1b2dc2c..baa9edde4 100644 --- a/crates/cli/src/commands/syn2mas.rs +++ b/crates/cli/src/commands/syn2mas.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, process::ExitCode}; +use std::{collections::HashMap, process::ExitCode, sync::atomic::Ordering, time::Duration}; use anyhow::Context; use camino::Utf8PathBuf; @@ -10,12 +10,18 @@ use mas_config::{ }; use mas_storage::SystemClock; use mas_storage_pg::MIGRATOR; +use opentelemetry::KeyValue; use rand::thread_rng; use sqlx::{Connection, Either, PgConnection, postgres::PgConnectOptions, types::Uuid}; -use syn2mas::{LockedMasDatabase, MasWriter, SynapseReader, synapse_config}; -use tracing::{Instrument, error, info_span, warn}; +use syn2mas::{ + LockedMasDatabase, MasWriter, Progress, ProgressStage, SynapseReader, synapse_config, +}; +use tracing::{Instrument, error, info, info_span, warn}; -use crate::util::{DatabaseConnectOptions, database_connection_from_config_with_options}; +use crate::{ + telemetry::METER, + util::{DatabaseConnectOptions, database_connection_from_config_with_options}, +}; /// The exit code used by `syn2mas check` and `syn2mas migrate` when there are /// errors preventing migration. @@ -248,7 +254,12 @@ impl Options { #[allow(clippy::disallowed_methods)] let mut rng = thread_rng(); - // TODO progress reporting + let progress = Progress::default(); + + let occasional_progress_logger_task = + tokio::spawn(occasional_progress_logger(progress.clone())); + let progress_telemetry_task = tokio::spawn(progress_telemetry(progress.clone())); + let mas_matrix = MatrixConfig::extract(figment)?; eprintln!("\n\n"); syn2mas::migrate( @@ -258,11 +269,75 @@ impl Options { &clock, &mut rng, provider_id_mappings, + &progress, ) .await?; + occasional_progress_logger_task.abort(); + progress_telemetry_task.abort(); + Ok(ExitCode::SUCCESS) } } } } + +/// Logs progress every 30 seconds, as a lightweight alternative to a progress +/// bar. For most deployments, the migration will not take 30 seconds so this +/// will not be relevant. In other cases, this will give the operator an idea of +/// what's going on. +async fn occasional_progress_logger(progress: Progress) { + loop { + tokio::time::sleep(Duration::from_secs(30)).await; + match &**progress.get_current_stage() { + ProgressStage::SettingUp => { + info!(name: "progress", "still setting up"); + } + ProgressStage::MigratingData { + entity, + migrated, + approx_count, + } => { + let migrated = migrated.load(Ordering::Relaxed); + #[allow(clippy::cast_precision_loss)] + let percent = (f64::from(migrated) / *approx_count as f64) * 100.0; + info!(name: "progress", "migrating {entity}: {migrated}/~{approx_count} (~{percent:.1}%)"); + } + ProgressStage::RebuildIndex { index_name } => { + info!(name: "progress", "still waiting for rebuild of index {index_name}"); + } + ProgressStage::RebuildConstraint { constraint_name } => { + info!(name: "progress", "still waiting for rebuild of constraint {constraint_name}"); + } + } + } +} + +/// Reports migration progress as OpenTelemetry metrics +async fn progress_telemetry(progress: Progress) { + let migrated_data_counter = METER + .u64_gauge("migrated_data") + .with_description("How many entities have been migrated so far") + .build(); + let max_data_counter = METER + .u64_gauge("max_data") + .with_description("How many entities of the given type exist (approximate)") + .build(); + + loop { + tokio::time::sleep(Duration::from_secs(10)).await; + if let ProgressStage::MigratingData { + entity, + migrated, + approx_count, + } = &**progress.get_current_stage() + { + let metrics_kv = [KeyValue::new("entity", *entity)]; + let migrated = migrated.load(Ordering::Relaxed); + migrated_data_counter.record(u64::from(migrated), &metrics_kv); + max_data_counter.record(*approx_count, &metrics_kv); + } else { + // not sure how to map other stages + } + } +} diff --git a/crates/syn2mas/src/lib.rs b/crates/syn2mas/src/lib.rs index d0d1162fb..5164c0dc7 100644 --- a/crates/syn2mas/src/lib.rs +++ b/crates/syn2mas/src/lib.rs @@ -7,6 +7,7 @@ mod mas_writer; mod synapse_reader; mod migration; +mod progress; type RandomState = rustc_hash::FxBuildHasher; type HashMap = rustc_hash::FxHashMap; @@ -14,6 +15,7 @@ type HashMap = rustc_hash::FxHashMap; pub use self::{ mas_writer::{MasWriter, checks::mas_pre_migration_checks, locking::LockedMasDatabase}, migration::migrate, + progress::{Progress, ProgressStage}, synapse_reader::{ SynapseReader, checks::{ diff --git a/crates/syn2mas/src/migration.rs b/crates/syn2mas/src/migration.rs index 5a648f3b4..596bb8735 100644 --- a/crates/syn2mas/src/migration.rs +++ b/crates/syn2mas/src/migration.rs @@ -11,7 +11,14 @@ //! This module does not implement any of the safety checks that should be run //! *before* the migration. -use std::{pin::pin, time::Instant}; +use std::{ + pin::pin, + sync::{ + Arc, + atomic::{AtomicU32, Ordering}, + }, + time::Instant, +}; use chrono::{DateTime, Utc}; use compact_str::CompactString; @@ -32,6 +39,7 @@ use crate::{ MasNewEmailThreepid, MasNewUnsupportedThreepid, MasNewUpstreamOauthLink, MasNewUser, MasNewUserPassword, MasWriteBuffer, MasWriter, }, + progress::{Progress, ProgressStage}, synapse_reader::{ self, ExtractLocalpartError, FullUserId, SynapseAccessToken, SynapseDevice, SynapseExternalId, SynapseRefreshableTokenPair, SynapseThreepid, SynapseUser, @@ -147,6 +155,7 @@ pub async fn migrate( clock: &dyn Clock, rng: &mut impl RngCore, provider_id_mapping: std::collections::HashMap, + progress: &Progress, ) -> Result<(), Error> { let counts = synapse.count_rows().await.into_synapse("counting users")?; @@ -162,14 +171,58 @@ pub async fn migrate( provider_id_mapping, }; - let (mas, state) = migrate_users(&mut synapse, mas, state, rng).await?; - let (mas, state) = migrate_threepids(&mut synapse, mas, rng, state).await?; - let (mas, state) = migrate_external_ids(&mut synapse, mas, rng, state).await?; + let migrated_counter = Arc::new(AtomicU32::new(0)); + progress.set_current_stage(ProgressStage::MigratingData { + entity: "users", + migrated: migrated_counter.clone(), + approx_count: counts.users as u64, + }); + let (mas, state) = migrate_users(&mut synapse, mas, state, rng, migrated_counter).await?; + + let migrated_counter = Arc::new(AtomicU32::new(0)); + progress.set_current_stage(ProgressStage::MigratingData { + entity: "threepids", + migrated: migrated_counter.clone(), + approx_count: counts.users as u64, + }); + let (mas, state) = migrate_threepids(&mut synapse, mas, rng, state, &migrated_counter).await?; + + let migrated_counter = Arc::new(AtomicU32::new(0)); + progress.set_current_stage(ProgressStage::MigratingData { + entity: "external_ids", + migrated: migrated_counter.clone(), + approx_count: counts.users as u64, + }); let (mas, state) = - migrate_unrefreshable_access_tokens(&mut synapse, mas, clock, rng, state).await?; + migrate_external_ids(&mut synapse, mas, rng, state, &migrated_counter).await?; + + let migrated_counter = Arc::new(AtomicU32::new(0)); + progress.set_current_stage(ProgressStage::MigratingData { + entity: "unrefreshable_access_tokens", + migrated: migrated_counter.clone(), + approx_count: counts.users as u64, + }); let (mas, state) = - migrate_refreshable_token_pairs(&mut synapse, mas, clock, rng, state).await?; - let (mas, _state) = migrate_devices(&mut synapse, mas, rng, state).await?; + migrate_unrefreshable_access_tokens(&mut synapse, mas, clock, rng, state, migrated_counter) + .await?; + + let migrated_counter = Arc::new(AtomicU32::new(0)); + progress.set_current_stage(ProgressStage::MigratingData { + entity: "refreshable_token_pairs", + migrated: migrated_counter.clone(), + approx_count: counts.users as u64, + }); + let (mas, state) = + migrate_refreshable_token_pairs(&mut synapse, mas, clock, rng, state, &migrated_counter) + .await?; + + let migrated_counter = Arc::new(AtomicU32::new(0)); + progress.set_current_stage(ProgressStage::MigratingData { + entity: "devices", + migrated: migrated_counter.clone(), + approx_count: counts.users as u64, + }); + let (mas, _state) = migrate_devices(&mut synapse, mas, rng, state, migrated_counter).await?; synapse .finish() @@ -189,6 +242,7 @@ async fn migrate_users( mut mas: MasWriter, mut state: MigrationState, rng: &mut impl RngCore, + progress_counter: Arc, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); @@ -261,6 +315,8 @@ async fn migrate_users( .await .into_mas("writing password")?; } + + progress_counter.fetch_add(1, Ordering::Relaxed); } user_buffer @@ -304,6 +360,7 @@ async fn migrate_threepids( mut mas: MasWriter, rng: &mut impl RngCore, state: MigrationState, + progress_counter: &AtomicU32, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); @@ -365,6 +422,8 @@ async fn migrate_threepids( .await .into_mas("writing unsupported threepid")?; } + + progress_counter.fetch_add(1, Ordering::Relaxed); } email_buffer @@ -394,6 +453,7 @@ async fn migrate_external_ids( mut mas: MasWriter, rng: &mut impl RngCore, state: MigrationState, + progress_counter: &AtomicU32, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); @@ -447,6 +507,8 @@ async fn migrate_external_ids( ) .await .into_mas("failed to write upstream link")?; + + progress_counter.fetch_add(1, Ordering::Relaxed); } write_buffer @@ -476,6 +538,7 @@ async fn migrate_devices( mut mas: MasWriter, rng: &mut impl RngCore, mut state: MigrationState, + progress_counter: Arc, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); @@ -563,6 +626,8 @@ async fn migrate_devices( ) .await .into_mas("writing compat sessions")?; + + progress_counter.fetch_add(1, Ordering::Relaxed); } write_buffer @@ -605,6 +670,7 @@ async fn migrate_unrefreshable_access_tokens( clock: &dyn Clock, rng: &mut impl RngCore, mut state: MigrationState, + progress_counter: Arc, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); @@ -704,6 +770,8 @@ async fn migrate_unrefreshable_access_tokens( ) .await .into_mas("writing compat access tokens")?; + + progress_counter.fetch_add(1, Ordering::Relaxed); } write_buffer .finish(&mut mas) @@ -749,6 +817,7 @@ async fn migrate_refreshable_token_pairs( clock: &dyn Clock, rng: &mut impl RngCore, mut state: MigrationState, + progress_counter: &AtomicU32, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); @@ -830,6 +899,8 @@ async fn migrate_refreshable_token_pairs( ) .await .into_mas("writing compat refresh tokens")?; + + progress_counter.fetch_add(1, Ordering::Relaxed); } access_token_write_buffer diff --git a/crates/syn2mas/src/progress.rs b/crates/syn2mas/src/progress.rs new file mode 100644 index 000000000..f2c86602a --- /dev/null +++ b/crates/syn2mas/src/progress.rs @@ -0,0 +1,54 @@ +use std::sync::{Arc, atomic::AtomicU32}; + +use arc_swap::ArcSwap; + +/// Tracker for the progress of the migration +/// +/// Cloning this struct intuitively gives a 'handle' to the same counters, +/// which means it can be shared between tasks/threads. +#[derive(Clone)] +pub struct Progress { + current_stage: Arc>, +} + +impl Progress { + /// Sets the current stage of progress. + /// + /// This is probably not cheap enough to use for every individual row, + /// so use of atomic integers for the fields that will be updated is + /// recommended. + #[inline] + pub fn set_current_stage(&self, stage: ProgressStage) { + self.current_stage.store(Arc::new(stage)); + } + + /// Returns the current stage of progress. + #[inline] + #[must_use] + pub fn get_current_stage(&self) -> arc_swap::Guard> { + self.current_stage.load() + } +} + +impl Default for Progress { + fn default() -> Self { + Self { + current_stage: Arc::new(ArcSwap::new(Arc::new(ProgressStage::SettingUp))), + } + } +} + +pub enum ProgressStage { + SettingUp, + MigratingData { + entity: &'static str, + migrated: Arc, + approx_count: u64, + }, + RebuildIndex { + index_name: String, + }, + RebuildConstraint { + constraint_name: String, + }, +}