fix: recover client control connection safely

2026-08-13 06:25:51 +00:00
parent eabbda3c3a
commit 62a0f24437
5 changed files with 228 additions and 53 deletions
@@ -110,6 +110,123 @@ class DaemonTransportTests(unittest.IsolatedAsyncioTestCase):
                ):
                    await asyncio.wait_for(daemon._connection(), 1)

+    async def test_failed_writer_ends_connection_without_waiting_for_heartbeat(self):
+        """A send failure must immediately reach the reconnect supervisor."""
+
+        async def control(websocket):
+            registration = decode(await websocket.recv())
+            response = new_envelope()
+            response.correlation_id = registration.message_id
+            response.register_response.status = (
+                client_pb2.REGISTRATION_STATUS_ACCEPTED
+            )
+            response.register_response.negotiated_version.major = 1
+            await websocket.send(encode(response))
+            await websocket.wait_closed()
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            token = root / "token"
+            token.write_text("shared-secret", encoding="utf-8")
+            os.chmod(token, 0o600)
+            async with serve(control, "127.0.0.1", 0, ping_interval=None) as server:
+                port = server.sockets[0].getsockname()[1]
+                service = ServiceConfig("http://local", PurePosixPath("/api"), root)
+                config = ClientConfig(
+                    "cache-1", "Cache 1", "cache",
+                    f"ws://127.0.0.1:{port}", token,
+                    root / "state.db", root / "backups", service, service,
+                )
+                probe = FilesystemProbe(root, True, True, True, True, True)
+                daemon = ArchiveClientDaemon(config, [probe, probe], [])
+                await asyncio.to_thread(daemon.store.initialize)
+
+                async def failed_writer(websocket, outbound):
+                    raise OSError("simulated broken socket")
+
+                daemon._writer = failed_writer
+                with self.assertRaisesRegex(RuntimeError, "writer failed"):
+                    await asyncio.wait_for(daemon._connection(), 1)
+
+    async def test_full_outbound_queue_aborts_connection_instead_of_blocking_heartbeats(self):
+        """Bulk output cannot indefinitely block the receive/heartbeat loop."""
+
+        async def control(websocket):
+            registration = decode(await websocket.recv())
+            response = new_envelope()
+            response.correlation_id = registration.message_id
+            response.register_response.status = (
+                client_pb2.REGISTRATION_STATUS_ACCEPTED
+            )
+            response.register_response.negotiated_version.major = 1
+            await websocket.send(encode(response))
+            for sequence in range(1, 102):
+                heartbeat = new_envelope()
+                heartbeat.heartbeat.sequence = sequence
+                await websocket.send(encode(heartbeat))
+            await websocket.wait_closed()
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            token = root / "token"
+            token.write_text("shared-secret", encoding="utf-8")
+            os.chmod(token, 0o600)
+            async with serve(control, "127.0.0.1", 0, ping_interval=None) as server:
+                port = server.sockets[0].getsockname()[1]
+                service = ServiceConfig("http://local", PurePosixPath("/api"), root)
+                config = ClientConfig(
+                    "cache-1", "Cache 1", "cache",
+                    f"ws://127.0.0.1:{port}", token,
+                    root / "state.db", root / "backups", service, service,
+                    ConnectionConfig(outbound_enqueue_timeout=0.01),
+                )
+                probe = FilesystemProbe(root, True, True, True, True, True)
+                daemon = ArchiveClientDaemon(config, [probe, probe], [])
+                await asyncio.to_thread(daemon.store.initialize)
+
+                async def stopped_writer(websocket, outbound):
+                    await asyncio.Event().wait()
+
+                daemon._writer = stopped_writer
+                with self.assertRaisesRegex(RuntimeError, "outbound queue is blocked"):
+                    await asyncio.wait_for(daemon._connection(), 2)
+
+    async def test_reconnect_resume_skips_historical_job_commands(self):
+        """Registration reconciliation, not command replay, recovers job state."""
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            token = root / "token"
+            token.write_text("shared-secret", encoding="utf-8")
+            os.chmod(token, 0o600)
+            service = ServiceConfig("http://local", PurePosixPath("/api"), root)
+            config = ClientConfig(
+                "cache-1", "Cache 1", "cache", "ws://control", token,
+                root / "state.db", root / "backups", service, service,
+            )
+            probe = FilesystemProbe(root, True, True, True, True, True)
+            daemon = ArchiveClientDaemon(config, [probe, probe], [])
+            await asyncio.to_thread(daemon.store.initialize)
+            command = control_pb2.Command(command_id=str(uuid4()))
+            command.execute_step.job_id = str(uuid4())
+            command.execute_step.expected_last_event_sequence = 1
+            acknowledgement = control_pb2.CommandAck(
+                command_id=command.command_id,
+                status=control_pb2.COMMAND_ACK_STATUS_ACCEPTED,
+            )
+            await asyncio.to_thread(
+                daemon.store.accept_command,
+                command.command_id,
+                encode_message(command),
+                encode_message(acknowledgement),
+            )
+            daemon.jobs = Mock()
+            outbound = asyncio.Queue()
+            tasks = set()
+            await daemon._resume_commands(outbound, tasks)
+            self.assertEqual(tasks, set())
+            self.assertTrue(outbound.empty())
+
    async def test_eviction_assignment_and_steps_are_admitted(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)