From 1efcbcf2ba808b4a68b00a9851c3a02227f9b1f9 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Sun, 15 Mar 2026 13:54:49 -0700 Subject: [PATCH] =?UTF-8?q?perf:=20parallelize=20=5Fcollect=5Fmcp=5Fstatus?= =?UTF-8?q?=20and=20=5Fnotify=5Fnodes=5Fmcp=5Freload=20wi=E2=80=A6=20(#73)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf: parallelize _collect_mcp_status and _notify_nodes_mcp_reload with asyncio.gather Both functions queried cluster nodes sequentially, making latency O(N × timeout). Use asyncio.gather to query all nodes concurrently, matching the existing admin_list_watches pattern. Also reuse the shared proxy_client instead of creating throwaway httpx clients per node, and add debug logging on MCP status fetch failures. * perf: bound node fan-out concurrency and improve debug logging Add _NODE_FAN_OUT_LIMIT (50) semaphore to all three gather fan-out sites (_collect_mcp_status, _notify_nodes_mcp_reload, admin_list_watches) to cap concurrent outbound connections below the httpx pool limit, leaving headroom for other proxy traffic at 1000-node scale. Add exc_info=True to all debug log calls for actionable diagnostics. * test: add unit tests for _collect_mcp_status and _notify_nodes_mcp_reload 11 tests covering success, non-200, missing URL, exceptions, empty cluster, and mixed multi-node scenarios for both fan-out helpers. --- tests/test_mcp_admin_api.py | 158 +++++++++++++++++++++++++++++++++++- turnstone/console/server.py | 85 +++++++++++-------- 2 files changed, 210 insertions(+), 33 deletions(-) diff --git a/tests/test_mcp_admin_api.py b/tests/test_mcp_admin_api.py index eafc734f..5ad62d91 100644 --- a/tests/test_mcp_admin_api.py +++ b/tests/test_mcp_admin_api.py @@ -5,7 +5,7 @@ from __future__ import annotations import json import uuid from typing import TYPE_CHECKING, Any -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from starlette.applications import Starlette @@ -19,6 +19,8 @@ if TYPE_CHECKING: from starlette.responses import Response from turnstone.console.server import ( + _collect_mcp_status, + _notify_nodes_mcp_reload, admin_create_mcp_server, admin_delete_mcp_server, admin_get_mcp_server, @@ -537,3 +539,157 @@ class TestPermission: def test_delete_without_permission(self, client_no_perm): r = client_no_perm.delete(f"/v1/api/admin/mcp-servers/{uuid.uuid4().hex}") assert r.status_code == 403 + + +# --------------------------------------------------------------------------- +# Unit tests for _collect_mcp_status / _notify_nodes_mcp_reload +# --------------------------------------------------------------------------- + + +def _fake_request(*nodes: dict[str, Any], proxy_client: Any = None) -> MagicMock: + """Build a minimal mock request with collector and proxy_client.""" + collector = MagicMock() + collector.get_nodes.return_value = (list(nodes), len(nodes)) + req = MagicMock() + req.app.state.collector = collector + req.app.state.proxy_client = proxy_client or AsyncMock() + req.app.state.proxy_token_mgr = None + req.app.state.proxy_auth_token = "tok" + return req + + +def _mock_resp(status_code: int = 200, json_data: Any = None) -> MagicMock: + """Build a mock httpx response (sync .json(), like the real thing).""" + resp = MagicMock() + resp.status_code = status_code + resp.json.return_value = json_data or {} + return resp + + +class TestCollectMcpStatus: + @pytest.mark.anyio + async def test_returns_servers_on_200(self): + resp = _mock_resp(200, {"servers": {"s1": {"status": "ok"}}}) + client = AsyncMock() + client.get.return_value = resp + req = _fake_request( + {"node_id": "n1", "server_url": "http://n1:8000"}, + proxy_client=client, + ) + result = await _collect_mcp_status(req) + assert result == {"n1": {"s1": {"status": "ok"}}} + + @pytest.mark.anyio + async def test_skips_non_200(self): + client = AsyncMock() + client.get.return_value = _mock_resp(503) + req = _fake_request( + {"node_id": "n1", "server_url": "http://n1:8000"}, + proxy_client=client, + ) + result = await _collect_mcp_status(req) + assert result == {} + + @pytest.mark.anyio + async def test_skips_nodes_without_url(self): + client = AsyncMock() + req = _fake_request( + {"node_id": "n1", "server_url": ""}, + {"node_id": "n2"}, + proxy_client=client, + ) + result = await _collect_mcp_status(req) + assert result == {} + client.get.assert_not_called() + + @pytest.mark.anyio + async def test_handles_exception(self): + client = AsyncMock() + client.get.side_effect = ConnectionError("refused") + req = _fake_request( + {"node_id": "n1", "server_url": "http://n1:8000"}, + proxy_client=client, + ) + result = await _collect_mcp_status(req) + assert result == {} + + @pytest.mark.anyio + async def test_empty_cluster(self): + req = _fake_request() + result = await _collect_mcp_status(req) + assert result == {} + + @pytest.mark.anyio + async def test_multiple_nodes_mixed(self): + ok_resp = _mock_resp(200, {"servers": {"s1": {"status": "ok"}}}) + err_resp = _mock_resp(500) + + client = AsyncMock() + client.get.side_effect = [ok_resp, ConnectionError("down"), err_resp] + req = _fake_request( + {"node_id": "n1", "server_url": "http://n1:8000"}, + {"node_id": "n2", "server_url": "http://n2:8000"}, + {"node_id": "n3", "server_url": "http://n3:8000"}, + proxy_client=client, + ) + result = await _collect_mcp_status(req) + assert result == {"n1": {"s1": {"status": "ok"}}} + + +class TestNotifyNodesMcpReload: + @pytest.mark.anyio + async def test_returns_json_on_success(self): + client = AsyncMock() + client.post.return_value = _mock_resp(200, {"reloaded": 3}) + req = _fake_request( + {"node_id": "n1", "server_url": "http://n1:8000"}, + proxy_client=client, + ) + result = await _notify_nodes_mcp_reload(req) + assert result == {"n1": {"reloaded": 3}} + + @pytest.mark.anyio + async def test_skips_nodes_without_url(self): + client = AsyncMock() + req = _fake_request( + {"node_id": "n1", "server_url": ""}, + proxy_client=client, + ) + result = await _notify_nodes_mcp_reload(req) + assert result == {} + client.post.assert_not_called() + + @pytest.mark.anyio + async def test_records_error_on_exception(self): + client = AsyncMock() + client.post.side_effect = ConnectionError("refused") + req = _fake_request( + {"node_id": "n1", "server_url": "http://n1:8000"}, + proxy_client=client, + ) + result = await _notify_nodes_mcp_reload(req) + assert "n1" in result + assert "error" in result["n1"] + assert "refused" in result["n1"]["error"] + + @pytest.mark.anyio + async def test_empty_cluster(self): + req = _fake_request() + result = await _notify_nodes_mcp_reload(req) + assert result == {} + + @pytest.mark.anyio + async def test_multiple_nodes_mixed(self): + client = AsyncMock() + client.post.side_effect = [ + _mock_resp(200, {"reloaded": 2}), + TimeoutError("timeout"), + ] + req = _fake_request( + {"node_id": "n1", "server_url": "http://n1:8000"}, + {"node_id": "n2", "server_url": "http://n2:8000"}, + proxy_client=client, + ) + result = await _notify_nodes_mcp_reload(req) + assert result["n1"] == {"reloaded": 2} + assert "error" in result["n2"] diff --git a/turnstone/console/server.py b/turnstone/console/server.py index 0988558d..0a5b63bf 100644 --- a/turnstone/console/server.py +++ b/turnstone/console/server.py @@ -1492,23 +1492,29 @@ async def admin_list_watches(request: Request) -> JSONResponse: nodes, _ = collector.get_nodes(limit=500) client: httpx.AsyncClient = request.app.state.proxy_client headers = _proxy_auth_headers(request) + sem = asyncio.Semaphore(_NODE_FAN_OUT_LIMIT) async def _fetch_node(node: dict[str, Any]) -> list[dict[str, Any]]: server_url = (node.get("server_url") or "").rstrip("/") if not server_url: return [] - try: - resp = await client.get(f"{server_url}/v1/api/watches", headers=headers) - if resp.status_code == 200: - data = resp.json() - watches: list[dict[str, Any]] = data.get("watches", []) - # Tag each watch with node_id in case the server omits it - for w in watches: - if not w.get("node_id"): - w["node_id"] = node["node_id"] - return watches - except Exception: - log.debug("Failed to fetch watches from node %s", node.get("node_id")) + async with sem: + try: + resp = await client.get(f"{server_url}/v1/api/watches", headers=headers) + if resp.status_code == 200: + data = resp.json() + watches: list[dict[str, Any]] = data.get("watches", []) + # Tag each watch with node_id in case the server omits it + for w in watches: + if not w.get("node_id"): + w["node_id"] = node["node_id"] + return watches + except Exception: + log.debug( + "Failed to fetch watches from node %s", + node.get("node_id"), + exc_info=True, + ) return [] tasks = [_fetch_node(n) for n in nodes] @@ -1524,6 +1530,11 @@ async def admin_list_watches(request: Request) -> JSONResponse: _VALID_WATCH_ID = re.compile(r"^[a-fA-F0-9]+$") +# Max concurrent outbound requests when fanning out to cluster nodes. +# Sized below the default httpx pool limit (100) to leave headroom for +# other proxy traffic (UI proxying, SSE streams, etc.). +_NODE_FAN_OUT_LIMIT = 50 + async def admin_cancel_watch(request: Request) -> Response: """POST /v1/api/admin/watches/{watch_id}/cancel — proxy cancel to the owning node.""" @@ -3140,24 +3151,30 @@ async def _collect_mcp_status( """Query all nodes for MCP status. Returns {node_id: {server_name: status}}.""" collector: ClusterCollector = request.app.state.collector nodes, _ = collector.get_nodes(sort_by="activity", limit=1000, offset=0) - result: dict[str, dict[str, dict[str, Any]]] = {} - for node in nodes: + client: httpx.AsyncClient = request.app.state.proxy_client + headers = _proxy_auth_headers(request) + sem = asyncio.Semaphore(_NODE_FAN_OUT_LIMIT) + + async def _fetch(node: dict[str, Any]) -> tuple[str, dict[str, dict[str, Any]] | None]: node_id = node.get("node_id", "") url = node.get("server_url", "") if not url: - continue - try: - headers = _proxy_auth_headers(request) - async with httpx.AsyncClient(timeout=httpx.Timeout(10)) as client: + return node_id, None + async with sem: + try: resp = await client.get( f"{url.rstrip('/')}/v1/api/_internal/mcp-status", headers=headers, + timeout=10, ) if resp.status_code == 200: - result[node_id] = resp.json().get("servers", {}) - except Exception: - pass - return result + return node_id, resp.json().get("servers", {}) + except Exception: + log.debug("Failed to fetch MCP status from node %s", node_id, exc_info=True) + return node_id, None + + results = await asyncio.gather(*[_fetch(n) for n in nodes]) + return {nid: servers for nid, servers in results if servers is not None} async def admin_list_mcp_servers(request: Request) -> JSONResponse: @@ -3485,25 +3502,29 @@ async def _notify_nodes_mcp_reload(request: Request) -> dict[str, Any]: """Tell all nodes to re-read the mcp_servers DB table and reconcile.""" collector: ClusterCollector = request.app.state.collector nodes, _ = collector.get_nodes(sort_by="activity", limit=1000, offset=0) - results: dict[str, Any] = {} + client: httpx.AsyncClient = request.app.state.proxy_client + headers = _proxy_auth_headers(request) + sem = asyncio.Semaphore(_NODE_FAN_OUT_LIMIT) - for node in nodes: + async def _notify(node: dict[str, Any]) -> tuple[str, Any]: node_id = node.get("node_id", "") url = node.get("server_url", "") if not url: - continue - try: - headers = _proxy_auth_headers(request) - async with httpx.AsyncClient(timeout=httpx.Timeout(30)) as client: + return node_id, None + async with sem: + try: resp = await client.post( f"{url.rstrip('/')}/v1/api/_internal/mcp-reload", headers=headers, + timeout=30, ) - results[node_id] = resp.json() - except Exception as exc: - results[node_id] = {"error": str(exc)} + return node_id, resp.json() + except Exception as exc: + log.debug("Failed to notify node %s for MCP reload", node_id, exc_info=True) + return node_id, {"error": str(exc)} - return results + results = await asyncio.gather(*[_notify(n) for n in nodes]) + return {nid: data for nid, data in results if data is not None} async def admin_mcp_reload(request: Request) -> JSONResponse: