Count entities from the reading end

Author: sandhose

crates/syn2mas/src/synapse_reader/mod.rs

@@ -12,10 +12,13 @@ use std::fmt::Display;
 
 use chrono::{DateTime, Utc};
 use futures_util::{Stream, TryStreamExt};
+use opentelemetry::KeyValue;
 use sqlx::{Acquire, FromRow, PgConnection, Postgres, Transaction, Type, query};
 use thiserror::Error;
 use thiserror_ext::ContextInto;
 
+use crate::telemetry::READER_ENTITY_COUNT;
+
 pub mod checks;
 pub mod config;
 
@@ -432,6 +435,7 @@ impl<'conn> SynapseReader<'conn> {
     /// Reads Synapse users, excluding application service users (which do not
     /// need to be migrated), from the database.
     pub fn read_users(&mut self) -> impl Stream<Item = Result<SynapseUser, Error>> + '_ {
+        let kv = [KeyValue::new("entity", "users")];
         sqlx::query_as(
             "
             SELECT
@@ -440,12 +444,14 @@ impl<'conn> SynapseReader<'conn> {
             ",
         )
         .fetch(&mut *self.txn)
+        .inspect_ok(move |_| READER_ENTITY_COUNT.add(1, &kv))
         .map_err(|err| err.into_database("reading Synapse users"))
     }
 
     /// Reads threepids (such as e-mail and phone number associations) from
     /// Synapse.
     pub fn read_threepids(&mut self) -> impl Stream<Item = Result<SynapseThreepid, Error>> + '_ {
+        let kv = [KeyValue::new("entity", "threepids")];
         sqlx::query_as(
             "
             SELECT
@@ -454,13 +460,15 @@ impl<'conn> SynapseReader<'conn> {
             ",
         )
         .fetch(&mut *self.txn)
+        .inspect_ok(move |_| READER_ENTITY_COUNT.add(1, &kv))
         .map_err(|err| err.into_database("reading Synapse threepids"))
     }
 
     /// Read associations between Synapse users and external identity providers
     pub fn read_user_external_ids(
         &mut self,
     ) -> impl Stream<Item = Result<SynapseExternalId, Error>> + '_ {
+        let kv = [KeyValue::new("entity", "external_ids")];
         sqlx::query_as(
             "
             SELECT
@@ -469,13 +477,15 @@ impl<'conn> SynapseReader<'conn> {
             ",
         )
         .fetch(&mut *self.txn)
+        .inspect_ok(move |_| READER_ENTITY_COUNT.add(1, &kv))
         .map_err(|err| err.into_database("reading Synapse user external IDs"))
     }
 
     /// Reads devices from the Synapse database.
     /// Does not include so-called 'hidden' devices, which are just a mechanism
     /// for storing various signing keys shared between the real devices.
     pub fn read_devices(&mut self) -> impl Stream<Item = Result<SynapseDevice, Error>> + '_ {
+        let kv = [KeyValue::new("entity", "devices")];
         sqlx::query_as(
             "
             SELECT
@@ -485,6 +495,7 @@ impl<'conn> SynapseReader<'conn> {
             ",
         )
         .fetch(&mut *self.txn)
+        .inspect_ok(move |_| READER_ENTITY_COUNT.add(1, &kv))
         .map_err(|err| err.into_database("reading Synapse devices"))
     }
 
@@ -500,6 +511,7 @@ impl<'conn> SynapseReader<'conn> {
     pub fn read_unrefreshable_access_tokens(
         &mut self,
     ) -> impl Stream<Item = Result<SynapseAccessToken, Error>> + '_ {
+        let kv = [KeyValue::new("entity", "nonrefreshable_access_tokens")];
         sqlx::query_as(
             "
             SELECT
@@ -517,6 +529,7 @@ impl<'conn> SynapseReader<'conn> {
             ",
         )
         .fetch(&mut *self.txn)
+        .inspect_ok(move |_| READER_ENTITY_COUNT.add(1, &kv))
         .map_err(|err| err.into_database("reading Synapse access tokens"))
     }
 
@@ -532,6 +545,7 @@ impl<'conn> SynapseReader<'conn> {
     pub fn read_refreshable_token_pairs(
         &mut self,
     ) -> impl Stream<Item = Result<SynapseRefreshableTokenPair, Error>> + '_ {
+        let kv = [KeyValue::new("entity", "refreshable_token_pairs")];
         sqlx::query_as(
             "
             SELECT
@@ -544,6 +558,7 @@ impl<'conn> SynapseReader<'conn> {
             ",
         )
         .fetch(&mut *self.txn)
+        .inspect_ok(move |_| READER_ENTITY_COUNT.add(1, &kv))
         .map_err(|err| err.into_database("reading Synapse refresh tokens"))
     }
 }

crates/syn2mas/src/telemetry.rs

@@ -53,6 +53,13 @@ pub static WRITER_WAIT_TIME: LazyLock<Histogram<u64>> = LazyLock::new(|| {
         .build()
 });
 
+pub static READER_ENTITY_COUNT: LazyLock<Counter<u64>> = LazyLock::new(|| {
+    METER
+        .u64_counter("syn2mas.reader.entity_count")
+        .with_description("Number of entities read from the reader")
+        .build()
+});
+
 /// Attribute key for syn2mas.entity metrics representing what entity.
 pub const K_ENTITY: &str = "entity";