From a76ec207b3a091630c8ea74fc8b6a85f2d78cf2a Mon Sep 17 00:00:00 2001
From: Paul Makles <paulmakles@gmail.com>
Date: Fri, 29 Jan 2021 21:56:37 +0000
Subject: [PATCH] Add Prometheus metrics.

---
 Cargo.lock            | 26 +++++++++++++++++++++++++-
 Cargo.toml            |  7 ++++---
 src/main.rs           | 21 ++++++++++++++++-----
 src/util/variables.rs |  3 ++-
 4 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8dc2be3..b0fded2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1971,6 +1971,20 @@ dependencies = [
  "yansi",
 ]
 
+[[package]]
+name = "prometheus"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c8425533e7122f0c3cc7a37e6244b16ad3a2cc32ae7ac6276e2a75da0d9c200d"
+dependencies = [
+ "cfg-if 1.0.0",
+ "fnv",
+ "lazy_static",
+ "parking_lot 0.11.1",
+ "regex",
+ "thiserror",
+]
+
 [[package]]
 name = "quick-error"
 version = "1.2.3"
@@ -2351,7 +2365,7 @@ dependencies = [
 
 [[package]]
 name = "revolt"
-version = "0.3.2"
+version = "0.3.3"
 dependencies = [
  "async-std",
  "async-tungstenite",
@@ -2377,6 +2391,7 @@ dependencies = [
  "rocket",
  "rocket_contrib",
  "rocket_cors",
+ "rocket_prometheus",
  "serde",
  "serde_json",
  "snafu",
@@ -2493,6 +2508,15 @@ dependencies = [
  "version_check",
 ]
 
+[[package]]
+name = "rocket_prometheus"
+version = "0.7.0"
+source = "git+https://github.com/insertish/rocket_prometheus?rev=3d825aedb42793246c306a81fe67c5b187948983#3d825aedb42793246c306a81fe67c5b187948983"
+dependencies = [
+ "prometheus",
+ "rocket",
+]
+
 [[package]]
 name = "rust-argon2"
 version = "0.8.3"
diff --git a/Cargo.toml b/Cargo.toml
index 7b62b25..80fbcb6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "revolt"
-version = "0.3.2"
+version = "0.3.3"
 authors = ["Paul Makles <paulmakles@gmail.com>"]
 edition = "2018"
 
@@ -8,15 +8,16 @@ edition = "2018"
 
 [dependencies]
 futures = "0.3.8"
-many-to-many = "0.1.2"
 impl_ops = "0.1.1"
+many-to-many = "0.1.2"
 ctrlc = { version = "3.0", features = ["termination"] }
+async-std = { version = "1.8.0", features = ["tokio02", "attributes"] }
 async-tungstenite = { version = "0.10.0", features = ["async-std-runtime"] }
 rauth = { git = "https://gitlab.insrt.uk/insert/rauth", rev = "d48b0087cc546d2c9ef53f27d0693c8378172dfe" }
-async-std = { version = "1.8.0", features = ["tokio02", "attributes"] }
 
 hive_pubsub = { version = "0.4.3", features = ["mongo"] }
 rocket_cors = { git = "https://github.com/insertish/rocket_cors", branch = "master" }
+rocket_prometheus = { git = "https://github.com/insertish/rocket_prometheus", rev = "3d825aedb42793246c306a81fe67c5b187948983" }
 rocket_contrib = { git = "https://github.com/SergioBenitez/Rocket", rev = "031948c1daaa146128d8a435be116476f2adde00" }
 rocket = { git = "https://github.com/SergioBenitez/Rocket", rev = "031948c1daaa146128d8a435be116476f2adde00", default-features = false }
 mongodb = { version = "1.1.1", features = ["tokio-runtime"], default-features = false }
diff --git a/src/main.rs b/src/main.rs
index 690a8ab..efe439b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,13 +18,14 @@ pub mod notifications;
 pub mod routes;
 pub mod util;
 
-use chrono::Duration;
-use futures::join;
 use log::info;
+use futures::join;
+use chrono::Duration;
 use rauth::auth::Auth;
-use rauth::options::{EmailVerification, Options, SMTP};
 use rocket_cors::AllowedOrigins;
-use util::variables::{PUBLIC_URL, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME, USE_EMAIL};
+use rocket_prometheus::PrometheusMetrics;
+use rauth::options::{EmailVerification, Options, SMTP};
+use util::variables::{PUBLIC_URL, SMTP_FROM, SMTP_HOST, SMTP_PASSWORD, SMTP_USERNAME, USE_EMAIL, USE_PROMETHEUS};
 
 #[async_std::main]
 async fn main() {
@@ -80,7 +81,17 @@ async fn launch_web() {
             }),
     );
 
-    routes::mount(rocket::ignite())
+    let mut rocket = rocket::ignite();
+
+    if *USE_PROMETHEUS {
+        info!("Enabled Prometheus metrics!");
+        let prometheus = PrometheusMetrics::new();
+        rocket = rocket
+            .attach(prometheus.clone())
+            .mount("/metrics", prometheus);
+    }
+
+    routes::mount(rocket)
         .mount("/", rocket_cors::catch_all_options_routes())
         .mount("/auth", rauth::routes::routes())
         .manage(auth)
diff --git a/src/util/variables.rs b/src/util/variables.rs
index 7630e33..cbf01cc 100644
--- a/src/util/variables.rs
+++ b/src/util/variables.rs
@@ -19,7 +19,7 @@ lazy_static! {
         env::var("REVOLT_WS_HOST").unwrap_or_else(|_| "0.0.0.0:9000".to_string());
 
     // Application Flags
-    pub static ref DISABLE_REGISTRATION: bool = env::var("REVOLT_DISABLE_REGISTRATION").map_or(false, |v| v == "*1");
+    pub static ref DISABLE_REGISTRATION: bool = env::var("REVOLT_DISABLE_REGISTRATION").map_or(false, |v| v == "1");
     pub static ref USE_EMAIL: bool = env::var("REVOLT_USE_EMAIL_VERIFICATION").map_or(
         env::var("REVOLT_SMTP_HOST").is_ok()
             && env::var("REVOLT_SMTP_USERNAME").is_ok()
@@ -28,6 +28,7 @@ lazy_static! {
         |v| v == *"1"
     );
     pub static ref USE_HCAPTCHA: bool = env::var("REVOLT_HCAPTCHA_KEY").is_ok();
+    pub static ref USE_PROMETHEUS: bool = env::var("REVOLT_ENABLE_PROMETHEUS").map_or(false, |v| v == "1");
 
     // SMTP Settings
     pub static ref SMTP_HOST: String =
-- 
GitLab