feat(pds): Prometheus metrics + structured logging for rsky-pds

rsky-pds/Cargo.toml

@@ -41,6 +41,8 @@ lexicon_cid = { workspace = true }
 lru = "0.14"
 mailchecker = "6.0.1"
 mailgun-rs = "0.1.10"
+metrics = "0.24"
+metrics-exporter-prometheus = { version = "0.17", default-features = false }
 rand = { workspace = true }
 rsky-oauth = { path = "../rsky-oauth", version = "0.3.0" }
 rand_core = { workspace = true }
@@ -71,7 +73,7 @@ time = "^0.3.36"
 tokio = { workspace = true }
 toml = "0.8.12"
 tracing = "0.1.41"
-tracing-subscriber = "0.3.19"
+tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 url = "2.5.2"
 ws = { package = "rocket_ws", version = "0.1.1" }
 

rsky-pds/src/apis/com/atproto/repo/apply_writes.rs

@@ -123,6 +123,13 @@ async fn inner_apply_writes(
                 commit.commit_data.rev,
             )
             .await?;
+        for write in &writes {
+            crate::metrics::record_repo_write(match write {
+                PreparedWrite::Create(_) => "create",
+                PreparedWrite::Update(_) => "update",
+                PreparedWrite::Delete(_) => "delete",
+            });
+        }
         // The lexicon declares a JSON object output; returning an empty body
         // instead makes a client that requires JSON treat a successful write as
         // failed and retry it, duplicating records.

rsky-pds/src/apis/com/atproto/repo/create_record.rs

@@ -100,6 +100,7 @@ async fn inner_create_record(
         account_manager
             .update_repo_root(did, commit.commit_data.cid, commit.commit_data.rev)
             .await?;
+        crate::metrics::record_repo_write("create");
 
         Ok(CreateRecordOutput {
             uri: write.uri.clone(),

rsky-pds/src/apis/com/atproto/repo/delete_record.rs

@@ -85,6 +85,7 @@ async fn inner_delete_record(
             account_manager
                 .update_repo_root(did, commit.commit_data.cid, commit.commit_data.rev)
                 .await?;
+            crate::metrics::record_repo_write("delete");
 
             Ok(())
         }

rsky-pds/src/apis/com/atproto/repo/put_record.rs

@@ -112,6 +112,11 @@ async fn inner_put_record(
             account_manager
                 .update_repo_root(did, commit.commit_data.cid, commit.commit_data.rev)
                 .await?;
+            crate::metrics::record_repo_write(match &write {
+                PreparedWrite::Create(_) => "create",
+                PreparedWrite::Update(_) => "update",
+                PreparedWrite::Delete(_) => "delete",
+            });
         }
         Ok(PutRecordOutput {
             uri: write.uri().to_string(),

rsky-pds/src/apis/com/atproto/repo/upload_blob.rs

@@ -78,6 +78,8 @@ async fn inner_upload_blob(
             .await?;
     }
 
+    crate::metrics::record_blob_upload(blobref.get_size().unwrap_or(0).max(0) as u64);
+
     Ok(BlobOutput {
         blob: Blob {
             r#type: Some("blob".to_string()),

rsky-pds/src/apis/com/atproto/server/create_session.rs

@@ -111,7 +111,13 @@ pub async fn create_session(
 ) -> Result<Json<CreateSessionOutput>, ApiError> {
     // @TODO: Add rate limiting
     match inner_create_session(body, account_manager).await {
-        Ok(res) => Ok(Json(res)),
-        Err(error) => Err(error),
+        Ok(res) => {
+            crate::metrics::record_login(true);
+            Ok(Json(res))
+        }
+        Err(error) => {
+            crate::metrics::record_login(false);
+            Err(error)
+        }
     }
 }

rsky-pds/src/apis/com/atproto/server/get_service_auth.rs

@@ -71,7 +71,10 @@ pub async fn get_service_auth(
     actor_store: &State<ActorStore>,
 ) -> Result<Json<GetServiceAuthOutput>, ApiError> {
     match inner_get_service_auth(aud, exp, lxm, auth, actor_store).await {
-        Ok(token) => Ok(Json(GetServiceAuthOutput { token })),
+        Ok(token) => {
+            crate::metrics::record_service_token_issued();
+            Ok(Json(GetServiceAuthOutput { token }))
+        }
         Err(error) => {
             tracing::error!("Internal Error: {error}");
             Err(ApiError::RuntimeError)

rsky-pds/src/apis/com/atproto/sync/subscribe_repos.rs

@@ -23,6 +23,25 @@ use std::time::SystemTime;
 use tokio::time::{interval, Duration as TokioDuration};
 use ws::Message;
 
+/// Tracks the `pds_firehose_subscribers` gauge for the lifetime of a single
+/// subscribeRepos connection: incremented on connect, decremented on drop
+/// (covers every exit path -- normal completion, an early `return`, a
+/// `break`, or the client simply disconnecting) so the count can never leak.
+struct FirehoseSubscriberGuard;
+
+impl FirehoseSubscriberGuard {
+    fn new() -> Self {
+        crate::metrics::record_firehose_subscriber_connected();
+        FirehoseSubscriberGuard
+    }
+}
+
+impl Drop for FirehoseSubscriberGuard {
+    fn drop(&mut self) {
+        crate::metrics::record_firehose_subscriber_disconnected();
+    }
+}
+
 fn get_backfill_limit(ms: u64) -> String {
     let system_time = SystemTime::now();
     let mut dt: DateTime<UtcOffset> = system_time.into();
@@ -45,6 +64,7 @@ pub async fn subscribe_repos<'a>(
     ws: ws::WebSocket,
 ) -> ws::Stream!['a] {
     ws::Stream! { ws =>
+        let _firehose_subscriber_guard = FirehoseSubscriberGuard::new();
         let sequencer_lock = sequencer.sequencer.read().await.clone();
         let mut outbox = Outbox::new(
             sequencer_lock.clone(),

rsky-pds/src/lib.rs

@@ -22,6 +22,7 @@ pub mod handle;
 pub mod image;
 pub mod lexicon;
 pub mod mailer;
+pub mod metrics;
 pub mod models;
 pub mod oauth;
 pub mod oauth_scope;
@@ -311,6 +312,8 @@ pub async fn build_rocket(rocket_cfg: Option<RocketConfig>) -> Rocket<Build> {
 
     let shield = Shield::default().enable(NoSniff::Enable);
 
+    let metrics_handle = crate::metrics::install_recorder();
+
     rocket::custom(figment)
         .mount(
             "/",
@@ -319,6 +322,7 @@ pub async fn build_rocket(rocket_cfg: Option<RocketConfig>) -> Rocket<Build> {
                 robots,
                 health,
                 health_live,
+                crate::metrics::metrics_route,
                 com::atproto::admin::delete_account::delete_account,
                 com::atproto::admin::disable_account_invites::disable_account_invites,
                 com::atproto::admin::disable_invite_codes::disable_invite_codes,
@@ -449,6 +453,8 @@ pub async fn build_rocket(rocket_cfg: Option<RocketConfig>) -> Rocket<Build> {
         .attach(CORS)
         .attach(oauth::OAuthHeaders)
         .attach(shield)
+        .attach(crate::metrics::XrpcMetrics)
+        .manage(metrics_handle)
         .manage(sequencer)
         .manage(blobstore_factory)
         .manage(id_resolver)

rsky-pds/src/main.rs

@@ -1,12 +1,19 @@
 use rsky_pds::build_rocket;
+use tracing_subscriber::fmt::Layer;
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::util::SubscriberInitExt;
+use tracing_subscriber::EnvFilter;
 
 #[rocket::main]
 async fn main() {
     let _ = &*rsky_pds::context::PDS_REPO_SIGNING_KEYPAIR;
     let _ = &*rsky_pds::auth_verifier::PDS_JWT_KEYPAIR;
     let _ = &*rsky_pds::apis::com::atproto::server::PDS_PLC_ROTATION_KEYPAIR;
 
-    let subscriber = tracing_subscriber::FmtSubscriber::new();
-    tracing::subscriber::set_global_default(subscriber).unwrap();
+    tracing_subscriber::registry()
+        .with(EnvFilter::from_default_env())
+        .with(Layer::new())
+        .init();
+
     let _ = build_rocket(None).await.launch().await;
 }