server, requestutil, resource_manager: reduce request body buffering

pkg/mcs/resourcemanager/metadataapi/config_service.go

@@ -181,13 +181,19 @@ func decodeResourceGroup(body io.Reader, group *rmpb.ResourceGroup) error {
 	if err != nil {
 		return err
 	}
-	legacyJSON, rawKeyspaceID, err := splitResourceGroupJSON(data)
-	if err != nil {
-		return err
+	var keyspaceID resourceGroupKeyspaceIDJSON
+	legacyGroup := resourceGroupJSON{
+		resourceGroupWithoutKeyspaceID: (*resourceGroupWithoutKeyspaceID)(group),
+		KeyspaceID:                     resourceGroupKeyspaceIDField{value: &keyspaceID},
+		KeyspaceIDCamel:                resourceGroupKeyspaceIDField{value: &keyspaceID},
 	}
 	// Keep the legacy encoding/json behavior for all existing ResourceGroup
 	// fields. In particular, it matches JSON field names case-insensitively.
-	if err := json.Unmarshal(legacyJSON, group); err != nil {
+	if err := json.Unmarshal(data, &legacyGroup); err != nil {
+		var keyspaceIDErr *resourceGroupKeyspaceIDError
+		if errors.As(err, &keyspaceIDErr) {
+			return keyspaceIDErr
+		}
 		// The updated ResourceGroup contains a protobuf oneof, so clients may
 		// serialize the whole message as protobuf JSON. Retry strictly to
 		// accept enum names and quoted 64-bit integers without silently
@@ -196,67 +202,60 @@ func decodeResourceGroup(body io.Reader, group *rmpb.ResourceGroup) error {
 		if protoErr := (&jsonpb.Unmarshaler{}).Unmarshal(bytes.NewReader(data), group); protoErr != nil {
 			return fmt.Errorf("invalid resource group JSON: legacy JSON: %v; protobuf JSON: %w", err, protoErr)
 		}
-		return validateResourceGroupKeyspaceID(group, rawKeyspaceID)
+		return validateResourceGroupKeyspaceID(group, keyspaceID.raw)
 	}
-	if rawKeyspaceID != nil {
-		keyspaceID, err := decodeKeyspaceIDJSON(rawKeyspaceID)
+	if keyspaceID.raw != nil {
+		decodedKeyspaceID, err := decodeKeyspaceIDJSON(keyspaceID.raw)
 		if err != nil {
 			return err
 		}
-		group.KeyspaceId = keyspaceID
+		group.KeyspaceId = decodedKeyspaceID
 	}
-	return validateResourceGroupKeyspaceID(group, rawKeyspaceID)
+	return validateResourceGroupKeyspaceID(group, keyspaceID.raw)
 }
 
-func splitResourceGroupJSON(data []byte) ([]byte, json.RawMessage, error) {
-	if isJSONNull(data) {
-		return data, nil, nil
+// resourceGroupWithoutKeyspaceID lets encoding/json decode every legacy field
+// directly. The explicit fields in resourceGroupJSON shadow KeyspaceId, whose
+// protobuf oneof still requires compatibility handling.
+type resourceGroupWithoutKeyspaceID rmpb.ResourceGroup
+
+type resourceGroupJSON struct {
+	*resourceGroupWithoutKeyspaceID
+	KeyspaceID      resourceGroupKeyspaceIDField `json:"keyspace_id"`
+	KeyspaceIDCamel resourceGroupKeyspaceIDField `json:"keyspaceId"`
+}
+
+type resourceGroupKeyspaceIDJSON struct {
+	raw json.RawMessage
+}
+
+type resourceGroupKeyspaceIDField struct {
+	value *resourceGroupKeyspaceIDJSON
+}
+
+// UnmarshalJSON records and normalizes a ResourceGroup keyspace ID without
+// making a second copy of the rest of the request body.
+func (f *resourceGroupKeyspaceIDField) UnmarshalJSON(data []byte) error {
+	if f.value.raw != nil {
+		return &resourceGroupKeyspaceIDError{errors.New("keyspace_id must be set only once")}
 	}
-	// KeyspaceIDValue became a protobuf oneof, which encoding/json cannot decode.
-	// Remove it from the legacy payload and decode it separately with jsonpb.
-	fields, err := decodeJSONObjectFields(data)
+	raw, err := normalizeKeyspaceIDJSON(data)
 	if err != nil {
-		return nil, nil, err
-	}
-	legacyFields := make([]jsonObjectField, 0, len(fields))
-	var rawKeyspaceID json.RawMessage
-	for _, field := range fields {
-		if !isKeyspaceIDJSONField(field.name) {
-			legacyFields = append(legacyFields, field)
-			continue
-		}
-		if rawKeyspaceID != nil {
-			return nil, nil, errors.New("keyspace_id must be set only once")
-		}
-		rawKeyspaceID, err = normalizeKeyspaceIDJSON(field.value)
-		if err != nil {
-			return nil, nil, err
-		}
+		return &resourceGroupKeyspaceIDError{err}
 	}
-	return marshalJSONObjectFields(legacyFields), rawKeyspaceID, nil
+	f.value.raw = raw
+	return nil
+}
+
+type resourceGroupKeyspaceIDError struct {
+	error
 }
 
 type jsonObjectField struct {
 	name  string
 	value json.RawMessage
 }
 
-func marshalJSONObjectFields(fields []jsonObjectField) []byte {
-	var buffer bytes.Buffer
-	buffer.WriteByte('{')
-	for i, field := range fields {
-		if i > 0 {
-			buffer.WriteByte(',')
-		}
-		name, _ := json.Marshal(field.name)
-		buffer.Write(name)
-		buffer.WriteByte(':')
-		buffer.Write(field.value)
-	}
-	buffer.WriteByte('}')
-	return buffer.Bytes()
-}
-
 func decodeJSONObjectFields(data []byte) ([]jsonObjectField, error) {
 	var object map[string]json.RawMessage
 	if err := json.Unmarshal(data, &object); err != nil {
@@ -293,10 +292,6 @@ func decodeJSONObjectFields(data []byte) ([]jsonObjectField, error) {
 	return fields, nil
 }
 
-func isKeyspaceIDJSONField(name string) bool {
-	return strings.EqualFold(name, "keyspace_id") || strings.EqualFold(name, "keyspaceId")
-}
-
 func isJSONNull(data []byte) bool {
 	return bytes.Equal(bytes.TrimSpace(data), []byte("null"))
 }

pkg/mcs/resourcemanager/metadataapi/config_service_test.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/gin-gonic/gin"
@@ -162,6 +163,22 @@ func TestConfigServiceGroupCRUDAndErrorCodes(t *testing.T) {
 		re.Equal(int64(987), storedGroup.RUSettings.RU.Settings.BurstLimit)
 	}
 
+	largeJobType := strings.Repeat("x", 1<<20)
+	largeResourceGroupBody, err := json.Marshal(&rmpb.ResourceGroup{
+		Name: "large_resource_group",
+		BackgroundSettings: &rmpb.BackgroundSettings{
+			JobTypes: []string{largeJobType},
+		},
+	})
+	re.NoError(err)
+	re.Greater(len(largeResourceGroupBody), 1<<20)
+	for _, method := range []string{http.MethodPost, http.MethodPut} {
+		resp = doRawResourceGroupRequest(handler, method, largeResourceGroupBody)
+		re.Equal(http.StatusOK, resp.Code, resp.Body.String())
+		largeResourceGroup := store.groups[groupKey(constant.NullKeyspaceID, "large_resource_group")]
+		re.Equal(largeJobType, largeResourceGroup.Background.JobTypes[0])
+	}
+
 	store.addErr = errors.New("add failed")
 	resp = doJSONRequest(re, handler, http.MethodPost, "/resource-manager/api/v1/config/group", group)
 	re.Equal(http.StatusInternalServerError, resp.Code)

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,13 @@ 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() {
+	if !rm.svr.GetServiceMiddlewarePersistOptions().IsAuditEnabled() &&
+		!rm.svr.GetServiceMiddlewarePersistOptions().IsRateLimitEnabled() {
 		next(w, r)
 		return
 	}
 
-	requestInfo := requestutil.GetRequestInfo(r)
+	requestInfo := requestutil.GetRequestInfoWithoutBody(r)
 	r = r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo))
 
 	failpoint.Inject("addRequestInfoMiddleware", func() {
@@ -79,6 +81,38 @@ func (rm *requestInfoMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Reques
 	next(w, r)
 }
 
+func captureRequestBodyForAudit(
+	r *http.Request,
+	requestInfo *requestutil.RequestInfo,
+	labels *audit.BackendLabels,
+) bool {
+	if requestInfo.BodyParam != "" || !auditNeedsRequestBody(labels) {
+		return false
+	}
+	requestInfo.CaptureBody(r)
+	return true
+}
+
+func auditNeedsRequestBody(labels *audit.BackendLabels) bool {
+	return labels != nil && slices.Contains(labels.Labels, audit.LocalLogLabel)
+}
+
+func prepareRequestForAudit(
+	r *http.Request,
+	requestInfo requestutil.RequestInfo,
+	hasRequestInfo bool,
+	labels *audit.BackendLabels,
+) *http.Request {
+	contextNeedsUpdate := !hasRequestInfo
+	if captureRequestBodyForAudit(r, &requestInfo, labels) {
+		contextNeedsUpdate = true
+	}
+	if contextNeedsUpdate {
+		return r.WithContext(requestutil.WithRequestInfo(r.Context(), requestInfo))
+	}
+	return r
+}
+
 type clusterMiddleware struct {
 	s                         *server.Server
 	rd                        *render.Render
@@ -185,14 +219,14 @@ func (s *auditMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request, next
 
 	requestInfo, ok := requestutil.RequestInfoFrom(r.Context())
 	if !ok {
-		requestInfo = requestutil.GetRequestInfo(r)
+		requestInfo = requestutil.GetRequestInfoWithoutBody(r)
 	}
-
 	labels := s.svr.GetServiceAuditBackendLabels(requestInfo.ServiceLabel)
 	if labels == nil {
 		next(w, r)
 		return
 	}
+	r = prepareRequestForAudit(r, requestInfo, ok, labels)
 
 	beforeNextBackends := make([]audit.Backend, 0)
 	afterNextBackends := make([]audit.Backend, 0)
@@ -234,7 +268,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