From 2068db3c40a01d130a90ccdf11170e7043da7bc3 Mon Sep 17 00:00:00 2001 From: Rerowros Date: Sun, 9 Aug 2026 08:41:08 +0400 Subject: [PATCH 1/3] Preserve lifecycle lease when node stop fails --- PasarGuardNodeBridge/controller.py | 14 ++++-- PasarGuardNodeBridge/grpclib.py | 17 +++---- PasarGuardNodeBridge/rest.py | 11 ++--- tests/test_stop_lifecycle.py | 72 ++++++++++++++++++++++++++++++ 4 files changed, 91 insertions(+), 23 deletions(-) create mode 100644 tests/test_stop_lifecycle.py diff --git a/PasarGuardNodeBridge/controller.py b/PasarGuardNodeBridge/controller.py index cd79d73..e6c0ab0 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,15 @@ 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 + 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..e6c7254 100644 --- a/PasarGuardNodeBridge/grpclib.py +++ b/PasarGuardNodeBridge/grpclib.py @@ -200,19 +200,14 @@ async def stop(self, timeout: int | None = None) -> None: 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._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..210d9d6 100644 --- a/PasarGuardNodeBridge/rest.py +++ b/PasarGuardNodeBridge/rest.py @@ -224,15 +224,10 @@ async def stop(self, timeout: int | None = None) -> None: async with self._node_lock: 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._make_request(method="PUT", endpoint="stop", timeout=timeout) + 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..7bfd3fc --- /dev/null +++ b/tests/test_stop_lifecycle.py @@ -0,0 +1,72 @@ +import asyncio +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock + +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.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.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, {}) + + 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() + + +if __name__ == "__main__": + unittest.main() From 6389b06dfa22d090cc370f44fd0d3ce9d29b40ca Mon Sep 17 00:00:00 2001 From: Rerowros Date: Sun, 9 Aug 2026 17:22:38 +0400 Subject: [PATCH 2/3] fix: preserve retryable node stop state on failure --- PasarGuardNodeBridge/controller.py | 2 ++ PasarGuardNodeBridge/grpclib.py | 3 +-- PasarGuardNodeBridge/rest.py | 3 +-- tests/test_stop_lifecycle.py | 19 ++++++++++++++++++- 4 files changed, 22 insertions(+), 5 deletions(-) diff --git a/PasarGuardNodeBridge/controller.py b/PasarGuardNodeBridge/controller.py index e6c0ab0..470a936 100644 --- a/PasarGuardNodeBridge/controller.py +++ b/PasarGuardNodeBridge/controller.py @@ -364,6 +364,8 @@ async def _stop_lifecycle_heartbeat(self, lease: LifecycleLease) -> None: 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 e6c7254..048a9b8 100644 --- a/PasarGuardNodeBridge/grpclib.py +++ b/PasarGuardNodeBridge/grpclib.py @@ -198,13 +198,12 @@ 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() - 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._stop_lifecycle_heartbeat(lease) diff --git a/PasarGuardNodeBridge/rest.py b/PasarGuardNodeBridge/rest.py index 210d9d6..50d81cc 100644 --- a/PasarGuardNodeBridge/rest.py +++ b/PasarGuardNodeBridge/rest.py @@ -222,9 +222,8 @@ 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() - await self._make_request(method="PUT", endpoint="stop", timeout=timeout) + await self.disconnect() await self._release_lifecycle_lease(lease, LifecycleStatus.STOPPED, desired=LifecycleStatus.STOPPED) except BaseException: await self._stop_lifecycle_heartbeat(lease) diff --git a/tests/test_stop_lifecycle.py b/tests/test_stop_lifecycle.py index 7bfd3fc..4ebaebb 100644 --- a/tests/test_stop_lifecycle.py +++ b/tests/test_stop_lifecycle.py @@ -1,7 +1,7 @@ import asyncio import unittest from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock from PasarGuardNodeBridge.controller import Health, NodeAPIError from PasarGuardNodeBridge.grpclib import Node as GrpcNode @@ -12,12 +12,14 @@ 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()) @@ -38,6 +40,7 @@ async def _assert_failed_stop_keeps_lease(self, node) -> None: ) 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() @@ -67,6 +70,20 @@ async def test_grpc_stop_propagates_error_without_releasing_lease(self): 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) + lease = await coordinator.try_acquire(node.node_id, node.worker_id, LifecycleOperation.STOP, 30) + self.assertIsNotNone(lease) + heartbeat = asyncio.get_running_loop().create_future() + heartbeat.set_exception(RuntimeError("heartbeat failed")) + node._lifecycle_heartbeat_tasks[lease.token] = heartbeat + + await node._stop_lifecycle_heartbeat(lease) + + node.logger.exception.assert_called_once() + if __name__ == "__main__": unittest.main() From c8152523d5fafb295080b001e7bd738ee92c6587 Mon Sep 17 00:00:00 2001 From: Rerowros Date: Sun, 9 Aug 2026 18:24:01 +0400 Subject: [PATCH 3/3] test: cover heartbeat failure during node stop --- tests/test_stop_lifecycle.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/test_stop_lifecycle.py b/tests/test_stop_lifecycle.py index 4ebaebb..63b9143 100644 --- a/tests/test_stop_lifecycle.py +++ b/tests/test_stop_lifecycle.py @@ -74,14 +74,26 @@ 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) - heartbeat = asyncio.get_running_loop().create_future() - heartbeat.set_exception(RuntimeError("heartbeat failed")) - node._lifecycle_heartbeat_tasks[lease.token] = heartbeat + 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) - await node._stop_lifecycle_heartbeat(lease) + 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()