From 318ffde0506b9fb6cff1ebb5755da59e78d8e8a8 Mon Sep 17 00:00:00 2001
From: Katalina Okano <git@kat.cafe>
Date: Sat, 26 Nov 2022 01:37:37 -0500
Subject: [PATCH] add health checks to gql

---
 services/api/src/graphql/mod.rs   |  8 ++-
 services/api/src/graphql/types.rs | 82 +++++++++++++++++++++++++++++++
 services/websocket/src/main.rs    |  8 +--
 3 files changed, 93 insertions(+), 5 deletions(-)

diff --git a/services/api/src/graphql/mod.rs b/services/api/src/graphql/mod.rs
index 6d7683b..dc35abb 100644
--- a/services/api/src/graphql/mod.rs
+++ b/services/api/src/graphql/mod.rs
@@ -1,7 +1,7 @@
 use crate::redispool::RedisPool;
 
-use self::types::World;
-use juniper::{graphql_object, FieldResult, ID};
+use self::types::{Health, World};
+use juniper::{graphql_object, meta::Field, FieldResult, ID};
 use rocket::response::content::RawHtml;
 
 pub mod types;
@@ -60,6 +60,10 @@ impl Query {
             id: ID::from(id.to_string()),
         })
     }
+
+    fn health() -> FieldResult<Health> {
+        Ok(Health {})
+    }
 }
 
 pub struct Context {
diff --git a/services/api/src/graphql/types.rs b/services/api/src/graphql/types.rs
index 979092e..4f5b218 100644
--- a/services/api/src/graphql/types.rs
+++ b/services/api/src/graphql/types.rs
@@ -275,3 +275,85 @@ impl Classes {
             .await
     }
 }
+
+#[derive(juniper::GraphQLEnum)]
+enum WebsocketState {
+    #[graphql(
+        description = "Using Nanite Systems manifold. This is the best possible running state."
+    )]
+    Primary,
+
+    #[graphql(
+        description = "Using backup Daybreak Games manifold. This means the primary socket hasn't recieved events for at least 60 seconds."
+    )]
+    Backup,
+
+    #[graphql(description = "Both event processors are down. This is bad.")]
+    Down,
+}
+
+#[derive(juniper::GraphQLEnum)]
+enum UpDown {
+    #[graphql(description = "Checks have passed.")]
+    Up,
+
+    #[graphql(description = "Checks have failed. This is bad.")]
+    Down,
+}
+
+pub struct Health {}
+
+impl Health {
+    async fn get_heartbeat(context: &super::Context, pair: &str) -> WebsocketState {
+        let mut con = (*context).con.get().await.unwrap();
+
+        let res: Result<i32, _> = cmd("GET")
+            .arg(format!("heartbeat:{}:primary", pair))
+            .query_async(&mut *con)
+            .await;
+        match res {
+            Ok(_) => WebsocketState::Primary,
+            Err(_) => {
+                let res: Result<i32, _> = cmd("GET")
+                    .arg(format!("heartbeat:{}:backup", pair))
+                    .query_async(&mut con)
+                    .await;
+                match res {
+                    Ok(_) => WebsocketState::Backup,
+                    Err(_) => WebsocketState::Down,
+                }
+            }
+        }
+    }
+}
+
+#[graphql_object(context = super::Context)]
+#[graphql(description = "Saerro's self-checks. Down is universally bad.")]
+impl Health {
+    #[graphql(description = "Checks PC event processors for its running state.")]
+    async fn pc(context: &super::Context) -> WebsocketState {
+        Health::get_heartbeat(context, "pc").await
+    }
+
+    #[graphql(description = "Checks PS4 US event processors for its running state.")]
+    async fn ps4us(context: &super::Context) -> WebsocketState {
+        Health::get_heartbeat(context, "ps4us").await
+    }
+
+    #[graphql(description = "Checks PS4 EU event processors for its running state.")]
+    async fn ps4eu(context: &super::Context) -> WebsocketState {
+        Health::get_heartbeat(context, "ps4eu").await
+    }
+
+    #[graphql(description = "Is our datastore working?")]
+    async fn redis(context: &super::Context) -> UpDown {
+        let mut con = (*context).con.get().await.unwrap();
+
+        let res: Result<String, _> = cmd("PING").query_async(&mut con).await;
+
+        match res {
+            Ok(_) => UpDown::Up,
+            Err(_) => UpDown::Down,
+        }
+    }
+}
diff --git a/services/websocket/src/main.rs b/services/websocket/src/main.rs
index 35e8bd0..92ce41a 100644
--- a/services/websocket/src/main.rs
+++ b/services/websocket/src/main.rs
@@ -120,14 +120,16 @@ async fn track_class(class_event: ClassEvent) {
 fn should_process_event() -> bool {
     let mut con = REDIS_CLIENT.get_connection().unwrap();
     let role: String = ROLE.parse().unwrap();
-    let heartbeat_key = format!("heartbeat:{}", PAIR.to_string());
+
+    let heartbeat_key = format!("heartbeat:{}:{}", PAIR.to_string(), role);
+    let _: () = con.set_ex(heartbeat_key, "1", 60).unwrap();
 
     if role == "primary" {
-        let _: () = con.set_ex(heartbeat_key, "1", 60).unwrap();
         return false;
     }
 
-    match con.get(heartbeat_key) {
+    let primary_heartbeat_key = format!("heartbeat:{}:primary", PAIR.to_string());
+    match con.get(primary_heartbeat_key) {
         Ok(1) => true,
         _ => false,
     }