refactor(hermes): state->wormhole/metrics downcasting

Author: Reisen

apps/hermes/src/api.rs

@@ -27,7 +27,7 @@ mod ws;
 pub struct ApiState<S = State> {
     pub state:   Arc<S>,
     pub ws:      Arc<ws::WsState>,
-    pub metrics: Arc<metrics_middleware::Metrics>,
+    pub metrics: Arc<metrics_middleware::ApiMetrics>,
 }
 
 /// Manually implement `Clone` as the derive macro will try and slap `Clone` on
@@ -49,7 +49,7 @@ impl ApiState<State> {
         requester_ip_header_name: String,
     ) -> Self {
         Self {
-            metrics: Arc::new(metrics_middleware::Metrics::new(state.clone())),
+            metrics: Arc::new(metrics_middleware::ApiMetrics::new(state.clone())),
             ws: Arc::new(ws::WsState::new(
                 ws_whitelist,
                 requester_ip_header_name,

apps/hermes/src/api/metrics_middleware.rs

@@ -1,6 +1,6 @@
 use {
     super::ApiState,
-    crate::state::State as AppState,
+    crate::state::metrics::Metrics,
     axum::{
         extract::{
             MatchedPath,
@@ -22,13 +22,19 @@ use {
     tokio::time::Instant,
 };
 
-pub struct Metrics {
+pub struct ApiMetrics {
     pub requests:  Family<Labels, Counter>,
     pub latencies: Family<Labels, Histogram>,
 }
 
-impl Metrics {
-    pub fn new(state: Arc<AppState>) -> Self {
+impl ApiMetrics {
+    pub fn new<S>(state: Arc<S>) -> Self
+    where
+        S: Metrics,
+        S: Send,
+        S: Sync,
+        S: 'static,
+    {
         let new = Self {
             requests:  Family::default(),
             latencies: Family::new_with_constructor(|| {
@@ -46,15 +52,21 @@ impl Metrics {
             let latencies = new.latencies.clone();
 
             tokio::spawn(async move {
-                let mut metrics_registry = state.metrics_registry.write().await;
-
-                metrics_registry.register("api_requests", "Total number of API requests", requests);
+                Metrics::register(
+                    &*state,
+                    ("api_requests", "Total number of API requests", requests),
+                )
+                .await;
 
-                metrics_registry.register(
-                    "api_request_latency_seconds",
-                    "API request latency in seconds",
-                    latencies,
-                );
+                Metrics::register(
+                    &*state,
+                    (
+                        "api_request_latency_seconds",
+                        "API request latency in seconds",
+                        latencies,
+                    ),
+                )
+                .await;
             });
         }
 
@@ -80,21 +92,18 @@ pub async fn track_metrics<B>(
     } else {
         req.uri().path().to_owned()
     };
-    let method = req.method().clone();
 
+    let method = req.method().clone();
     let response = next.run(req).await;
-
     let latency = start.elapsed().as_secs_f64();
     let status = response.status().as_u16();
-
     let labels = Labels {
         method: method.to_string(),
         path,
         status,
     };
 
     api_state.metrics.requests.get_or_create(&labels).inc();
-
     api_state
         .metrics
         .latencies

apps/hermes/src/api/ws.rs

@@ -12,6 +12,7 @@ use {
             AggregationEvent,
             RequestTime,
         },
+        metrics::Metrics,
         State,
     },
     anyhow::{
@@ -115,12 +116,18 @@ pub struct Labels {
     pub status:      Status,
 }
 
-pub struct Metrics {
+pub struct WsMetrics {
     pub interactions: Family<Labels, Counter>,
 }
 
-impl Metrics {
-    pub fn new(state: Arc<State>) -> Self {
+impl WsMetrics {
+    pub fn new<S>(state: Arc<S>) -> Self
+    where
+        S: Metrics,
+        S: Send,
+        S: Sync,
+        S: 'static,
+    {
         let new = Self {
             interactions: Family::default(),
         };
@@ -129,11 +136,15 @@ impl Metrics {
             let interactions = new.interactions.clone();
 
             tokio::spawn(async move {
-                state.metrics_registry.write().await.register(
-                    "ws_interactions",
-                    "Total number of websocket interactions",
-                    interactions,
-                );
+                Metrics::register(
+                    &*state,
+                    (
+                        "ws_interactions",
+                        "Total number of websocket interactions",
+                        interactions,
+                    ),
+                )
+                .await;
             });
         }
 
@@ -146,7 +157,7 @@ pub struct WsState {
     pub bytes_limit_whitelist:    Vec<IpNet>,
     pub rate_limiter:             DefaultKeyedRateLimiter<IpAddr>,
     pub requester_ip_header_name: String,
-    pub metrics:                  Metrics,
+    pub metrics:                  WsMetrics,
 }
 
 impl WsState {
@@ -158,7 +169,7 @@ impl WsState {
             ))),
             bytes_limit_whitelist: whitelist,
             requester_ip_header_name,
-            metrics: Metrics::new(state.clone()),
+            metrics: WsMetrics::new(state.clone()),
         }
     }
 }

apps/hermes/src/metrics_server.rs

@@ -5,7 +5,10 @@
 use {
     crate::{
         config::RunOptions,
-        state::State as AppState,
+        state::{
+            metrics::Metrics,
+            State as AppState,
+        },
     },
     anyhow::Result,
     axum::{
@@ -15,7 +18,6 @@ use {
         routing::get,
         Router,
     },
-    prometheus_client::encoding::text::encode,
     std::sync::Arc,
 };
 
@@ -43,13 +45,7 @@ pub async fn run(opts: RunOptions, state: Arc<AppState>) -> Result<()> {
 }
 
 pub async fn metrics(State(state): State<Arc<AppState>>) -> impl IntoResponse {
-    let registry = state.metrics_registry.read().await;
-    let mut buffer = String::new();
-
-    // Should not fail if the metrics are valid and there is memory available
-    // to write to the buffer.
-    encode(&mut buffer, &registry).unwrap();
-
+    let buffer = Metrics::encode(&*state).await;
     (
         [(
             header::CONTENT_TYPE,

apps/hermes/src/network/pythnet.rs

@@ -7,7 +7,6 @@ use {
         api::types::PriceFeedMetadata,
         config::RunOptions,
         network::wormhole::{
-            update_guardian_set,
             BridgeData,
             GuardianSet,
             GuardianSetData,
@@ -22,6 +21,7 @@ use {
                 Aggregates,
                 Update,
             },
+            wormhole::Wormhole,
             State,
         },
     },
@@ -215,11 +215,14 @@ pub async fn run(store: Arc<State>, pythnet_ws_endpoint: String) -> Result<!> {
 /// This method performs the necessary work to pull down the bridge state and associated guardian
 /// sets from a deployed Wormhole contract. Note that we only fetch the last two accounts due to
 /// the fact that during a Wormhole upgrade, there will only be messages produces from those two.
-async fn fetch_existing_guardian_sets(
-    state: Arc<State>,
+async fn fetch_existing_guardian_sets<S>(
+    state: Arc<S>,
     pythnet_http_endpoint: String,
     wormhole_contract_addr: Pubkey,
-) -> Result<()> {
+) -> Result<()>
+where
+    S: Wormhole,
+{
     let client = RpcClient::new(pythnet_http_endpoint.to_string());
     let bridge = fetch_bridge_data(&client, &wormhole_contract_addr).await?;
 
@@ -233,7 +236,7 @@ async fn fetch_existing_guardian_sets(
         "Retrieved Current GuardianSet.",
     );
 
-    update_guardian_set(&state, bridge.guardian_set_index, current).await;
+    Wormhole::update_guardian_set(&*state, bridge.guardian_set_index, current).await;
 
     // If there are more than one guardian set, we want to fetch the previous one as well as it
     // may still be in transition phase if a guardian upgrade has just occurred.
@@ -251,7 +254,7 @@ async fn fetch_existing_guardian_sets(
             "Retrieved Previous GuardianSet.",
         );
 
-        update_guardian_set(&state, bridge.guardian_set_index - 1, previous).await;
+        Wormhole::update_guardian_set(&*state, bridge.guardian_set_index - 1, previous).await;
     }
 
     Ok(())