transport: Pool read buffers used by the HTTP/2 framer

credentials/alts/internal/conn/record.go

@@ -27,8 +27,9 @@ import (
 	"net"
 
 	core "google.golang.org/grpc/credentials/alts/internal"
-	"google.golang.org/grpc/internal/mem"
+	imem "google.golang.org/grpc/internal/mem"
 	"google.golang.org/grpc/internal/transport/readyreader"
+	"google.golang.org/grpc/mem"
 )
 
 // ALTSRecordCrypto is the interface for gRPC ALTS record protocol.
@@ -75,16 +76,19 @@ const (
 
 var (
 	protocols    = make(map[string]ALTSRecordFunc)
-	writeBufPool *mem.BinaryTieredBufferPool
+	writeBufPool *imem.BinaryTieredBufferPool
 	// readBufPool pools buffers of at least `altsReadBufferInitialSize` size.
 	// Since the read buffer size is slightly larger than 32KB, using a regular
 	// BinaryTieredBufferPool results in allocating buffers of almost double the
 	// required length.
-	readBufPool = mem.NewDirtySimplePool()
+	readBufPool = imem.NewDirtySimplePool()
+
+	// Compile-time check to ensure conn implements ReadyReader.
+	_ readyreader.Reader = &conn{}
 )
 
 func init() {
-	pool, err := mem.NewDirtyBinaryTieredBufferPool(
+	pool, err := imem.NewDirtyBinaryTieredBufferPool(
 		8,
 		12, // Go page size, 4KB
 		14, // 16KB (max HTTP/2 frame size used by gRPC)
@@ -126,7 +130,8 @@ type conn struct {
 	// nextFrame stores the next frame (in protected buffer) info.
 	nextFrame []byte
 	// overhead is the calculated overhead of each frame.
-	overhead int
+	overhead  int
+	constPool constBufferPool // stored as a field to avoid heap allocations.
 }
 
 // NewConn creates a new secure channel instance given the other party role and
@@ -163,21 +168,38 @@ func NewConn(c net.Conn, side core.Side, recordProtocol string, key []byte, prot
 	return altsConn, nil
 }
 
+type constBufferPool struct {
+	buffer []byte
+}
+
+func (p *constBufferPool) Get(int) *[]byte {
+	return &p.buffer
+}
+
+func (p *constBufferPool) Put(*[]byte) {}
+
 // Read reads and decrypts a frame from the underlying connection, and copies the
 // decrypted payload into b. If the size of the payload is greater than len(b),
 // Read retains the remaining bytes in an internal buffer, and subsequent calls
 // to Read will read from this buffer until it is exhausted.
 func (p *conn) Read(b []byte) (n int, err error) {
+	p.constPool.buffer = b
+	_, n, err = p.ReadOnReady(len(b), &p.constPool)
+	return n, err
+}
+
+func (p *conn) ReadOnReady(bufSize int, pool mem.BufferPool) (*[]byte, int, error) {
 	if len(p.buf) == 0 {
 		var framedMsg []byte
 		var protected []byte
 		if p.protectedHandle != nil {
 			protected = *p.protectedHandle
 			protected = protected[:cap(protected)]
 		}
+		var err error
 		framedMsg, p.nextFrame, err = ParseFramedMsg(p.nextFrame, altsRecordLengthLimit)
 		if err != nil {
-			return 0, err
+			return nil, 0, err
 		}
 		// Check whether the next frame to be decrypted has been
 		// completely received yet.
@@ -217,40 +239,42 @@ func (p *conn) Read(b []byte) (n int, err error) {
 				// Connection was idle, need to re-allocate the read buffer.
 				newBuf, nRead, err := p.reader.ReadOnReady(altsReadBufferInitialSize, readBufPool)
 				if err != nil {
-					return 0, err
+					return nil, 0, err
 				}
 				p.protectedHandle = newBuf
 				protected = (*newBuf)[:nRead]
 			} else {
 				nRead, err := p.Conn.Read(protected[len(protected):cap(protected)])
 				if err != nil {
-					return 0, err
+					return nil, 0, err
 				}
 				protected = protected[:len(protected)+nRead]
 			}
 			framedMsg, p.nextFrame, err = ParseFramedMsg(protected, altsRecordLengthLimit)
 			if err != nil {
-				return 0, err
+				return nil, 0, err
 			}
 		}
 		// Now we have a complete frame, decrypted it.
 		msg := framedMsg[MsgLenFieldSize:]
 		msgType := binary.LittleEndian.Uint32(msg[:msgTypeFieldSize])
 		if msgType&0xff != altsRecordMsgType {
-			return 0, fmt.Errorf("received frame with incorrect message type %v, expected lower byte %v",
+			return nil, 0, fmt.Errorf("received frame with incorrect message type %v, expected lower byte %v",
 				msgType, altsRecordMsgType)
 		}
 		ciphertext := msg[msgTypeFieldSize:]
 
 		// Decrypt directly into the buffer, avoiding a copy from p.buf if
 		// possible.
-		if len(b) >= len(ciphertext) {
-			dec, err := p.crypto.Decrypt(b[:0], ciphertext)
+		if bufSize >= len(ciphertext) {
+			allocatedBuf := pool.Get(bufSize)
+			dec, err := p.crypto.Decrypt((*allocatedBuf)[:0], ciphertext)
 			if err != nil {
-				return 0, err
+				pool.Put(allocatedBuf)
+				return nil, 0, err
 			}
 			p.dropProtectedIfEmtpy()
-			return len(dec), nil
+			return allocatedBuf, len(dec), nil
 		}
 		// Decrypt requires that if the dst and ciphertext alias, they
 		// must alias exactly. Code here used to use msg[:0], but msg
@@ -261,14 +285,15 @@ func (p *conn) Read(b []byte) (n int, err error) {
 		// check: https://golang.org/pkg/crypto/cipher/#AEAD.
 		p.buf, err = p.crypto.Decrypt(ciphertext[:0], ciphertext)
 		if err != nil {
-			return 0, err
+			return nil, 0, err
 		}
 	}
 
-	n = copy(b, p.buf)
+	allocatedBuf := pool.Get(bufSize)
+	n := copy(*allocatedBuf, p.buf)
 	p.buf = p.buf[n:]
 	p.dropProtectedIfEmtpy()
-	return n, nil
+	return allocatedBuf, n, nil
 }
 
 func (p *conn) dropProtectedIfEmtpy() {

internal/envconfig/envconfig.go

@@ -142,6 +142,16 @@ var (
 	//
 	// TODO: In release v1.82.0, env var will be enabled by default.
 	Enable8KBDefaultHeaderListSize = boolFromEnv("GRPC_GO_EXPERIMENTAL_ENABLE_8KB_DEFAULT_HEADER_LIST_SIZE", false)
+
+	// EnableHTTPFramerReadBufferPooling enables the use of the
+	// readyreader.Reader interface to perform non-memory-pinning reads,
+	// provided the underlying net.Conn supports it. This reduces memory usage
+	// when subchannels are idle.
+	//
+	// 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.
+	EnableHTTPFramerReadBufferPooling = boolFromEnv("GRPC_GO_EXPERIMENTAL_HTTP_FRAMER_READ_BUFFER_POOLING", true)
 )
 
 func boolFromEnv(envVar string, def bool) bool {

internal/transport/http_util.go

@@ -36,6 +36,9 @@ import (
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/hpack"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/internal/envconfig"
+	imem "google.golang.org/grpc/internal/mem"
+	"google.golang.org/grpc/internal/transport/readyreader"
 	"google.golang.org/grpc/mem"
 )
 
@@ -296,15 +299,15 @@ func decodeGrpcMessageUnchecked(msg string) string {
 }
 
 type bufWriter struct {
-	pool      *sync.Pool
+	pool      *imem.SimpleBufferPool
 	buf       []byte
 	offset    int
 	batchSize int
 	conn      io.Writer
 	err       error
 }
 
-func newBufWriter(conn io.Writer, batchSize int, pool *sync.Pool) *bufWriter {
+func newBufWriter(conn io.Writer, batchSize int, pool *imem.SimpleBufferPool) *bufWriter {
 	w := &bufWriter{
 		batchSize: batchSize,
 		conn:      conn,
@@ -326,7 +329,7 @@ func (w *bufWriter) Write(b []byte) (int, error) {
 		return n, toIOError(err)
 	}
 	if w.buf == nil {
-		b := w.pool.Get().(*[]byte)
+		b := w.pool.Get(w.batchSize)
 		w.buf = *b
 	}
 	written := 0
@@ -407,22 +410,33 @@ type framer struct {
 	errDetail error
 }
 
-var writeBufferPoolMap = make(map[int]*sync.Pool)
-var writeBufferMutex sync.Mutex
+var ioBufferPoolMap = make(map[int]*imem.SimpleBufferPool)
+var ioBufferMutex sync.Mutex
+
+func bufferedReader(r io.Reader, bufSize int) io.Reader {
+	if bufSize <= 0 {
+		return r
+	}
+	if !envconfig.EnableHTTPFramerReadBufferPooling {
+		return bufio.NewReaderSize(r, bufSize)
+	}
+	if rr := readyreader.NewNonBlocking(r); rr != nil {
+		readPool := getIOBufferPool(bufSize)
+		return readyreader.NewBuffered(rr, bufSize, readPool)
+	}
+	return bufio.NewReaderSize(r, bufSize)
+}
 
 func newFramer(conn io.ReadWriter, writeBufferSize, readBufferSize int, sharedWriteBuffer bool, maxHeaderListSize uint32, memPool mem.BufferPool) *framer {
 	if writeBufferSize < 0 {
 		writeBufferSize = 0
 	}
-	var r io.Reader = conn
-	if readBufferSize > 0 {
-		r = bufio.NewReaderSize(r, readBufferSize)
-	}
-	var pool *sync.Pool
+	r := bufferedReader(conn, readBufferSize)
+	var writePool *imem.SimpleBufferPool
 	if sharedWriteBuffer {
-		pool = getWriteBufferPool(writeBufferSize)
+		writePool = getIOBufferPool(writeBufferSize)
 	}
-	w := newBufWriter(conn, writeBufferSize, pool)
+	w := newBufWriter(conn, writeBufferSize, writePool)
 	f := &framer{
 		writer: w,
 		fr:     http2.NewFramer(w, r),
@@ -578,20 +592,15 @@ func (df *parsedDataFrame) Header() http2.FrameHeader {
 	return df.FrameHeader
 }
 
-func getWriteBufferPool(size int) *sync.Pool {
-	writeBufferMutex.Lock()
-	defer writeBufferMutex.Unlock()
-	pool, ok := writeBufferPoolMap[size]
+func getIOBufferPool(size int) *imem.SimpleBufferPool {
+	ioBufferMutex.Lock()
+	defer ioBufferMutex.Unlock()
+	pool, ok := ioBufferPoolMap[size]
 	if ok {
 		return pool
 	}
-	pool = &sync.Pool{
-		New: func() any {
-			b := make([]byte, size)
-			return &b
-		},
-	}
-	writeBufferPoolMap[size] = pool
+	pool = imem.NewDirtySimplePool()
+	ioBufferPoolMap[size] = pool
 	return pool
 }
 

internal/transport/http_util_test.go

@@ -19,6 +19,7 @@
 package transport
 
 import (
+	"bufio"
 	"bytes"
 	"errors"
 	"fmt"
@@ -31,6 +32,9 @@ import (
 	"time"
 
 	"golang.org/x/net/http2"
+	"google.golang.org/grpc/internal/envconfig"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/internal/transport/readyreader"
 	"google.golang.org/grpc/mem"
 )
 
@@ -259,7 +263,7 @@ func (s) TestWriteBadConnection(t *testing.T) {
 	// Configure the bufWriter with a batchsize that results in data being flushed
 	// to the underlying conn, midway through Write().
 	writeBufferSize := (len(data) - 1) / 2
-	writer := newBufWriter(&badNetworkConn{}, writeBufferSize, getWriteBufferPool(writeBufferSize))
+	writer := newBufWriter(&badNetworkConn{}, writeBufferSize, getIOBufferPool(writeBufferSize))
 
 	errCh := make(chan error, 1)
 	go func() {
@@ -413,3 +417,74 @@ func (s) TestFramer_ParseDataFrame(t *testing.T) {
 		})
 	}
 }
+
+type testReadyReader struct {
+	readyreader.Reader
+}
+
+func (t *testReadyReader) Read([]byte) (int, error) {
+	return 0, io.EOF
+}
+
+func (s) TestBufferedReader(t *testing.T) {
+	normalReader := bytes.NewReader(nil)
+
+	tests := []struct {
+		name          string
+		reader        io.Reader
+		bufSize       int
+		enablePooling bool
+		wantTypeOf    any
+	}{
+		{
+			name:          "bufSize_0",
+			reader:        normalReader,
+			bufSize:       0,
+			enablePooling: true,
+			wantTypeOf:    (*bytes.Reader)(nil),
+		},
+		{
+			name:          "env_var_disabled_normal_reader",
+			reader:        normalReader,
+			bufSize:       10,
+			enablePooling: false,
+			wantTypeOf:    (*bufio.Reader)(nil),
+		},
+		{
+			name:          "env_var_disabled_ready_reader",
+			reader:        &testReadyReader{},
+			bufSize:       10,
+			enablePooling: false,
+			wantTypeOf:    (*bufio.Reader)(nil),
+		},
+		{
+			name:          "env_var_enabled_normal_reader",
+			reader:        normalReader,
+			bufSize:       10,
+			enablePooling: true,
+			wantTypeOf:    (*bufio.Reader)(nil),
+		},
+		{
+			name:          "env_var_enabled_ready_reader",
+			reader:        &testReadyReader{},
+			bufSize:       10,
+			enablePooling: true,
+			wantTypeOf:    readyreader.NewBuffered(nil, 10, mem.DefaultBufferPool()),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			testutils.SetEnvConfig(t, &envconfig.EnableHTTPFramerReadBufferPooling, tt.enablePooling)
+
+			got := bufferedReader(tt.reader, tt.bufSize)
+
+			gotType := reflect.TypeOf(got)
+			wantType := reflect.TypeOf(tt.wantTypeOf)
+
+			if gotType != wantType {
+				t.Errorf("bufferedReader() type = %v, want %v", gotType, wantType)
+			}
+		})
+	}
+}