diff --git a/packages/bigquery-magics/bigquery_magics/graph_server.py b/packages/bigquery-magics/bigquery_magics/graph_server.py index fac2bb281122..74d3e937e173 100644 --- a/packages/bigquery-magics/bigquery_magics/graph_server.py +++ b/packages/bigquery-magics/bigquery_magics/graph_server.py @@ -297,6 +297,7 @@ def __init__(self): self.port = None self.url = None self._server = None + self._server_ready = threading.Event() def build_route(self, endpoint): """ @@ -317,6 +318,7 @@ class ThreadedTCPServer(socketserver.TCPServer): with ThreadedTCPServer(("127.0.0.1", self.port), GraphServerHandler) as httpd: self._server = httpd + self._server_ready.set() self._server.serve_forever() def init(self): @@ -327,9 +329,11 @@ def init(self): self.port = portpicker.pick_unused_port() self.url = f"{GraphServer.host}:{self.port}" + self._server_ready.clear() server_thread = threading.Thread(target=self._start_server) server_thread.start() + self._server_ready.wait() return server_thread def stop_server(self): diff --git a/packages/bigquery-magics/tests/unit/test_graph_server.py b/packages/bigquery-magics/tests/unit/test_graph_server.py index 28f4ba58f8ec..2f4d10283560 100644 --- a/packages/bigquery-magics/tests/unit/test_graph_server.py +++ b/packages/bigquery-magics/tests/unit/test_graph_server.py @@ -13,6 +13,7 @@ # limitations under the License. import json +import time import unittest from unittest import mock @@ -747,3 +748,55 @@ def test_convert_schema_shared_label(): labels = {label["name"]: label for label in result["labels"]} assert "Person" in labels assert set(labels["Person"]["propertyDeclarationNames"]) == {"id", "name"} + + +def test_stop_server_immediately_after_init(): + """Verifies that GraphServer.init() waits until the underlying TCP server + is initialized before returning, preventing a race condition with stop_server(). + """ + # Mock portpicker to return port 0 (OS-assigned ephemeral port) so this test + # runs in all unit test environments, even when optional dependencies like + # `spanner-graph-notebook` and `portpicker` are not installed. + mock_portpicker = mock.MagicMock() + mock_portpicker.pick_unused_port.return_value = 0 + + with mock.patch.dict("sys.modules", {"portpicker": mock_portpicker}): + server = graph_server.GraphServer() + server_thread = server.init() + + try: + # 1. Immediately call stop_server() after init() without making any HTTP + # requests. If init() returns before `self._server` is assigned in + # _start_server(), `if self._server:` inside stop_server() evaluates to + # False, making stop_server() a no-op. + server.stop_server() + + # 2. Wait up to 1.0 second for the background server thread to exit. + # When properly synchronized, init() blocks until `self._server` is + # assigned, so stop_server() calls `self._server.shutdown()` and the + # thread exits almost immediately. If the shutdown signal is missed, + # the background thread remains blocked in `serve_forever()`. Using + # `timeout=1.0` prevents `join()` from hanging indefinitely. + server_thread.join(timeout=1.0) + + # 3. Assert that the server thread actually terminated. + assert not server_thread.is_alive(), ( + "GraphServer thread is still alive after stop_server(); " + "stop_server() likely raced with _start_server() before self._server was set." + ) + finally: + # 4. Cleanup in case the assertion above fails: + # Because `server_thread` is a non-daemon thread, Python's interpreter + # shutdown (`threading._shutdown()`) will hang at the end of the pytest + # run if `server_thread` is left running. + if server_thread.is_alive(): # pragma: NO COVER + # Poll with a short `time.sleep()` (releasing the GIL so the background + # thread can finish binding the socket and setting `server._server`) + # up to a bounded deadline. + deadline = time.monotonic() + 5.0 + while server._server is None and time.monotonic() < deadline: + time.sleep(0.01) + + # Shut down the server so the background thread exits cleanly. + server.stop_server() + server_thread.join(timeout=2.0)