changes no tests

Author: arjan-bal

internal/transport/http2_client.go

@@ -335,7 +335,7 @@ func NewHTTP2Client(connectCtx, ctx context.Context, addr resolver.Address, opts
 		writerDone:            make(chan struct{}),
 		goAway:                make(chan struct{}),
 		keepaliveDone:         make(chan struct{}),
-		framer:                newFramer(conn, writeBufSize, readBufSize, opts.SharedWriteBuffer, maxHeaderListSize),
+		framer:                newFramer(conn, writeBufSize, readBufSize, opts.SharedWriteBuffer, maxHeaderListSize, opts.BufferPool),
 		fc:                    &trInFlow{limit: uint32(icwz)},
 		scheme:                scheme,
 		activeStreams:         make(map[uint32]*ClientStream),
@@ -1177,7 +1177,12 @@ func (t *http2Client) updateFlowControl(n uint32) {
 	})
 }
 
-func (t *http2Client) handleData(f *http2.DataFrame) {
+func (t *http2Client) handleData(f *parsedDataFrame) {
+	defer func() {
+		if f.data != nil {
+			f.data.Free()
+		}
+	}()
 	size := f.Header().Length
 	var sendBDPPing bool
 	if t.bdpEst != nil {
@@ -1221,22 +1226,15 @@ func (t *http2Client) handleData(f *http2.DataFrame) {
 			t.closeStream(s, io.EOF, true, http2.ErrCodeFlowControl, status.New(codes.Internal, err.Error()), nil, false)
 			return
 		}
+		dataLen := f.data.Len()
 		if f.Header().Flags.Has(http2.FlagDataPadded) {
-			if w := s.fc.onRead(size - uint32(len(f.Data()))); w > 0 {
+			if w := s.fc.onRead(size - uint32(dataLen)); w > 0 {
 				t.controlBuf.put(&outgoingWindowUpdate{s.id, w})
 			}
 		}
-		// TODO(bradfitz, zhaoq): A copy is required here because there is no
-		// guarantee f.Data() is consumed before the arrival of next frame.
-		// Can this copy be eliminated?
-		if len(f.Data()) > 0 {
-			pool := t.bufferPool
-			if pool == nil {
-				// Note that this is only supposed to be nil in tests. Otherwise, stream is
-				// always initialized with a BufferPool.
-				pool = mem.DefaultBufferPool()
-			}
-			s.write(recvMsg{buffer: mem.Copy(f.Data(), pool)})
+		if dataLen > 0 {
+			s.write(recvMsg{buffer: f.data})
+			f.data = nil
 		}
 	}
 	// The server has closed the stream without sending trailers.  Record that
@@ -1656,6 +1654,13 @@ func (t *http2Client) reader(errCh chan<- error) {
 		}
 	}()
 
+	pool := t.bufferPool
+	if pool == nil {
+		// Note that this is only supposed to be nil in tests. Otherwise, stream
+		// is always initialized with a BufferPool.
+		pool = mem.DefaultBufferPool()
+	}
+
 	if err := t.readServerPreface(); err != nil {
 		errCh <- err
 		return
@@ -1668,7 +1673,7 @@ func (t *http2Client) reader(errCh chan<- error) {
 	// loop to keep reading incoming messages on this transport.
 	for {
 		t.controlBuf.throttle()
-		frame, err := t.framer.fr.ReadFrame()
+		frame, err := t.framer.readFrame()
 		if t.keepaliveEnabled {
 			atomic.StoreInt64(&t.lastRead, time.Now().UnixNano())
 		}
@@ -1701,7 +1706,7 @@ func (t *http2Client) reader(errCh chan<- error) {
 		switch frame := frame.(type) {
 		case *http2.MetaHeadersFrame:
 			t.operateHeaders(frame)
-		case *http2.DataFrame:
+		case *parsedDataFrame:
 			t.handleData(frame)
 		case *http2.RSTStreamFrame:
 			t.handleRSTStream(frame)

internal/transport/http2_server.go

@@ -169,7 +169,8 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
 	if config.MaxHeaderListSize != nil {
 		maxHeaderListSize = *config.MaxHeaderListSize
 	}
-	framer := newFramer(conn, writeBufSize, readBufSize, config.SharedWriteBuffer, maxHeaderListSize)
+
+	framer := newFramer(conn, writeBufSize, readBufSize, config.SharedWriteBuffer, maxHeaderListSize, config.BufferPool)
 	// Send initial settings as connection preface to client.
 	isettings := []http2.Setting{{
 		ID:  http2.SettingMaxFrameSize,
@@ -668,10 +669,16 @@ func (t *http2Server) HandleStreams(ctx context.Context, handle func(*ServerStre
 		close(t.readerDone)
 		<-t.loopyWriterDone
 	}()
+	pool := t.bufferPool
+	if pool == nil {
+		// Note that this is only supposed to be nil in tests. Otherwise, stream
+		// is always initialized with a BufferPool.
+		pool = mem.DefaultBufferPool()
+	}
 	for {
 		t.controlBuf.throttle()
-		frame, err := t.framer.fr.ReadFrame()
 		atomic.StoreInt64(&t.lastRead, time.Now().UnixNano())
+		frame, err := t.framer.readFrame()
 		if err != nil {
 			if se, ok := err.(http2.StreamError); ok {
 				if t.logger.V(logLevel) {
@@ -707,7 +714,7 @@ func (t *http2Server) HandleStreams(ctx context.Context, handle func(*ServerStre
 				})
 				continue
 			}
-		case *http2.DataFrame:
+		case *parsedDataFrame:
 			t.handleData(frame)
 		case *http2.RSTStreamFrame:
 			t.handleRSTStream(frame)
@@ -788,7 +795,12 @@ func (t *http2Server) updateFlowControl(n uint32) {
 
 }
 
-func (t *http2Server) handleData(f *http2.DataFrame) {
+func (t *http2Server) handleData(f *parsedDataFrame) {
+	defer func() {
+		if f.data != nil {
+			f.data.Free()
+		}
+	}()
 	size := f.Header().Length
 	var sendBDPPing bool
 	if t.bdpEst != nil {
@@ -833,22 +845,15 @@ func (t *http2Server) handleData(f *http2.DataFrame) {
 			t.closeStream(s, true, http2.ErrCodeFlowControl, false)
 			return
 		}
+		dataLen := f.data.Len()
 		if f.Header().Flags.Has(http2.FlagDataPadded) {
-			if w := s.fc.onRead(size - uint32(len(f.Data()))); w > 0 {
+			if w := s.fc.onRead(size - uint32(dataLen)); w > 0 {
 				t.controlBuf.put(&outgoingWindowUpdate{s.id, w})
 			}
 		}
-		// TODO(bradfitz, zhaoq): A copy is required here because there is no
-		// guarantee f.Data() is consumed before the arrival of next frame.
-		// Can this copy be eliminated?
-		if len(f.Data()) > 0 {
-			pool := t.bufferPool
-			if pool == nil {
-				// Note that this is only supposed to be nil in tests. Otherwise, stream is
-				// always initialized with a BufferPool.
-				pool = mem.DefaultBufferPool()
-			}
-			s.write(recvMsg{buffer: mem.Copy(f.Data(), pool)})
+		if dataLen > 0 {
+			s.write(recvMsg{buffer: f.data})
+			f.data = nil
 		}
 	}
 	if f.StreamEnded() {

internal/transport/http_util.go

@@ -37,6 +37,7 @@ import (
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/hpack"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/mem"
 )
 
 const (
@@ -388,15 +389,34 @@ func toIOError(err error) error {
 	return ioError{error: err}
 }
 
+type parsedDataFrame struct {
+	http2.FrameHeader
+	data mem.Buffer
+}
+
+func (df *parsedDataFrame) StreamEnded() bool {
+	return df.FrameHeader.Flags.Has(http2.FlagDataEndStream)
+}
+
 type framer struct {
 	writer *bufWriter
 	fr     *http2.Framer
+	reader io.Reader
+	// Cached data frame to avoid heap allocations.
+	dataFrame parsedDataFrame
+	pool      mem.BufferPool
 }
 
 var writeBufferPoolMap = make(map[int]*sync.Pool)
 var writeBufferMutex sync.Mutex
 
-func newFramer(conn net.Conn, writeBufferSize, readBufferSize int, sharedWriteBuffer bool, maxHeaderListSize uint32) *framer {
+func newFramer(conn net.Conn, writeBufferSize, readBufferSize int, sharedWriteBuffer bool, maxHeaderListSize uint32, memPool mem.BufferPool) *framer {
+	if memPool == nil {
+		// Note that this is only supposed to be nil in tests. Otherwise, stream
+		// is always initialized with a BufferPool.
+		memPool = mem.DefaultBufferPool()
+	}
+
 	if writeBufferSize < 0 {
 		writeBufferSize = 0
 	}
@@ -412,6 +432,8 @@ func newFramer(conn net.Conn, writeBufferSize, readBufferSize int, sharedWriteBu
 	f := &framer{
 		writer: w,
 		fr:     http2.NewFramer(w, r),
+		reader: r,
+		pool:   memPool,
 	}
 	f.fr.SetMaxReadFrameSize(http2MaxFrameLen)
 	// Opt-in to Frame reuse API on framer to reduce garbage.
@@ -422,6 +444,73 @@ func newFramer(conn net.Conn, writeBufferSize, readBufferSize int, sharedWriteBu
 	return f
 }
 
+func (f *framer) readFrame() (any, error) {
+	fh, err := f.fr.ReadFrameHeader()
+	if err != nil {
+		return nil, err
+	}
+	if fh.Type == http2.FrameData {
+		err = f.readDataFrame(fh, f.pool, &f.dataFrame)
+		return &f.dataFrame, err
+	} else {
+		return f.fr.ReadFrameForHeader(fh)
+	}
+}
+
+// readDataFrame reads and parses a data frame from the underlying io.Reader.
+// Frames aren't safe to read from after a subsequent call to ReadFrame.
+func (f *framer) readDataFrame(fh http2.FrameHeader, pool mem.BufferPool, df *parsedDataFrame) (err error) {
+	if fh.StreamID == 0 {
+		// DATA frames MUST be associated with a stream. If a
+		// DATA frame is received whose stream identifier
+		// field is 0x0, the recipient MUST respond with a
+		// connection error (Section 5.4.1) of type
+		// PROTOCOL_ERROR.
+		return fmt.Errorf("DATA frame with stream ID 0")
+	}
+	payload := pool.Get(int(fh.Length))
+	defer func() {
+		if err != nil {
+			f.pool.Put(payload)
+		}
+	}()
+	if fh.Flags.Has(http2.FlagDataPadded) {
+		if fh.Length == 0 {
+			return io.ErrUnexpectedEOF
+		}
+		// This initial 1-byte read can be inefficient for unbuffered readers,
+		// but it allows the rest of the payload to be read directly to the
+		// start of the destination slice. This makes it easy to return the
+		// original slice back to the buffer pool.
+		if _, err := io.ReadFull(f.reader, (*payload)[:1]); err != nil {
+			return err
+		}
+		padSize := (*payload)[0]
+		*payload = (*payload)[:len(*payload)-1]
+		if _, err := io.ReadFull(f.reader, *payload); err != nil {
+			return err
+		}
+		if int(padSize) > len(*payload) {
+			// If the length of the padding is greater than the
+			// length of the frame payload, the recipient MUST
+			// treat this as a connection error.
+			// Filed: https://github.com/http2/http2-spec/issues/610
+			return fmt.Errorf("pad size larger than data payload")
+		}
+		*payload = (*payload)[:len(*payload)-int(padSize)]
+	} else if _, err := io.ReadFull(f.reader, *payload); err != nil {
+		return err
+	}
+
+	df.FrameHeader = fh
+	df.data = mem.NewBuffer(payload, pool)
+	return nil
+}
+
+func (df *parsedDataFrame) Header() http2.FrameHeader {
+	return df.FrameHeader
+}
+
 func getWriteBufferPool(size int) *sync.Pool {
 	writeBufferMutex.Lock()
 	defer writeBufferMutex.Unlock()

internal/transport/keepalive_test.go

@@ -40,6 +40,7 @@ import (
 	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/syscall"
 	"google.golang.org/grpc/keepalive"
+	"google.golang.org/grpc/mem"
 	"google.golang.org/grpc/testdata"
 )
 
@@ -192,7 +193,7 @@ func (s) TestKeepaliveServerClosesUnresponsiveClient(t *testing.T) {
 	if n, err := conn.Write(clientPreface); err != nil || n != len(clientPreface) {
 		t.Fatalf("conn.Write(clientPreface) failed: n=%v, err=%v", n, err)
 	}
-	framer := newFramer(conn, defaultWriteBufSize, defaultReadBufSize, false, 0)
+	framer := newFramer(conn, defaultWriteBufSize, defaultReadBufSize, false, 0, mem.DefaultBufferPool())
 	if err := framer.fr.WriteSettings(http2.Setting{}); err != nil {
 		t.Fatal("framer.WriteSettings(http2.Setting{}) failed:", err)
 	}