From d5989ab49ae6663e32037415fd591386aa906301 Mon Sep 17 00:00:00 2001
From: Chris Smith <1979423+chris13524@users.noreply.github.com>
Date: Tue, 12 Dec 2023 11:03:43 -0500
Subject: [PATCH] fix: add timing traces (#160)

---
 justfile                         | 11 +++++++++++
 src/handlers/identity/resolve.rs | 23 +++++++++++------------
 src/stores/keys.rs               |  6 ++++++
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/justfile b/justfile
index 25be31f..c883ed4 100644
--- a/justfile
+++ b/justfile
@@ -42,6 +42,17 @@ fmt:
   @echo '==> Reformatting code'
   cargo +nightly fmt
 
+fmt-imports:
+  #!/bin/bash
+  set -euo pipefail
+
+  if command -v cargo-fmt >/dev/null; then
+    echo '==> Running rustfmt'
+    cargo +nightly fmt -- --config group_imports=StdExternalCrate,imports_granularity=One
+  else
+    echo '==> rustfmt not found in PATH, skipping'
+  fi
+
 # Build docker image
 build-docker:
   @echo '=> Build keys-server docker image'
diff --git a/src/handlers/identity/resolve.rs b/src/handlers/identity/resolve.rs
index 2503f0c..b0c24f8 100644
--- a/src/handlers/identity/resolve.rs
+++ b/src/handlers/identity/resolve.rs
@@ -13,19 +13,15 @@ use {
     serde::{Deserialize, Serialize},
     serde_json::{json, Value},
     std::sync::Arc,
+    tracing::instrument,
     validator::Validate,
 };
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Debug, Validate)]
 pub struct ResolveIdentityPayload {
     #[serde(rename = "publicKey")]
-    public_key: String,
-}
-
-#[derive(Validate, Debug)]
-pub struct ResolveIdentityParams {
     #[validate(custom = "validate_identity_key")]
-    identity_key: String,
+    public_key: String,
 }
 
 #[derive(Serialize)]
@@ -39,15 +35,14 @@ impl From<ResolveIdentityResponse> for Value {
     }
 }
 
+#[instrument(name = "resolve_handler", skip(state))]
 pub async fn handler(
     State(state): State<Arc<AppState>>,
-    Query(payload): Query<ResolveIdentityPayload>,
+    Query(params): Query<ResolveIdentityPayload>,
 ) -> error::Result<Response> {
-    let params = ResolveIdentityParams {
-        identity_key: payload.public_key,
-    };
     info!("Handling - Resolve identity with params: {:?}", params);
 
+    info!("Timing - Validating params - Start");
     params.validate().map_err(|error| {
         info!(
             "Failure - Resolve identity with params: {:?}, error: {:?}",
@@ -55,10 +50,12 @@ pub async fn handler(
         );
         error
     })?;
+    info!("Timing - Validating params - End");
 
+    info!("Timing - get_cacao_by_identity_key - Start");
     let cacao = state
         .keys_persitent_storage
-        .get_cacao_by_identity_key(&params.identity_key)
+        .get_cacao_by_identity_key(&params.public_key)
         .await
         .map_err(|error| {
             warn!(
@@ -67,11 +64,13 @@ pub async fn handler(
             );
             error
         })?;
+    info!("Timing - get_cacao_by_identity_key - End");
 
     let response = ResolveIdentityResponse { cacao };
 
     info!("Success - Resolve identity with params: {:?}", params);
     increment_counter!(state.metrics, identity_resolved);
+    info!("Incremented counter");
 
     Ok(Response::new_success_with_value(
         StatusCode::OK,
diff --git a/src/stores/keys.rs b/src/stores/keys.rs
index 6b6a00c..98335f8 100644
--- a/src/stores/keys.rs
+++ b/src/stores/keys.rs
@@ -208,12 +208,15 @@ impl KeysPersistentStorage for MongoPersistentStorage {
     }
 
     async fn get_cacao_by_identity_key(&self, identity_key: &str) -> Result<Cacao, StoreError> {
+        info!("get_cacao_by_identity_key");
         let filter = doc! {
             "identities.identity_key": identity_key,
         };
 
+        info!("constructing not_found");
         let not_found = StoreError::NotFound("Identity key".to_string(), identity_key.to_string());
 
+        info!("find_one");
         match MongoKeys::find_one(&self.db, Some(filter), None).await? {
             Some(mongo_keys) => {
                 info!(
@@ -223,11 +226,14 @@ impl KeysPersistentStorage for MongoPersistentStorage {
                     mongo_keys
                 );
 
+                info!("Timing - find - Start");
                 let mongo_identity = mongo_keys
                     .identities
                     .into_iter()
                     .find(|i| i.identity_key == *identity_key)
                     .ok_or(not_found)?;
+                info!("Timing - find - End");
+
                 Ok(mongo_identity.cacao)
             }
             None => Err(not_found),