浏览代码

conn: always allocate 9 bytes for the header buf

If we are in protocol 2 mode we will allocate an 8 byte header buffer,
if the server does not support protocol 2 it will respond with a 9
byte header. Always allocate a 9 byte header but only use 8 if we are
in protocol 2 mode.
Chris Bannister 9 年之前
父节点
当前提交
bf3e78b459
共有 2 个文件被更改,包括 6 次插入10 次删除
  1. 2 8
      conn.go
  2. 4 2
      frame.go

+ 2 - 8
conn.go

@@ -127,7 +127,7 @@ type Conn struct {
 	timeout time.Duration
 	cfg     *ConnConfig
 
-	headerBuf []byte
+	headerBuf [maxFrameHeaderSize]byte
 
 	streams *streams.IDGenerator
 	mu      sync.RWMutex
@@ -175,11 +175,6 @@ func Connect(host *HostInfo, addr string, cfg *ConnConfig,
 		return nil, err
 	}
 
-	headerSize := 8
-	if cfg.ProtoVersion > protoVersion2 {
-		headerSize = 9
-	}
-
 	c := &Conn{
 		conn:         conn,
 		r:            bufio.NewReader(conn),
@@ -191,7 +186,6 @@ func Connect(host *HostInfo, addr string, cfg *ConnConfig,
 		errorHandler: errorHandler,
 		compressor:   cfg.Compressor,
 		auth:         cfg.Authenticator,
-		headerBuf:    make([]byte, headerSize),
 		quit:         make(chan struct{}),
 		session:      session,
 		streams:      streams.New(cfg.ProtoVersion),
@@ -445,7 +439,7 @@ func (c *Conn) recv() error {
 	}
 
 	// were just reading headers over and over and copy bodies
-	head, err := readHeader(c.r, c.headerBuf)
+	head, err := readHeader(c.r, c.headerBuf[:])
 	if err != nil {
 		return err
 	}

+ 4 - 2
frame.go

@@ -231,6 +231,8 @@ var (
 	ErrFrameTooBig = errors.New("frame length is bigger than the maximum allowed")
 )
 
+const maxFrameHeaderSize = 9
+
 func writeInt(p []byte, n int32) {
 	p[0] = byte(n >> 24)
 	p[1] = byte(n >> 16)
@@ -366,7 +368,7 @@ func readHeader(r io.Reader, p []byte) (head frameHeader, err error) {
 	head.flags = p[1]
 
 	if version > protoVersion2 {
-		if len(p) < 9 {
+		if len(p) != 9 {
 			return frameHeader{}, fmt.Errorf("not enough bytes to read header require 9 got: %d", len(p))
 		}
 
@@ -374,7 +376,7 @@ func readHeader(r io.Reader, p []byte) (head frameHeader, err error) {
 		head.op = frameOp(p[4])
 		head.length = int(readInt(p[5:]))
 	} else {
-		if len(p) < 8 {
+		if len(p) != 8 {
 			return frameHeader{}, fmt.Errorf("not enough bytes to read header require 8 got: %d", len(p))
 		}