From 62a0f2443745e293c5029b82e053ac1a0ec29ef2 Mon Sep 17 00:00:00 2001 From: Cabbagec Date: Thu, 13 Aug 2026 06:25:51 +0000 Subject: [PATCH] fix: recover client control connection safely --- Dockerfile | 1 + docs/deployment-and-usage.md | 1 + src/archive_clients/config.py | 7 ++ src/archive_clients/daemon.py | 155 ++++++++++++++++++++++------------ tests/test_daemon.py | 117 +++++++++++++++++++++++++ 5 files changed, 228 insertions(+), 53 deletions(-) diff --git a/Dockerfile b/Dockerfile index 1732729..be5cc9a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,6 +7,7 @@ RUN pip wheel --no-cache-dir --wheel-dir /wheels . FROM builder AS test RUN pip install --no-cache-dir /wheels/*.whl COPY tests ./tests +COPY scripts ./scripts CMD ["python", "-m", "unittest", "discover", "-s", "tests", "-v"] FROM python:3.11-slim diff --git a/docs/deployment-and-usage.md b/docs/deployment-and-usage.md index 9727aa4..d82d8fd 100644 --- a/docs/deployment-and-usage.md +++ b/docs/deployment-and-usage.md @@ -95,6 +95,7 @@ backup_dir = "/var/backups/archive-control" registration_timeout = "10s" # First response must arrive within this window. heartbeat_interval = "15s" # Server may negotiate a different effective value. offline_timeout = "45s" +outbound_enqueue_timeout = "15s" # Reconnect rather than wedging if sends stop draining. reconnect_initial = "1s" reconnect_max = "60s" # Retry forever, never wait longer than this. reconnect_reset_after = "60s" diff --git a/src/archive_clients/config.py b/src/archive_clients/config.py index 442f275..b748707 100644 --- a/src/archive_clients/config.py +++ b/src/archive_clients/config.py @@ -85,6 +85,7 @@ class ConnectionConfig: registration_timeout: float = 10 heartbeat_interval: float = 15 offline_timeout: float = 45 + outbound_enqueue_timeout: float = 15 reconnect_initial: float = 1 reconnect_max: float = 60 reconnect_reset_after: float = 60 @@ -249,6 +250,7 @@ def _connection(value: Any) -> ConnectionConfig: raise ConfigError("connection must be a table") _keys(value, { "registration_timeout", "heartbeat_interval", "offline_timeout", + "outbound_enqueue_timeout", "reconnect_initial", "reconnect_max", "reconnect_reset_after", "reconnect_jitter", }, "connection") @@ -256,6 +258,9 @@ def _connection(value: Any) -> ConnectionConfig: registration_timeout=_duration(value.get("registration_timeout", "10s")), heartbeat_interval=_duration(value.get("heartbeat_interval", "15s")), offline_timeout=_duration(value.get("offline_timeout", "45s")), + outbound_enqueue_timeout=_duration( + value.get("outbound_enqueue_timeout", "15s") + ), reconnect_initial=_duration(value.get("reconnect_initial", "1s")), reconnect_max=_duration(value.get("reconnect_max", "60s")), reconnect_reset_after=_duration(value.get("reconnect_reset_after", "60s")), @@ -265,6 +270,8 @@ def _connection(value: Any) -> ConnectionConfig: raise ConfigError("reconnect_initial cannot exceed reconnect_max") if result.offline_timeout <= result.heartbeat_interval: raise ConfigError("offline_timeout must exceed heartbeat_interval") + if result.outbound_enqueue_timeout <= 0: + raise ConfigError("outbound_enqueue_timeout must be positive") if not isinstance(result.reconnect_jitter, bool): raise ConfigError("reconnect_jitter must be boolean") return result diff --git a/src/archive_clients/daemon.py b/src/archive_clients/daemon.py index 6367132..c2dec05 100644 --- a/src/archive_clients/daemon.py +++ b/src/archive_clients/daemon.py @@ -229,16 +229,11 @@ class ArchiveClientDaemon: # detects the missing application heartbeats and reconnects. while True: try: - frame = await asyncio.wait_for( - websocket.recv(), - self.config.connection.offline_timeout, + frame = await self._next_frame( + websocket, writer ) except ConnectionClosedOK: return - except asyncio.TimeoutError as exc: - raise RuntimeError( - "control heartbeat timed out" - ) from exc await self._handle(decode(frame), outbound, command_tasks) finally: writer.cancel() @@ -323,6 +318,63 @@ class ArchiveClientDaemon: while True: await websocket.send(await outbound.get()) + async def _next_frame(self, websocket: Any, writer: asyncio.Task[None]): + """Receive one frame while supervising the connection writer. + + The previous receive-only wait let a failed writer go unnoticed until + the remote side happened to close or a heartbeat timeout elapsed. + Waiting for both directions makes a failed send an immediate, + reconnectable connection failure. + """ + + receiver = asyncio.create_task(websocket.recv()) + try: + done, _ = await asyncio.wait( + {receiver, writer}, + timeout=self.config.connection.offline_timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if writer in done: + if not receiver.done(): + receiver.cancel() + await asyncio.gather(receiver, return_exceptions=True) + if writer.cancelled(): + raise RuntimeError("control writer was cancelled") + error = writer.exception() + if error is not None: + raise RuntimeError("control writer failed") from error + raise RuntimeError("control writer stopped") + if receiver not in done: + receiver.cancel() + await asyncio.gather(receiver, return_exceptions=True) + raise RuntimeError("control heartbeat timed out") + return receiver.result() + except BaseException: + if not receiver.done(): + receiver.cancel() + await asyncio.gather(receiver, return_exceptions=True) + raise + + async def _enqueue( + self, outbound: asyncio.Queue[str], message: str + ) -> None: + """Bound connection-local backpressure so a dead writer cannot wedge I/O. + + A job event is durable before it is sent, so dropping this particular + connection after a bounded wait is safe: reconciliation/redelivery on + the next session will recover it. In contrast, indefinitely waiting + for a full queue can prevent heartbeat acknowledgements from being + read or sent, which prevents the reconnect supervisor from running. + """ + + try: + await asyncio.wait_for( + outbound.put(message), + timeout=self.config.connection.outbound_enqueue_timeout, + ) + except asyncio.TimeoutError as exc: + raise RuntimeError("control outbound queue is blocked") from exc + async def _handle( self, envelope: Any, @@ -334,13 +386,20 @@ class ArchiveClientDaemon: response = new_envelope() response.correlation_id = envelope.message_id response.heartbeat_ack.sequence = envelope.heartbeat.sequence - await outbound.put(encode(response)) + await self._enqueue(outbound, encode(response)) elif payload == "command": await self._accept_command(envelope, outbound, command_tasks) elif payload == "protocol_error": logger.warning( "control_reported_protocol_error", - extra={"error_code": envelope.protocol_error.error.code}, + extra={ + "error_code": envelope.protocol_error.error.code, + "error_detail": envelope.protocol_error.error.message, + "offending_message_id": ( + envelope.protocol_error.offending_message_id + ), + "retryable": envelope.protocol_error.error.retryable, + }, ) async def _accept_command( @@ -398,7 +457,7 @@ class ArchiveClientDaemon: response = new_envelope() response.correlation_id = envelope.message_id response.command_ack.CopyFrom(acknowledgement) - await outbound.put(encode(response)) + await self._enqueue(outbound, encode(response)) if ( accepted is not None @@ -419,7 +478,7 @@ class ArchiveClientDaemon: job_snapshot.last_event_sequence = int( row["last_event_sequence"] ) - await outbound.put(encode(snapshot)) + await self._enqueue(outbound, encode(snapshot)) elif ( accepted is not None and accepted_for_execution @@ -472,7 +531,6 @@ class ArchiveClientDaemon: outbound: asyncio.Queue[str], command_tasks: set[asyncio.Task[None]], ) -> None: - job_commands: list[control_pb2.Command] = [] for row in await asyncio.to_thread(self.store.list_accepted_commands): acknowledgement = decode_message( str(row["acknowledgement_json"]), control_pb2.CommandAck() @@ -484,28 +542,14 @@ class ArchiveClientDaemon: ) if command.WhichOneof("payload") == "ensure_route": self._schedule_route_command( - command, "", outbound, command_tasks + command, "", outbound, command_tasks, restore_ready=True ) - elif ( - command.WhichOneof("payload") in { - "assign_job", "execute_step", "cancel_job" - } - and self.jobs is not None - ): - if command.WhichOneof("payload") == "cancel_job": - self.jobs.request_cancel(command.cancel_job.job_id) - job_commands.append(command) - if job_commands: - task = asyncio.create_task( - self._resume_job_commands(job_commands, outbound), - name="resume-job-commands", - ) - command_tasks.add(task) - task.add_done_callback( - lambda completed: self._command_finished( - completed, command_tasks - ) - ) + # Job commands are intentionally not replayed here. A command + # acknowledgement is durable on both sides; the control daemon + # redelivers an unacknowledged command, while registration + # reconciliation requests snapshots for active jobs. Replaying + # every historical command re-emits overlapping event ranges and + # can flood the single connection's bounded outbound queue. async def _resume_route_commands( self, @@ -604,10 +648,13 @@ class ArchiveClientDaemon: response.correlation_id = correlation_id response.job_event.CopyFrom(event) future = asyncio.run_coroutine_threadsafe( - outbound.put(encode(response)), loop + self._enqueue(outbound, encode(response)), loop ) try: - future.result(timeout=30) + future.result( + timeout=self.config.connection.outbound_enqueue_timeout + + 1 + ) except Exception: # The event is already durable in the client DB and will # be replayed after reconnect. @@ -627,7 +674,7 @@ class ArchiveClientDaemon: response = new_envelope() response.correlation_id = correlation_id response.job_event.CopyFrom(event) - await outbound.put(encode(response)) + await self._enqueue(outbound, encode(response)) def _schedule_route_command( self, @@ -635,12 +682,15 @@ class ArchiveClientDaemon: correlation_id: str, outbound: asyncio.Queue[str], command_tasks: set[asyncio.Task[None]], + restore_ready: bool = False, ) -> None: if command.command_id in self._active_route_commands: return self._active_route_commands.add(command.command_id) task = asyncio.create_task( - self._execute_route(command, correlation_id, outbound), + self._execute_route( + command, correlation_id, outbound, restore_ready=restore_ready + ), name=f"route-{command.ensure_route.route.route_id}", ) command_tasks.add(task) @@ -671,13 +721,14 @@ class ArchiveClientDaemon: response = new_envelope() response.correlation_id = correlation_id response.inventory_chunk.CopyFrom(chunk) - await outbound.put(encode(response)) + await self._enqueue(outbound, encode(response)) async def _execute_route( self, command: Any, correlation_id: str, outbound: asyncio.Queue[str], + restore_ready: bool = False, ) -> None: assert self.routes is not None spec = command.ensure_route.route @@ -696,21 +747,19 @@ class ArchiveClientDaemon: response.route_update.CopyFrom( decode_message(str(row["update_json"]), control_pb2.RouteUpdate()) ) - await outbound.put(encode(response)) + await self._enqueue(outbound, encode(response)) if attempt["state"] == "ready": - # Route readiness is durable, but this process-local lookup is - # not. Reconfigure (which validates the existing Syncthing folder - # and recreates its local directory if necessary) before replaying - # the stored READY updates. This is essential after a daemon or - # mount restart: a previously ready route may otherwise point at a - # path that is no longer present, and a later source-stage command - # would fail with a bare ENOENT. - configured = await asyncio.to_thread( - self.routes.configure, - spec, - time.monotonic() + spec.setup_timeout_seconds, - ) - self._known_route_paths[spec.route_id] = configured.local_path + if restore_ready: + # Route readiness is durable, but this process-local lookup is + # not. Reconfigure (which validates the existing Syncthing + # folder and recreates its local directory if necessary) after + # a daemon restart, before replaying stored READY updates. + configured = await asyncio.to_thread( + self.routes.configure, + spec, + time.monotonic() + spec.setup_timeout_seconds, + ) + self._known_route_paths[spec.route_id] = configured.local_path return if attempt["state"] == "failed": return @@ -863,7 +912,7 @@ class ArchiveClientDaemon: response = new_envelope() response.correlation_id = correlation_id response.route_update.CopyFrom(update) - await outbound.put(encode(response)) + await self._enqueue(outbound, encode(response)) def _local_route( self, spec: route_pb2.EnsureRouteSpec, state: int diff --git a/tests/test_daemon.py b/tests/test_daemon.py index 408d2ee..1157e69 100644 --- a/tests/test_daemon.py +++ b/tests/test_daemon.py @@ -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)