http2: optimize buffer allocation in transport

We have identified a high memory usage problem in our production service, which utilizes Traefik as a gRPC proxy. This service handles a substantial volume of gRPC bi-directional streaming requests that can persist for extended periods, spanning many days. Currently, there exists only a single global buffer pool in the http2 package. The allocated buffers, regardless of their sizes, are shared among requests with vastly different characteristics. For instance, gRPC streaming requests typically require smaller buffer sizes and occupy buffers for significant durations. Conversely, general HTTP requests may necessitate larger buffer sizes but only retain them temporarily. Unfortunately, the large buffers allocated by HTTP requests are can be chosen for subsequent gRPC streaming requests, resulting in numerous large buffers being unable to be recycled.

In our production environment, which processes approximately 1 million gRPC streaming requests, memory usage can soar to an excessive 800 GiB. This is a substantial waste of resources.

To address this challenge, we propose implementing a multi-layered buffer pool mechanism. This mechanism allows requests with varying characteristics to select buffers of appropriate sizes, optimizing resource allocation and recycling.

Change-Id: I834f7c08d90fd298aac7971ad45dc1a36251788b
GitHub-Last-Rev: 4771976
GitHub-Pull-Request: #182
Reviewed-on: https://go-review.googlesource.com/c/net/+/508415
Run-TryBot: Damien Neil <dneil@google.com>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
Auto-Submit: Damien Neil <dneil@google.com>
Reviewed-by: Damien Neil <dneil@google.com>

http2/transport.go

@@ -19,6 +19,7 @@ import (
 	"io/fs"
 	"log"
 	"math"
+	"math/bits"
 	mathrand "math/rand"
 	"net"
 	"net/http"
@@ -1680,7 +1681,27 @@ func (cs *clientStream) frameScratchBufferLen(maxFrameSize int) int {
 	return int(n) // doesn't truncate; max is 512K
 }
 
-var bufPool sync.Pool // of *[]byte
+// Seven bufPools manage different frame sizes. This helps to avoid scenarios where long-running
+// streaming requests using small frame sizes occupy large buffers initially allocated for prior
+// requests needing big buffers. The size ranges are as follows:
+// {0 KB, 16 KB], {16 KB, 32 KB], {32 KB, 64 KB], {64 KB, 128 KB], {128 KB, 256 KB],
+// {256 KB, 512 KB], {512 KB, infinity}
+// In practice, the maximum scratch buffer size should not exceed 512 KB due to
+// frameScratchBufferLen(maxFrameSize), thus the "infinity pool" should never be used.
+// It exists mainly as a safety measure, for potential future increases in max buffer size.
+var bufPools [7]sync.Pool // of *[]byte
+func bufPoolIndex(size int) int {
+	if size <= 16384 {
+		return 0
+	}
+	size -= 1
+	bits := bits.Len(uint(size))
+	index := bits - 14
+	if index >= len(bufPools) {
+		return len(bufPools) - 1
+	}
+	return index
+}
 
 func (cs *clientStream) writeRequestBody(req *http.Request) (err error) {
 	cc := cs.cc
@@ -1698,12 +1719,13 @@ func (cs *clientStream) writeRequestBody(req *http.Request) (err error) {
 	// Scratch buffer for reading into & writing from.
 	scratchLen := cs.frameScratchBufferLen(maxFrameSize)
 	var buf []byte
-	if bp, ok := bufPool.Get().(*[]byte); ok && len(*bp) >= scratchLen {
-		defer bufPool.Put(bp)
+	index := bufPoolIndex(scratchLen)
+	if bp, ok := bufPools[index].Get().(*[]byte); ok && len(*bp) >= scratchLen {
+		defer bufPools[index].Put(bp)
 		buf = *bp
 	} else {
 		buf = make([]byte, scratchLen)
-		defer bufPool.Put(&buf)
+		defer bufPools[index].Put(&buf)
 	}
 
 	var sawEOF bool