From 53669437d58be705716d733c7a5030af6d382d0d Mon Sep 17 00:00:00 2001 From: Blizzard Date: Wed, 24 Jun 2026 15:40:50 +0800 Subject: [PATCH] =?UTF-8?q?perf(tools):=20=E5=B7=A5=E5=85=B7=E6=9C=8D?= =?UTF-8?q?=E5=8A=A1=E5=8D=95=E5=AE=9E=E4=BE=8B=E5=B9=B6=E5=8F=91=EF=BC=88?= =?UTF-8?q?ServeTool=20=E5=8D=8F=E7=A8=8B=E5=8C=96=EF=BC=89+=20mcp-py=20?= =?UTF-8?q?=E5=90=8C=E6=94=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 垂直扩：单个 mcp 实例内每个工具调用分发到独立 goroutine/task 并发处理， 不再受 NATS 单订阅回调串行所限；配合队列组多副本即「单实例并发 × 副本数」横向扩。 - bus.ServeTool：回调立即起 goroutine 并返回（不阻塞 NATS 投递），信号量限并发 （MCP_TOOL_CONCURRENCY，默认 16），handler panic → 回错误结果避免调用方干等超时。 - mcp-py mcp_gateway：_on_call 改 asyncio.create_task 派发 + Semaphore 限并发（同默认 16）。 - 测试 TestConcurrentToolServe：slow 工具阻塞时 fast 工具仍返回（旧串行下会超时）。 验证：单测通过；mcp-go/mcp-py live 重启工具就绪，工具往返正常（memory/history/kb_search 均正常响应，无 panic）。全模块 build+vet+test 全绿。 注：下游共享后端（Milvus/Neo4j/PG/嵌入端点）是横向扩的最终上限，届时扩这些而非 mcp 实例。 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/sundynix_mcp_py/mcp_gateway.py | 45 ++++++++------ sundynix-shared/bus/bus.go | 59 +++++++++++++------ sundynix-shared/bus/bus_e2e_test.go | 52 ++++++++++++++++ 3 files changed, 118 insertions(+), 38 deletions(-) diff --git a/sundynix-mcp-py/src/sundynix_mcp_py/mcp_gateway.py b/sundynix-mcp-py/src/sundynix_mcp_py/mcp_gateway.py index e4be237..0801dff 100644 --- a/sundynix-mcp-py/src/sundynix_mcp_py/mcp_gateway.py +++ b/sundynix-mcp-py/src/sundynix_mcp_py/mcp_gateway.py @@ -47,6 +47,8 @@ class McpGateway: self.interpreter = CodeInterpreter() # Docker 隔离沙箱 self._nc: nats.NATS | None = None self._tools: dict[str, callable] = {} + # 单实例工具并发上限（垂直扩）；队列组多副本再叠加水平扩。 + self._sem = asyncio.Semaphore(int(os.getenv("MCP_TOOL_CONCURRENCY", "16"))) async def register_tools(self) -> None: """注册工具名 → 处理协程。""" @@ -77,25 +79,30 @@ class McpGateway: await asyncio.Event().wait() async def _on_call(self, msg) -> None: - """解析 ToolCall → 路由到工具 → Respond ToolResult。""" - try: - req = json.loads(msg.data) - except Exception as e: # noqa: BLE001 - await self._reply(msg, ok=False, error=f"bad tool call: {e}") - return - tool = req.get("tool", "") - task_id = req.get("task_id", "") - args = req.get("args") or {} - log.info("[mcp_py] tool=%s task=%s args=%s", tool, task_id, args) - fn = self._tools.get(tool) - if fn is None: - await self._reply(msg, ok=False, error=f"unknown tool: {tool}") - return - try: - content = await fn(args) - await self._reply(msg, ok=True, content=content) - except Exception as e: # noqa: BLE001 - await self._reply(msg, ok=False, error=str(e)) + """立即派发到独立 task 并发处理，不阻塞订阅消息循环（nats-py 单订阅回调本是串行 await）。""" + asyncio.create_task(self._handle(msg)) + + async def _handle(self, msg) -> None: + """解析 ToolCall → 路由到工具 → Respond ToolResult（限并发）。""" + async with self._sem: + try: + req = json.loads(msg.data) + except Exception as e: # noqa: BLE001 + await self._reply(msg, ok=False, error=f"bad tool call: {e}") + return + tool = req.get("tool", "") + task_id = req.get("task_id", "") + args = req.get("args") or {} + log.info("[mcp_py] tool=%s task=%s args=%s", tool, task_id, args) + fn = self._tools.get(tool) + if fn is None: + await self._reply(msg, ok=False, error=f"unknown tool: {tool}") + return + try: + content = await fn(args) + await self._reply(msg, ok=True, content=content) + except Exception as e: # noqa: BLE001 + await self._reply(msg, ok=False, error=str(e)) @staticmethod async def _reply(msg, *, ok: bool, content: str = "", error: str = "") -> None: diff --git a/sundynix-shared/bus/bus.go b/sundynix-shared/bus/bus.go index a33b201..5b0e1ad 100644 --- a/sundynix-shared/bus/bus.go +++ b/sundynix-shared/bus/bus.go @@ -233,27 +233,48 @@ func (b *Bus) CallTool(ctx context.Context, subject string, call *contract.ToolC // ToolHandler 处理一次工具调用并返回结果。 type ToolHandler func(ctx context.Context, call *contract.ToolCall) *contract.ToolResult -// ServeTool 以队列组订阅工具主题（可用通配 sundynix.tools.go.>）， -// 对每个请求调用 h 并 Respond，队列组内多副本自动负载均衡。 -// 返回的 unsub 用于退订。 +// toolConcurrency 返回单实例工具并发处理上限（env MCP_TOOL_CONCURRENCY，默认 16）。 +func toolConcurrency() int { + if v := os.Getenv("MCP_TOOL_CONCURRENCY"); v != "" { + if n, err := strconv.Atoi(v); err == nil && n > 0 { + return n + } + } + return 16 +} + +// ServeTool 以队列组订阅工具主题（可用通配 sundynix.tools.go.>），队列组内多副本自动负载均衡（水平扩）。 +// 单实例内每个请求分发到独立 goroutine 并发处理（垂直扩，上限 MCP_TOOL_CONCURRENCY）—— +// 回调立即返回不阻塞 NATS 投递；信号量约束同时执行的 handler 数。返回的 unsub 用于退订。 func (b *Bus) ServeTool(subject, queue string, h ToolHandler) (unsub func() error, err error) { + sem := make(chan struct{}, toolConcurrency()) sub, err := b.nc.QueueSubscribe(subject, queue, func(m *nats.Msg) { - var call contract.ToolCall - if err := json.Unmarshal(m.Data, &call); err != nil { - respond(m, &contract.ToolResult{OK: false, Error: "bad tool call: " + err.Error()}) - return - } - // 还原上游链路并开服务端 span（成为 dispatcher tool.call span 的子节点）。 - mctx := extractTrace(context.Background(), m.Header) - mctx, span := tracer().Start(mctx, "tool.serve "+call.Tool, - trace.WithSpanKind(trace.SpanKindServer), - trace.WithAttributes(attribute.String("sundynix.tool", call.Tool))) - res := h(mctx, &call) - if res != nil { - span.SetAttributes(attribute.Bool("sundynix.tool.ok", res.OK)) - } - span.End() - respond(m, res) + // 立即起 goroutine 并返回，让 NATS 继续投递下一条（核心 NATS 单订阅回调本是串行）。 + go func() { + sem <- struct{}{} // 限并发：超额则在此短暂排队 + defer func() { + <-sem + if r := recover(); r != nil { // handler panic → 回错误结果，调用方不必干等超时 + respond(m, &contract.ToolResult{OK: false, Error: fmt.Sprintf("tool panic: %v", r)}) + } + }() + var call contract.ToolCall + if err := json.Unmarshal(m.Data, &call); err != nil { + respond(m, &contract.ToolResult{OK: false, Error: "bad tool call: " + err.Error()}) + return + } + // 还原上游链路并开服务端 span（成为 dispatcher tool.call span 的子节点）。 + mctx := extractTrace(context.Background(), m.Header) + mctx, span := tracer().Start(mctx, "tool.serve "+call.Tool, + trace.WithSpanKind(trace.SpanKindServer), + trace.WithAttributes(attribute.String("sundynix.tool", call.Tool))) + defer span.End() + res := h(mctx, &call) + if res != nil { + span.SetAttributes(attribute.Bool("sundynix.tool.ok", res.OK)) + } + respond(m, res) + }() }) if err != nil { return nil, fmt.Errorf("serve tool %s: %w", subject, err) diff --git a/sundynix-shared/bus/bus_e2e_test.go b/sundynix-shared/bus/bus_e2e_test.go index 80cbcbe..c5f6bde 100644 --- a/sundynix-shared/bus/bus_e2e_test.go +++ b/sundynix-shared/bus/bus_e2e_test.go @@ -254,3 +254,55 @@ func TestConcurrentConsume(t *testing.T) { t.Fatal("B 被阻塞的 A 卡住——并发消费未生效（仍是串行）") } } + +// TestConcurrentToolServe 验证单实例工具服务并发：一个慢工具调用不阻塞另一个快调用。 +// 旧的串行 ServeTool 下本测试会超时失败。 +func TestConcurrentToolServe(t *testing.T) { + url := startEmbeddedNATS(t) + + srv, err := bus.Connect(url) + if err != nil { + t.Fatalf("server connect: %v", err) + } + defer srv.Close() + + blockSlow := make(chan struct{}) + unsub, err := srv.ServeTool(contract.SubjectToolsGoAll, contract.QueueToolsGo, + func(_ context.Context, call *contract.ToolCall) *contract.ToolResult { + if call.Tool == "slow" { + <-blockSlow // 模拟慢工具（如 RAG 嵌入） + } + return &contract.ToolResult{OK: true, Content: call.Tool} + }) + if err != nil { + t.Fatalf("serve tool: %v", err) + } + defer unsub() + defer close(blockSlow) + + cli, err := bus.Connect(url) + if err != nil { + t.Fatalf("client connect: %v", err) + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + // 先发 slow（会阻塞在 handler 里），再发 fast——fast 应在 slow 仍阻塞时返回。 + go func() { + _, _ = cli.CallTool(ctx, contract.ToolSubjectGo("slow"), &contract.ToolCall{Tool: "slow"}) + }() + time.Sleep(300 * time.Millisecond) + + fctx, fcancel := context.WithTimeout(ctx, 3*time.Second) + defer fcancel() + res, err := cli.CallTool(fctx, contract.ToolSubjectGo("fast"), &contract.ToolCall{Tool: "fast"}) + if err != nil { + t.Fatalf("fast 工具被慢工具阻塞——单实例并发未生效: %v", err) + } + if res.Content != "fast" { + t.Fatalf("unexpected result: %q", res.Content) + } + t.Log("✓ 单实例工具并发生效：slow 仍阻塞时 fast 已返回") +}