Cleanups

internal/pool/conn.go

@@ -14,7 +14,9 @@ type Conn struct {
 	netConn net.Conn
 
 	Rd *proto.Reader
-	WB *proto.WriteBuffer
+	wb *proto.WriteBuffer
+
+	concurrentReadWrite bool
 
 	Inited bool
 	usedAt atomic.Value
@@ -26,7 +28,7 @@ func NewConn(netConn net.Conn) *Conn {
 	}
 	buf := proto.NewElasticBufReader(netConn)
 	cn.Rd = proto.NewReader(buf)
-	cn.WB = proto.NewWriteBuffer(buf)
+	cn.wb = proto.NewWriteBuffer()
 	cn.SetUsedAt(time.Now())
 	return cn
 }
@@ -76,6 +78,27 @@ 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) PrepareWriteBuffer() *proto.WriteBuffer {
+	if !cn.concurrentReadWrite {
+		cn.wb.ResetBuffer(cn.Rd.Buffer())
+	}
+	cn.wb.Reset()
+	return cn.wb
+}
+
+func (cn *Conn) FlushWriteBuffer(wb *proto.WriteBuffer) error {
+	_, err := cn.netConn.Write(wb.Bytes())
+	if !cn.concurrentReadWrite {
+		cn.Rd.ResetBuffer(wb.Buffer())
+	}
+	return err
+}
+
 func (cn *Conn) Close() error {
 	return cn.netConn.Close()
 }

internal/proto/reader.go

@@ -44,6 +44,14 @@ func (r *Reader) Reset(rd io.Reader) {
 	r.src.Reset(rd)
 }
 
+func (r *Reader) Buffer() []byte {
+	return r.src.Buffer()
+}
+
+func (r *Reader) ResetBuffer(buf []byte) {
+	r.src.ResetBuffer(buf)
+}
+
 func (r *Reader) Bytes() []byte {
 	return r.src.Bytes()
 }

internal/proto/write_buffer.go

@@ -7,14 +7,11 @@ import (
 )
 
 type WriteBuffer struct {
-	rb  *ElasticBufReader
 	buf []byte
 }
 
-func NewWriteBuffer(rb *ElasticBufReader) *WriteBuffer {
-	return &WriteBuffer{
-		rb: rb,
-	}
+func NewWriteBuffer() *WriteBuffer {
+	return &WriteBuffer{}
 }
 
 func (w *WriteBuffer) Len() int {
@@ -25,26 +22,16 @@ func (w *WriteBuffer) Bytes() []byte {
 	return w.buf
 }
 
-func (w *WriteBuffer) AllocBuffer() {
-	w.rb = nil
-	w.buf = make([]byte, defaultBufSize)
-}
-
 func (w *WriteBuffer) Reset() {
-	if w.rb != nil {
-		w.buf = w.rb.Buffer()[:0]
-	} else {
-		w.buf = w.buf[:0]
-	}
+	w.buf = w.buf[:0]
 }
 
-func (w *WriteBuffer) Flush() []byte {
-	b := w.buf
-	if w.rb != nil {
-		w.rb.ResetBuffer(w.buf[:cap(w.buf)])
-		w.buf = nil
-	}
-	return b
+func (w *WriteBuffer) Buffer() []byte {
+	return w.buf[:cap(w.buf)]
+}
+
+func (w *WriteBuffer) ResetBuffer(buf []byte) {
+	w.buf = buf[:0]
 }
 
 func (w *WriteBuffer) Append(args []interface{}) error {

internal/proto/write_buffer_test.go

@@ -1,7 +1,6 @@
 package proto_test
 
 import (
-	"strings"
 	"testing"
 	"time"
 
@@ -15,7 +14,7 @@ var _ = Describe("WriteBuffer", func() {
 	var buf *proto.WriteBuffer
 
 	BeforeEach(func() {
-		buf = proto.NewWriteBuffer(proto.NewElasticBufReader(strings.NewReader("")))
+		buf = proto.NewWriteBuffer()
 	})
 
 	It("should reset", func() {
@@ -54,7 +53,7 @@ var _ = Describe("WriteBuffer", func() {
 })
 
 func BenchmarkWriteBuffer_Append(b *testing.B) {
-	buf := proto.NewWriteBuffer(proto.NewElasticBufReader(strings.NewReader("")))
+	buf := proto.NewWriteBuffer()
 	args := []interface{}{"hello", "world", "foo", "bar"}
 
 	for i := 0; i < b.N; i++ {