buffer reads

perf is now roughly the same as pgbouncer

lib/zap/conn.go

@@ -7,21 +7,22 @@ import (
 	"net"
 )
 
+const pktBufSize = 4096
+
 type Conn struct {
 	conn net.Conn
 	w    io.Writer
-	r    io.Reader
 
-	buffers net.Buffers
+	writeBuf net.Buffers
 
-	byteBuf [1]byte
+	pktBuf  [pktBufSize]byte
+	readBuf []byte
 }
 
 func WrapNetConn(conn net.Conn) *Conn {
 	return &Conn{
 		conn: conn,
 		w:    conn,
-		r:    conn,
 	}
 }
 
@@ -29,7 +30,6 @@ func (T *Conn) EnableSSLClient(config *tls.Config) error {
 	sslConn := tls.Client(T.conn, config)
 	T.conn = sslConn
 	T.w = sslConn
-	T.r = sslConn
 	return sslConn.Handshake()
 }
 
@@ -37,29 +37,61 @@ func (T *Conn) EnableSSLServer(config *tls.Config) error {
 	sslConn := tls.Server(T.conn, config)
 	T.conn = sslConn
 	T.w = sslConn
-	T.r = sslConn
 	return sslConn.Handshake()
 }
 
 func (T *Conn) flush() error {
-	if len(T.buffers) == 0 {
+	if len(T.writeBuf) == 0 {
 		return nil
 	}
 
-	_, err := T.buffers.WriteTo(T.w)
-	T.buffers = T.buffers[0:]
+	_, err := T.writeBuf.WriteTo(T.w)
+	T.writeBuf = T.writeBuf[0:]
 	return err
 }
 
+func (T *Conn) read(buf []byte) (n int, err error) {
+	for {
+		if len(T.readBuf) > 0 {
+			cn := copy(buf, T.readBuf)
+			buf = buf[cn:]
+			T.readBuf = T.readBuf[cn:]
+			n += cn
+		}
+
+		if len(buf) == 0 {
+			return
+		}
+
+		if len(buf) > len(T.pktBuf) {
+			var rn int
+			rn, err = T.conn.Read(buf)
+			n += rn
+			if err != nil {
+				return
+			}
+			buf = buf[rn:]
+		} else {
+			var rn int
+			rn, err = T.conn.Read(T.pktBuf[:])
+			if err != nil {
+				return
+			}
+			T.readBuf = T.pktBuf[:rn]
+		}
+	}
+}
+
 func (T *Conn) ReadByte() (byte, error) {
 	if err := T.flush(); err != nil {
 		return 0, err
 	}
-	_, err := io.ReadFull(T.r, T.byteBuf[:])
+	var b [1]byte
+	_, err := T.read(b[:])
 	if err != nil {
 		return 0, err
 	}
-	return T.byteBuf[0], nil
+	return b[0], nil
 }
 
 func (T *Conn) ReadPacket(typed bool) (Packet, error) {
@@ -68,12 +100,12 @@ func (T *Conn) ReadPacket(typed bool) (Packet, error) {
 	}
 	packet := NewPacket(0)
 	if typed {
-		_, err := io.ReadFull(T.r, packet)
+		_, err := T.read(packet)
 		if err != nil {
 			return nil, err
 		}
 	} else {
-		_, err := io.ReadFull(T.r, packet[1:])
+		_, err := T.read(packet[1:])
 		if err != nil {
 			return nil, err
 		}
@@ -81,7 +113,7 @@ func (T *Conn) ReadPacket(typed bool) (Packet, error) {
 
 	length := binary.BigEndian.Uint32(packet[1:])
 	packet = packet.Grow(int(length) - 4)
-	_, err := io.ReadFull(T.r, packet.Payload())
+	_, err := T.read(packet.Payload())
 	if err != nil {
 		return nil, err
 	}
@@ -89,12 +121,12 @@ func (T *Conn) ReadPacket(typed bool) (Packet, error) {
 }
 
 func (T *Conn) WriteByte(b byte) error {
-	T.buffers = append(T.buffers, []byte{b})
+	T.writeBuf = append(T.writeBuf, []byte{b})
 	return nil
 }
 
 func (T *Conn) WritePacket(packet Packet) error {
-	T.buffers = append(T.buffers, packet.Bytes())
+	T.writeBuf = append(T.writeBuf, packet.Bytes())
 	return nil
 }