[512] Graceful Shutdown Handlers (#809)

* feat: implement graceful shutdown handlers for Django and Celery

Handle SIGTERM/SIGINT by draining in-flight HTTP requests, waiting for Celery tasks (30s timeout), closing DB connections, and logging the shutdown reason. Closes #512.

Co-authored-by: Cursor <cursoragent@cursor.com>

* test: patch shutdown.connections in close_db test

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: add deployment shutdown timeouts and fix DB close test

Set gunicorn --graceful-timeout and Celery --stop-timeout to 30s in
docker-compose and k8s manifests. Patch soroscan.shutdown.connections in tests.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix: remove unused MagicMock import for ruff

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: Junirezz

django-backend/soroscan/asgi.py

@@ -48,3 +48,7 @@
         ),
     }
 )
+
+from soroscan.shutdown import register_shutdown_handlers  # noqa: E402
+
+register_shutdown_handlers()

django-backend/soroscan/celery.py

@@ -6,7 +6,13 @@
 import time
 
 from celery import Celery
-from celery.signals import task_failure, task_postrun, task_prerun
+from celery.signals import (
+    task_failure,
+    task_postrun,
+    task_prerun,
+    worker_shutdown,
+    worker_shutting_down,
+)
 
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "soroscan.settings")
 
@@ -63,6 +69,20 @@ def record_celery_task_failure(sender, exception, **kwargs):
     ).inc()
 
 
+@worker_shutting_down.connect
+def _celery_worker_shutting_down(sender, sig=None, how=None, exitcode=None, **kwargs):
+    from soroscan.shutdown import on_celery_worker_shutting_down
+
+    on_celery_worker_shutting_down(sig=sig, how=how, exitcode=exitcode, **kwargs)
+
+
+@worker_shutdown.connect
+def _celery_worker_shutdown(sender, **kwargs):
+    from soroscan.shutdown import on_celery_worker_shutdown
+
+    on_celery_worker_shutdown(**kwargs)
+
+
 @app.task(bind=True, ignore_result=True)
 def debug_task(self):
     print(f"Request: {self.request!r}")

django-backend/soroscan/ingest/apps.py

@@ -9,5 +9,7 @@ class IngestConfig(AppConfig):
     def ready(self):
         import soroscan.ingest.signals  # noqa: F401 — registers signal handlers
         from soroscan.operational_metrics import register_operational_collector
+        from soroscan.shutdown import register_shutdown_handlers
 
         register_operational_collector()
+        register_shutdown_handlers()

django-backend/soroscan/ingest/tests/test_graceful_shutdown.py

@@ -0,0 +1,170 @@
+import json
+import signal
+import threading
+import time
+from unittest.mock import patch
+
+import pytest
+from django.http import HttpResponse
+from django.test import RequestFactory, override_settings
+
+from soroscan.middleware import GracefulShutdownMiddleware
+from soroscan.shutdown import (
+    begin_shutdown,
+    close_database_connections,
+    end_request,
+    in_flight_request_count,
+    is_shutting_down,
+    on_celery_worker_shutdown,
+    on_celery_worker_shutting_down,
+    perform_graceful_shutdown,
+    register_shutdown_handlers,
+    reset_shutdown_state,
+    try_begin_request,
+    wait_for_in_flight_requests,
+)
+
+
+@pytest.fixture(autouse=True)
+def _clean_shutdown_state():
+    reset_shutdown_state()
+    yield
+    reset_shutdown_state()
+
+
+@pytest.fixture
+def rf():
+    return RequestFactory()
+
+
+def test_try_begin_request_tracks_in_flight():
+    assert try_begin_request() is True
+    assert in_flight_request_count() == 1
+    end_request()
+    assert in_flight_request_count() == 0
+
+
+def test_try_begin_request_rejects_when_shutting_down():
+    begin_shutdown("test shutdown")
+    assert try_begin_request() is False
+    assert in_flight_request_count() == 0
+
+
+def test_wait_for_in_flight_requests_completes():
+    try_begin_request()
+
+    def finish_request():
+        time.sleep(0.05)
+        end_request()
+
+    thread = threading.Thread(target=finish_request)
+    thread.start()
+    assert wait_for_in_flight_requests(timeout=2) is True
+    thread.join()
+
+
+def test_wait_for_in_flight_requests_times_out():
+    try_begin_request()
+    assert wait_for_in_flight_requests(timeout=0.05) is False
+    end_request()
+
+
+@override_settings(SHUTDOWN_TIMEOUT_SECONDS=30)
+def test_perform_graceful_shutdown_waits_and_closes_db():
+    try_begin_request()
+    closed = threading.Event()
+
+    def finish_request():
+        time.sleep(0.05)
+        end_request()
+
+    thread = threading.Thread(target=finish_request)
+    thread.start()
+
+    with patch("soroscan.shutdown.close_database_connections", side_effect=lambda: closed.set()):
+        with patch("soroscan.shutdown.sys.exit") as mock_exit:
+            perform_graceful_shutdown("SIGTERM test", exit_process=True)
+
+    thread.join()
+    assert closed.is_set()
+    assert is_shutting_down()
+    mock_exit.assert_called_once_with(0)
+
+
+@override_settings(SHUTDOWN_TIMEOUT_SECONDS=0.05)
+def test_perform_graceful_shutdown_logs_timeout():
+    try_begin_request()
+
+    with patch("soroscan.shutdown.close_database_connections"):
+        with patch("soroscan.shutdown.sys.exit"):
+            with patch("soroscan.shutdown.logger") as mock_logger:
+                perform_graceful_shutdown("SIGTERM test", exit_process=True)
+
+    assert mock_logger.warning.called
+    end_request()
+
+
+def test_perform_graceful_shutdown_is_idempotent():
+    with patch("soroscan.shutdown.close_database_connections") as mock_close:
+        with patch("soroscan.shutdown.sys.exit"):
+            perform_graceful_shutdown("first", exit_process=False)
+            perform_graceful_shutdown("second", exit_process=False)
+
+    mock_close.assert_called_once()
+
+
+def test_register_shutdown_handlers_installs_sigterm():
+    register_shutdown_handlers()
+    handler = signal.getsignal(signal.SIGTERM)
+    assert handler is not None
+    assert handler.__name__ == "_handle_shutdown_signal"
+
+
+def test_middleware_allows_requests_when_healthy(rf):
+    request = rf.get("/api/health/")
+    middleware = GracefulShutdownMiddleware(lambda req: HttpResponse("OK", status=200))
+    response = middleware(request)
+    assert response.status_code == 200
+
+
+def test_middleware_rejects_requests_during_shutdown(rf):
+    begin_shutdown("SIGTERM")
+    request = rf.get("/api/health/")
+    middleware = GracefulShutdownMiddleware(lambda req: HttpResponse("OK", status=200))
+    response = middleware(request)
+    assert response.status_code == 503
+    data = json.loads(response.content)
+    assert data["error"] == "Server is shutting down"
+
+
+def test_middleware_tracks_in_flight_during_request(rf):
+    in_flight_during_request = None
+
+    def slow_response(req):
+        nonlocal in_flight_during_request
+        in_flight_during_request = in_flight_request_count()
+        return HttpResponse("OK", status=200)
+
+    request = rf.get("/api/events/")
+    middleware = GracefulShutdownMiddleware(slow_response)
+    middleware(request)
+
+    assert in_flight_during_request == 1
+    assert in_flight_request_count() == 0
+
+
+def test_close_database_connections():
+    with patch("soroscan.shutdown.connections") as mock_connections:
+        close_database_connections()
+    mock_connections.close_all.assert_called_once()
+
+
+def test_celery_worker_shutting_down_sets_state():
+    on_celery_worker_shutting_down(sig="SIGTERM", how="warm", exitcode=0)
+    assert is_shutting_down()
+
+
+def test_celery_worker_shutdown_closes_db():
+    with patch("soroscan.shutdown.close_database_connections") as mock_close:
+        on_celery_worker_shutdown()
+    mock_close.assert_called_once()

