feat(backend): add WebSocket check (codemod-com#1052)

* websocket healthcheck * added checking wether service is up and bump version * fixed import errors * fixed import errors --------- Co-authored-by: mirai2k <[email protected]>
mohebifar · Jul 17, 2024 · 68fe224 · 68fe224
1 parent 82a2f5e
commit 68fe224
Show file tree

Hide file tree

Showing 3 changed files with 981 additions and 909 deletions.
diff --git a/apps/backend/package.json b/apps/backend/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@codemod-com/backend",
-  "version": "0.0.133",
+  "version": "0.0.134",
   "scripts": {
     "build": "tsc && node esbuild.config.js",
     "start": "node build/index.js",
@@ -22,6 +22,7 @@
     "@types/pg": "catalog:",
     "@types/semver": "7.5.8",
     "@types/supertest": "catalog:",
+    "@types/ws": "^8.5.11",
     "dotenv-cli": "catalog:",
     "esbuild": "^0.19.12",
     "esbuild-plugin-copy": "catalog:",
@@ -63,6 +64,7 @@
     "semver": "7.6.0",
     "tar": "^6.2.0",
     "valibot": "catalog:",
+    "ws": "^8.18.0",
     "zod": "catalog:"
   },
   "type": "module",

diff --git a/apps/backend/src/cron.ts b/apps/backend/src/cron.ts
@@ -1,11 +1,11 @@
 import { prisma } from "@codemod-com/database";
 import { WebClient } from "@slack/web-api";
+import axios from "axios";
 import { CronJob } from "cron";
+import WebSocket from "ws";
 import { PostHogService } from "./services/PostHogService";
 import { environment } from "./util";
 
-import axios from "axios";
-
 // TODO: Move crons into independent CronService
 
 const cleanupLoginIntentsCron = new CronJob(
@@ -67,57 +67,100 @@ const syncDatabaseWithPosthogDataCron = new CronJob(
   false, // start
 );
 
+const services: Array<{
+  name: string;
+  url: string;
+  type: "http" | "websocket";
+  available: boolean;
+}> = [
+  {
+    name: "Backend API",
+    url: process.env.BACKEND_API_URL ?? "",
+    type: "http",
+    available: true,
+  },
+  {
+    name: "Auth Service",
+    url: process.env.AUTH_SERVICE_URL ?? "",
+    type: "http",
+    available: true,
+  },
+  {
+    name: "ModGPT Service",
+    url: process.env.MODGPT_SERVICE_URL ?? "",
+    type: "http",
+    available: true,
+  },
+  {
+    name: "Codemod AI Service",
+    url: process.env.CODEMOD_AI_SERVICE_URL ?? "",
+    type: "websocket",
+    available: true,
+  },
+  {
+    name: "Run Service",
+    url: process.env.RUN_SERVICE_URL ?? "",
+    type: "http",
+    available: true,
+  },
+];
+
 const systemHealthCheckCron = new CronJob(
-  "*/10 * * * *",
+  "*/10 * * * * *",
   async () => {
-    const token = environment.SLACK_TOKEN;
-    const channel = environment.SLACK_CHANNEL;
+    const token = process.env.SLACK_TOKEN ?? "";
+    const channel = process.env.SLACK_CHANNEL ?? "";
+    const web = new WebClient(token);
 
-    try {
-      const web = new WebClient(token);
+    await Promise.all(
+      services.map(async (service) => {
+        try {
+          if (service.type === "http") {
+            const response = await axios.get(service.url);
 
-      const services: Array<{ name: string; url: string }> = [
-        {
-          name: "Backend API",
-          url: environment.BACKEND_API_URL ?? "",
-        },
-        {
-          name: "Auth Service",
-          url: environment.AUTH_SERVICE_URL ?? "",
-        },
-        {
-          name: "ModGPT Service",
-          url: environment.MODGPT_SERVICE_URL ?? "",
-        },
-        {
-          name: "Codemod AI Service",
-          url: environment.CODEMOD_AI_SERVICE_URL ?? "",
-        },
-        {
-          name: "Run Service",
-          url: environment.RUN_SERVICE_URL ?? "",
-        },
-      ];
+            if (response.status === 200 && service.available === false) {
+              await web.chat.postMessage({
+                channel: channel,
+                text: `${service.name} is now up.`,
+              });
+              service.available = true;
+            }
+          } else if (service.type === "websocket") {
+            await new Promise((resolve, reject) => {
+              const ws = new WebSocket(service.url);
 
-      for (const service of services) {
-        try {
-          const { status } = await axios.get(service.url);
-          if (status !== 200) {
-            throw new Error("Service did not respond with OK");
+              ws.on("open", async () => {
+                if (service.available === false) {
+                  await web.chat.postMessage({
+                    channel: channel,
+                    text: `${service.name} is now up.`,
+                  });
+                }
+                service.available = true;
+
+                ws.close();
+                resolve(true);
+              });
+
+              ws.on("error", (error) => {
+                reject(
+                  new Error(`WebSocket connection error: ${error.message}`),
+                );
+              });
+            });
           }
         } catch (error) {
-          console.error(`${service.name} is down`, error);
+          if (service.available === true) {
+            await web.chat.postMessage({
+              channel: channel,
+              text: `${service.name} is down. Error: ${error}`,
+            });
 
-          await web.chat.postMessage({
-            channel: channel,
-            text: `${service.name} is down. Error: ${error}`,
-          });
+            service.available = false;
+          }
         }
-      }
-    } catch (err) {
-      console.error("Failed to check system health.");
-      console.error((err as Error).message);
-    }
+      }),
+    );
   }, // onTick
   null, // onComplete
   false, // start