server, requestutil, resource_manager: reduce request body buffering

pkg/utils/requestutil/request_info.go

@@ -15,11 +15,11 @@
 package requestutil
 
 import (
-	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/tikv/pd/pkg/utils/apiutil"
@@ -47,6 +47,14 @@ func (info *RequestInfo) String() string {
 
 // GetRequestInfo returns request info needed from http.Request
 func GetRequestInfo(r *http.Request) RequestInfo {
+	info := GetRequestInfoWithoutBody(r)
+	info.CaptureBody(r)
+	return info
+}
+
+// GetRequestInfoWithoutBody returns request info without consuming or buffering
+// the request body.
+func GetRequestInfoWithoutBody(r *http.Request) RequestInfo {
 	ip, port := apiutil.GetIPPortFromHTTPRequest(r)
 	return RequestInfo{
 		ServiceLabel:   apiutil.GetRouteName(r),
@@ -55,11 +63,16 @@ func GetRequestInfo(r *http.Request) RequestInfo {
 		IP:             ip,
 		Port:           port,
 		URLParam:       getURLParam(r),
-		BodyParam:      getBodyParam(r),
 		StartTimeStamp: time.Now().Unix(),
 	}
 }
 
+// CaptureBody consumes the request body into BodyParam and restores an
+// equivalent body for the handler.
+func (info *RequestInfo) CaptureBody(r *http.Request) {
+	info.BodyParam = getBodyParam(r)
+}
+
 func getURLParam(r *http.Request) string {
 	buf, err := json.Marshal(r.URL.Query())
 	if err != nil {
@@ -72,10 +85,12 @@ func getBodyParam(r *http.Request) string {
 	if r.Body == nil {
 		return ""
 	}
-	// http request body is a io.Reader between bytes.Reader and strings.Reader, it only has EOF error
 	buf, _ := io.ReadAll(r.Body)
 	r.Body.Close()
+	// Restore the body from BodyParam so both views share one long-lived backing
+	// store. Restoring from buf would retain a second complete copy until the
+	// handler finishes.
 	bodyParam := string(buf)
-	r.Body = io.NopCloser(bytes.NewBuffer(buf))
+	r.Body = io.NopCloser(strings.NewReader(bodyParam))
 	return bodyParam
 }

pkg/utils/requestutil/request_info_test.go

@@ -0,0 +1,73 @@
+// Copyright 2026 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// 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.
+
+package requestutil
+
+import (
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type trackingReadCloser struct {
+	io.Reader
+	read  bool
+	close bool
+}
+
+func (r *trackingReadCloser) Read(p []byte) (int, error) {
+	r.read = true
+	return r.Reader.Read(p)
+}
+
+func (r *trackingReadCloser) Close() error {
+	r.close = true
+	return nil
+}
+
+func TestGetRequestInfoWithoutBodyDoesNotConsumeBody(t *testing.T) {
+	re := require.New(t)
+	body := &trackingReadCloser{Reader: strings.NewReader("request-body")}
+	req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1/test", body)
+	re.NoError(err)
+
+	info := GetRequestInfoWithoutBody(req)
+	re.Empty(info.BodyParam)
+	re.False(body.read)
+	re.False(body.close)
+
+	data, err := io.ReadAll(req.Body)
+	re.NoError(err)
+	re.Equal("request-body", string(data))
+}
+
+func TestGetRequestInfoRestoresExactBody(t *testing.T) {
+	re := require.New(t)
+	const requestBody = "request-body-\x00-\xff"
+	body := &trackingReadCloser{Reader: strings.NewReader(requestBody)}
+	req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1/test", body)
+	re.NoError(err)
+
+	info := GetRequestInfo(req)
+	re.Equal(requestBody, info.BodyParam)
+	re.True(body.read)
+	re.True(body.close)
+
+	data, err := io.ReadAll(req.Body)
+	re.NoError(err)
+	re.Equal([]byte(requestBody), data)
+}

server/api/middleware.go

@@ -17,6 +17,7 @@ package api
 import (
 	"context"
 	"net/http"
+	"slices"
 	"time"
 
 	"github.com/unrolled/render"
@@ -59,12 +60,15 @@ func newRequestInfoMiddleware(s *server.Server) negroni.Handler {
 }
 
 func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
-	if !rm.svr.GetServiceMiddlewarePersistOptions().IsAuditEnabled() && !rm.svr.GetServiceMiddlewarePersistOptions().IsRateLimitEnabled() {
+	auditEnabled := rm.svr.GetServiceMiddlewarePersistOptions().IsAuditEnabled()
+	if !auditEnabled && !rm.svr.GetServiceMiddlewarePersistOptions().IsRateLimitEnabled() {
 		next(w, r)
 		return
 	}
 
-	requestInfo := requestutil.GetRequestInfo(r)
+	requestInfo := requestutil.GetRequestInfoWithoutBody(r)
+	labels := rm.svr.GetServiceAuditBackendLabels(requestInfo.ServiceLabel)
+	captureRequestBodyForAudit(r, &requestInfo, auditEnabled, labels)
 	r = r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo))
 
 	failpoint.Inject("addRequestInfoMiddleware", func() {
@@ -79,6 +83,21 @@ func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Reques
 	next(w, r)
 }
 
+func captureRequestBodyForAudit(
+	r *http.Request,
+	requestInfo *requestutil.RequestInfo,
+	auditEnabled bool,
+	labels *audit.BackendLabels,
+) {
+	if auditEnabled && auditNeedsRequestBody(labels) {
+		requestInfo.CaptureBody(r)
+	}
+}
+
+func auditNeedsRequestBody(labels *audit.BackendLabels) bool {
+	return labels != nil && slices.Contains(labels.Labels, audit.LocalLogLabel)
+}
+
 type clusterMiddleware struct {
 	s                         *server.Server
 	rd                        *render.Render
@@ -184,15 +203,20 @@ func (s *auditMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next
 	}
 
 	requestInfo, ok := requestutil.RequestInfoFrom(r.Context())
+	serviceLabel := requestInfo.ServiceLabel
 	if !ok {
-		requestInfo = requestutil.GetRequestInfo(r)
+		serviceLabel = apiutil.GetRouteName(r)
 	}
-
-	labels := s.svr.GetServiceAuditBackendLabels(requestInfo.ServiceLabel)
+	labels := s.svr.GetServiceAuditBackendLabels(serviceLabel)
 	if labels == nil {
 		next(w, r)
 		return
 	}
+	if !ok {
+		requestInfo = requestutil.GetRequestInfoWithoutBody(r)
+		captureRequestBodyForAudit(r, &requestInfo, true, labels)
+		r = r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo))
+	}
 
 	beforeNextBackends := make([]audit.Backend, 0)
 	afterNextBackends := make([]audit.Backend, 0)
@@ -234,7 +258,7 @@ func (s *rateLimitMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request,
 	}
 	requestInfo, ok := requestutil.RequestInfoFrom(r.Context())
 	if !ok {
-		requestInfo = requestutil.GetRequestInfo(r)
+		requestInfo = requestutil.GetRequestInfoWithoutBody(r)
 	}
 
 	// There is no need to check whether rateLimiter is nil. CreateServer ensures that it is created

server/api/middleware_test.go

@@ -0,0 +1,77 @@
+// Copyright 2026 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// 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.
+
+package api
+
+import (
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/tikv/pd/pkg/audit"
+	"github.com/tikv/pd/pkg/utils/requestutil"
+)
+
+func TestCaptureRequestBodyForAuditReadsBodyOnlyForLocalLog(t *testing.T) {
+	testCases := []struct {
+		name         string
+		auditEnabled bool
+		labels       *audit.BackendLabels
+		expectBody   bool
+	}{
+		{
+			name:         "audit-disabled",
+			auditEnabled: false,
+			labels:       &audit.BackendLabels{Labels: []string{audit.LocalLogLabel}},
+		},
+		{
+			name:         "prometheus-only",
+			auditEnabled: true,
+			labels:       &audit.BackendLabels{Labels: []string{audit.PrometheusHistogram}},
+		},
+		{
+			name:         "no-audit-backend",
+			auditEnabled: true,
+		},
+		{
+			name:         "local-log",
+			auditEnabled: true,
+			labels:       &audit.BackendLabels{Labels: []string{audit.LocalLogLabel, audit.PrometheusHistogram}},
+			expectBody:   true,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			re := require.New(t)
+			req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1/test", strings.NewReader("request-body"))
+			re.NoError(err)
+
+			info := requestutil.GetRequestInfoWithoutBody(req)
+			captureRequestBodyForAudit(req, &info, testCase.auditEnabled, testCase.labels)
+			if testCase.expectBody {
+				re.Equal("request-body", info.BodyParam)
+			} else {
+				re.Empty(info.BodyParam)
+			}
+
+			data, err := io.ReadAll(req.Body)
+			re.NoError(err)
+			re.Equal("request-body", string(data))
+		})
+	}
+}