Make proto/parser an internal package

internal/proto/proto.go

@@ -0,0 +1,122 @@
+package proto
+
+import (
+	"encoding"
+	"fmt"
+	"strconv"
+
+	"gopkg.in/redis.v4/internal/errors"
+)
+
+const (
+	ErrorReply  = '-'
+	StatusReply = '+'
+	IntReply    = ':'
+	StringReply = '$'
+	ArrayReply  = '*'
+)
+
+const defaultBufSize = 4096
+
+var errScanNil = errors.RedisError("redis: Scan(nil)")
+
+func Scan(b []byte, val interface{}) error {
+	switch v := val.(type) {
+	case nil:
+		return errScanNil
+	case *string:
+		*v = string(b)
+		return nil
+	case *[]byte:
+		*v = b
+		return nil
+	case *int:
+		var err error
+		*v, err = strconv.Atoi(string(b))
+		return err
+	case *int8:
+		n, err := strconv.ParseInt(string(b), 10, 8)
+		if err != nil {
+			return err
+		}
+		*v = int8(n)
+		return nil
+	case *int16:
+		n, err := strconv.ParseInt(string(b), 10, 16)
+		if err != nil {
+			return err
+		}
+		*v = int16(n)
+		return nil
+	case *int32:
+		n, err := strconv.ParseInt(string(b), 10, 32)
+		if err != nil {
+			return err
+		}
+		*v = int32(n)
+		return nil
+	case *int64:
+		n, err := strconv.ParseInt(string(b), 10, 64)
+		if err != nil {
+			return err
+		}
+		*v = n
+		return nil
+	case *uint:
+		n, err := strconv.ParseUint(string(b), 10, 64)
+		if err != nil {
+			return err
+		}
+		*v = uint(n)
+		return nil
+	case *uint8:
+		n, err := strconv.ParseUint(string(b), 10, 8)
+		if err != nil {
+			return err
+		}
+		*v = uint8(n)
+		return nil
+	case *uint16:
+		n, err := strconv.ParseUint(string(b), 10, 16)
+		if err != nil {
+			return err
+		}
+		*v = uint16(n)
+		return nil
+	case *uint32:
+		n, err := strconv.ParseUint(string(b), 10, 32)
+		if err != nil {
+			return err
+		}
+		*v = uint32(n)
+		return nil
+	case *uint64:
+		n, err := strconv.ParseUint(string(b), 10, 64)
+		if err != nil {
+			return err
+		}
+		*v = n
+		return nil
+	case *float32:
+		n, err := strconv.ParseFloat(string(b), 32)
+		if err != nil {
+			return err
+		}
+		*v = float32(n)
+		return err
+	case *float64:
+		var err error
+		*v, err = strconv.ParseFloat(string(b), 64)
+		return err
+	case *bool:
+		*v = len(b) == 1 && b[0] == '1'
+		return nil
+	default:
+		if bu, ok := val.(encoding.BinaryUnmarshaler); ok {
+			return bu.UnmarshalBinary(b)
+		}
+		err := fmt.Errorf(
+			"redis: can't unmarshal %T (consider implementing BinaryUnmarshaler)", val)
+		return err
+	}
+}

internal/proto/proto_test.go

@@ -0,0 +1,13 @@
+package proto_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+func TestGinkgoSuite(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "proto")
+}

internal/proto/reader.go

