Add metric for number of heartbeat tokens

src/metrics.rs

@@ -4,10 +4,12 @@
 //! to allow exposting it on a private network only
 //! independently of the main service.
 
+use std::sync::atomic::AtomicI64;
 use std::sync::Arc;
 
 use prometheus_client::encoding::text::encode;
 use prometheus_client::metrics::counter::Counter;
+use prometheus_client::metrics::gauge::Gauge;
 use prometheus_client::registry::Registry;
 
 use anyhow::Result;
@@ -16,28 +18,47 @@ use anyhow::Result;
 pub struct Metrics {
     pub registry: Registry,
 
+    /// Number of successfully sent visible notifications.
     pub direct_notifications_total: Counter,
+
+    /// Number of tokens registered for heartbeat notifications.
+    pub heartbeat_token_count: Gauge<i64, AtomicI64>,
 }
 
 impl Metrics {
     pub fn new() -> Self {
         let mut registry = Registry::default();
+
         let direct_notifications_total = Counter::default();
         registry.register(
             "direct_notifications",
             "Number of direct notifications",
             direct_notifications_total.clone(),
         );
+
+        let heartbeat_token_count = Gauge::<i64, AtomicI64>::default();
+        registry.register(
+            "heartbeat_token_count",
+            "Number of tokens registered for heartbeat notifications",
+            heartbeat_token_count.clone(),
+        );
+
         Self {
             registry,
             direct_notifications_total,
+            heartbeat_token_count,
         }
     }
 
     /// Counts direct notification.
     pub fn inc_direct_notification(&self) {
         self.direct_notifications_total.inc();
     }
+
+    /// Sets number of tokens registered for heartbeat notifications.
+    pub fn set_heartbeat_token_count(&self, value: usize) {
+        self.heartbeat_token_count.set(value as i64);
+    }
 }
 
 type State = Arc<Metrics>;

src/notifier.rs

@@ -6,10 +6,12 @@ use anyhow::Result;
 use async_std::prelude::*;
 use log::*;
 
+use crate::metrics::Metrics;
 use crate::state::State;
 
 pub async fn start(state: State, interval: std::time::Duration) -> Result<()> {
     let db = state.db();
+    let metrics = state.metrics();
     let production_client = state.production_client();
     let sandbox_client = state.sandbox_client();
     let topic = state.topic();
@@ -20,19 +22,20 @@ pub async fn start(state: State, interval: std::time::Duration) -> Result<()> {
     );
 
     // first wakeup on startup
-    wakeup(db, production_client, sandbox_client, topic).await;
+    wakeup(db, metrics, production_client, sandbox_client, topic).await;
 
     // create interval
     let mut interval = async_std::stream::interval(interval);
     while interval.next().await.is_some() {
-        wakeup(db, production_client, sandbox_client, topic).await;
+        wakeup(db, metrics, production_client, sandbox_client, topic).await;
     }
 
     Ok(())
 }
 
 async fn wakeup(
     db: &sled::Db,
+    metrics: &Metrics,
     production_client: &Client,
     sandbox_client: &Client,
     topic: Option<&str>,
@@ -46,6 +49,7 @@ async fn wakeup(
         .collect::<Vec<_>>();
 
     info!("sending notifications to {} devices", tokens.len());
+    metrics.set_heartbeat_token_count(tokens.len());
 
     for device_token in tokens {
         info!("notify: {}", device_token);