Cleanup buffers manipulation

internal/pool/conn.go

@@ -8,15 +8,20 @@ import (
 	"github.com/go-redis/redis/internal/proto"
 )
 
+func makeBuffer() []byte {
+	const defaulBufSize = 4096
+	return make([]byte, defaulBufSize)
+}
+
 var noDeadline = time.Time{}
 
 type Conn struct {
 	netConn net.Conn
 
-	Rd proto.Reader
-	wb *proto.WriteBuffer
-
-	concurrentReadWrite bool
+	buf      []byte
+	rd       proto.Reader
+	rdLocked bool
+	wb       *proto.WriteBuffer
 
 	InitedAt time.Time
 	pooled   bool
@@ -26,9 +31,9 @@ type Conn struct {
 func NewConn(netConn net.Conn) *Conn {
 	cn := &Conn{
 		netConn: netConn,
+		buf:     makeBuffer(),
 	}
-	buf := proto.NewElasticBufReader(netConn)
-	cn.Rd = proto.NewReader(buf)
+	cn.rd = proto.NewReader(proto.NewElasticBufReader(netConn))
 	cn.wb = proto.NewWriteBuffer()
 	cn.SetUsedAt(time.Now())
 	return cn
@@ -44,27 +49,25 @@ func (cn *Conn) SetUsedAt(tm time.Time) {
 
 func (cn *Conn) SetNetConn(netConn net.Conn) {
 	cn.netConn = netConn
-	cn.Rd.Reset(netConn)
+	cn.rd.Reset(netConn)
 }
 
-func (cn *Conn) SetReadTimeout(timeout time.Duration) {
+func (cn *Conn) setReadTimeout(timeout time.Duration) error {
 	now := time.Now()
 	cn.SetUsedAt(now)
 	if timeout > 0 {
-		cn.netConn.SetReadDeadline(now.Add(timeout))
-	} else {
-		cn.netConn.SetReadDeadline(noDeadline)
+		return cn.netConn.SetReadDeadline(now.Add(timeout))
 	}
+	return cn.netConn.SetReadDeadline(noDeadline)
 }
 
-func (cn *Conn) SetWriteTimeout(timeout time.Duration) {
+func (cn *Conn) setWriteTimeout(timeout time.Duration) error {
 	now := time.Now()
 	cn.SetUsedAt(now)
 	if timeout > 0 {
-		cn.netConn.SetWriteDeadline(now.Add(timeout))
-	} else {
-		cn.netConn.SetWriteDeadline(noDeadline)
+		return cn.netConn.SetWriteDeadline(now.Add(timeout))
 	}
+	return cn.netConn.SetWriteDeadline(noDeadline)
 }
 
 func (cn *Conn) Write(b []byte) (int, error) {
@@ -75,28 +78,43 @@ func (cn *Conn) RemoteAddr() net.Addr {
 	return cn.netConn.RemoteAddr()
 }
 
-func (cn *Conn) EnableConcurrentReadWrite() {
-	cn.concurrentReadWrite = true
-	cn.wb.ResetBuffer(make([]byte, 4096))
+func (cn *Conn) LockReaderBuffer() {
+	cn.rdLocked = true
+	cn.rd.ResetBuffer(makeBuffer())
 }
 
-func (cn *Conn) PrepareWriteBuffer() *proto.WriteBuffer {
-	if cn.concurrentReadWrite {
-		cn.wb.Reset()
-	} else {
-		cn.wb.ResetBuffer(cn.Rd.Buffer())
-	}
-	return cn.wb
-}
+func (cn *Conn) WithReader(timeout time.Duration, fn func(rd proto.Reader) error) error {
+	_ = cn.setReadTimeout(timeout)
 
-func (cn *Conn) FlushWriteBuffer(wb *proto.WriteBuffer) error {
-	_, err := cn.netConn.Write(wb.Bytes())
-	if !cn.concurrentReadWrite {
-		cn.Rd.ResetBuffer(wb.Buffer())
+	if !cn.rdLocked {
+		cn.rd.ResetBuffer(cn.buf)
 	}
+
+	err := fn(cn.rd)
+
+	if !cn.rdLocked {
+		cn.buf = cn.rd.Buffer()
+	}
+
 	return err
 }
 
+func (cn *Conn) WithWriter(timeout time.Duration, fn func(wb *proto.WriteBuffer) error) error {
+	_ = cn.setWriteTimeout(timeout)
+
+	cn.wb.ResetBuffer(cn.buf)
+
+	firstErr := fn(cn.wb)
+
+	_, err := cn.netConn.Write(cn.wb.Bytes())
+	cn.buf = cn.wb.Buffer()
+	if err != nil && firstErr == nil {
+		firstErr = err
+	}
+
+	return firstErr
+}
+
 func (cn *Conn) Close() error {
 	return cn.netConn.Close()
 }