diff --git a/Cargo.lock b/Cargo.lock index 9215ca52f..71588b9b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6108,6 +6108,8 @@ dependencies = [ "mas-config", "mas-storage", "mas-storage-pg", + "opentelemetry", + "opentelemetry-semantic-conventions", "rand", "rand_chacha", "rustc-hash 2.1.1", diff --git a/crates/cli/src/commands/syn2mas.rs b/crates/cli/src/commands/syn2mas.rs index baa9edde4..b75a02175 100644 --- a/crates/cli/src/commands/syn2mas.rs +++ b/crates/cli/src/commands/syn2mas.rs @@ -10,7 +10,6 @@ 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::{ @@ -18,10 +17,7 @@ use syn2mas::{ }; use tracing::{Instrument, error, info, info_span, warn}; -use crate::{ - telemetry::METER, - util::{DatabaseConnectOptions, database_connection_from_config_with_options}, -}; +use crate::util::{DatabaseConnectOptions, database_connection_from_config_with_options}; /// The exit code used by `syn2mas check` and `syn2mas migrate` when there are /// errors preventing migration. @@ -258,7 +254,6 @@ impl Options { 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"); @@ -274,7 +269,6 @@ impl Options { .await?; occasional_progress_logger_task.abort(); - progress_telemetry_task.abort(); Ok(ExitCode::SUCCESS) } @@ -312,32 +306,3 @@ async fn occasional_progress_logger(progress: Progress) { } } } - -/// 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/Cargo.toml b/crates/syn2mas/Cargo.toml index d48f875e0..909c9fa88 100644 --- a/crates/syn2mas/Cargo.toml +++ b/crates/syn2mas/Cargo.toml @@ -35,6 +35,9 @@ ulid = { workspace = true, features = ["uuid"] } mas-config.workspace = true mas-storage.workspace = true +opentelemetry.workspace = true +opentelemetry-semantic-conventions.workspace = true + [dev-dependencies] mas-storage-pg.workspace = true diff --git a/crates/syn2mas/src/lib.rs b/crates/syn2mas/src/lib.rs index 5164c0dc7..0fd91ac79 100644 --- a/crates/syn2mas/src/lib.rs +++ b/crates/syn2mas/src/lib.rs @@ -8,6 +8,7 @@ mod synapse_reader; mod migration; mod progress; +mod telemetry; type RandomState = rustc_hash::FxBuildHasher; type HashMap = rustc_hash::FxHashMap; diff --git a/crates/syn2mas/src/migration.rs b/crates/syn2mas/src/migration.rs index 596bb8735..d56f1d487 100644 --- a/crates/syn2mas/src/migration.rs +++ b/crates/syn2mas/src/migration.rs @@ -24,6 +24,7 @@ use chrono::{DateTime, Utc}; use compact_str::CompactString; use futures_util::{SinkExt, StreamExt as _, TryFutureExt, TryStreamExt as _}; use mas_storage::Clock; +use opentelemetry::{KeyValue, metrics::Counter}; use rand::{RngCore, SeedableRng}; use thiserror::Error; use thiserror_ext::ContextInto; @@ -44,6 +45,11 @@ use crate::{ self, ExtractLocalpartError, FullUserId, SynapseAccessToken, SynapseDevice, SynapseExternalId, SynapseRefreshableTokenPair, SynapseThreepid, SynapseUser, }, + telemetry::{ + K_ENTITY, METER, V_ENTITY_DEVICES, V_ENTITY_EXTERNAL_IDS, + V_ENTITY_NONREFRESHABLE_ACCESS_TOKENS, V_ENTITY_REFRESHABLE_TOKEN_PAIRS, + V_ENTITY_THREEPIDS, V_ENTITY_USERS, + }, }; #[derive(Debug, Error, ContextInto)] @@ -147,7 +153,7 @@ struct MigrationState { /// /// - An underlying database access error, either to MAS or to Synapse. /// - Invalid data in the Synapse database. -#[allow(clippy::implicit_hasher)] +#[allow(clippy::implicit_hasher, clippy::too_many_lines)] pub async fn migrate( mut synapse: SynapseReader<'_>, mas: MasWriter, @@ -159,6 +165,45 @@ pub async fn migrate( ) -> Result<(), Error> { let counts = synapse.count_rows().await.into_synapse("counting users")?; + let approx_total_counter = METER + .u64_counter("syn2mas.entity.approx_total") + .with_description("Approximate number of entities of this type to be migrated") + .build(); + let migrated_otel_counter = METER + .u64_counter("syn2mas.entity.migrated") + .with_description("Number of entities of this type that have been migrated so far") + .build(); + + approx_total_counter.add( + counts.users as u64, + &[KeyValue::new(K_ENTITY, V_ENTITY_USERS)], + ); + approx_total_counter.add( + counts.devices as u64, + &[KeyValue::new(K_ENTITY, V_ENTITY_DEVICES)], + ); + approx_total_counter.add( + counts.threepids as u64, + &[KeyValue::new(K_ENTITY, V_ENTITY_THREEPIDS)], + ); + approx_total_counter.add( + counts.external_ids as u64, + &[KeyValue::new(K_ENTITY, V_ENTITY_EXTERNAL_IDS)], + ); + // assume 1 refreshable access token per refresh token. + let approx_nonrefreshable_access_tokens = counts.access_tokens - counts.refresh_tokens; + approx_total_counter.add( + approx_nonrefreshable_access_tokens as u64, + &[KeyValue::new( + K_ENTITY, + V_ENTITY_NONREFRESHABLE_ACCESS_TOKENS, + )], + ); + approx_total_counter.add( + counts.refresh_tokens as u64, + &[KeyValue::new(K_ENTITY, V_ENTITY_REFRESHABLE_TOKEN_PAIRS)], + ); + let state = MigrationState { server_name, // We oversize the hashmaps, as the estimates are innaccurate, and we would like to avoid @@ -173,56 +218,101 @@ pub async fn migrate( let migrated_counter = Arc::new(AtomicU32::new(0)); progress.set_current_stage(ProgressStage::MigratingData { - entity: "users", + entity: V_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 (mas, state) = migrate_users( + &mut synapse, + mas, + state, + rng, + migrated_counter, + migrated_otel_counter.clone(), + ) + .await?; let migrated_counter = Arc::new(AtomicU32::new(0)); progress.set_current_stage(ProgressStage::MigratingData { - entity: "threepids", + entity: V_ENTITY_THREEPIDS, migrated: migrated_counter.clone(), - approx_count: counts.users as u64, + approx_count: counts.threepids as u64, }); - let (mas, state) = migrate_threepids(&mut synapse, mas, rng, state, &migrated_counter).await?; + let (mas, state) = migrate_threepids( + &mut synapse, + mas, + rng, + state, + &migrated_counter, + migrated_otel_counter.clone(), + ) + .await?; let migrated_counter = Arc::new(AtomicU32::new(0)); progress.set_current_stage(ProgressStage::MigratingData { - entity: "external_ids", + entity: V_ENTITY_EXTERNAL_IDS, migrated: migrated_counter.clone(), - approx_count: counts.users as u64, + approx_count: counts.external_ids as u64, }); - let (mas, state) = - migrate_external_ids(&mut synapse, mas, rng, state, &migrated_counter).await?; + let (mas, state) = migrate_external_ids( + &mut synapse, + mas, + rng, + state, + &migrated_counter, + migrated_otel_counter.clone(), + ) + .await?; let migrated_counter = Arc::new(AtomicU32::new(0)); progress.set_current_stage(ProgressStage::MigratingData { - entity: "unrefreshable_access_tokens", + entity: V_ENTITY_NONREFRESHABLE_ACCESS_TOKENS, migrated: migrated_counter.clone(), - approx_count: counts.users as u64, + approx_count: (counts.access_tokens - counts.refresh_tokens) as u64, }); - let (mas, state) = - migrate_unrefreshable_access_tokens(&mut synapse, mas, clock, rng, state, migrated_counter) - .await?; + let (mas, state) = migrate_unrefreshable_access_tokens( + &mut synapse, + mas, + clock, + rng, + state, + migrated_counter, + migrated_otel_counter.clone(), + ) + .await?; let migrated_counter = Arc::new(AtomicU32::new(0)); progress.set_current_stage(ProgressStage::MigratingData { - entity: "refreshable_token_pairs", + entity: V_ENTITY_REFRESHABLE_TOKEN_PAIRS, migrated: migrated_counter.clone(), - approx_count: counts.users as u64, + approx_count: counts.refresh_tokens as u64, }); - let (mas, state) = - migrate_refreshable_token_pairs(&mut synapse, mas, clock, rng, state, &migrated_counter) - .await?; + let (mas, state) = migrate_refreshable_token_pairs( + &mut synapse, + mas, + clock, + rng, + state, + &migrated_counter, + migrated_otel_counter.clone(), + ) + .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, + approx_count: counts.devices as u64, }); - let (mas, _state) = migrate_devices(&mut synapse, mas, rng, state, migrated_counter).await?; + let (mas, _state) = migrate_devices( + &mut synapse, + mas, + rng, + state, + migrated_counter, + migrated_otel_counter.clone(), + ) + .await?; synapse .finish() @@ -243,8 +333,10 @@ async fn migrate_users( mut state: MigrationState, rng: &mut impl RngCore, progress_counter: Arc, + migrated_otel_counter: Counter, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); + let otel_kv = [KeyValue::new(K_ENTITY, V_ENTITY_USERS)]; let (tx, mut rx) = tokio::sync::mpsc::channel::(10 * 1024 * 1024); @@ -316,6 +408,7 @@ async fn migrate_users( .into_mas("writing password")?; } + migrated_otel_counter.add(1, &otel_kv); progress_counter.fetch_add(1, Ordering::Relaxed); } @@ -361,8 +454,10 @@ async fn migrate_threepids( rng: &mut impl RngCore, state: MigrationState, progress_counter: &AtomicU32, + migrated_otel_counter: Counter, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); + let otel_kv = [KeyValue::new(K_ENTITY, V_ENTITY_THREEPIDS)]; let mut email_buffer = MasWriteBuffer::new(&mas, MasWriter::write_email_threepids); let mut unsupported_buffer = MasWriteBuffer::new(&mas, MasWriter::write_unsupported_threepids); @@ -423,6 +518,7 @@ async fn migrate_threepids( .into_mas("writing unsupported threepid")?; } + migrated_otel_counter.add(1, &otel_kv); progress_counter.fetch_add(1, Ordering::Relaxed); } @@ -454,8 +550,10 @@ async fn migrate_external_ids( rng: &mut impl RngCore, state: MigrationState, progress_counter: &AtomicU32, + migrated_otel_counter: Counter, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); + let otel_kv = [KeyValue::new(K_ENTITY, V_ENTITY_EXTERNAL_IDS)]; let mut write_buffer = MasWriteBuffer::new(&mas, MasWriter::write_upstream_oauth_links); let mut extids_stream = pin!(synapse.read_user_external_ids()); @@ -508,6 +606,7 @@ async fn migrate_external_ids( .await .into_mas("failed to write upstream link")?; + migrated_otel_counter.add(1, &otel_kv); progress_counter.fetch_add(1, Ordering::Relaxed); } @@ -539,8 +638,10 @@ async fn migrate_devices( rng: &mut impl RngCore, mut state: MigrationState, progress_counter: Arc, + migrated_otel_counter: Counter, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); + let otel_kv = [KeyValue::new(K_ENTITY, V_ENTITY_DEVICES)]; let (tx, mut rx) = tokio::sync::mpsc::channel(10 * 1024 * 1024); @@ -627,6 +728,7 @@ async fn migrate_devices( .await .into_mas("writing compat sessions")?; + migrated_otel_counter.add(1, &otel_kv); progress_counter.fetch_add(1, Ordering::Relaxed); } @@ -671,8 +773,13 @@ async fn migrate_unrefreshable_access_tokens( rng: &mut impl RngCore, mut state: MigrationState, progress_counter: Arc, + migrated_otel_counter: Counter, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); + let otel_kv = [KeyValue::new( + K_ENTITY, + V_ENTITY_NONREFRESHABLE_ACCESS_TOKENS, + )]; let (tx, mut rx) = tokio::sync::mpsc::channel(10 * 1024 * 1024); @@ -771,6 +878,7 @@ async fn migrate_unrefreshable_access_tokens( .await .into_mas("writing compat access tokens")?; + migrated_otel_counter.add(1, &otel_kv); progress_counter.fetch_add(1, Ordering::Relaxed); } write_buffer @@ -818,8 +926,10 @@ async fn migrate_refreshable_token_pairs( rng: &mut impl RngCore, mut state: MigrationState, progress_counter: &AtomicU32, + migrated_otel_counter: Counter, ) -> Result<(MasWriter, MigrationState), Error> { let start = Instant::now(); + let otel_kv = [KeyValue::new(K_ENTITY, V_ENTITY_REFRESHABLE_TOKEN_PAIRS)]; let mut token_stream = pin!(synapse.read_refreshable_token_pairs()); let mut access_token_write_buffer = @@ -900,6 +1010,7 @@ async fn migrate_refreshable_token_pairs( .await .into_mas("writing compat refresh tokens")?; + migrated_otel_counter.add(1, &otel_kv); progress_counter.fetch_add(1, Ordering::Relaxed); } diff --git a/crates/syn2mas/src/synapse_reader/mod.rs b/crates/syn2mas/src/synapse_reader/mod.rs index 7bce256d9..68746ecf8 100644 --- a/crates/syn2mas/src/synapse_reader/mod.rs +++ b/crates/syn2mas/src/synapse_reader/mod.rs @@ -268,6 +268,10 @@ const TABLES_TO_LOCK: &[&str] = &[ pub struct SynapseRowCounts { pub users: usize, pub devices: usize, + pub threepids: usize, + pub external_ids: usize, + pub access_tokens: usize, + pub refresh_tokens: usize, } pub struct SynapseReader<'c> { @@ -367,7 +371,62 @@ impl<'conn> SynapseReader<'conn> { .try_into() .unwrap_or(usize::MAX); - Ok(SynapseRowCounts { users, devices }) + let threepids = sqlx::query_scalar::<_, i64>( + " + SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'user_threepids'::regclass; + " + ) + .fetch_one(&mut *self.txn) + .await + .into_database("estimating count of threepids")? + .max(0) + .try_into() + .unwrap_or(usize::MAX); + + let access_tokens = sqlx::query_scalar::<_, i64>( + " + SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'access_tokens'::regclass; + " + ) + .fetch_one(&mut *self.txn) + .await + .into_database("estimating count of access tokens")? + .max(0) + .try_into() + .unwrap_or(usize::MAX); + + let refresh_tokens = sqlx::query_scalar::<_, i64>( + " + SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'refresh_tokens'::regclass; + " + ) + .fetch_one(&mut *self.txn) + .await + .into_database("estimating count of refresh tokens")? + .max(0) + .try_into() + .unwrap_or(usize::MAX); + + let external_ids = sqlx::query_scalar::<_, i64>( + " + SELECT reltuples::bigint AS estimate FROM pg_class WHERE oid = 'user_external_ids'::regclass; + " + ) + .fetch_one(&mut *self.txn) + .await + .into_database("estimating count of external IDs")? + .max(0) + .try_into() + .unwrap_or(usize::MAX); + + Ok(SynapseRowCounts { + users, + devices, + threepids, + external_ids, + access_tokens, + refresh_tokens, + }) } /// Reads Synapse users, excluding application service users (which do not diff --git a/crates/syn2mas/src/telemetry.rs b/crates/syn2mas/src/telemetry.rs new file mode 100644 index 000000000..5c1c0a54a --- /dev/null +++ b/crates/syn2mas/src/telemetry.rs @@ -0,0 +1,32 @@ +use std::sync::LazyLock; + +use opentelemetry::{InstrumentationScope, metrics::Meter}; +use opentelemetry_semantic_conventions as semcov; + +static SCOPE: LazyLock = LazyLock::new(|| { + InstrumentationScope::builder(env!("CARGO_PKG_NAME")) + .with_version(env!("CARGO_PKG_VERSION")) + .with_schema_url(semcov::SCHEMA_URL) + .build() +}); + +pub static METER: LazyLock = + LazyLock::new(|| opentelemetry::global::meter_with_scope(SCOPE.clone())); + +/// Attribute key for syn2mas.entity metrics representing what entity. +pub const K_ENTITY: &str = "entity"; + +/// Attribute value for syn2mas.entity metrics representing users. +pub const V_ENTITY_USERS: &str = "users"; +/// Attribute value for syn2mas.entity metrics representing devices. +pub const V_ENTITY_DEVICES: &str = "devices"; +/// Attribute value for syn2mas.entity metrics representing threepids. +pub const V_ENTITY_THREEPIDS: &str = "threepids"; +/// Attribute value for syn2mas.entity metrics representing external IDs. +pub const V_ENTITY_EXTERNAL_IDS: &str = "external_ids"; +/// Attribute value for syn2mas.entity metrics representing non-refreshable +/// access token entities. +pub const V_ENTITY_NONREFRESHABLE_ACCESS_TOKENS: &str = "nonrefreshable_access_tokens"; +/// Attribute value for syn2mas.entity metrics representing refreshable +/// access/refresh token pairs. +pub const V_ENTITY_REFRESHABLE_TOKEN_PAIRS: &str = "refreshable_token_pairs";