Repository: PrefectHQ/prefect

Issue:
Orion server ignores SIGTERM termination signal from docker/k8s
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar issue and didn't find it.
- [X] I searched the Prefect documentation for this issue.
- [X] I checked that this issue is related to Prefect and not one of its dependencies.

### Bug summary

A running orion process in docker/k8s does not respect SIGTERM

### Reproduction

- start a single-pod Deployment in k8s with command `["prefect", "orion", "start", "--host", "0.0.0.0", "--log-level", "INFO", "--port", "4200"]`
- pod goes in state Running
- pod starts logging heartbeats (INFO)
- scale Deployment to 0
- pod goes in state Terminating (k8s sends SIGTERM)
- pod keeps logging heartbeats like nothing happened
- k8s kills the pod after `terminationGracePeriodSeconds: 60`


### Error

```python3
# delete k8s Pod, logs continue like normal
12:57:00.421 | INFO | prefect.orion.services.recentdeploymentsscheduler - Scheduled 0 runs.
12:57:00.422 | INFO | prefect.orion.services.marklateruns - Finished monitoring for late runs.
12:57:05.422 | INFO | prefect.orion.services.recentdeploymentsscheduler - Scheduled 0 runs.
12:57:05.422 | INFO | prefect.orion.services.marklateruns - Finished monitoring for late runs.
12:57:10.422 | INFO | prefect.orion.services.recentdeploymentsscheduler - Scheduled 0 runs.
12:57:10.422 | INFO | prefect.orion.services.marklateruns - Finished monitoring for late runs.
12:57:15.359 | INFO | prefect.orion.services.scheduler - Scheduled 0 runs.
...
# log stream disconnects
```


### Versions

```Text
Version:             2.7.3
API version:         0.8.4
Python version:      3.11.1
Git commit:          26c145b3
Built:               Fri, Dec 16, 2022 12:08 PM
OS/Arch:             linux/x86_64
Profile:             default
Server type:         ephemeral
Server:
  Database:          postgresql
```


### Additional context

```
docker pull prefecthq/prefect:2.7.3-python3.11
```


Discussion:
Hm we convert SIGINT to SIGTERM at https://github.com/PrefectHQ/prefect/blob/main/src/prefect/cli/orion.py#L148, perhaps this as simple as adding SIGTERM forwarding to the subprocess.

Resolving patch:
diff --git a/src/prefect/cli/server.py b/src/prefect/cli/server.py
--- a/src/prefect/cli/server.py
+++ b/src/prefect/cli/server.py
@@ -32,7 +32,7 @@
     PREFECT_UI_ENABLED,
 )
 from prefect.utilities.asyncutils import run_sync_in_worker_thread
