perf(mcp-proxy): reduce response body allocations (#3174)

* perf(mcp-proxy): reduce response body allocations

* perf(mcp-proxy): preallocate responses up to 10 MiB

Author: wklken

src/mcp-proxy/pkg/infra/proxy/proxy.go

@@ -20,6 +20,7 @@
 package proxy
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"encoding/json"
@@ -56,6 +57,9 @@ import (
 	"mcp_proxy/pkg/util"
 )
 
+// maxResponseBodyPreallocateSize limits only the initial allocation; response bodies are always read to EOF.
+const maxResponseBodyPreallocateSize = 10 << 20
+
 // sharedTransport 是所有 tool call 共用的 HTTP Transport，避免每次调用创建新连接池。
 // 通过 InitSharedTransport 从配置初始化，参数可在 config.yaml 的 mcpServer.transport 段调整。
 var (
@@ -1123,7 +1127,10 @@ func genToolHandler(toolApiConfig *ToolConfig, serverName string, rawResponseEna
 					var bodyBytes []byte
 					if response.Body() != nil {
 						var e error
-						bodyBytes, e = io.ReadAll(response.Body())
+						bodyBytes, e = readResponseBody(
+							response.Body(),
+							response.GetHeader("Content-Length"),
+						)
 						if e != nil {
 							return nil, e
 						}
@@ -1327,6 +1334,17 @@ func recordToolCallMetrics(
 	}
 }
 
+func readResponseBody(body io.Reader, contentLength string) ([]byte, error) {
+	size, err := strconv.ParseInt(contentLength, 10, 64)
+	if err != nil || size <= 0 || size > maxResponseBodyPreallocateSize {
+		return io.ReadAll(body)
+	}
+
+	buffer := bytes.NewBuffer(make([]byte, 0, int(size)+bytes.MinRead))
+	_, err = buffer.ReadFrom(body)
+	return buffer.Bytes(), err
+}
+
 // logToolCall logs MCP protocol-level request/response for tools/call.
 // This is needed because tools/call handler is invoked inside callTool(),
 // which does not go through the MCP middleware chain.

src/mcp-proxy/pkg/infra/proxy/proxy_benchmark_test.go

@@ -58,7 +58,7 @@ var (
 func BenchmarkGenToolHandlerLargeJSONResponse(b *testing.B) {
 	initBenchmarkRuntime(b)
 
-	for _, size := range []int{64 << 10, 1 << 20} {
+	for _, size := range []int{64 << 10, 1 << 20, 10 << 20} {
 		size := size
 		b.Run(strconv.Itoa(size)+"B", func(b *testing.B) {
 			responseBody := buildBenchmarkJSONBody(size)
@@ -286,10 +286,10 @@ func BenchmarkMCPToolResultWireEncode(b *testing.B) {
 	}
 }
 
-// BenchmarkReadLargeResponseBody compares the current unhinted io.ReadAll path with
-// the allocation headroom available when an upstream Content-Length can be trusted.
+// BenchmarkReadLargeResponseBody compares the previous unhinted io.ReadAll path with
+// the production response reader when an upstream Content-Length is available.
 func BenchmarkReadLargeResponseBody(b *testing.B) {
-	for _, size := range []int{64 << 10, 1 << 20} {
+	for _, size := range []int{64 << 10, 1 << 20, 10 << 20} {
 		size := size
 		b.Run(strconv.Itoa(size)+"B", func(b *testing.B) {
 			body := buildBenchmarkJSONBody(size)
@@ -307,16 +307,16 @@ func BenchmarkReadLargeResponseBody(b *testing.B) {
 				}
 			})
 
-			b.Run("content-length-preallocated", func(b *testing.B) {
+			b.Run("production-content-length", func(b *testing.B) {
 				b.ReportAllocs()
 				b.SetBytes(int64(len(body)))
 				for i := 0; i < b.N; i++ {
 					reader := benchmarkChunkReader{body: body}
-					buffer := bytes.NewBuffer(make([]byte, 0, len(body)+bytes.MinRead))
-					if _, err := buffer.ReadFrom(&reader); err != nil {
-						b.Fatalf("read preallocated response body: %v", err)
+					readBody, err := readResponseBody(&reader, strconv.Itoa(len(body)))
+					if err != nil {
+						b.Fatalf("read response body: %v", err)
 					}
-					benchmarkResponseBody = buffer.Bytes()
+					benchmarkResponseBody = readBody
 				}
 			})
 		})

src/mcp-proxy/pkg/infra/proxy/response_body_test.go

@@ -0,0 +1,82 @@
+/*
+ * TencentBlueKing is pleased to support the open source community by making
+ * 蓝鲸智云 - API 网关(BlueKing - APIGateway) available.
+ * Copyright (C) Tencent. All rights reserved.
+ * Licensed under the MIT License (the "License"); you may not use this file except
+ * in compliance with the License. You may obtain a copy of the License at
+ *
+ *     http://opensource.org/licenses/MIT
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under
+ * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * We undertake not to change the open source license (MIT license) applicable
+ * to the current version of the project delivered to anyone in the future.
+ */
+
+package proxy
+
+import (
+	"errors"
+	"io"
+	"strconv"
+	"strings"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("readResponseBody", func() {
+	DescribeTable("reads the complete response body regardless of the Content-Length hint",
+		func(contentLength, body string) {
+			result, err := readResponseBody(strings.NewReader(body), contentLength)
+
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal([]byte(body)))
+		},
+		Entry("with an exact length", "13", "complete-body"),
+		Entry("without a length", "", "complete-body"),
+		Entry("with an invalid length", "not-a-number", "complete-body"),
+		Entry("with a zero length", "0", "complete-body"),
+		Entry("with a negative length", "-1", "complete-body"),
+		Entry("with a length above the preallocation limit",
+			strconv.Itoa(maxResponseBodyPreallocateSize+1), "complete-body"),
+		Entry("with an underestimated length", "3", "complete-body"),
+		Entry("with an overestimated length", "1024", "complete-body"),
+	)
+
+	It("returns bytes read before a reader error and preserves the error", func() {
+		readErr := errors.New("response read failed")
+		reader := &errorResponseReader{
+			body: []byte("partial-body"),
+			err:  readErr,
+		}
+
+		result, err := readResponseBody(reader, strconv.Itoa(len(reader.body)))
+
+		Expect(result).To(Equal([]byte("partial-body")))
+		Expect(err).To(MatchError(readErr))
+	})
+})
+
+type errorResponseReader struct {
+	body []byte
+	err  error
+}
+
+func (r *errorResponseReader) Read(p []byte) (int, error) {
+	if len(r.body) == 0 {
+		return 0, r.err
+	}
+
+	n := copy(p, r.body)
+	r.body = r.body[n:]
+	if len(r.body) == 0 {
+		return n, r.err
+	}
+	return n, nil
+}
+
+var _ io.Reader = (*errorResponseReader)(nil)