@@ -0,0 +1,264 @@
+package proto
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"io"
+	"strconv"
+
+	ierrors "gopkg.in/redis.v4/internal/errors"
+)
+
+type MultiBulkParse func(*Reader, int64) (interface{}, error)
+
+var errEmptyReply = errors.New("redis: reply is empty")
+
+type Reader struct {
+	src *bufio.Reader
+	buf []byte
+}
+
+func NewReader(rd io.Reader) *Reader {
+	return &Reader{
+		src: bufio.NewReader(rd),
+		buf: make([]byte, 0, defaultBufSize),
+	}
+}
+
+func (p *Reader) PeekBuffered() []byte {
+	if n := p.src.Buffered(); n != 0 {
+		b, _ := p.src.Peek(n)
+		return b
+	}
+	return nil
+}
+
+func (p *Reader) ReadN(n int) ([]byte, error) {
+	// grow internal buffer, if necessary
+	if d := n - cap(p.buf); d > 0 {
+		p.buf = p.buf[:cap(p.buf)]
+		p.buf = append(p.buf, make([]byte, d)...)
+	} else {
+		p.buf = p.buf[:n]
+	}
+
+	_, err := io.ReadFull(p.src, p.buf)
+	return p.buf, err
+}
+
+func (p *Reader) ReadLine() ([]byte, error) {
+	line, isPrefix, err := p.src.ReadLine()
+	if err != nil {
+		return nil, err
+	}
+	if isPrefix {
+		return nil, bufio.ErrBufferFull
+	}
+	if len(line) == 0 {
+		return nil, errEmptyReply
+	}
+	if isNilReply(line) {
+		return nil, ierrors.Nil
+	}
+	return line, nil
+}
+
+func (p *Reader) ReadReply(m MultiBulkParse) (interface{}, error) {
+	line, err := p.ReadLine()
+	if err != nil {
+		return nil, err
+	}
+
+	switch line[0] {
+	case ErrorReply:
+		return nil, parseErrorValue(line)
+	case StatusReply:
+		return parseStatusValue(line)
+	case IntReply:
+		return parseIntValue(line)
+	case StringReply:
+		return p.parseBytesValue(line)
+	case ArrayReply:
+		n, err := parseArrayLen(line)
+		if err != nil {
+			return nil, err
+		}
+		return m(p, n)
+	}
+	return nil, fmt.Errorf("redis: can't parse %.100q", line)
+}
+
+func (p *Reader) ReadIntReply() (int64, error) {
+	line, err := p.ReadLine()
+	if err != nil {
+		return 0, err
+	}
+	switch line[0] {
+	case ErrorReply:
+		return 0, parseErrorValue(line)
+	case IntReply:
+		return parseIntValue(line)
+	default:
+		return 0, fmt.Errorf("redis: can't parse int reply: %.100q", line)
+	}
+}
+
+func (p *Reader) ReadBytesReply() ([]byte, error) {
+	line, err := p.ReadLine()
+	if err != nil {
+		return nil, err
+	}
+	switch line[0] {
+	case ErrorReply:
+		return nil, parseErrorValue(line)
+	case StringReply:
+		return p.parseBytesValue(line)
+	case StatusReply:
+		return parseStatusValue(line)
+	default:
+		return nil, fmt.Errorf("redis: can't parse string reply: %.100q", line)
+	}
+}
+
+func (p *Reader) ReadStringReply() (string, error) {
+	b, err := p.ReadBytesReply()
+	if err != nil {
+		return "", err
+	}
+	return string(b), nil
+}
+
+func (p *Reader) ReadFloatReply() (float64, error) {
+	s, err := p.ReadStringReply()
+	if err != nil {
+		return 0, err
+	}
+	return strconv.ParseFloat(s, 64)
+}
+
+func (p *Reader) ReadArrayReply(m MultiBulkParse) (interface{}, error) {
+	line, err := p.ReadLine()
+	if err != nil {
+		return nil, err
+	}
+	switch line[0] {
+	case ErrorReply:
+		return nil, parseErrorValue(line)
+	case ArrayReply:
+		n, err := parseArrayLen(line)
+		if err != nil {
+			return nil, err
+		}
+		return m(p, n)
+	default:
+		return nil, fmt.Errorf("redis: can't parse array reply: %.100q", line)
+	}
+}
+
+func (p *Reader) ReadArrayLen() (int64, error) {
+	line, err := p.ReadLine()
+	if err != nil {
+		return 0, err
+	}
+	switch line[0] {
+	case ErrorReply:
+		return 0, parseErrorValue(line)
+	case ArrayReply:
+		return parseArrayLen(line)
+	default:
+		return 0, fmt.Errorf("redis: can't parse array reply: %.100q", line)
+	}
+}
+
+func (p *Reader) ReadScanReply() ([]string, uint64, error) {
+	n, err := p.ReadArrayLen()
+	if err != nil {
+		return nil, 0, err
+	}
+	if n != 2 {
+		return nil, 0, fmt.Errorf("redis: got %d elements in scan reply, expected 2", n)
+	}
+
+	s, err := p.ReadStringReply()
+	if err != nil {
+		return nil, 0, err
+	}
+
+	cursor, err := strconv.ParseUint(s, 10, 64)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	n, err = p.ReadArrayLen()
+	if err != nil {
+		return nil, 0, err
+	}
+
+	keys := make([]string, n)
+	for i := int64(0); i < n; i++ {
+		key, err := p.ReadStringReply()
+		if err != nil {
+			return nil, 0, err
+		}
+		keys[i] = key
+	}
+
+	return keys, cursor, err
+}
+
+func (p *Reader) parseBytesValue(line []byte) ([]byte, error) {
+	if isNilReply(line) {
+		return nil, ierrors.Nil
+	}
+
+	replyLen, err := strconv.Atoi(string(line[1:]))
+	if err != nil {
+		return nil, err
+	}
+
+	b, err := p.ReadN(replyLen + 2)
+	if err != nil {
+		return nil, err
+	}
+	return b[:replyLen], nil
+}
+
+// --------------------------------------------------------------------
+
+func formatInt(n int64) string {
+	return strconv.FormatInt(n, 10)
+}
+
+func formatUint(u uint64) string {
+	return strconv.FormatUint(u, 10)
+}
+
+func formatFloat(f float64) string {
+	return strconv.FormatFloat(f, 'f', -1, 64)
+}
+
+func isNilReply(b []byte) bool {
+	return len(b) == 3 &&
+		(b[0] == StringReply || b[0] == ArrayReply) &&
+		b[1] == '-' && b[2] == '1'
+}
+
+func parseErrorValue(line []byte) error {
+	return ierrors.RedisError(string(line[1:]))
+}
+
+func parseStatusValue(line []byte) ([]byte, error) {
+	return line[1:], nil
+}
+
+func parseIntValue(line []byte) (int64, error) {
+	return strconv.ParseInt(string(line[1:]), 10, 64)
+}
+
+func parseArrayLen(line []byte) (int64, error) {
+	if isNilReply(line) {
+		return 0, ierrors.Nil
+	}
+	return parseIntValue(line)
+}

