瀏覽代碼

Handle RESP strings longer than bufio.Reader size

Gary Burd 7 年之前
父節點
當前提交
02dc2736db
共有 2 個文件被更改,包括 50 次插入1 次删除
  1. 12 1
      redis/conn.go
  2. 38 0
      redis/conn_test.go

+ 12 - 1
redis/conn.go

@@ -427,10 +427,21 @@ func (pe protocolError) Error() string {
 	return fmt.Sprintf("redigo: %s (possible server error or unsupported concurrent read by application)", string(pe))
 }
 
+// readLine reads a line of input from the RESP stream.
 func (c *conn) readLine() ([]byte, error) {
+	// To avoid allocations, attempt to read the line using ReadSlice. This
+	// call typically succeeds. The known case where the call fails is when
+	// reading the output from the MONITOR command.
 	p, err := c.br.ReadSlice('\n')
 	if err == bufio.ErrBufferFull {
-		return nil, protocolError("long response line")
+		// The line does not fit in the bufio.Reader's buffer. Fall back to
+		// allocating a buffer for the line.
+		buf := append([]byte{}, p...)
+		for err == bufio.ErrBufferFull {
+			p, err = c.br.ReadSlice('\n')
+			buf = append(buf, p...)
+		}
+		p = buf
 	}
 	if err != nil {
 		return nil, err

+ 38 - 0
redis/conn_test.go

@@ -169,6 +169,10 @@ var readTests = []struct {
 		"+PONG\r\n",
 		"PONG",
 	},
+	{
+		"+OK\n\n", // no \r
+		errorSentinel,
+	},
 	{
 		"@OK\r\n",
 		errorSentinel,
@@ -206,6 +210,11 @@ var readTests = []struct {
 		[]interface{}{[]byte("foo"), nil, []byte("bar")},
 	},
 
+	{
+		// "" is not a valid length
+		"$\r\nfoobar\r\n",
+		errorSentinel,
+	},
 	{
 		// "x" is not a valid length
 		"$x\r\nfoobar\r\n",
@@ -216,6 +225,11 @@ var readTests = []struct {
 		"$-2\r\n",
 		errorSentinel,
 	},
+	{
+		// ""  is not a valid integer
+		":\r\n",
+		errorSentinel,
+	},
 	{
 		// "x"  is not a valid integer
 		":x\r\n",
@@ -258,6 +272,30 @@ func TestRead(t *testing.T) {
 	}
 }
 
+func TestReadString(t *testing.T) {
+	// n is value of bufio.defaultBufSize
+	const n = 4096
+
+	// Test read string lengths near bufio.Reader buffer boundaries.
+	testRanges := [][2]int{{0, 64}, {n - 64, n + 64}, {2*n - 64, 2*n + 64}}
+
+	p := make([]byte, 2*n+64)
+	for i := range p {
+		p[i] = byte('a' + i%26)
+	}
+	s := string(p)
+
+	for _, r := range testRanges {
+		for i := r[0]; i < r[1]; i++ {
+			c, _ := redis.Dial("", "", dialTestConn("+"+s[:i]+"\r\n", nil))
+			actual, err := c.Receive()
+			if err != nil || actual != s[:i] {
+				t.Fatalf("Receive(string len %d) -> err=%v, equal=%v", i, err, actual != s[:i])
+			}
+		}
+	}
+}
+
 var testCommands = []struct {
 	args     []interface{}
 	expected interface{}