From 6b6f4e39e780d4bd6d0d1e569c74c60840189a10 Mon Sep 17 00:00:00 2001
From: Nick Muerdter <12112+GUI@users.noreply.github.com>
Date: Fri, 19 Jul 2024 11:16:34 -0600
Subject: [PATCH] Add health check to ensure trafficserver is responding.

---
 src/api-umbrella/http-api/health.lua          | 26 ++++++++++++++-----
 .../etc/trafficserver/remap.config.etlua      |  6 +++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/api-umbrella/http-api/health.lua b/src/api-umbrella/http-api/health.lua
index 06ad9b685..99e009a76 100644
--- a/src/api-umbrella/http-api/health.lua
+++ b/src/api-umbrella/http-api/health.lua
@@ -13,6 +13,7 @@ local function status_response(quick)
     status = "red",
     details = {
       apis_config = "red",
+      cache_server = "red"
     },
   }
 
@@ -21,8 +22,22 @@ local function status_response(quick)
     response["details"]["apis_config"] = "green"
   end
 
+  local httpc = http.new()
+  httpc:set_timeout(3000)
+
+  local res, err = httpc:request_uri("http://127.0.0.1:" .. config["trafficserver"]["port"] .. "/_trafficserver-health/nocache/1", {
+    headers = {
+      ["Host"] = "api-umbrella-trafficserver-health.internal",
+    },
+  })
+  if err then
+    ngx.log(ngx.ERR, "failed to fetch web app: ", err)
+  elseif res.status == 200 then
+    response["details"]["cache_server"] = "green"
+  end
+
   if quick then
-    if response["details"]["apis_config"] == "green" then
+    if response["details"]["apis_config"] == "green" and response["details"]["cache_server"] == "green" then
       response["status"] = "green"
     end
 
@@ -33,11 +48,8 @@ local function status_response(quick)
   response["details"]["analytics_db_setup"] = "red"
   response["details"]["web_app"] = "red"
 
-  local httpc = http.new()
-  httpc:set_timeout(3000)
-
   -- Check the health of the OpenSearch cluster
-  local res, err = opensearch_query("/_cluster/health")
+  res, err = opensearch_query("/_cluster/health")
   if err then
     ngx.log(ngx.ERR, "failed to fetch cluster health from opensearch: ", err)
   elseif res.body_json then
@@ -65,9 +77,9 @@ local function status_response(quick)
   -- content), OpenSearch seems to get stuck in the yellow status, even though
   -- everything appears operational (but then it becomes green once content
   -- starts indexing).
-  if response["details"]["apis_config"] == "green" and (response["details"]["analytics_db"] == "yellow" or response["details"]["analytics_db"] == "green") and response["details"]["analytics_db_setup"] == "green" and response["details"]["web_app"] == "green" then
+  if response["details"]["apis_config"] == "green" and response["details"]["cache_server"] == "green" and (response["details"]["analytics_db"] == "yellow" or response["details"]["analytics_db"] == "green") and response["details"]["analytics_db_setup"] == "green" and response["details"]["web_app"] == "green" then
     response["status"] = "green"
-  elseif response["details"]["apis_config"] == "green" then
+  elseif response["details"]["apis_config"] == "green" and response["details"]["cache_server"] == "green" then
     response["status"] = "yellow"
   end
 
diff --git a/templates/etc/trafficserver/remap.config.etlua b/templates/etc/trafficserver/remap.config.etlua
index 80ee634fa..f9cde2fe8 100644
--- a/templates/etc/trafficserver/remap.config.etlua
+++ b/templates/etc/trafficserver/remap.config.etlua
@@ -1 +1,7 @@
+# A fake endpoint that just returns a 200 OK response. We'll use this to check
+# if Trafficserver is up and responding (since Envoy may not be up if that's
+# hosted separately).
+map http://api-umbrella-trafficserver-health.internal/_trafficserver-health/ http://127.0.0.1 @plugin=generator.so
+
+# Route all requests to Envoy for proxying to underlying API backends.
 map / <%- config["envoy"]["scheme"] %>://<%- config["envoy"]["host"] %>:<%- config["envoy"]["port"] %>