internal/proto/reader_test.go

@@ -0,0 +1,86 @@
+package proto_test
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+	"gopkg.in/redis.v4/internal/proto"
+)
+
+var _ = Describe("Reader", func() {
+
+	It("should read n bytes", func() {
+		data, err := proto.NewReader(strings.NewReader("ABCDEFGHIJKLMNO")).ReadN(10)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(len(data)).To(Equal(10))
+		Expect(string(data)).To(Equal("ABCDEFGHIJ"))
+
+		data, err = proto.NewReader(strings.NewReader(strings.Repeat("x", 8192))).ReadN(6000)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(len(data)).To(Equal(6000))
+	})
+
+	It("should read lines", func() {
+		p := proto.NewReader(strings.NewReader("$5\r\nhello\r\n"))
+
+		data, err := p.ReadLine()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(string(data)).To(Equal("$5"))
+
+		data, err = p.ReadLine()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(string(data)).To(Equal("hello"))
+	})
+
+})
+
+func BenchmarkReader_ParseReply_Status(b *testing.B) {
+	benchmarkParseReply(b, "+OK\r\n", nil, false)
+}
+
+func BenchmarkReader_ParseReply_Int(b *testing.B) {
+	benchmarkParseReply(b, ":1\r\n", nil, false)
+}
+
+func BenchmarkReader_ParseReply_Error(b *testing.B) {
+	benchmarkParseReply(b, "-Error message\r\n", nil, true)
+}
+
+func BenchmarkReader_ParseReply_String(b *testing.B) {
+	benchmarkParseReply(b, "$5\r\nhello\r\n", nil, false)
+}
+
+func BenchmarkReader_ParseReply_Slice(b *testing.B) {
+	benchmarkParseReply(b, "*2\r\n$5\r\nhello\r\n$5\r\nworld\r\n", multiBulkParse, false)
+}
+
+func benchmarkParseReply(b *testing.B, reply string, m proto.MultiBulkParse, wanterr bool) {
+	buf := &bytes.Buffer{}
+	for i := 0; i < b.N; i++ {
+		buf.WriteString(reply)
+	}
+	p := proto.NewReader(buf)
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		_, err := p.ReadReply(m)
+		if !wanterr && err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func multiBulkParse(p *proto.Reader, n int64) (interface{}, error) {
+	vv := make([]interface{}, 0, n)
+	for i := int64(0); i < n; i++ {
+		v, err := p.ReadReply(multiBulkParse)
+		if err != nil {
+			return nil, err
+		}
+		vv = append(vv, v)
+	}
+	return vv, nil
+}

internal/proto/writebuffer.go

@@ -0,0 +1,101 @@
+package proto
+
+import (
+	"encoding"
+	"fmt"
+	"strconv"
+)
+
+type WriteBuffer struct{ b []byte }
+
+func NewWriteBuffer() *WriteBuffer {
+	return &WriteBuffer{
+		b: make([]byte, 0, defaultBufSize),
+	}
+}
+
+func (w *WriteBuffer) Len() int      { return len(w.b) }
+func (w *WriteBuffer) Bytes() []byte { return w.b }
+func (w *WriteBuffer) Reset()        { w.b = w.b[:0] }
+
+func (w *WriteBuffer) Append(args []interface{}) error {
+	w.b = append(w.b, ArrayReply)
+	w.b = strconv.AppendUint(w.b, uint64(len(args)), 10)
+	w.b = append(w.b, '\r', '\n')
+
+	for _, arg := range args {
+		if err := w.append(arg); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (w *WriteBuffer) append(val interface{}) error {
+	switch v := val.(type) {
+	case nil:
+		w.AppendString("")
+	case string:
+		w.AppendString(v)
+	case []byte:
+		w.AppendBytes(v)
+	case int:
+		w.AppendString(formatInt(int64(v)))
+	case int8:
+		w.AppendString(formatInt(int64(v)))
+	case int16:
+		w.AppendString(formatInt(int64(v)))
+	case int32:
+		w.AppendString(formatInt(int64(v)))
+	case int64:
+		w.AppendString(formatInt(v))
+	case uint:
+		w.AppendString(formatUint(uint64(v)))
+	case uint8:
+		w.AppendString(formatUint(uint64(v)))
+	case uint16:
+		w.AppendString(formatUint(uint64(v)))
+	case uint32:
+		w.AppendString(formatUint(uint64(v)))
+	case uint64:
+		w.AppendString(formatUint(v))
+	case float32:
+		w.AppendString(formatFloat(float64(v)))
+	case float64:
+		w.AppendString(formatFloat(v))
+	case bool:
+		if v {
+			w.AppendString("1")
+		} else {
+			w.AppendString("0")
+		}
+	default:
+		if bm, ok := val.(encoding.BinaryMarshaler); ok {
+			bb, err := bm.MarshalBinary()
+			if err != nil {
+				return err
+			}
+			w.AppendBytes(bb)
+		} else {
+			return fmt.Errorf(
+				"redis: can't marshal %T (consider implementing encoding.BinaryMarshaler)", val)
+		}
+	}
+	return nil
+}
+
+func (w *WriteBuffer) AppendString(s string) {
+	w.b = append(w.b, StringReply)
+	w.b = strconv.AppendUint(w.b, uint64(len(s)), 10)
+	w.b = append(w.b, '\r', '\n')
+	w.b = append(w.b, s...)
+	w.b = append(w.b, '\r', '\n')
+}
+
+func (w *WriteBuffer) AppendBytes(p []byte) {
+	w.b = append(w.b, StringReply)
+	w.b = strconv.AppendUint(w.b, uint64(len(p)), 10)
+	w.b = append(w.b, '\r', '\n')
+	w.b = append(w.b, p...)
+	w.b = append(w.b, '\r', '\n')
+}

internal/proto/writebuffer_test.go

@@ -0,0 +1,62 @@
+package proto_test
+
+import (
+	"testing"
+	"time"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+	"gopkg.in/redis.v4/internal/proto"
+)
+
+var _ = Describe("WriteBuffer", func() {
+	var buf *proto.WriteBuffer
+
+	BeforeEach(func() {
+		buf = proto.NewWriteBuffer()
+	})
+
+	It("should reset", func() {
+		buf.AppendString("string")
+		Expect(buf.Len()).To(Equal(12))
+		buf.Reset()
+		Expect(buf.Len()).To(Equal(0))
+	})
+
+	It("should append args", func() {
+		err := buf.Append([]interface{}{
+			"string",
+			12,
+			34.56,
+			[]byte{'b', 'y', 't', 'e', 's'},
+			true,
+			nil,
+		})
+		Expect(err).NotTo(HaveOccurred())
+		Expect(buf.Bytes()).To(Equal([]byte("*6\r\n" +
+			"$6\r\nstring\r\n" +
+			"$2\r\n12\r\n" +
+			"$5\r\n34.56\r\n" +
+			"$5\r\nbytes\r\n" +
+			"$1\r\n1\r\n" +
+			"$0\r\n" +
+			"\r\n")))
+	})
+
+	It("should append marshalable args", func() {
+		err := buf.Append([]interface{}{time.Unix(1414141414, 0)})
+		Expect(err).NotTo(HaveOccurred())
+		Expect(buf.Len()).To(Equal(26))
+	})
+
+})
+
+func BenchmarkWriteBuffer_Append(b *testing.B) {
+	buf := proto.NewWriteBuffer()
+	args := []interface{}{"hello", "world", "foo", "bar"}
+
+	for i := 0; i < b.N; i++ {
+		buf.Append(args)
+		buf.Reset()
+	}
+}