Merge pull request #65 from weedge/master

fix: []byte nil resp && feature: add WriteBulkFrom with bufio.Writer for Conn interface WriteBulkFrom func

Author: tidwall

redcon.go

@@ -113,6 +113,8 @@ type Conn interface {
 	PeekPipeline() []Command
 	// NetConn returns the base net.Conn connection
 	NetConn() net.Conn
+	// WriteBulkFrom write bulk from io.Reader, size n
+	WriteBulkFrom(n int64, rb io.Reader)
 }
 
 // NewServer returns a new Redcon server configured on "tcp" network net.
@@ -494,6 +496,9 @@ func (c *conn) PeekPipeline() []Command {
 func (c *conn) NetConn() net.Conn {
 	return c.conn
 }
+func (c *conn) WriteBulkFrom(n int64, rb io.Reader) {
+	c.wr.WriteBulkFrom(n, rb)
+}
 
 // BaseWriter returns the underlying connection writer, if any
 func BaseWriter(c Conn) *Writer {
@@ -589,13 +594,27 @@ type Writer struct {
 	w   io.Writer
 	b   []byte
 	err error
+
+	// buff use io buffer write to w(io.Writer)
+	// for io.Copy r(io.Reader) to w(io.Writer)
+	buff *bufio.Writer
 }
 
 // NewWriter creates a new RESP writer.
 func NewWriter(wr io.Writer) *Writer {
 	return &Writer{
-		w: wr,
+		w:    wr,
+		buff: bufio.NewWriter(wr),
+	}
+}
+
+func (w *Writer) WriteBulkFrom(n int64, r io.Reader) {
+	if w != nil && w.err != nil {
+		return
 	}
+	w.buff.Write(appendPrefix(w.b, '$', n))
+	io.Copy(w.buff, r)
+	w.buff.Write([]byte{'\r', '\n'})
 }
 
 // WriteNull writes a null to the client
@@ -656,6 +675,10 @@ func (w *Writer) SetBuffer(raw []byte) {
 
 // Flush writes all unflushed Write* calls to the underlying writer.
 func (w *Writer) Flush() error {
+	if w.buff != nil {
+		w.buff.Flush()
+	}
+
 	if w.err != nil {
 		return w.err
 	}

redcon_test.go

@@ -365,6 +365,35 @@ func testServerNetwork(t *testing.T, network, laddr string) {
 	<-done
 }
 
+func TestConnImpl(t *testing.T) {
+	var i interface{} = &conn{}
+	if _, ok := i.(Conn); !ok {
+		t.Fatalf("conn does not implement Conn interface")
+	}
+}
+
+func TestWriteBulkFrom(t *testing.T) {
+	wbuf := &bytes.Buffer{}
+	wr := NewWriter(wbuf)
+	rbuf := &bytes.Buffer{}
+	testStr := "hello world"
+	rbuf.WriteString(testStr)
+	wr.WriteBulkFrom(int64(len(testStr)), rbuf)
+	wr.Flush()
+	if wbuf.String() != fmt.Sprintf("$%d\r\n%s\r\n", len(testStr), testStr) {
+		t.Fatal("failed")
+	}
+	wbuf.Reset()
+	testStr1 := "hi world"
+	rbuf.WriteString(testStr1)
+	wr.WriteBulkFrom(int64(len(testStr1)), rbuf)
+	wr.Flush()
+	if wbuf.String() != fmt.Sprintf("$%d\r\n%s\r\n", len(testStr1), testStr1) {
+		t.Fatal("failed")
+	}
+	wbuf.Reset()
+}
+
 func TestWriter(t *testing.T) {
 	buf := &bytes.Buffer{}
 	wr := NewWriter(buf)

resp.go

@@ -556,18 +556,19 @@ type Marshaler interface {
 }
 
 // AppendAny appends any type to valid Redis type.
-//   nil             -> null
-//   error           -> error (adds "ERR " when first word is not uppercase)
-//   string          -> bulk-string
-//   numbers         -> bulk-string
-//   []byte          -> bulk-string
-//   bool            -> bulk-string ("0" or "1")
-//   slice           -> array
-//   map             -> array with key/value pairs
-//   SimpleString    -> string
-//   SimpleInt       -> integer
-//   Marshaler       -> raw bytes
-//   everything-else -> bulk-string representation using fmt.Sprint()
+//
+//	nil             -> null
+//	error           -> error (adds "ERR " when first word is not uppercase)
+//	string          -> bulk-string
+//	numbers         -> bulk-string
+//	[]byte          -> bulk-string
+//	bool            -> bulk-string ("0" or "1")
+//	slice           -> array
+//	map             -> array with key/value pairs
+//	SimpleString    -> string
+//	SimpleInt       -> integer
+//	Marshaler       -> raw bytes
+//	everything-else -> bulk-string representation using fmt.Sprint()
 func AppendAny(b []byte, v interface{}) []byte {
 	switch v := v.(type) {
 	case SimpleString:
@@ -583,7 +584,11 @@ func AppendAny(b []byte, v interface{}) []byte {
 	case string:
 		b = AppendBulkString(b, v)
 	case []byte:
-		b = AppendBulk(b, v)
+		if v == nil {
+			b = AppendNull(b)
+		} else {
+			b = AppendBulk(b, v)
+		}
 	case bool:
 		if v {
 			b = AppendBulkString(b, "1")