Merge pull request from go-redis/fix/limit-alloc

Limit allocation.
This commit is contained in:
Vladimir Mihailenco 2016-11-19 14:14:48 +02:00 committed by GitHub
commit c985686167
3 changed files with 41 additions and 16 deletions

View File

@ -9,6 +9,8 @@ import (
"gopkg.in/redis.v5/internal"
)
const bytesAllocLimit = 1024 * 1024 // 1mb
const errEmptyReply = internal.RedisError("redis: reply is empty")
type MultiBulkParse func(*Reader, int64) (interface{}, error)
@ -34,16 +36,7 @@ func (p *Reader) PeekBuffered() []byte {
}
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
return readN(p.src, p.buf, n)
}
func (p *Reader) ReadLine() ([]byte, error) {
@ -225,6 +218,36 @@ func (p *Reader) readBytesValue(line []byte) ([]byte, error) {
// --------------------------------------------------------------------
func readN(r io.Reader, b []byte, n int) ([]byte, error) {
if n == 0 && b == nil {
return make([]byte, 0), nil
}
if cap(b) >= n {
b = b[:n]
_, err := io.ReadFull(r, b)
return b, err
}
b = b[:cap(b)]
pos := 0
for pos < n {
diff := n - len(b)
if diff > bytesAllocLimit {
diff = bytesAllocLimit
}
b = append(b, make([]byte, diff)...)
nn, err := io.ReadFull(r, b[pos:])
if err != nil {
return nil, err
}
pos += nn
}
return b, nil
}
func formatInt(n int64) string {
return strconv.FormatInt(n, 10)
}

View File

@ -104,7 +104,9 @@ var _ = Describe("races", func() {
})
It("should handle big vals in Get", func() {
bigVal := string(bytes.Repeat([]byte{'*'}, 1<<17)) // 128kb
C, N = 4, 100
bigVal := bytes.Repeat([]byte{'*'}, 1<<17) // 128kb
err := client.Set("key", bigVal, 0).Err()
Expect(err).NotTo(HaveOccurred())
@ -115,7 +117,7 @@ var _ = Describe("races", func() {
perform(C, func(id int) {
for i := 0; i < N; i++ {
got, err := client.Get("key").Result()
got, err := client.Get("key").Bytes()
Expect(err).NotTo(HaveOccurred())
Expect(got).To(Equal(bigVal))
}
@ -124,7 +126,8 @@ var _ = Describe("races", func() {
It("should handle big vals in Set", func() {
C, N = 4, 100
bigVal := string(bytes.Repeat([]byte{'*'}, 1<<17)) // 128kb
bigVal := bytes.Repeat([]byte{'*'}, 1<<17) // 128kb
perform(C, func(id int) {
for i := 0; i < N; i++ {

View File

@ -185,7 +185,7 @@ var _ = Describe("Client", func() {
})
It("should handle big vals", func() {
bigVal := string(bytes.Repeat([]byte{'*'}, 1<<17)) // 128kb
bigVal := bytes.Repeat([]byte{'*'}, 2e6)
err := client.Set("key", bigVal, 0).Err()
Expect(err).NotTo(HaveOccurred())
@ -194,9 +194,8 @@ var _ = Describe("Client", func() {
Expect(client.Close()).To(BeNil())
client = redis.NewClient(redisOptions())
got, err := client.Get("key").Result()
got, err := client.Get("key").Bytes()
Expect(err).NotTo(HaveOccurred())
Expect(len(got)).To(Equal(len(bigVal)))
Expect(got).To(Equal(bigVal))
})