From 2f636403752749faaaaa9673fd9684547a954616 Mon Sep 17 00:00:00 2001
From: Mackenzie Halliday <mackenzie.halliday@gmail.com>
Date: Thu, 21 Nov 2024 09:38:36 -0500
Subject: [PATCH] Squash #2041 - Update Health Check to search for PID under
 dd-trace

---
 .../dev/vaec-celery-task-definition.json      | 13 ++++++++++++-
 scripts/check_celery.sh                       | 19 ++++++++++++-------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/cd/application-deployment/dev/vaec-celery-task-definition.json b/cd/application-deployment/dev/vaec-celery-task-definition.json
index aba3c83614..620413baa4 100644
--- a/cd/application-deployment/dev/vaec-celery-task-definition.json
+++ b/cd/application-deployment/dev/vaec-celery-task-definition.json
@@ -106,7 +106,18 @@
         "sh",
         "-c",
         "/app/scripts/run_celery.sh"
-      ]
+      ],
+      "healthCheck": {
+        "command": [
+          "CMD",
+          "sh",
+          "-c",
+          "./scripts/check_celery.sh"
+        ],
+        "interval": 30,
+        "retries": 5,
+        "timeout": 10
+      }
     },
     {
       "name": "datadog-agent",
diff --git a/scripts/check_celery.sh b/scripts/check_celery.sh
index f4c2bf88e7..09d8fb5a5c 100755
--- a/scripts/check_celery.sh
+++ b/scripts/check_celery.sh
@@ -3,14 +3,19 @@
 set -e
 
 function get_celery_pids {
-  # get the PIDs of the process whose parent is the main celery process, saved in celery.pid
-  # print only pid and their command, get the ones with "celery" in their name
-  # and keep only these PIDs
+  # First, get the PID from the celery.pid file
+  MAIN_PID=$(cat /tmp/celery.pid)
+
+  # Check if the main process is ddtrace-run or Celery directly
+  if pstree -p ${MAIN_PID} | grep -q 'ddtrace'; then
+    # If ddtrace-run is present, navigate to the child process
+    APP_PIDS=$(pstree -p ${MAIN_PID} | sed 's/.*-ddtrace(\([0-9]*\)).*-celery(\([0-9]*\)).*/\2/')
+  else
+    # If no ddtrace-run, assume the main process is Celery
+    APP_PIDS=$(pstree -p ${MAIN_PID} | sed 's/.*-celery(\([0-9]*\)).*/\1/')
+  fi
 
-  set +o pipefail # so grep returning no matches does not premature fail pipe
-  APP_PIDS=$(pstree -p `cat /tmp/celery.pid` | sed 's/\(.*\)-celery(\(\d*\))/\2/')
   echo "Here are the APP_PIDS: ${APP_PIDS}"
-  set -o pipefail # pipefail should be set everywhere else
 }
 
 function ensure_celery_is_running {
@@ -20,7 +25,7 @@ function ensure_celery_is_running {
   fi
 
   for APP_PID in ${APP_PIDS}; do
-      kill -0 ${APP_PID} 2&>/dev/null || return 1
+      kill -0 ${APP_PID} 2>/dev/null || return 1
   done
 }