feat: gRPC supports reusing write buffer for each connection (#1918)

Author: DMwangnima

client/option.go

@@ -547,6 +547,28 @@ func WithGRPCKeepaliveParams(kp grpc.ClientKeepalive) Option {
 	}}
 }
 
+// WithGRPCReuseWriteBuffer enables write buffer reusing across connection lifecycle.
+// When enabled, a fixed-size buffer allocation (default 64KB) is no longer assigned to each connection.
+// Instead, buffers are allocated on-demand from buffer pool and put back to buffer pool after each flush.
+// Trading CPU overhead (malloc/free) for reduced memory usage.
+//
+// Use cases:
+//   - Scenarios with thousands of idle/low-traffic connections
+//   - Memory-constrained environments
+//
+// Trade-offs:
+//   - Reduces memory: ~2*WriteBufferSize per connection
+//   - Increases CPU: buffer pool malloc/free overhead on each write=>flush cycle
+//   - Slightly higher GC pressure
+//
+// This feature is disabled by default. (optimized for throughput)
+func WithGRPCReuseWriteBuffer(cfg grpc.ReuseWriteBufferConfig) Option {
+	return Option{F: func(o *client.Options, di *utils.Slice) {
+		di.Push(fmt.Sprintf("WithGRPCReuseWriteBuffer(%+v)", cfg))
+		o.GRPCConnectOpts.ReuseWriteBufferConfig = cfg
+	}}
+}
+
 // WithWarmingUp forces the client to do some warm-ups at the end of the initialization.
 func WithWarmingUp(wuo *warmup.ClientOption) Option {
 	return Option{F: func(o *client.Options, di *utils.Slice) {

client/option_test.go

@@ -755,6 +755,17 @@ func TestWithGRPCTLSConfig(t *testing.T) {
 	test.Assert(t, opts.GRPCConnectOpts != nil)
 }
 
+func TestWithGRPCReuseWriteBuffer(t *testing.T) {
+	opts := client.NewOptions([]client.Option{})
+	test.Assert(t, !opts.GRPCConnectOpts.ReuseWriteBufferConfig.Enable, opts.GRPCConnectOpts)
+	opts = client.NewOptions([]client.Option{WithGRPCReuseWriteBuffer(grpc.ReuseWriteBufferConfig{Enable: true})})
+	test.Assert(t, opts.GRPCConnectOpts.ReuseWriteBufferConfig.Enable, opts.GRPCConnectOpts)
+	opts = client.NewOptions([]client.Option{WithGRPCReuseWriteBuffer(grpc.ReuseWriteBufferConfig{Enable: false})})
+	test.Assert(t, !opts.GRPCConnectOpts.ReuseWriteBufferConfig.Enable, opts.GRPCConnectOpts)
+	opts = client.NewOptions([]client.Option{WithGRPCReuseWriteBuffer(grpc.ReuseWriteBufferConfig{})})
+	test.Assert(t, !opts.GRPCConnectOpts.ReuseWriteBufferConfig.Enable, opts.GRPCConnectOpts)
+}
+
 func TestTailOption(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()

pkg/remote/trans/nphttp2/grpc/controlbuf.go

@@ -610,7 +610,7 @@ func (l *loopyWriter) run(remoteAddr string) (err error) {
 			}
 			if gosched {
 				gosched = false
-				if l.framer.writer.offset < minBatchSize {
+				if l.framer.writer.GetOffset() < minBatchSize {
 					runtime.Gosched()
 					continue hasdata
 				}

pkg/remote/trans/nphttp2/grpc/framer.go

@@ -21,6 +21,7 @@ import (
 	"net"
 
 	"github.com/bytedance/gopkg/lang/dirtmake"
+	"github.com/bytedance/gopkg/lang/mcache"
 	"github.com/cloudwego/netpoll"
 	"golang.org/x/net/http2/hpack"
 
@@ -30,17 +31,17 @@ import (
 type framer struct {
 	*grpcframe.Framer
 	reader netpoll.Reader
-	writer *bufWriter
+	writer bufWriter
 }
 
-func newFramer(conn net.Conn, writeBufferSize, readBufferSize, maxHeaderListSize uint32) *framer {
+func newFramer(conn net.Conn, writeBufferSize, readBufferSize, maxHeaderListSize uint32, reuseCfg ReuseWriteBufferConfig) *framer {
 	var r netpoll.Reader
 	if npConn, ok := conn.(interface{ Reader() netpoll.Reader }); ok {
 		r = npConn.Reader()
 	} else {
 		r = netpoll.NewReader(conn)
 	}
-	w := newBufWriter(conn, int(writeBufferSize))
+	w := newBufWriter(conn, int(writeBufferSize), reuseCfg)
 	fr := &framer{
 		reader: r,
 		writer: w,
@@ -55,23 +56,41 @@ func newFramer(conn net.Conn, writeBufferSize, readBufferSize, maxHeaderListSize
 	return fr
 }
 
-type bufWriter struct {
-	writer    io.Writer
-	buf       []byte
-	offset    int
-	batchSize int
-	err       error
+type ReuseWriteBufferConfig struct {
+	Enable bool
+}
+
+type bufWriter interface {
+	Write(b []byte) (n int, err error)
+	Flush() error
+	GetOffset() int
 }
 
-func newBufWriter(writer io.Writer, batchSize int) *bufWriter {
-	return &bufWriter{
-		buf:       dirtmake.Bytes(batchSize*2, batchSize*2),
+func newBufWriter(writer io.Writer, batchSize int, reuseCfg ReuseWriteBufferConfig) bufWriter {
+	w := &keepBufWriter{
 		batchSize: batchSize,
 		writer:    writer,
 	}
+
+	if !reuseCfg.Enable {
+		// Using pre-allocated memory dedicated to each connection
+		w.buf = dirtmake.Bytes(batchSize*2, batchSize*2)
+		return w
+	}
+
+	return &reuseBufWriter{w}
 }
 
-func (w *bufWriter) Write(b []byte) (n int, err error) {
+// keepBufWriter pre-allocates batchSize * 2 buf and keeps it resident throughout the connection lifecycle.
+type keepBufWriter struct {
+	writer    io.Writer
+	buf       []byte
+	offset    int
+	batchSize int
+	err       error
+}
+
+func (w *keepBufWriter) Write(b []byte) (n int, err error) {
 	if w.err != nil {
 		return 0, w.err
 	}
@@ -84,13 +103,23 @@ func (w *bufWriter) Write(b []byte) (n int, err error) {
 		w.offset += nn
 		n += nn
 		if w.offset >= w.batchSize {
-			err = w.Flush()
+			if err = w.flushAllocatedBuffer(); err != nil {
+				return n, err
+			}
 		}
 	}
 	return n, err
 }
 
-func (w *bufWriter) Flush() error {
+func (w *keepBufWriter) Flush() error {
+	return w.flushAllocatedBuffer()
+}
+
+func (w *keepBufWriter) GetOffset() int {
+	return w.offset
+}
+
+func (w *keepBufWriter) flushAllocatedBuffer() error {
 	if w.err != nil {
 		return w.err
 	}
@@ -101,3 +130,47 @@ func (w *bufWriter) Flush() error {
 	w.offset = 0
 	return w.err
 }
+
+// During the Write=>Write=> … =>Write=>Flush cycle,
+// the first Write allocates batchSize * 2 buf from mcache, and the final Flush returns it.
+type reuseBufWriter struct {
+	*keepBufWriter
+}
+
+func (w *reuseBufWriter) Write(b []byte) (n int, err error) {
+	if w.err != nil {
+		return 0, w.err
+	}
+	if w.batchSize == 0 { // buffer has been disabled.
+		return w.writer.Write(b)
+	}
+	if w.buf == nil {
+		w.buf = mcache.Malloc(w.batchSize*2, w.batchSize*2)
+	}
+	for len(b) > 0 {
+		nn := copy(w.buf[w.offset:], b)
+		b = b[nn:]
+		w.offset += nn
+		n += nn
+		if w.offset >= w.batchSize {
+			if err = w.flushAllocatedBuffer(); err != nil {
+				return n, err
+			}
+		}
+	}
+	return n, err
+}
+
+func (w *reuseBufWriter) Flush() error {
+	err := w.flushAllocatedBuffer()
+	// Reuse only when err is nil, to prevent the underlying connection from still holding buf after a Write failure.
+	if err == nil && w.buf != nil {
+		mcache.Free(w.buf)
+		w.buf = nil
+	}
+	return err
+}
+
+func (w *reuseBufWriter) GetOffset() int {
+	return w.offset
+}

pkg/remote/trans/nphttp2/grpc/http2_client.go

@@ -186,7 +186,7 @@ func newHTTP2Client(ctx context.Context, conn net.Conn, opts ConnectOptions,
 		readerDone:            make(chan struct{}),
 		writerDone:            make(chan struct{}),
 		goAway:                make(chan struct{}),
-		framer:                newFramer(conn, writeBufSize, readBufSize, maxHeaderListSize),
+		framer:                newFramer(conn, writeBufSize, readBufSize, maxHeaderListSize, opts.ReuseWriteBufferConfig),
 		fc:                    &trInFlow{limit: icwz},
 		activeStreams:         make(map[uint32]*Stream),
 		kp:                    kp,

pkg/remote/trans/nphttp2/grpc/http2_server.go

@@ -148,7 +148,7 @@ func newHTTP2Server(ctx context.Context, conn net.Conn, config *ServerConfig) (_
 		maxHeaderListSize = *config.MaxHeaderListSize
 	}
 
-	framer := newFramer(conn, config.WriteBufferSize, config.ReadBufferSize, maxHeaderListSize)
+	framer := newFramer(conn, config.WriteBufferSize, config.ReadBufferSize, maxHeaderListSize, config.ReuseWriteBufferConfig)
 	// Send initial settings as connection preface to client.
 	isettings := []http2.Setting{{
 		ID:  http2.SettingMaxFrameSize,