Cherry-pick #9331 to v1.84.x (#9334)

Original PR: #9331

RELEASE NOTES: N/A

Author: arjan-bal

internal/envconfig/envconfig.go

@@ -150,7 +150,7 @@ var (
 	// throttling limit if unforeseen issues arise, and it will be removed in a
 	// future release.
 	//
-	// TODO: Remove this env var once v1.83.0 is release.
+	// TODO: Remove this env var once v1.83.0 is released.
 	ControlBufferThrottleLimit = uint64FromEnv("GRPC_GO_EXPERIMENTAL_CONTROL_BUFFER_THROTTLE_LIMIT", 100, 1, 10000)
 
 	// ALTSMaxFrameSize is the maximum frame size for ALTS in bytes.
@@ -159,6 +159,16 @@ var (
 	// 512KiB to match altsWriteBufferMaxSize in
 	// credentials/alts/internal/conn/record.go).
 	ALTSMaxFrameSize = uint64FromEnv("GRPC_GO_EXPERIMENTAL_ALTS_MAX_FRAME_SIZE", 4096, 4096, 512*1024)
+
+	// EnableReceiveBufferCompaction enables the compaction of data buffers
+	// to reduce the number of buffers in the receive buffer.
+	//
+	// This environment variable serves as an escape hatch to disable the
+	// feature if unforeseen issues arise, and it will be removed in a future
+	// release.
+	//
+	// TODO: Remove this env var once v1.85.0 is released.
+	EnableReceiveBufferCompaction = boolFromEnv("GRPC_GO_EXPERIMENTAL_ENABLE_RECEIVE_BUFFER_COMPACTION", true)
 )
 
 func boolFromEnv(envVar string, def bool) bool {

internal/mem/buffer_pool.go

@@ -26,12 +26,26 @@ import (
 	"slices"
 	"sort"
 	"sync"
+
+	"google.golang.org/grpc/internal"
 )
 
 const (
 	goPageSize = 4 * 1024 // 4KiB. N.B. this must be a power of 2.
 )
 
+var (
+	// BufferPoolingThreshold is the minimum size of a buffer that can be pooled.
+	// This is used to determine whether to pool buffers or allocate them directly.
+	BufferPoolingThreshold = 1 << 10
+)
+
+func init() {
+	internal.SetBufferPoolingThresholdForTesting = func(threshold int) {
+		BufferPoolingThreshold = threshold
+	}
+}
+
 var uintSize = bits.UintSize // use a variable for mocking during tests.
 
 // bufferPool is a copy of the public bufferPool interface used to avoid

internal/transport/handler_server.go

@@ -424,7 +424,7 @@ func (ht *serverHandlerTransport) HandleStreams(ctx context.Context, startStream
 		st:               ht,
 		headerWireLength: 0, // won't have access to header wire length until golang/go#18997.
 	}
-	s.Stream.buf.init()
+	s.Stream.buf.init(ht.bufferPool)
 	s.readRequester = s
 	s.trReader = transportReader{
 		reader:        recvBufferReader{ctx: s.ctx, ctxDone: s.ctx.Done(), recv: &s.buf},

internal/transport/http2_client.go

@@ -500,7 +500,7 @@ func (t *http2Client) newStream(ctx context.Context, callHdr *CallHdr, handler s
 		headerChan:   make(chan struct{}),
 		statsHandler: handler,
 	}
-	s.Stream.buf.init()
+	s.Stream.buf.init(t.bufferPool)
 	s.Stream.wq.init(defaultWriteQuota, s.done)
 	s.readRequester = s
 	// The client side stream context should have exactly the same life cycle with the user provided context.

internal/transport/http2_server.go

@@ -407,7 +407,7 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 		st:               t,
 		headerWireLength: int(frame.Header().Length),
 	}
-	s.Stream.buf.init()
+	s.Stream.buf.init(t.bufferPool)
 	var (
 		// if false, content-type was missing or invalid
 		isGRPC      = false

internal/transport/transport.go

@@ -30,11 +30,14 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"golang.org/x/net/http2"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/envconfig"
+	imem "google.golang.org/grpc/internal/mem"
 	"google.golang.org/grpc/internal/transport/internal"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/mem"
@@ -45,7 +48,30 @@ import (
 	"google.golang.org/grpc/tap"
 )
 
-const logLevel = 2
+const (
+	logLevel = 2
+	// recvMsgSize estimates the memory overhead of a recvMsg in the backlog.
+	// It accounts for the recvMsg struct itself and the slice header of the
+	// underlying buffer's data.
+	recvMsgSize = int(unsafe.Sizeof(recvMsg{}) + unsafe.Sizeof([]byte{}))
+
+	// utilizationFactor controls when we consider memory utilization acceptable.
+	// When backlogHeapSize / payloadSize <= utilizationFactor (meaning at least
+	// 50% of the heap memory is actual payload data), compaction is skipped.
+	utilizationFactor = 2
+)
+
+var (
+	// compactionThreshold is approx 57KB (on 64-bit systems). It allows
+	// accumulating up to 1024 1-byte payloads before triggering compaction.
+	//
+	// Because individual payloads <= 1024 bytes are allocated on the heap
+	// outside mem.BufferPool, waiting for at least 1024 bytes to accumulate
+	// ensures that compaction coalesces those small heap allocations into a
+	// single large buffer from mem.BufferPool, enabling buffer reuse while
+	// avoiding frequent copying for small bursts of frames.
+	compactionThreshold = imem.BufferPoolingThreshold * (recvMsgSize + 1)
+)
 
 func init() {
 	internal.TimeNowFunc = func() int64 { return time.Now().UnixNano() }
@@ -71,23 +97,31 @@ type recvBuffer struct {
 	c       chan recvMsg
 	mu      sync.Mutex
 	backlog []recvMsg
-	err     error
+	// uncompactedSuffixLen tracks the number of consecutive data messages at
+	// the tail of backlog that have not been compacted.
+	uncompactedSuffixLen int
+	// uncompactedBytes tracks the total payload bytes across the trailing
+	// uncompactedSuffixLen messages.
+	uncompactedBytes int
+	err              error
+	bufPool          mem.BufferPool
 }
 
 // init allows a recvBuffer to be initialized in-place, which is useful
 // for resetting a buffer or for avoiding a heap allocation when the buffer
 // is embedded in another struct.
-func (b *recvBuffer) init() {
+func (b *recvBuffer) init(pool mem.BufferPool) {
 	b.c = make(chan recvMsg, 1)
+	b.bufPool = pool
 }
 
 func (b *recvBuffer) put(r recvMsg) {
 	b.mu.Lock()
+	defer b.mu.Unlock()
 	if b.err != nil {
 		// drop the buffer on the floor. Since b.err is not nil, any subsequent reads
 		// will always return an error, making this buffer inaccessible.
 		r.buffer.Free()
-		b.mu.Unlock()
 		// An error had occurred earlier, don't accept more
 		// data or errors.
 		return
@@ -96,20 +130,84 @@ func (b *recvBuffer) put(r recvMsg) {
 	if len(b.backlog) == 0 {
 		select {
 		case b.c <- r:
-			b.mu.Unlock()
 			return
 		default:
 		}
 	}
 	b.backlog = append(b.backlog, r)
-	b.mu.Unlock()
+	b.compactBacklogLocked(r)
+}
+
+func (b *recvBuffer) compactBacklogLocked(r recvMsg) {
+	if !envconfig.EnableReceiveBufferCompaction {
+		return
+	}
+	if r.buffer == nil {
+		b.uncompactedBytes = 0
+		b.uncompactedSuffixLen = 0
+		return
+	}
+
+	b.uncompactedSuffixLen++
+	b.uncompactedBytes += r.buffer.Len()
+	backlogHeapSize := b.uncompactedSuffixLen*recvMsgSize + b.uncompactedBytes
+
+	// If the memory overhead is less than 50% of the heap usage (e.g., because
+	// a large DATA frame arrived), the average message size in the suffix is
+	// large enough that memory bloat is not a concern. Reset suffix tracking.
+	if backlogHeapSize <= utilizationFactor*b.uncompactedBytes {
+		b.uncompactedBytes = 0
+		b.uncompactedSuffixLen = 0
+		return
+	}
+	// Avoid compacting too frequently for short bursts of small frames.
+	// Wait until we have accumulated at least ~1024 small messages (~57 KB).
+	if backlogHeapSize <= compactionThreshold {
+		// Still can accumulate more payloads.
+		return
+	}
+
+	// Since the memory utilization is less than 50%, the average payload size
+	// of each recvMsg must be less than recvMsgSize (approx 56 bytes).
+	// In the worst case for bytes copied (where the average payload is just
+	// below recvMsgSize), compaction will occur once every:
+	//   compactionThreshold / (recvMsgSize + avg_payload) = ~520 messages,
+	// copying ~29KB of data.
+
+	start := 0
+	newBuf := b.bufPool.Get(b.uncompactedBytes)
+	startIdx := len(b.backlog) - b.uncompactedSuffixLen
+
+	for i := startIdx; i < len(b.backlog); i++ {
+		m := b.backlog[i]
+		b.backlog[i] = recvMsg{}
+		start += copy((*newBuf)[start:], m.buffer.ReadOnlyData())
+		m.buffer.Free()
+	}
+	b.backlog[startIdx] = recvMsg{
+		buffer: mem.NewBuffer(newBuf, b.bufPool),
+	}
+	b.backlog = b.backlog[:startIdx+1]
+	// After compaction, the suffix is replaced with a single message containing
+	// the combined payload. The new utilization is close to 1.0 (overhead of
+	// one recvMsg relative to the large compacted payload), which is well
+	// below the utilization factor of 2.
+	b.uncompactedBytes = 0
+	b.uncompactedSuffixLen = 0
 }
 
 func (b *recvBuffer) load() {
 	b.mu.Lock()
 	if len(b.backlog) > 0 {
 		select {
 		case b.c <- b.backlog[0]:
+			// backlog[0] is only part of the tracked uncompacted suffix if the
+			// entire backlog currently consists of the suffix. If an earlier
+			// compaction or reset occurred, backlog[0] is already compacted.
+			if envconfig.EnableReceiveBufferCompaction && b.uncompactedSuffixLen == len(b.backlog) {
+				b.uncompactedSuffixLen--
+				b.uncompactedBytes -= b.backlog[0].buffer.Len()
+			}
 			b.backlog[0] = recvMsg{}
 			b.backlog = b.backlog[1:]
 		default: