feat: add archive client foundation
This commit is contained in:
@@ -0,0 +1,198 @@
|
||||
"""Reconnectable Archive Control client transport."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from websockets.asyncio.client import connect
|
||||
|
||||
from archive_clients.config import ClientConfig
|
||||
from archive_clients.probes import FilesystemProbe
|
||||
from archive_clients.protocol import (
|
||||
decode,
|
||||
decode_message,
|
||||
encode,
|
||||
encode_message,
|
||||
new_envelope,
|
||||
)
|
||||
from archive_clients.state import ClientStore, CommandConflict
|
||||
from archive_control.v1 import client_pb2, common_pb2, control_pb2, job_pb2
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ArchiveClientDaemon:
|
||||
def __init__(self, config: ClientConfig, probes: list[FilesystemProbe]):
|
||||
self.config = config
|
||||
self.probes = probes
|
||||
self.store = ClientStore(config.state_db)
|
||||
|
||||
async def run(self) -> None:
|
||||
await asyncio.to_thread(self.store.initialize)
|
||||
delay = self.config.connection.reconnect_initial
|
||||
while True:
|
||||
started = time.monotonic()
|
||||
try:
|
||||
await self._connection()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("control connection ended: %s", type(exc).__name__)
|
||||
if time.monotonic() - started >= self.config.connection.reconnect_reset_after:
|
||||
delay = self.config.connection.reconnect_initial
|
||||
wait = random.uniform(0, delay) if self.config.connection.reconnect_jitter else delay
|
||||
await asyncio.sleep(wait)
|
||||
delay = min(delay * 2, self.config.connection.reconnect_max)
|
||||
|
||||
async def _connection(self) -> None:
|
||||
async with connect(
|
||||
self.config.control_endpoint, ping_interval=None, compression=None,
|
||||
max_size=1024 * 1024,
|
||||
) as websocket:
|
||||
registration = self._registration()
|
||||
await websocket.send(encode(registration))
|
||||
registration.register_request.shared_token = ""
|
||||
response = decode(await asyncio.wait_for(
|
||||
websocket.recv(), self.config.connection.registration_timeout
|
||||
))
|
||||
if response.WhichOneof("payload") != "register_response":
|
||||
raise RuntimeError("control did not answer registration")
|
||||
if response.correlation_id != registration.message_id:
|
||||
raise RuntimeError("registration response correlation mismatch")
|
||||
if response.register_response.status != client_pb2.REGISTRATION_STATUS_ACCEPTED:
|
||||
raise RuntimeError("control rejected registration")
|
||||
if response.register_response.negotiated_version.major != 1:
|
||||
raise RuntimeError("control negotiated an unsupported protocol version")
|
||||
outbound: asyncio.Queue[str] = asyncio.Queue(maxsize=100)
|
||||
writer = asyncio.create_task(self._writer(websocket, outbound))
|
||||
try:
|
||||
async for frame in websocket:
|
||||
await self._handle(decode(frame), outbound)
|
||||
finally:
|
||||
writer.cancel()
|
||||
await asyncio.gather(writer, return_exceptions=True)
|
||||
|
||||
def _registration(self):
|
||||
envelope = new_envelope()
|
||||
request = envelope.register_request
|
||||
request.protocol_version.CopyFrom(envelope.protocol_version)
|
||||
request.client.client_id = self.config.client_id
|
||||
request.client.display_name = self.config.display_name
|
||||
request.client.role = (
|
||||
common_pb2.CLIENT_ROLE_ARCHIVE
|
||||
if self.config.role == "archive" else common_pb2.CLIENT_ROLE_CACHE
|
||||
)
|
||||
request.connection_instance_id = str(uuid.uuid4())
|
||||
request.shared_token = self.config.read_shared_token()
|
||||
request.capabilities.max_envelope_bytes = 1024 * 1024
|
||||
request.capabilities.syncthing_advertised_addresses.extend(
|
||||
self.config.syncthing.advertised_addresses
|
||||
)
|
||||
for root_name, probe in zip(
|
||||
("qbittorrent", "syncthing"), self.probes, strict=True
|
||||
):
|
||||
filesystem = request.capabilities.filesystems.add()
|
||||
filesystem.root_name = root_name
|
||||
filesystem.readable = probe.readable
|
||||
filesystem.writable = probe.writable
|
||||
filesystem.hard_link = probe.hard_link
|
||||
filesystem.sparse_files = probe.sparse_files
|
||||
if all(probe.sparse_files for probe in self.probes):
|
||||
request.capabilities.features.append(client_pb2.CLIENT_FEATURE_SPARSE_FILES)
|
||||
for cursor in self.store.list_active_job_cursors():
|
||||
active = request.active_jobs.add()
|
||||
active.job_id = str(cursor["job_id"])
|
||||
active.job_revision = int(cursor["revision"])
|
||||
active.last_event_sequence = int(cursor["last_event_sequence"])
|
||||
active.state = job_pb2.JobState.Value(str(cursor["state"]))
|
||||
active.committed = bool(cursor["committed"])
|
||||
return envelope
|
||||
|
||||
async def _writer(
|
||||
self, websocket: Any, outbound: asyncio.Queue[str]
|
||||
) -> None:
|
||||
while True:
|
||||
await websocket.send(await outbound.get())
|
||||
|
||||
async def _handle(
|
||||
self, envelope: Any, outbound: asyncio.Queue[str]
|
||||
) -> None:
|
||||
payload = envelope.WhichOneof("payload")
|
||||
if payload == "heartbeat":
|
||||
response = new_envelope()
|
||||
response.correlation_id = envelope.message_id
|
||||
response.heartbeat_ack.sequence = envelope.heartbeat.sequence
|
||||
await outbound.put(encode(response))
|
||||
elif payload == "command":
|
||||
await self._accept_command(envelope, outbound)
|
||||
elif payload == "protocol_error":
|
||||
logger.warning("control reported protocol error code=%s", envelope.protocol_error.error.code)
|
||||
|
||||
async def _accept_command(
|
||||
self, envelope: Any, outbound: asyncio.Queue[str]
|
||||
) -> None:
|
||||
command = envelope.command
|
||||
acknowledgement = self._initial_acknowledgement(command)
|
||||
accepted = None
|
||||
try:
|
||||
accepted = await asyncio.to_thread(
|
||||
self.store.accept_command,
|
||||
command.command_id,
|
||||
encode_message(command),
|
||||
encode_message(acknowledgement),
|
||||
)
|
||||
if accepted.duplicate:
|
||||
acknowledgement = decode_message(
|
||||
accepted.acknowledgement_json, control_pb2.CommandAck()
|
||||
)
|
||||
if (
|
||||
acknowledgement.status
|
||||
== control_pb2.COMMAND_ACK_STATUS_ACCEPTED
|
||||
):
|
||||
acknowledgement.status = (
|
||||
control_pb2.COMMAND_ACK_STATUS_DUPLICATE
|
||||
)
|
||||
except CommandConflict:
|
||||
acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_REJECTED
|
||||
acknowledgement.error.code = common_pb2.ERROR_CODE_CONFLICT
|
||||
acknowledgement.error.message = "command ID content conflict"
|
||||
response = new_envelope()
|
||||
response.correlation_id = envelope.message_id
|
||||
response.command_ack.CopyFrom(acknowledgement)
|
||||
await outbound.put(encode(response))
|
||||
|
||||
if (
|
||||
accepted is not None
|
||||
and not accepted.duplicate
|
||||
and acknowledgement.status
|
||||
== control_pb2.COMMAND_ACK_STATUS_ACCEPTED
|
||||
and command.WhichOneof("payload") == "request_job_snapshot"
|
||||
):
|
||||
snapshot = new_envelope()
|
||||
snapshot.correlation_id = envelope.message_id
|
||||
snapshot.client_state_snapshot.snapshot_id = str(uuid.uuid4())
|
||||
snapshot.client_state_snapshot.observed_at.CopyFrom(snapshot.sent_at)
|
||||
await outbound.put(encode(snapshot))
|
||||
|
||||
@staticmethod
|
||||
def _initial_acknowledgement(command: Any) -> control_pb2.CommandAck:
|
||||
acknowledgement = control_pb2.CommandAck(command_id=command.command_id)
|
||||
if not command.command_id or command.WhichOneof("payload") is None:
|
||||
acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_REJECTED
|
||||
acknowledgement.error.code = common_pb2.ERROR_CODE_INVALID_ARGUMENT
|
||||
acknowledgement.error.message = "command ID and payload are required"
|
||||
elif command.WhichOneof("payload") == "request_job_snapshot":
|
||||
acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_ACCEPTED
|
||||
else:
|
||||
acknowledgement.status = control_pb2.COMMAND_ACK_STATUS_REJECTED
|
||||
acknowledgement.error.code = common_pb2.ERROR_CODE_UNSUPPORTED
|
||||
acknowledgement.error.message = (
|
||||
"command is not supported by this client build"
|
||||
)
|
||||
return acknowledgement
|
||||
Reference in New Issue
Block a user