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()