cloudwego · ppzqh · Jun 6, 2025
@@ -20,6 +20,8 @@ import (
 	bytes2 "bytes"
 	"net"
 	"time"
+
+	"github.com/cloudwego/gopkg/bufiox"
 )
 
 var _ net.Conn = &Conn{}
@@ -111,3 +113,12 @@ func NewIOConn() *Conn {
 		},
 	}
 }
+
+type MockConnWithBufioxReader struct {
+	net.Conn
+	BufioxReader bufiox.Reader
+}
+
+func (c *MockConnWithBufioxReader) Reader() bufiox.Reader {
+	return c.BufioxReader
+}
@@ -23,6 +23,7 @@ files=(
 ../../pkg/proxy/proxy.go proxy/proxy.go proxy
 ../../pkg/utils/sharedticker.go utils/sharedticker.go utils
 ../../../netpoll/connection.go netpoll/connection.go netpoll
+../../../gopkg/bufiox/bufreader.go bufiox/bufreader.go bufiox
 $GOROOT/src/net/net.go net/net.go net
 )
 

@@ -20,6 +20,8 @@ import (
 	"io"
 	"net"
 
+	"github.com/cloudwego/gopkg/bufiox"
+
 	"github.com/bytedance/gopkg/lang/dirtmake"
 	"github.com/cloudwego/netpoll"
 	"golang.org/x/net/http2/hpack"
@@ -29,16 +31,24 @@ import (
 
 type framer struct {
 	*grpcframe.Framer
-	reader netpoll.Reader
+	reader bufiox.Reader
 	writer *bufWriter
 }
 
 func newFramer(conn net.Conn, writeBufferSize, readBufferSize, maxHeaderListSize uint32) *framer {
-	var r netpoll.Reader
-	if npConn, ok := conn.(interface{ Reader() netpoll.Reader }); ok {
-		r = npConn.Reader()
+	var r bufiox.Reader
+	// Initialize a bufiox.Reader based on the connection:
+	// 1. If the connection's reader is a `bufiox.Reader`, use it directly.
+	// 2. If the connection's reader is a `netpoll.Reader`, wrap it in a `netpollBufioxReader`.
+	// 3. Otherwise, create a `bufiox.DefaultReader` with the connection.
+	if bc, ok := conn.(interface{ Reader() bufiox.Reader }); ok {
+		r = bc.Reader()
 	} else {
-		r = netpoll.NewReader(conn)
+		if npConn, ok := conn.(interface{ Reader() netpoll.Reader }); ok {
+			r = &netpollBufioxReader{Reader: npConn.Reader()}
+		} else {
+			r = bufiox.NewDefaultReader(conn)
+		}
 	}
 	w := newBufWriter(conn, int(writeBufferSize))
 	fr := &framer{
@@ -101,3 +111,47 @@ func (w *bufWriter) Flush() error {
 	w.offset = 0
 	return w.err
 }
+
+var _ bufiox.Reader = &netpollBufioxReader{}
+
+// netpollBufioxReader implements bufiox.Reader with netpoll.Reader
+type netpollBufioxReader struct {
+	netpoll.Reader
+	readLen int
+}
+
+func (r *netpollBufioxReader) Next(n int) (p []byte, err error) {
+	p, err = r.Reader.Next(n)
+	if err != nil {
+		return nil, err
+	}
+	r.readLen += len(p)
+	return p, nil
+}
+
+func (r *netpollBufioxReader) ReadBinary(bs []byte) (n int, err error) {
+	p, err := r.Next(len(bs))
+	if err != nil {
+		return 0, err
+	}
+	n = copy(bs, p)
+	return n, nil
+}
+
+func (r *netpollBufioxReader) Skip(n int) (err error) {
+	err = r.Reader.Skip(n)
+	if err != nil {
+		return err
+	}
+	r.readLen += n
+	return nil
+}
+
+func (r *netpollBufioxReader) ReadLen() (n int) {
+	return r.readLen
+}
+
+func (r *netpollBufioxReader) Release(e error) (err error) {
+	r.readLen = 0
+	return r.Reader.Release()
+}
@@ -0,0 +1,114 @@
+/*
+ * Copyright 2025 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package grpc
+
+import (
+	"io"
+	"net"
+	"testing"
+
+	"github.com/cloudwego/gopkg/bufiox"
+	"github.com/golang/mock/gomock"
+
+	mocksnetpoll "github.com/cloudwego/kitex/internal/mocks/netpoll"
+	"github.com/cloudwego/kitex/internal/test"
+)
+
+type mockConnWithBufioxReader struct {
+	net.Conn
+	bufioxReader bufiox.Reader
+}
+
+func (c *mockConnWithBufioxReader) Reader() bufiox.Reader {
+	return c.bufioxReader
+}
+
+func TestNewFramer(t *testing.T) {
+	// conn without bufiox reader
+	var conn net.Conn
+	conn = &mockConn{}
+	fr := newFramer(conn, 0, 0, 0)
+	_, ok := fr.reader.(*bufiox.DefaultReader)
+	test.Assert(t, ok)
+
+	// conn with bufiox reader
+	reader := &bufiox.DefaultReader{}
+	conn = &mockConnWithBufioxReader{bufioxReader: reader}
+	fr = newFramer(conn, 0, 0, 0)
+	test.Assert(t, fr.reader == reader)
+
+	// netpoll conn
+	conn = &mockNetpollConn{}
+	fr = newFramer(conn, 0, 0, 0)
+	_, ok = fr.reader.(*netpollBufioxReader)
+	test.Assert(t, ok)
+}
+
+func TestNetpollBufioxReader(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var (
+		validReadLength = 5
+		eofReadLength   = 8
+		validSkipLength = 5
+		eofSkipLength   = 3
+		expectedString  = "Hello"
+	)
+
+	mockReader := mocksnetpoll.NewMockReader(ctrl)
+	mockReader.EXPECT().Next(validReadLength).Return([]byte(expectedString), nil).Times(2) // readBinary also calls reader.Next
+	mockReader.EXPECT().Next(eofReadLength).Return(nil, io.EOF).Times(1)
+	mockReader.EXPECT().Skip(validSkipLength).Return(nil).Times(1)
+	mockReader.EXPECT().Skip(eofSkipLength).Return(io.EOF).Times(1)
+	mockReader.EXPECT().Release().Return(nil).Times(1)
+	reader := &netpollBufioxReader{
+		Reader: mockReader,
+	}
+	currReadLen := 0
+
+	// Next
+	p, err := reader.Next(validReadLength)
+	test.Assert(t, err == nil)
+	test.Assert(t, "Hello" == string(p))
+	test.Assert(t, validReadLength == reader.ReadLen())
+	currReadLen = validReadLength
+	p, err = reader.Next(8)
+	test.Assert(t, io.EOF == err)
+	test.Assert(t, p == nil)
+
+	// ReadBinary
+	buf := make([]byte, validReadLength)
+	n, err := reader.ReadBinary(buf)
+	test.Assert(t, err == nil)
+	test.Assert(t, validReadLength == n)
+	test.Assert(t, "Hello" == string(buf))
+	test.Assert(t, currReadLen+validReadLength == reader.ReadLen())
+	currReadLen = reader.ReadLen()
+
+	// Skip
+	err = reader.Skip(validReadLength)
+	test.Assert(t, err == nil)
+	test.Assert(t, currReadLen+validReadLength == reader.ReadLen())
+	err = reader.Skip(3)
+	test.Assert(t, io.EOF == err)
+
+	// Release
+	err = reader.Release(nil)
+	test.Assert(t, err == nil)
+	test.Assert(t, reader.ReadLen() == 0)
+}