diff --git a/PasarGuardNodeBridge/controller.py b/PasarGuardNodeBridge/controller.py index cd79d73..470a936 100644 --- a/PasarGuardNodeBridge/controller.py +++ b/PasarGuardNodeBridge/controller.py @@ -339,10 +339,7 @@ async def _release_lifecycle_lease( node_version: str = "", core_version: str = "", ) -> None: - heartbeat = self._lifecycle_heartbeat_tasks.pop(lease.token, None) - if heartbeat is not None: - heartbeat.cancel() - await heartbeat + await self._stop_lifecycle_heartbeat(lease) if observed is None: await self._lifecycle_coordinator.release(lease) @@ -359,6 +356,17 @@ async def _release_lifecycle_lease( ) await self._lifecycle_coordinator.release(lease, state) + async def _stop_lifecycle_heartbeat(self, lease: LifecycleLease) -> None: + heartbeat = self._lifecycle_heartbeat_tasks.pop(lease.token, None) + if heartbeat is not None: + heartbeat.cancel() + try: + await heartbeat + except asyncio.CancelledError: + pass + except Exception: + self.logger.exception("[%s] Lifecycle heartbeat failed during cleanup", self.name) + async def get_lifecycle_state(self) -> NodeLifecycleState | None: return await self._lifecycle_coordinator.get_state(self.node_id) diff --git a/PasarGuardNodeBridge/grpclib.py b/PasarGuardNodeBridge/grpclib.py index a9eb7d0..048a9b8 100644 --- a/PasarGuardNodeBridge/grpclib.py +++ b/PasarGuardNodeBridge/grpclib.py @@ -198,21 +198,15 @@ async def stop(self, timeout: int | None = None) -> None: lease = await self._acquire_lifecycle_lease(LifecycleOperation.STOP) try: async with self._node_lock: - await self.disconnect() - - try: - await self._handle_grpc_request( - method=self._client.Stop, - request=service.Empty(), - timeout=timeout, - ) - except Exception: - pass - await self._release_lifecycle_lease( - lease, LifecycleStatus.STOPPED, desired=LifecycleStatus.STOPPED + await self._handle_grpc_request( + method=self._client.Stop, + request=service.Empty(), + timeout=timeout, ) + await self.disconnect() + await self._release_lifecycle_lease(lease, LifecycleStatus.STOPPED, desired=LifecycleStatus.STOPPED) except BaseException: - await self._release_lifecycle_lease(lease, LifecycleStatus.BROKEN, desired=LifecycleStatus.STOPPED) + await self._stop_lifecycle_heartbeat(lease) raise finally: await self._json_client.close() diff --git a/PasarGuardNodeBridge/rest.py b/PasarGuardNodeBridge/rest.py index 9011aa6..50d81cc 100644 --- a/PasarGuardNodeBridge/rest.py +++ b/PasarGuardNodeBridge/rest.py @@ -222,17 +222,11 @@ async def stop(self, timeout: int | None = None) -> None: lease = await self._acquire_lifecycle_lease(LifecycleOperation.STOP) try: async with self._node_lock: + await self._make_request(method="PUT", endpoint="stop", timeout=timeout) await self.disconnect() - - try: - await self._make_request(method="PUT", endpoint="stop", timeout=timeout) - except Exception: - pass - await self._release_lifecycle_lease( - lease, LifecycleStatus.STOPPED, desired=LifecycleStatus.STOPPED - ) + await self._release_lifecycle_lease(lease, LifecycleStatus.STOPPED, desired=LifecycleStatus.STOPPED) except BaseException: - await self._release_lifecycle_lease(lease, LifecycleStatus.BROKEN, desired=LifecycleStatus.STOPPED) + await self._stop_lifecycle_heartbeat(lease) raise finally: await self._client.close() diff --git a/tests/test_stop_lifecycle.py b/tests/test_stop_lifecycle.py new file mode 100644 index 0000000..63b9143 --- /dev/null +++ b/tests/test_stop_lifecycle.py @@ -0,0 +1,101 @@ +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +from PasarGuardNodeBridge.controller import Health, NodeAPIError +from PasarGuardNodeBridge.grpclib import Node as GrpcNode +from PasarGuardNodeBridge.rest import Node as RestNode +from PasarGuardNodeBridge.storage import InMemoryNodeLifecycleCoordinator, LifecycleOperation, LifecycleStatus + + +class StopLifecycleTests(unittest.IsolatedAsyncioTestCase): + def _configure_node(self, node, coordinator: InMemoryNodeLifecycleCoordinator) -> None: + node.node_id = "node-1" + node.name = "node-1" + node.worker_id = "worker-1" + node._default_timeout = 10 + node._node_lock = asyncio.Lock() + node._lifecycle_coordinator = coordinator + node._lifecycle_lease_seconds = 30 + node._lifecycle_heartbeat_tasks = {} + node.logger = Mock() + node.get_health = AsyncMock(return_value=Health.HEALTHY) + node.disconnect = AsyncMock() + node._json_client = SimpleNamespace(close=AsyncMock()) + + async def _assert_failed_stop_keeps_lease(self, node) -> None: + with self.assertRaises(NodeAPIError) as error: + await node.stop() + + self.assertEqual(error.exception.code, 503) + state = await node._lifecycle_coordinator.get_state(node.node_id) + self.assertEqual(state.observed, LifecycleStatus.STOPPING) + + competing_lease = await node._lifecycle_coordinator.try_acquire( + node.node_id, + "worker-2", + LifecycleOperation.START, + 30, + ) + self.assertIsNone(competing_lease) + self.assertEqual(node._lifecycle_heartbeat_tasks, {}) + node.disconnect.assert_not_awaited() + + async def test_rest_stop_propagates_error_without_releasing_lease(self): + coordinator = InMemoryNodeLifecycleCoordinator() + node = RestNode.__new__(RestNode) + self._configure_node(node, coordinator) + node._client = SimpleNamespace(close=AsyncMock()) + node._make_request = AsyncMock(side_effect=NodeAPIError(503, "REST stop failed")) + + await self._assert_failed_stop_keeps_lease(node) + + node._make_request.assert_awaited_once_with(method="PUT", endpoint="stop", timeout=10) + node._client.close.assert_awaited_once() + node._json_client.close.assert_awaited_once() + + async def test_grpc_stop_propagates_error_without_releasing_lease(self): + coordinator = InMemoryNodeLifecycleCoordinator() + node = GrpcNode.__new__(GrpcNode) + self._configure_node(node, coordinator) + node._client = SimpleNamespace(Stop=AsyncMock()) + node._handle_grpc_request = AsyncMock(side_effect=NodeAPIError(503, "gRPC stop failed")) + + await self._assert_failed_stop_keeps_lease(node) + + node._handle_grpc_request.assert_awaited_once() + request = node._handle_grpc_request.await_args.kwargs + self.assertIs(request["method"], node._client.Stop) + self.assertEqual(request["timeout"], 10) + node._json_client.close.assert_awaited_once() + + async def test_failed_heartbeat_does_not_mask_stop_error(self): + coordinator = InMemoryNodeLifecycleCoordinator() + node = RestNode.__new__(RestNode) + self._configure_node(node, coordinator) + node._client = SimpleNamespace(close=AsyncMock()) + node._make_request = AsyncMock(side_effect=NodeAPIError(503, "REST stop failed")) + + async def heartbeat_that_fails_during_cleanup(): + try: + await asyncio.Event().wait() + except asyncio.CancelledError as exc: + raise RuntimeError("heartbeat failed") from exc + + lease = await coordinator.try_acquire(node.node_id, node.worker_id, LifecycleOperation.STOP, 30) + self.assertIsNotNone(lease) + node._acquire_lifecycle_lease = AsyncMock(return_value=lease) + node._lifecycle_heartbeat_tasks[lease.token] = asyncio.create_task(heartbeat_that_fails_during_cleanup()) + await asyncio.sleep(0) + + with self.assertRaises(NodeAPIError) as error: + await node.stop() + + self.assertEqual(error.exception.code, 503) + self.assertEqual(error.exception.detail, "REST stop failed") + node.logger.exception.assert_called_once() + + +if __name__ == "__main__": + unittest.main()