Optimize reading big values

internal/proto/elastic_reader.go

@@ -183,3 +183,23 @@ func (b *ElasticBufReader) grow(n int) {
 		b.buf = append(b.buf, make([]byte, d)...)
 	}
 }
+
+func (b *ElasticBufReader) Read(p []byte) (n int, err error) {
+	if len(p) == 0 {
+		return 0, b.readErr()
+	}
+
+	if b.r != b.w {
+		// copy as much as we can
+		n = copy(p, b.buf[b.r:b.w])
+		b.r += n
+		return n, nil
+	}
+
+	if b.err != nil {
+		return 0, b.readErr()
+	}
+
+	n, b.err = b.rd.Read(p)
+	return n, b.readErr()
+}

internal/proto/reader.go

@@ -54,10 +54,6 @@ func (r Reader) Bytes() []byte {
 	return r.src.Bytes()
 }
 
-func (r Reader) ReadN(n int) ([]byte, error) {
-	return r.src.ReadN(n)
-}
-
 func (r Reader) ReadLine() ([]byte, error) {
 	line, err := r.src.ReadLine()
 	if err != nil {
@@ -82,11 +78,11 @@ func (r Reader) ReadReply(m MultiBulkParse) (interface{}, error) {
 	case ErrorReply:
 		return nil, ParseErrorReply(line)
 	case StatusReply:
-		return parseTmpStatusReply(line), nil
+		return string(line[1:]), nil
 	case IntReply:
 		return util.ParseInt(line[1:], 10, 64)
 	case StringReply:
-		return r.readTmpBytesReply(line)
+		return r.readStringReply(line)
 	case ArrayReply:
 		n, err := parseArrayLen(line)
 		if err != nil {
@@ -112,47 +108,42 @@ func (r Reader) ReadIntReply() (int64, error) {
 	}
 }
 
-func (r Reader) ReadTmpBytesReply() ([]byte, error) {
+func (r Reader) ReadString() (string, error) {
 	line, err := r.ReadLine()
-	if err != nil {
-		return nil, err
-	}
-	switch line[0] {
-	case ErrorReply:
-		return nil, ParseErrorReply(line)
-	case StringReply:
-		return r.readTmpBytesReply(line)
-	case StatusReply:
-		return parseTmpStatusReply(line), nil
-	default:
-		return nil, fmt.Errorf("redis: can't parse string reply: %.100q", line)
-	}
-}
-
-func (r Reader) ReadBytesReply() ([]byte, error) {
-	b, err := r.ReadTmpBytesReply()
-	if err != nil {
-		return nil, err
-	}
-	cp := make([]byte, len(b))
-	copy(cp, b)
-	return cp, nil
-}
-
-func (r Reader) ReadStringReply() (string, error) {
-	b, err := r.ReadTmpBytesReply()
 	if err != nil {
 		return "", err
 	}
-	return string(b), nil
+	switch line[0] {
+	case ErrorReply:
+		return "", ParseErrorReply(line)
+	case StringReply:
+		return r.readStringReply(line)
+	case StatusReply:
+		return string(line[1:]), nil
+	case IntReply:
+		return string(line[1:]), nil
+	default:
+		return "", fmt.Errorf("redis: can't parse reply=%.100q reading string", line)
+	}
 }
 
-func (r Reader) ReadFloatReply() (float64, error) {
-	b, err := r.ReadTmpBytesReply()
-	if err != nil {
-		return 0, err
+func (r Reader) readStringReply(line []byte) (string, error) {
+	if isNilReply(line) {
+		return "", Nil
 	}
-	return util.ParseFloat(b, 64)
+
+	replyLen, err := strconv.Atoi(string(line[1:]))
+	if err != nil {
+		return "", err
+	}
+
+	b := make([]byte, replyLen+2)
+	_, err = io.ReadFull(r.src, b)
+	if err != nil {
+		return "", err
+	}
+
+	return util.BytesToString(b[:replyLen]), nil
 }
 
 func (r Reader) ReadArrayReply(m MultiBulkParse) (interface{}, error) {
@@ -210,7 +201,7 @@ func (r Reader) ReadScanReply() ([]string, uint64, error) {
 
 	keys := make([]string, n)
 	for i := int64(0); i < n; i++ {
-		key, err := r.ReadStringReply()
+		key, err := r.ReadString()
 		if err != nil {
 			return nil, 0, err
 		}
@@ -220,7 +211,48 @@ func (r Reader) ReadScanReply() ([]string, uint64, error) {
 	return keys, cursor, err
 }
 
-func (r Reader) readTmpBytesReply(line []byte) ([]byte, error) {
+func (r Reader) ReadInt() (int64, error) {
+	b, err := r.readTmpBytesReply()
+	if err != nil {
+		return 0, err
+	}
+	return util.ParseInt(b, 10, 64)
+}
+
+func (r Reader) ReadUint() (uint64, error) {
+	b, err := r.readTmpBytesReply()
+	if err != nil {
+		return 0, err
+	}
+	return util.ParseUint(b, 10, 64)
+}
+
+func (r Reader) ReadFloatReply() (float64, error) {
+	b, err := r.readTmpBytesReply()
+	if err != nil {
+		return 0, err
+	}
+	return util.ParseFloat(b, 64)
+}
+
+func (r Reader) readTmpBytesReply() ([]byte, error) {
+	line, err := r.ReadLine()
+	if err != nil {
+		return nil, err
+	}
+	switch line[0] {
+	case ErrorReply:
+		return nil, ParseErrorReply(line)
+	case StringReply:
+		return r._readTmpBytesReply(line)
+	case StatusReply:
+		return line[1:], nil
+	default:
+		return nil, fmt.Errorf("redis: can't parse string reply: %.100q", line)
+	}
+}
+
+func (r Reader) _readTmpBytesReply(line []byte) ([]byte, error) {
 	if isNilReply(line) {
 		return nil, Nil
 	}
@@ -230,29 +262,13 @@ func (r Reader) readTmpBytesReply(line []byte) ([]byte, error) {
 		return nil, err
 	}
 
-	b, err := r.ReadN(replyLen + 2)
+	b, err := r.src.ReadN(replyLen + 2)
 	if err != nil {
 		return nil, err
 	}
 	return b[:replyLen], nil
 }
 
-func (r Reader) ReadInt() (int64, error) {
-	b, err := r.ReadTmpBytesReply()
-	if err != nil {
-		return 0, err
-	}
-	return util.ParseInt(b, 10, 64)
-}
-
-func (r Reader) ReadUint() (uint64, error) {
-	b, err := r.ReadTmpBytesReply()
-	if err != nil {
-		return 0, err
-	}
-	return util.ParseUint(b, 10, 64)
-}
-
 func isNilReply(b []byte) bool {
 	return len(b) == 3 &&
 		(b[0] == StringReply || b[0] == ArrayReply) &&
@@ -263,10 +279,6 @@ func ParseErrorReply(line []byte) error {
 	return RedisError(string(line[1:]))
 }
 
-func parseTmpStatusReply(line []byte) []byte {
-	return line[1:]
-}
-
 func parseArrayLen(line []byte) (int64, error) {
 	if isNilReply(line) {
 		return 0, Nil

internal/proto/reader_test.go

@@ -6,42 +6,12 @@ import (
 	"testing"
 
 	"github.com/go-redis/redis/internal/proto"
-
-	. "github.com/onsi/ginkgo"
-	. "github.com/onsi/gomega"
 )
 
 func newReader(s string) proto.Reader {
 	return proto.NewReader(proto.NewElasticBufReader(strings.NewReader(s)))
 }
 
-var _ = Describe("Reader", func() {
-
-	It("should read n bytes", func() {
-		data, err := newReader("ABCDEFGHIJKLMNO").ReadN(10)
-		Expect(err).NotTo(HaveOccurred())
-		Expect(len(data)).To(Equal(10))
-		Expect(string(data)).To(Equal("ABCDEFGHIJ"))
-
-		data, err = newReader(strings.Repeat("x", 8192)).ReadN(6000)
-		Expect(err).NotTo(HaveOccurred())
-		Expect(len(data)).To(Equal(6000))
-	})
-
-	It("should read lines", func() {
-		r := newReader("$5\r\nhello\r\n")
-
-		data, err := r.ReadLine()
-		Expect(err).NotTo(HaveOccurred())
-		Expect(string(data)).To(Equal("$5"))
-
-		data, err = r.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)
 }