From cd98275030dd090b73eac4e60e3ffd3f1c5f22ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Far=C3=ADas=20Santana?= <tomas@tomasfarias.dev>
Date: Thu, 10 Oct 2024 13:17:27 +0100
Subject: [PATCH] fix(temporal): Handle SIGTERM correctly on worker cleanup
 (#25505)

---
 bin/temporal-django-worker        | 19 ++++++++++++++++---
 posthog/temporal/common/worker.py |  8 ++++----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/bin/temporal-django-worker b/bin/temporal-django-worker
index ab10a959a87..05fb98ba2be 100755
--- a/bin/temporal-django-worker
+++ b/bin/temporal-django-worker
@@ -2,8 +2,21 @@
 
 set -e
 
-trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT
+cleanup() {
+    echo "Stopping worker..."
+    if kill -0 "$worker_pid" >/dev/null 2>&1; then
+        kill "$worker_pid"
+    else
+        echo "Worker process is not running."
+    fi
+}
 
-python3 manage.py start_temporal_worker "$@"
+trap cleanup SIGINT SIGTERM EXIT
 
-wait
+python3 manage.py start_temporal_worker "$@" &
+
+worker_pid=$!
+
+wait $worker_pid
+
+cleanup
diff --git a/posthog/temporal/common/worker.py b/posthog/temporal/common/worker.py
index ac500b4e7ae..6d525b428bf 100644
--- a/posthog/temporal/common/worker.py
+++ b/posthog/temporal/common/worker.py
@@ -1,5 +1,5 @@
+import asyncio
 import signal
-import sys
 from datetime import timedelta
 
 from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
@@ -43,10 +43,10 @@ async def start_worker(
 
     # catch the TERM signal, and stop the worker gracefully
     # https://github.com/temporalio/sdk-python#worker-shutdown
-    async def signal_handler(sig, frame):
+    async def shutdown_worker():
         await worker.shutdown()
-        sys.exit(0)
 
-    signal.signal(signal.SIGTERM, signal_handler)
+    loop = asyncio.get_event_loop()
+    loop.add_signal_handler(signal.SIGTERM, lambda: asyncio.create_task(shutdown_worker()))
 
     await worker.run()