fix: Bound shutdown when an RPC handler is blocked

Author: marianogappa

cloudquery/sdk/serve/plugin.py

@@ -8,6 +8,7 @@
 import subprocess
 import tarfile
 import threading
+import time
 from concurrent import futures
 from pathlib import Path
 
@@ -26,6 +27,8 @@
 from cloudquery.sdk.schema import table
 
 _IS_WINDOWS = sys.platform == "win32"
+_STOP_GRACE_SECONDS = 5
+_FORCED_EXIT_DELAY_SECONDS = 5
 
 try:
     import colorama
@@ -409,11 +412,26 @@ def _serve(self, args):
     def _stop_on_termination_signals(self):
         if threading.current_thread() is not threading.main_thread():
             return
+        self._termination_requested = False
         for signum in (signal.SIGTERM, signal.SIGINT):
-            signal.signal(signum, self._stop_in_background)
+            signal.signal(signum, self._on_termination_signal)
 
-    def _stop_in_background(self, signum, frame):
-        threading.Thread(target=self.stop, daemon=True).start()
+    def _on_termination_signal(self, signum, frame):
+        if self._termination_requested:
+            _exit_now()
+        self._termination_requested = True
+        threading.Thread(target=self._stop_then_force_exit, daemon=True).start()
+
+    def _stop_then_force_exit(self):
+        self.stop().wait()
+        time.sleep(_FORCED_EXIT_DELAY_SECONDS)
+        _exit_now()
 
     def stop(self):
-        self._server.stop(5)
+        return self._server.stop(_STOP_GRACE_SECONDS)
+
+
+def _exit_now():
+    sys.stdout.flush()
+    sys.stderr.flush()
+    os._exit(0)

tests/serve/signals.py

@@ -1,18 +1,34 @@
 import signal
+import socket
 import subprocess
 import sys
 import textwrap
 import time
 
+import grpc
 import pytest
 
+from cloudquery.plugin_v3 import plugin_pb2, plugin_pb2_grpc
+
 SERVE_PLUGIN = """
 from cloudquery.sdk import serve
 from cloudquery.sdk.plugin import Plugin
 
 serve.PluginCommand(Plugin("test", "v0.0.1")).run(["serve", "--address", "127.0.0.1:0"])
 """
 
+SERVE_PLUGIN_WITH_BLOCKED_INIT = """
+import sys, threading
+from cloudquery.sdk import serve
+from cloudquery.sdk.plugin import Plugin
+
+class BlockedInitPlugin(Plugin):
+    def init(self, spec, no_connection=False):
+        threading.Event().wait()
+
+serve.PluginCommand(BlockedInitPlugin("test", "v0.0.1")).run(["serve", "--address", sys.argv[1]])
+"""
+
 SIGTERM_DURING_LOCK_RESTORE = """
 import os, signal, sys, threading
 from cloudquery.sdk import serve
@@ -38,11 +54,11 @@ def _acquire_restore(self, x):
 """
 
 
-def start_plugin(tmp_path, source):
+def start_plugin(tmp_path, source, *args):
     script = tmp_path / "plugin_main.py"
     script.write_text(textwrap.dedent(source))
     return subprocess.Popen(
-        [sys.executable, str(script)],
+        [sys.executable, str(script), *args],
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
         text=True,
@@ -70,3 +86,47 @@ def test_signal_while_serving_exits_zero(tmp_path, signum):
     process.send_signal(signum)
     output, _ = process.communicate(timeout=30)
     assert process.returncode == 0, output
+
+
+def free_address():
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return "127.0.0.1:%d" % sock.getsockname()[1]
+
+
+def start_plugin_with_blocked_init(tmp_path):
+    address = free_address()
+    process = start_plugin(tmp_path, SERVE_PLUGIN_WITH_BLOCKED_INIT, address)
+    wait_until_serving(process)
+    channel = grpc.insecure_channel(address)
+    init_call = plugin_pb2_grpc.PluginStub(channel).Init.future(
+        plugin_pb2.Init.Request(spec=b"")
+    )
+    time.sleep(0.5)
+    return process, channel, init_call
+
+
+def test_sigterm_with_blocked_handler_exits_zero_within_bound(tmp_path):
+    process, channel, init_call = start_plugin_with_blocked_init(tmp_path)
+    try:
+        process.send_signal(signal.SIGTERM)
+        output, _ = process.communicate(timeout=20)
+        assert process.returncode == 0, output
+    finally:
+        init_call.cancel()
+        channel.close()
+        process.kill()
+
+
+def test_second_signal_with_blocked_handler_exits_zero_immediately(tmp_path):
+    process, channel, init_call = start_plugin_with_blocked_init(tmp_path)
+    try:
+        process.send_signal(signal.SIGTERM)
+        time.sleep(0.5)
+        process.send_signal(signal.SIGINT)
+        output, _ = process.communicate(timeout=3)
+        assert process.returncode == 0, output
+    finally:
+        init_call.cancel()
+        channel.close()
+        process.kill()