fixup! syn2mas: Add progress reporting to log and to opentelemetry metrics

Add metrics directly within syn2mas, no background thread

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",

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
-        }
-    }
-}

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
 

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<K, V> = rustc_hash::FxHashMap<K, V>;

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<AtomicU32>,
+    migrated_otel_counter: Counter<u64>,
 ) -> 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::<SynapseUser>(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<u64>,
 ) -> 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<u64>,
 ) -> 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<AtomicU32>,
+    migrated_otel_counter: Counter<u64>,
 ) -> 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<AtomicU32>,
+    migrated_otel_counter: Counter<u64>,
 ) -> 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<u64>,
 ) -> 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);
     }
 

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

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<InstrumentationScope> = 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<Meter> =
+    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";