diff --git a/amber/src/main/python/core/runnables/heartbeat.py b/amber/src/main/python/core/runnables/heartbeat.py index 1e4c45837d4..941782f6998 100644 --- a/amber/src/main/python/core/runnables/heartbeat.py +++ b/amber/src/main/python/core/runnables/heartbeat.py @@ -79,22 +79,26 @@ def run(self) -> None: def _check_heartbeat(self) -> bool: """ - Attempt to connect to JVM on the specific port. If succeeds, it means the - socket is still available and the JVM is still alive. Otherwise, the JVM - might have been gone. + Attempt to connect to the JVM port. Returns True iff socket.create_connection succeeds; + a close() failure after a successful connection is logged but DOES NOT flip the return value. - :return: bool, indicating if the socket is available. + :return: bool, True if connect succeeded, False if connect failed. """ try: temp_socket = socket.create_connection( (self._parsed_server_host, self._parsed_server_port), timeout=1 ) - temp_socket.close() - return True except Exception as e: logger.warning(f"Server is down with exception: {e}") return False + try: + temp_socket.close() + except Exception as e: + logger.warning(f"Failed to close socket: {e}") + + return True + @overrides def stop(self): # clean up every process under the python worker diff --git a/amber/src/test/python/core/runnables/test_heartbeat.py b/amber/src/test/python/core/runnables/test_heartbeat.py index 76dc93071aa..0fe5b321f50 100644 --- a/amber/src/test/python/core/runnables/test_heartbeat.py +++ b/amber/src/test/python/core/runnables/test_heartbeat.py @@ -79,12 +79,7 @@ def test_returns_false_when_socket_connection_times_out(self): ): assert hb._check_heartbeat() is False - def test_returns_false_when_socket_close_raises(self): - # Pins the false-negative path: connect succeeds but the subsequent - # close() throws (e.g. broken pipe on a half-open socket). The bare - # `except Exception` in _check_heartbeat() catches it and reports - # "server down", and a regression that narrows that handler would be - # caught here. + def test_returns_true_when_connection_succeeds_but_socket_close_raises(self): hb = make_heartbeat() fake_sock = MagicMock() fake_sock.close.side_effect = OSError("close failed") @@ -92,7 +87,8 @@ def test_returns_false_when_socket_close_raises(self): "core.runnables.heartbeat.socket.create_connection", return_value=fake_sock, ): - assert hb._check_heartbeat() is False + assert hb._check_heartbeat() is True + fake_sock.close.assert_called_once() class TestRunEarlyExit: