feat: feature flag metrics in node

meilin96 · Feb 24, 2022 · 7bc5642 · 7bc5642
1 parent b8a7ce9
commit 7bc5642
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 14 deletions.
diff --git a/network/src/ledger.rs b/network/src/ledger.rs
@@ -31,7 +31,7 @@ use snarkos_environment::{
 use snarkos_storage::{storage::Storage, BlockLocators, LedgerState, MAXIMUM_LINEAR_BLOCK_LOCATORS};
 use snarkvm::dpc::prelude::*;
 
-#[cfg(feature = "prometheus")]
+#[cfg(any(feature = "test", feature = "prometheus"))]
 use snarkos_metrics as metrics;
 
 use anyhow::Result;
@@ -521,7 +521,7 @@ impl<N: Network, E: Environment> Ledger<N, E> {
                             self.canon.latest_block_hash()
                         );
 
-                        #[cfg(feature = "prometheus")]
+                        #[cfg(any(feature = "test", feature = "prometheus"))]
                         metrics::gauge!(metrics::blocks::HEIGHT, latest_block_height as f64);
 
                         // Update the timestamp of the last block increment.
@@ -574,7 +574,7 @@ impl<N: Network, E: Environment> Ledger<N, E> {
                 let latest_block_height = self.canon.latest_block_height();
                 info!("Ledger successfully reverted to block {}", latest_block_height);
 
-                #[cfg(feature = "prometheus")]
+                #[cfg(any(feature = "test", feature = "prometheus"))]
                 metrics::gauge!(metrics::blocks::HEIGHT, latest_block_height as f64);
 
                 // Update the last block update timestamp.

diff --git a/network/src/peers.rs b/network/src/peers.rs
@@ -18,7 +18,7 @@ use crate::{Data, DisconnectReason, LedgerReader, LedgerRouter, Message, Operato
 use snarkos_environment::Environment;
 use snarkvm::dpc::prelude::*;
 
-#[cfg(feature = "prometheus")]
+#[cfg(any(feature = "test", feature = "prometheus"))]
 use snarkos_metrics as metrics;
 
 use anyhow::Result;
@@ -516,7 +516,7 @@ impl<N: Network, E: Environment> Peers<N, E> {
                 // Remove an entry for this `Peer` in the candidate peers, if it exists.
                 self.candidate_peers.write().await.remove(&peer_ip);
 
-                #[cfg(feature = "prometheus")]
+                #[cfg(any(feature = "test", feature = "prometheus"))]
                 {
                     let number_of_connected_peers = self.number_of_connected_peers().await;
                     let number_of_candidate_peers = self.number_of_candidate_peers().await;
@@ -530,7 +530,7 @@ impl<N: Network, E: Environment> Peers<N, E> {
                 // Add an entry for this `Peer` in the candidate peers.
                 self.candidate_peers.write().await.insert(peer_ip);
 
-                #[cfg(feature = "prometheus")]
+                #[cfg(any(feature = "test", feature = "prometheus"))]
                 {
                     let number_of_connected_peers = self.number_of_connected_peers().await;
                     let number_of_candidate_peers = self.number_of_candidate_peers().await;
@@ -544,7 +544,7 @@ impl<N: Network, E: Environment> Peers<N, E> {
                 // Add an entry for this `Peer` in the restricted peers.
                 self.restricted_peers.write().await.insert(peer_ip, Instant::now());
 
-                #[cfg(feature = "prometheus")]
+                #[cfg(any(feature = "test", feature = "prometheus"))]
                 {
                     let number_of_connected_peers = self.number_of_connected_peers().await;
                     let number_of_restricted_peers = self.number_of_restricted_peers().await;
@@ -560,7 +560,7 @@ impl<N: Network, E: Environment> Peers<N, E> {
             PeersRequest::ReceivePeerResponse(peer_ips) => {
                 self.add_candidate_peers(peer_ips.iter()).await;
 
-                #[cfg(feature = "prometheus")]
+                #[cfg(any(feature = "test", feature = "prometheus"))]
                 {
                     let number_of_candidate_peers = self.number_of_candidate_peers().await;
                     metrics::gauge!(metrics::peers::CANDIDATE, number_of_candidate_peers as f64);
@@ -601,7 +601,7 @@ impl<N: Network, E: Environment> Peers<N, E> {
                     trace!("Outbound channel failed: {}", error);
                     self.connected_peers.write().await.remove(&peer);
 
-                    #[cfg(feature = "prometheus")]
+                    #[cfg(any(feature = "test", feature = "prometheus"))]
                     {
                         let number_of_connected_peers = self.number_of_connected_peers().await;
                         metrics::gauge!(metrics::peers::CONNECTED, number_of_connected_peers as f64);

diff --git a/snarkos/server.rs b/snarkos/server.rs
@@ -30,10 +30,12 @@ use snarkvm::prelude::*;
 
 #[cfg(feature = "rpc")]
 use snarkos_rpc::{initialize_rpc_server, RpcContext};
+
+#[cfg(any(feature = "test", feature = "prometheus"))]
+use snarkos_metrics as metrics;
+
 #[cfg(feature = "rpc")]
 use tokio::sync::RwLock;
-#[cfg(feature = "prometheus")]
-use snarkos_metrics as metrics;
 
 use anyhow::Result;
 use std::{net::SocketAddr, sync::Arc, time::Duration};
@@ -176,7 +178,7 @@ impl<N: Network, E: Environment> Server<N, E> {
         Self::initialize_notification(ledger.reader(), prover.clone(), address).await;
 
         // Initialise the metrics exporter.
-        #[cfg(feature = "prometheus")]
+        #[cfg(any(feature = "test", feature = "prometheus"))]
         Self::initialize_metrics(ledger.reader());
 
         Ok(Self {
@@ -425,9 +427,12 @@ impl<N: Network, E: Environment> Server<N, E> {
         let _ = handler.await;
     }
 
-    #[cfg(feature = "prometheus")]
+    #[cfg(any(feature = "test", feature = "prometheus"))]
     fn initialize_metrics(ledger: LedgerReader<N>) {
-        E::tasks().append(snarkos_metrics::initialize().expect("couldn't initialise the metrics"));
+        #[cfg(not(feature = "test"))]
+        if let Some(handler) = snarkos_metrics::initialize() {
+            E::tasks().append(handler);
+        }
 
         // Set the block height as it could already be non-zero.
         metrics::gauge!(metrics::blocks::HEIGHT, ledger.latest_block_height() as f64);