From dd5cfcf3c128e450ac044528a914d47f48aa256e Mon Sep 17 00:00:00 2001 From: Mariano Gappa Date: Tue, 3 Nov 2026 11:09:31 +0400 Subject: [PATCH 1/3] fix: Exit 0 on SIGTERM while serving --- cloudquery/sdk/serve/plugin.py | 12 ++++++ tests/serve/signals.py | 72 ++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+) create mode 100644 tests/serve/signals.py diff --git a/cloudquery/sdk/serve/plugin.py b/cloudquery/sdk/serve/plugin.py index 7ea4e88..2a9a2ff 100644 --- a/cloudquery/sdk/serve/plugin.py +++ b/cloudquery/sdk/serve/plugin.py @@ -4,8 +4,10 @@ import logging import os import shutil +import signal import subprocess import tarfile +import threading from concurrent import futures from pathlib import Path @@ -401,7 +403,17 @@ def _serve(self, args): self._server.add_insecure_port(args.address) logger.info("Starting server", address=args.address) self._server.start() + self._stop_on_termination_signals() self._server.wait_for_termination() + def _stop_on_termination_signals(self): + if threading.current_thread() is not threading.main_thread(): + return + for signum in (signal.SIGTERM, signal.SIGINT): + signal.signal(signum, self._stop_in_background) + + def _stop_in_background(self, signum, frame): + threading.Thread(target=self.stop, daemon=True).start() + def stop(self): self._server.stop(5) diff --git a/tests/serve/signals.py b/tests/serve/signals.py new file mode 100644 index 0000000..b26a56c --- /dev/null +++ b/tests/serve/signals.py @@ -0,0 +1,72 @@ +import signal +import subprocess +import sys +import textwrap +import time + +import pytest + +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"]) +""" + +SIGTERM_DURING_LOCK_RESTORE = """ +import os, signal, sys, threading +from cloudquery.sdk import serve +from cloudquery.sdk.plugin import Plugin + +def signal_handler(sig, frame): + sys.exit(0) + +signal.signal(signal.SIGTERM, signal_handler) + +_orig = threading.Condition._acquire_restore +_fired = threading.Event() + +def _acquire_restore(self, x): + if threading.current_thread() is threading.main_thread() and not _fired.is_set(): + _fired.set() + os.kill(os.getpid(), signal.SIGTERM) + return _orig(self, x) + +threading.Timer(1.0, lambda: setattr(threading.Condition, "_acquire_restore", _acquire_restore)).start() + +serve.PluginCommand(Plugin("test", "v0.0.1")).run(["serve", "--address", "127.0.0.1:0"]) +""" + + +def start_plugin(tmp_path, source): + script = tmp_path / "plugin_main.py" + script.write_text(textwrap.dedent(source)) + return subprocess.Popen( + [sys.executable, str(script)], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + +def wait_until_serving(process): + for line in process.stdout: + if "Starting server" in line: + time.sleep(0.5) + return + pytest.fail("plugin exited before it started serving") + + +def test_sigterm_during_lock_restore_exits_zero(tmp_path): + process = start_plugin(tmp_path, SIGTERM_DURING_LOCK_RESTORE) + output, _ = process.communicate(timeout=30) + assert process.returncode == 0, output + + +@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGINT]) +def test_signal_while_serving_exits_zero(tmp_path, signum): + process = start_plugin(tmp_path, SERVE_PLUGIN) + wait_until_serving(process) + process.send_signal(signum) + output, _ = process.communicate(timeout=30) + assert process.returncode == 0, output From afba5378e62a29b2da2176932914730afd9b46ac Mon Sep 17 00:00:00 2001 From: Mariano Gappa Date: Tue, 3 Nov 2026 14:04:09 +0400 Subject: [PATCH 2/3] fix: Bound shutdown when an RPC handler is blocked --- cloudquery/sdk/serve/plugin.py | 26 +++++++++++--- tests/serve/signals.py | 64 ++++++++++++++++++++++++++++++++-- 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/cloudquery/sdk/serve/plugin.py b/cloudquery/sdk/serve/plugin.py index 2a9a2ff..d91485d 100644 --- a/cloudquery/sdk/serve/plugin.py +++ b/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) diff --git a/tests/serve/signals.py b/tests/serve/signals.py index b26a56c..efeb5d9 100644 --- a/tests/serve/signals.py +++ b/tests/serve/signals.py @@ -1,11 +1,15 @@ 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 @@ -13,6 +17,18 @@ 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() From 320207b8eee68ca88a56e1879c5a3d242de7acf1 Mon Sep 17 00:00:00 2001 From: Mariano Gappa Date: Tue, 3 Nov 2026 15:58:11 +0400 Subject: [PATCH 3/3] fix: Keep output flushing from blocking forced exit --- cloudquery/sdk/serve/plugin.py | 14 ++++++- tests/serve/signals.py | 68 +++++++++++++++++++++++++++++++++- 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/cloudquery/sdk/serve/plugin.py b/cloudquery/sdk/serve/plugin.py index d91485d..852f237 100644 --- a/cloudquery/sdk/serve/plugin.py +++ b/cloudquery/sdk/serve/plugin.py @@ -29,6 +29,7 @@ _IS_WINDOWS = sys.platform == "win32" _STOP_GRACE_SECONDS = 5 _FORCED_EXIT_DELAY_SECONDS = 5 +_FORCED_EXIT_FLUSH_TIMEOUT_SECONDS = 1 try: import colorama @@ -432,6 +433,15 @@ def stop(self): def _exit_now(): - sys.stdout.flush() - sys.stderr.flush() + flusher = threading.Thread(target=_flush_output_streams, daemon=True) + flusher.start() + flusher.join(_FORCED_EXIT_FLUSH_TIMEOUT_SECONDS) os._exit(0) + + +def _flush_output_streams(): + for stream in (sys.stdout, sys.stderr): + try: + stream.flush() + except Exception: + pass diff --git a/tests/serve/signals.py b/tests/serve/signals.py index efeb5d9..06e2841 100644 --- a/tests/serve/signals.py +++ b/tests/serve/signals.py @@ -29,6 +29,32 @@ def init(self, spec, no_connection=False): serve.PluginCommand(BlockedInitPlugin("test", "v0.0.1")).run(["serve", "--address", sys.argv[1]]) """ +SERVE_PLUGIN_WITH_INIT_FLOODING_STDOUT = """ +import sys +from cloudquery.sdk import serve +from cloudquery.sdk.plugin import Plugin + +class FloodingInitPlugin(Plugin): + def init(self, spec, no_connection=False): + while True: + sys.stdout.write("x" * 1024) + +serve.PluginCommand(FloodingInitPlugin("test", "v0.0.1")).run(["serve", "--address", sys.argv[1]]) +""" + +SERVE_PLUGIN_WITH_INIT_BLOCKED_ON_BUFFERED_STDOUT = """ +import sys, threading +from cloudquery.sdk import serve +from cloudquery.sdk.plugin import Plugin + +class BufferedInitPlugin(Plugin): + def init(self, spec, no_connection=False): + sys.stdout.write("buffered and never flushed") + threading.Event().wait() + +serve.PluginCommand(BufferedInitPlugin("test", "v0.0.1")).run(["serve", "--address", sys.argv[1]]) +""" + SIGTERM_DURING_LOCK_RESTORE = """ import os, signal, sys, threading from cloudquery.sdk import serve @@ -94,9 +120,9 @@ def free_address(): return "127.0.0.1:%d" % sock.getsockname()[1] -def start_plugin_with_blocked_init(tmp_path): +def start_plugin_with_blocked_init(tmp_path, source=SERVE_PLUGIN_WITH_BLOCKED_INIT): address = free_address() - process = start_plugin(tmp_path, SERVE_PLUGIN_WITH_BLOCKED_INIT, address) + process = start_plugin(tmp_path, source, address) wait_until_serving(process) channel = grpc.insecure_channel(address) init_call = plugin_pb2_grpc.PluginStub(channel).Init.future( @@ -130,3 +156,41 @@ def test_second_signal_with_blocked_handler_exits_zero_immediately(tmp_path): init_call.cancel() channel.close() process.kill() + + +def send_termination_signals(process, second_signal): + process.send_signal(signal.SIGTERM) + if second_signal: + time.sleep(0.5) + process.send_signal(signal.SIGINT) + return 3 + return 20 + + +@pytest.mark.parametrize("second_signal", [False, True]) +def test_forced_exit_with_full_stdout_pipe_exits_zero(tmp_path, second_signal): + process, channel, init_call = start_plugin_with_blocked_init( + tmp_path, SERVE_PLUGIN_WITH_INIT_FLOODING_STDOUT + ) + try: + timeout = send_termination_signals(process, second_signal) + assert process.wait(timeout=timeout) == 0 + finally: + init_call.cancel() + channel.close() + process.kill() + + +@pytest.mark.parametrize("second_signal", [False, True]) +def test_forced_exit_with_closed_stdout_pipe_exits_zero(tmp_path, second_signal): + process, channel, init_call = start_plugin_with_blocked_init( + tmp_path, SERVE_PLUGIN_WITH_INIT_BLOCKED_ON_BUFFERED_STDOUT + ) + try: + process.stdout.close() + timeout = send_termination_signals(process, second_signal) + assert process.wait(timeout=timeout) == 0 + finally: + init_call.cancel() + channel.close() + process.kill()