-from prefect.utilities.processutils import kill_on_interrupt, run_process
+from prefect.utilities.processutils import run_process, setup_signal_handlers_server
 
 server_app = PrefectTyper(
     name="server",
@@ -170,7 +170,9 @@ async def start(
         # and will start shutting down subprocesses:
         # https://github.com/PrefectHQ/server/issues/2475
 
-        kill_on_interrupt(server_process_id, "the Prefect server", app.console.print)
+        setup_signal_handlers_server(
+            server_process_id, "the Prefect server", app.console.print
+        )
 
     app.console.print("Server stopped!")
 
diff --git a/src/prefect/utilities/processutils.py b/src/prefect/utilities/processutils.py
--- a/src/prefect/utilities/processutils.py
+++ b/src/prefect/utilities/processutils.py
@@ -5,6 +5,7 @@
 import sys
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from functools import partial
 from io import TextIOBase
 from typing import (
     IO,
@@ -311,28 +312,43 @@ async def stream_text(source: TextReceiveStream, sink: Optional[TextSink]):
             raise TypeError(f"Unsupported sink type {type(sink).__name__}")
 
 
-def kill_on_interrupt(pid: int, process_name: str, print_fn: Callable):
-    """Kill a process with the given `pid` when a SIGNINT is received."""
+def forward_signal_handler(
+    pid: int, signum: int, *signums: int, process_name: str, print_fn: Callable
+):
+    """Forward subsequent signum events (e.g. interrupts) to respective signums."""
+    current_signal, future_signals = signums[0], signums[1:]
+
+    def handler(*args):
+        print_fn(
+            f"Received {getattr(signum, 'name', signum)}. "
+            f"Sending {getattr(current_signal, 'name', current_signal)} to"
+            f" {process_name} (PID {pid})..."
+        )
+        os.kill(pid, current_signal)
+        if future_signals:
+            forward_signal_handler(
+                pid,
+                signum,
+                *future_signals,
+                process_name=process_name,
+                print_fn=print_fn,
+            )
 
-    # In a non-windows enviornment first interrupt with send a SIGTERM, then
-    # subsequent interrupts will send SIGKILL. In Windows we use
-    # CTRL_BREAK_EVENT as SIGTERM is useless:
-    # https://bugs.python.org/issue26350
-    if sys.platform == "win32":
+    # register current and future signal handlers
+    signal.signal(signum, handler)
 
-        def stop_process(*args):
-            print_fn(f"\nStopping {process_name}...")
-            os.kill(pid, signal.CTRL_BREAK_EVENT)
 
+def setup_signal_handlers_server(pid: int, process_name: str, print_fn: Callable):
+    """Handle interrupts of the server gracefully."""
+    setup_handler = partial(
+        forward_signal_handler, pid, process_name=process_name, print_fn=print_fn
+    )
+    # on Windows, use CTRL_BREAK_EVENT as SIGTERM is useless:
+    # https://bugs.python.org/issue26350
+    if sys.platform == "win32":
+        setup_handler(signal.SIGINT, signal.CTRL_BREAK_EVENT)
     else:
-
-        def stop_process(*args):
-            print_fn(f"\nStopping {process_name}...")
-            os.kill(pid, signal.SIGTERM)
-            signal.signal(signal.SIGINT, kill_process)
-
-        def kill_process(*args):
-            print_fn(f"\nKilling {process_name}...")
-            os.kill(pid, signal.SIGKILL)
-
-    signal.signal(signal.SIGINT, stop_process)
+        # first interrupt: SIGTERM, second interrupt: SIGKILL
+        setup_handler(signal.SIGINT, signal.SIGTERM, signal.SIGKILL)
+        # forward first SIGTERM directly, send SIGKILL on subsequent SIGTERM
+        setup_handler(signal.SIGTERM, signal.SIGTERM, signal.SIGKILL)


Test patch:
diff --git a/tests/cli/test_start_server.py b/tests/cli/test_start_server.py
new file mode 100644
--- /dev/null
+++ b/tests/cli/test_start_server.py
@@ -0,0 +1,189 @@
+import os
+import signal
+import sys
+import tempfile
+
+import anyio
+import httpx
+import pytest
+
+from prefect.settings import get_current_settings
+from prefect.testing.fixtures import is_port_in_use
+from prefect.utilities.processutils import open_process
+
+STARTUP_TIMEOUT = 20
+SHUTDOWN_TIMEOUT = 20
+
+
+@pytest.fixture(scope="function")
+async def server_process():
+    """
+    Runs an instance of the server. Requires a port from 2222-2229 to be available.
+    Uses the same database as the rest of the tests.
+    Yields:
+        The anyio.Process.
+    """
+
+    ports = list(range(2222, 2230))
+
+    while True:
+        try:
+            port = ports.pop()
+        except IndexError as exc:
+            raise RuntimeError("No ports available to run test API.") from exc
+
+        if not is_port_in_use(port):
+            break
+
+    out = tempfile.TemporaryFile()  # capture output for test assertions
+
+    # Will connect to the same database as normal test clients
+    async with open_process(
+        command=[
+            "prefect",
+            "server",
+            "start",
+            "--host",
+            "127.0.0.1",
+            "--port",
+            str(port),
+            "--log-level",
+            "INFO",
+        ],
+        stdout=out,
+        stderr=out,
+        env={**os.environ, **get_current_settings().to_environment_variables()},
+    ) as process:
+        process.out = out
+        api_url = f"http://localhost:{port}/api"
+
+        # Wait for the server to be ready
+        async with httpx.AsyncClient() as client:
+            response = None
+            with anyio.move_on_after(STARTUP_TIMEOUT):
+                while True:
+                    try:
+                        response = await client.get(api_url + "/health")
+                    except httpx.ConnectError:
+                        pass
+                    else:
+                        if response.status_code == 200:
+                            await anyio.sleep(0.5)  # extra sleep for less flakiness
+                            break
+                    await anyio.sleep(0.1)
+            if response:
+                response.raise_for_status()
+            if not response:
+                raise RuntimeError(
+                    "Timed out while attempting to connect to hosted test server."
+                )
+
+        # Yield to the consuming tests
+        yield process
+
+        # Then shutdown the process
+        try:
+            process.terminate()
+        except ProcessLookupError:
+            pass
+        out.close()
+
+
+class TestUvicornSignalForwarding:
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigint_sends_sigterm(self, server_process):
+        server_process.send_signal(signal.SIGINT)
+        with anyio.fail_after(SHUTDOWN_TIMEOUT):
+            await server_process.wait()
+        server_process.out.seek(0)
+        out = server_process.out.read().decode()
+
+        assert "Sending SIGTERM" in out, (
+            "When sending a SIGINT, the main process should send a SIGTERM to the"
+            f" uvicorn subprocess. Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigterm_sends_sigterm_directly(self, server_process):
+        server_process.send_signal(signal.SIGTERM)
+        with anyio.fail_after(SHUTDOWN_TIMEOUT):
+            await server_process.wait()
+        server_process.out.seek(0)
+        out = server_process.out.read().decode()
+
+        assert "Sending SIGTERM" in out, (
+            "When sending a SIGTERM, the main process should send a SIGTERM to the"
+            f" uvicorn subprocess. Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigint_sends_sigterm_then_sigkill(self, server_process):
+        server_process.send_signal(signal.SIGINT)
+        await anyio.sleep(0.001)  # some time needed for the recursive signal handler
+        server_process.send_signal(signal.SIGINT)
+        with anyio.fail_after(SHUTDOWN_TIMEOUT):
+            await server_process.wait()
+        server_process.out.seek(0)
+        out = server_process.out.read().decode()
+
+        assert (
+            # either the main PID is still waiting for shutdown, so forwards the SIGKILL
+            "Sending SIGKILL" in out
+            # or SIGKILL came too late, and the main PID is already closing
+            or "KeyboardInterrupt" in out
+            or "Server stopped!" in out
+        ), (
+            "When sending two SIGINT shortly after each other, the main process should"
+            " first send a SIGTERM and then a SIGKILL to the uvicorn subprocess."
+            f" Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigterm_sends_sigterm_then_sigkill(self, server_process):
+        server_process.send_signal(signal.SIGTERM)
+        await anyio.sleep(0.001)  # some time needed for the recursive signal handler
+        server_process.send_signal(signal.SIGTERM)
+        with anyio.fail_after(SHUTDOWN_TIMEOUT):
+            await server_process.wait()
+        server_process.out.seek(0)
+        out = server_process.out.read().decode()
+
+        assert (
+            # either the main PID is still waiting for shutdown, so forwards the SIGKILL
+            "Sending SIGKILL" in out
+            # or SIGKILL came too late, and the main PID is already closing
+            or "KeyboardInterrupt" in out
+            or "Server stopped!" in out
+        ), (
+            "When sending two SIGTERM shortly after each other, the main process should"
+            " first send a SIGTERM and then a SIGKILL to the uvicorn subprocess."
+            f" Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform != "win32",
+        reason="CTRL_BREAK_EVENT is only defined in Windows",
+    )
+    async def test_sends_ctrl_break_win32(self, server_process):
+        server_process.send_signal(signal.SIGINT)
+        with anyio.fail_after(SHUTDOWN_TIMEOUT):
+            await server_process.wait()
+        server_process.out.seek(0)
+        out = server_process.out.read().decode()
+
+        assert "Sending CTRL_BREAK_EVENT" in out, (
+            "When sending a SIGINT, the main process should send a CTRL_BREAK_EVENT to"
+            f" the uvicorn subprocess. Output:\n{out}"
+        )
diff --git a/tests/utilities/test_processutils.py b/tests/utilities/test_processutils.py
--- a/tests/utilities/test_processutils.py
+++ b/tests/utilities/test_processutils.py
@@ -1,4 +1,3 @@
-import signal
 import subprocess
 import sys
 from unittest import mock
@@ -6,7 +5,7 @@
 import pytest
 
 import prefect.utilities.processutils
-from prefect.utilities.processutils import kill_on_interrupt, open_process, run_process
+from prefect.utilities.processutils import open_process, run_process
 
 
 class TestRunProcess:
@@ -108,63 +107,3 @@ async def test_adds_ctrl_c_handler_to_win32_process_group(self, monkeypatch):
             creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
         )
         mock_set_console_ctrl_handler.assert_called_once_with(mock_ctrl_c_handler, 1)
-
-
-class TestKillOnInterrupt:
-    @pytest.mark.skipif(
-        sys.platform == "win32",
-        reason="SIGTERM is only used in non-Windows environments",
-    )
-    def test_sends_sigterm_then_sigkill(self, monkeypatch):
-        print_fn = mock.Mock()
-        os_kill = mock.Mock()
-        signal_receiver = mock.Mock()
-        monkeypatch.setattr("os.kill", os_kill)
-        monkeypatch.setattr("signal.signal", signal_receiver)
-
-        kill_on_interrupt(123, "My Process", print_fn)
-
-        # This is a bit convoluted, but the point of the `signal_receiver` mock
-        # is to try to avoid actually sending the SIGINT. Doing so causes
-        # pytest to receive the signal and cancel the test run. Note that this
-        # doesn't happen with `signal.raise_signal`, but that's unavalable in
-        # python 3.7.
-        assert signal_receiver.call_args_list[0][0][0] == signal.SIGINT
-        signal_handler = signal_receiver.call_args_list[0][0][1]
-        signal_receiver.reset_mock()
-
-        # Call SIGINT handler expect a SIGTERM
-        signal_handler()
-        print_fn.assert_called_once_with("\nStopping My Process...")
-        os_kill.assert_called_once_with(123, signal.SIGTERM)
-
-        # Reset mocks and send second SIGINT and expect SIGKILL
-        print_fn.reset_mock()
-        os_kill.reset_mock()
-
-        assert signal_receiver.call_args_list[0][0][0] == signal.SIGINT
-        signal_handler = signal_receiver.call_args_list[0][0][1]
-        signal_handler()
-
-        print_fn.assert_called_once_with("\nKilling My Process...")
-        os_kill.assert_called_once_with(123, signal.SIGKILL)
-
-    @pytest.mark.skipif(
-        sys.platform != "win32",
-        reason="CTRL_BREAK_EVENT is only defined in Windows",
-    )
-    def test_sends_ctrl_break_win32(self, monkeypatch):
-        print_fn = mock.Mock()
-        os_kill = mock.Mock()
-        signal_receiver = mock.Mock()
-        monkeypatch.setattr("os.kill", os_kill)
-        monkeypatch.setattr("signal.signal", signal_receiver)
-
-        kill_on_interrupt(123, "My Process", print_fn)
-
-        assert signal_receiver.call_args_list[0][0][0] == signal.SIGINT
-        signal_handler = signal_receiver.call_args_list[0][0][1]
-        signal_handler()
-
-        print_fn.assert_called_once_with("\nStopping My Process...")
-        os_kill.assert_called_once_with(123, signal.CTRL_BREAK_EVENT)