music-assistant-server

1"""
2Websocket API client for the Music Assistant performance benchmark.
3
4A minimal asyncio client for the MA websocket API with a single reader task that
5routes command results to their waiters and fans events out to subscribers, so
6events are never lost while a command result is being awaited.
7"""
8
9from __future__ import annotations
10
11import asyncio
12import contextlib
13import json
14import statistics
15import time
16import uuid
17from dataclasses import dataclass, field
18from typing import Any
19
20import aiohttp
21
22COMMAND_TIMEOUT = 300
23MAX_WS_MSG_SIZE = 64 * 1024 * 1024
24
25
26@dataclass
27class CommandResult:
28    """Outcome of a single websocket API command."""
29
30    result: Any
31    duration_seconds: float
32    payload_bytes: int
33    items: int
34
35
36@dataclass
37class _PendingCommand:
38    """Bookkeeping for an in-flight command awaiting its (partial) result(s)."""
39
40    future: asyncio.Future[CommandResult]
41    sent_at: float
42    payload_bytes: int = 0
43    items: int = 0
44    partial_results: list[Any] = field(default_factory=list)
45
46
47class PerfWsClient:
48    """Websocket API client with command latency/payload accounting and event capture."""
49
50    def __init__(self, port: int, token: str) -> None:
51        """
52        Initialize the client.
53
54        :param port: The (loopback) webserver port of the benchmark server.
55        :param token: A long-lived auth token for an admin user.
56        """
57        self.url = f"ws://127.0.0.1:{port}/ws"
58        self.token = token
59        self._session: aiohttp.ClientSession | None = None
60        self._ws: aiohttp.ClientWebSocketResponse | None = None
61        self._reader_task: asyncio.Task[None] | None = None
62        self._pending: dict[str, _PendingCommand] = {}
63        self._event_waiters: list[tuple[str, asyncio.Future[dict[str, Any]]]] = []
64
65    async def connect(self) -> None:
66        """Connect, consume the server-info message and authenticate."""
67        self._session = aiohttp.ClientSession()
68        self._ws = await self._session.ws_connect(self.url, max_msg_size=MAX_WS_MSG_SIZE)
69        # first message is the server info
70        await self._ws.receive_json()
71        self._reader_task = asyncio.create_task(self._reader())
72        result = await self.command("auth", {"token": self.token})
73        if not (isinstance(result.result, dict) and result.result.get("authenticated")):
74            raise RuntimeError(f"websocket auth failed: {result.result}")
75
76    async def close(self) -> None:
77        """Close the websocket connection."""
78        if self._reader_task:
79            self._reader_task.cancel()
80            with contextlib.suppress(asyncio.CancelledError):
81                await self._reader_task
82        if self._ws:
83            await self._ws.close()
84        if self._session:
85            await self._session.close()
86
87    async def command(
88        self, command: str, args: dict[str, Any] | None = None, timeout: float = COMMAND_TIMEOUT
89    ) -> CommandResult:
90        """
91        Send an API command and wait for its full (possibly chunked) result.
92
93        :param command: The API command to execute (e.g. music/tracks/library_items).
94        :param args: Optional arguments for the command.
95        :param timeout: Maximum seconds to wait for the (final) result.
96        """
97        assert self._ws is not None
98        message_id = uuid.uuid4().hex
99        payload: dict[str, Any] = {"message_id": message_id, "command": command}
100        if args:
101            payload["args"] = args
102        pending = _PendingCommand(
103            future=asyncio.get_running_loop().create_future(), sent_at=time.perf_counter()
104        )
105        self._pending[message_id] = pending
106        try:
107            await self._ws.send_str(json.dumps(payload))
108            return await asyncio.wait_for(pending.future, timeout=timeout)
109        finally:
110            self._pending.pop(message_id, None)
111
112    async def wait_for_event(self, event: str, timeout: float) -> dict[str, Any]:
113        """
114        Wait until the server pushes the given event type.
115
116        :param event: The event type value to wait for (e.g. music_sync_completed).
117        :param timeout: Maximum seconds to wait.
118        """
119        future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future()
120        waiter = (event, future)
121        self._event_waiters.append(waiter)
122        try:
123            return await asyncio.wait_for(future, timeout=timeout)
124        finally:
125            if waiter in self._event_waiters:
126                self._event_waiters.remove(waiter)
127
128    async def _reader(self) -> None:
129        """Route incoming messages to pending commands and event waiters."""
130        assert self._ws is not None
131        try:
132            async for msg in self._ws:
133                if msg.type != aiohttp.WSMsgType.TEXT:
134                    break
135                self._handle_message(json.loads(msg.data), len(msg.data.encode()))
136        finally:
137            # fail fast on a closed/errored connection instead of letting callers
138            # wait out their full command timeouts
139            error = ConnectionError("websocket connection closed")
140            for pending in self._pending.values():
141                if not pending.future.done():
142                    pending.future.set_exception(error)
143            for _, future in self._event_waiters:
144                if not future.done():
145                    future.set_exception(error)
146
147    def _handle_message(self, data: dict[str, Any], size_bytes: int) -> None:
148        """Dispatch a single incoming message to its pending command or event waiters."""
149        if event := data.get("event"):
150            for waiter_event, future in list(self._event_waiters):
151                if waiter_event == event and not future.done():
152                    future.set_result(data)
153            return
154        message_id = data.get("message_id")
155        if not message_id or not (pending := self._pending.get(message_id)):
156            return
157        pending.payload_bytes += size_bytes
158        if "error_code" in data:
159            if not pending.future.done():
160                pending.future.set_exception(
161                    RuntimeError(f"command failed: {data.get('details') or data['error_code']}")
162                )
163            return
164        result = data.get("result")
165        if isinstance(result, list):
166            pending.items += len(result)
167            pending.partial_results.extend(result)
168        if data.get("partial"):
169            return
170        if not pending.future.done():
171            final_result = result
172            if pending.partial_results and isinstance(result, list):
173                final_result = pending.partial_results
174            pending.future.set_result(
175                CommandResult(
176                    result=final_result,
177                    duration_seconds=time.perf_counter() - pending.sent_at,
178                    payload_bytes=pending.payload_bytes,
179                    items=pending.items,
180                )
181            )
182
183
184def summarize_latencies(durations_ms: list[float]) -> tuple[float, float]:
185    """Return (median, p95) for a list of latency samples in milliseconds."""
186    median = statistics.median(durations_ms)
187    if len(durations_ms) >= 20:
188        p95 = statistics.quantiles(durations_ms, n=20)[18]
189    else:
190        p95 = max(durations_ms)
191    return round(median, 1), round(p95, 1)
192