django-backend/soroscan/middleware.py

@@ -149,6 +149,26 @@ def __call__(self, request):
 
         return self.get_response(request)
 
+class GracefulShutdownMiddleware:
+    """Reject new requests during shutdown and track in-flight request count."""
+
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        from soroscan.shutdown import end_request, try_begin_request
+
+        if not try_begin_request():
+            return JsonResponse(
+                {"error": "Server is shutting down"},
+                status=503,
+            )
+        try:
+            return self.get_response(request)
+        finally:
+            end_request()
+
+
 class MaintenanceModeMiddleware:
     """Return 503 for all non-admin routes when MAINTENANCE_MODE=True."""
 

django-backend/soroscan/settings.py

@@ -105,6 +105,7 @@ def _load_software_version() -> str:
 MIDDLEWARE = [
     # PrometheusBeforeMiddleware must be first to capture all requests.
     "django_prometheus.middleware.PrometheusBeforeMiddleware",
+    "soroscan.middleware.GracefulShutdownMiddleware",
     "soroscan.monitoring.ErrorRateMetricsMiddleware",
     "soroscan.middleware.RequestBodySizeMiddleware",
     "soroscan.middleware.MaintenanceModeMiddleware",
@@ -296,6 +297,9 @@ def _load_software_version() -> str:
 CELERY_TASK_SERIALIZER = "json"
 CELERY_RESULT_SERIALIZER = "json"
 CELERY_TIMEZONE = TIME_ZONE
+# Graceful shutdown: wait up to 30s for active tasks after SIGTERM
+CELERY_WORKER_SOFT_SHUTDOWN_TIMEOUT = 30
+SHUTDOWN_TIMEOUT_SECONDS = env.int("SHUTDOWN_TIMEOUT_SECONDS", default=30)
 CELERY_TASK_ROUTES = {
     "ingest.tasks.ingest_latest_events": {"queue": "high_priority"},
     "ingest.tasks.dispatch_webhook": {"queue": "default"},

django-backend/soroscan/settings_test.py

@@ -36,6 +36,7 @@
 
 MIDDLEWARE = [
     "django_prometheus.middleware.PrometheusBeforeMiddleware",
+    "soroscan.middleware.GracefulShutdownMiddleware",
     "soroscan.middleware.RequestBodySizeMiddleware",
     "soroscan.middleware.MaintenanceModeMiddleware",
     "django.middleware.security.SecurityMiddleware",
@@ -148,6 +149,8 @@
 CELERY_ACCEPT_CONTENT = ["json"]
 CELERY_TASK_SERIALIZER = "json"
 CELERY_RESULT_SERIALIZER = "json"
+SHUTDOWN_TIMEOUT_SECONDS = 30
+CELERY_WORKER_SOFT_SHUTDOWN_TIMEOUT = 30
 CELERY_TIMEZONE = TIME_ZONE
 
 # Stellar / Soroban Configuration