Preserve lifecycle lease when node stop fails

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)
 

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()

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()

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()