transport: Fix reporting of bytes read while reading headers (#7660) (#7667)

arjan-bal · web-flow · commit 935f8cb5ac28 · 2024-09-24T23:15:46.000+05:30
diff --git a/internal/transport/transport.go b/internal/transport/transport.go
@@ -616,7 +616,7 @@ func (t *transportReader) ReadHeader(header []byte) (int, error) {
 		t.er = err
 		return 0, err
 	}
-	t.windowHandler(len(header))
+	t.windowHandler(n)
 	return n, nil
 }
 
diff --git a/internal/transport/transport_test.go b/internal/transport/transport_test.go
@@ -2845,3 +2845,34 @@ func (s) TestClientCloseReturnsEarlyWhenGoAwayWriteHangs(t *testing.T) {
 	isGreetingDone.Store(true)
 	ct.Close(errors.New("manually closed by client"))
 }
+
+// TestReadHeaderMultipleBuffers tests the stream when the gRPC headers are
+// split across multiple buffers. It verifies that the reporting of the
+// number of bytes read for flow control is correct.
+func (s) TestReadHeaderMultipleBuffers(t *testing.T) {
+	headerLen := 5
+	recvBuffer := newRecvBuffer()
+	recvBuffer.put(recvMsg{buffer: make(mem.SliceBuffer, 3)})
+	recvBuffer.put(recvMsg{buffer: make(mem.SliceBuffer, headerLen-3)})
+	bytesRead := 0
+	s := Stream{
+		requestRead: func(int) {},
+		trReader: &transportReader{
+			reader: &recvBufferReader{
+				recv: recvBuffer,
+			},
+			windowHandler: func(i int) {
+				bytesRead += i
+			},
+		},
+	}
+
+	header := make([]byte, headerLen)
+	err := s.ReadHeader(header)
+	if err != nil {
+		t.Fatalf("ReadHeader(%v) = %v", header, err)
+	}
+	if bytesRead != headerLen {
+		t.Errorf("bytesRead = %d, want = %d", bytesRead, headerLen)
+	}
+}

Original file line number	Diff line number	Diff line change
`@@ -616,7 +616,7 @@ func (t *transportReader) ReadHeader(header []byte) (int, error) {`
`616`	`616`	`t.er = err`
`617`	`617`	`return 0, err`
`618`	`618`	`}`
`619`		`- t.windowHandler(len(header))`
	`619`	`+ t.windowHandler(n)`
`620`	`620`	`return n, nil`
`621`	`621`	`}`
`622`	`622`