From 6d27c4366619001722b6bc878d37c2f3c3ee5b25 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Tue, 28 May 2019 18:05:52 +1000
Subject: [PATCH] Fix bug with metrics scraping

---
 beacon_node/http_server/src/key.rs     |  7 +++
 beacon_node/http_server/src/metrics.rs | 64 +++++++++++++++-----------
 2 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/beacon_node/http_server/src/key.rs b/beacon_node/http_server/src/key.rs
index 2cbe68cd0..b84c5f685 100644
--- a/beacon_node/http_server/src/key.rs
+++ b/beacon_node/http_server/src/key.rs
@@ -1,3 +1,4 @@
+use crate::metrics::LocalMetrics;
 use beacon_chain::{BeaconChain, BeaconChainTypes};
 use iron::typemap::Key;
 use prometheus::Registry;
@@ -17,3 +18,9 @@ pub struct MetricsRegistryKey;
 impl Key for MetricsRegistryKey {
     type Value = Registry;
 }
+
+pub struct LocalMetricsKey;
+
+impl Key for LocalMetricsKey {
+    type Value = LocalMetrics;
+}
diff --git a/beacon_node/http_server/src/metrics.rs b/beacon_node/http_server/src/metrics.rs
index 608f38efd..4c243ce71 100644
--- a/beacon_node/http_server/src/metrics.rs
+++ b/beacon_node/http_server/src/metrics.rs
@@ -1,12 +1,12 @@
 use crate::{
-    key::{BeaconChainKey, MetricsRegistryKey},
+    key::{BeaconChainKey, LocalMetricsKey, MetricsRegistryKey},
     map_persistent_err_to_500,
 };
 use beacon_chain::{BeaconChain, BeaconChainTypes};
 use iron::prelude::*;
 use iron::{status::Status, Handler, IronResult, Request, Response};
 use persistent::Read;
-use prometheus::{Encoder, IntCounter, Opts, Registry, TextEncoder};
+use prometheus::{Encoder, IntGauge, Opts, Registry, TextEncoder};
 use slot_clock::SlotClock;
 use std::sync::Arc;
 use types::Slot;
@@ -18,12 +18,42 @@ pub fn build_handler<T: BeaconChainTypes + 'static>(
 ) -> impl Handler {
     let mut chain = Chain::new(handle_metrics::<T>);
 
+    let local_metrics = LocalMetrics::new().unwrap();
+    local_metrics.register(&metrics_registry).unwrap();
+
     chain.link(Read::<BeaconChainKey<T>>::both(beacon_chain));
     chain.link(Read::<MetricsRegistryKey>::both(metrics_registry));
+    chain.link(Read::<LocalMetricsKey>::both(local_metrics));
 
     chain
 }
 
+pub struct LocalMetrics {
+    present_slot: IntGauge,
+    validator_count: IntGauge,
+}
+
+impl LocalMetrics {
+    pub fn new() -> Result<Self, prometheus::Error> {
+        Ok(Self {
+            present_slot: {
+                let opts = Opts::new("present_slot", "slot_at_time_of_scrape");
+                IntGauge::with_opts(opts)?
+            },
+            validator_count: {
+                let opts = Opts::new("validator_count", "number_of_validators");
+                IntGauge::with_opts(opts)?
+            },
+        })
+    }
+
+    pub fn register(&self, registry: &Registry) -> Result<(), prometheus::Error> {
+        registry.register(Box::new(self.present_slot.clone()))?;
+
+        Ok(())
+    }
+}
+
 /// Handle a request for Prometheus metrics.
 ///
 /// Returns a text string containing all metrics.
@@ -36,23 +66,19 @@ fn handle_metrics<T: BeaconChainTypes + 'static>(req: &mut Request) -> IronResul
         .get::<Read<MetricsRegistryKey>>()
         .map_err(map_persistent_err_to_500)?;
 
+    let local_metrics = req
+        .get::<Read<LocalMetricsKey>>()
+        .map_err(map_persistent_err_to_500)?;
+
     let present_slot = beacon_chain
         .slot_clock
         .present_slot()
         .unwrap_or_else(|_| None)
         .unwrap_or_else(|| Slot::new(0));
-    register_and_set_slot(&r, "present_slot", "slock_clock_reading", present_slot);
-
-    let best_slot = beacon_chain.head().beacon_block.slot;
-    register_and_set_slot(&r, "best_slot", "slot_of_block_at_head_of_chain", best_slot);
+    local_metrics.present_slot.set(present_slot.as_u64() as i64);
 
     let validator_count = beacon_chain.head().beacon_state.validator_registry.len();
-    register_and_set(
-        &r,
-        "validator_count",
-        "total_number_of_validators",
-        validator_count as i64,
-    );
+    local_metrics.validator_count.set(validator_count as i64);
 
     // Gather the metrics.
     let mut buffer = vec![];
@@ -64,17 +90,3 @@ fn handle_metrics<T: BeaconChainTypes + 'static>(req: &mut Request) -> IronResul
 
     Ok(Response::with((Status::Ok, prom_string)))
 }
-
-fn register_and_set(registry: &Registry, name: &str, help: &str, value: i64) {
-    let counter_opts = Opts::new(name, help);
-    let counter = IntCounter::with_opts(counter_opts).unwrap();
-    registry.register(Box::new(counter.clone())).unwrap();
-    counter.inc_by(value);
-}
-
-fn register_and_set_slot(registry: &Registry, name: &str, help: &str, slot: Slot) {
-    let counter_opts = Opts::new(name, help);
-    let counter = IntCounter::with_opts(counter_opts).unwrap();
-    registry.register(Box::new(counter.clone())).unwrap();
-    counter.inc_by(slot.as_u64() as i64);
-}