mirror of
https://github.com/redis/go-redis.git
synced 2025-07-29 17:41:15 +03:00
Use bufio.Writer
This commit is contained in:
Vladimir Mihailenco
@ -1,205 +0,0 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
)
|
||||
|
||||
const defaultBufSize = 4096
|
||||
|
||||
// ElasticBufReader is like bufio.Reader but instead of returning ErrBufferFull
|
||||
// it automatically grows the buffer.
|
||||
type ElasticBufReader struct {
|
||||
buf []byte
|
||||
rd io.Reader // reader provided by the client
|
||||
r, w int // buf read and write positions
|
||||
err error
|
||||
}
|
||||
|
||||
func NewElasticBufReader(rd io.Reader) *ElasticBufReader {
|
||||
return &ElasticBufReader{
|
||||
rd: rd,
|
||||
}
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) Reset(rd io.Reader) {
|
||||
b.rd = rd
|
||||
b.r, b.w = 0, 0
|
||||
b.err = nil
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) Buffer() []byte {
|
||||
return b.buf
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) ResetBuffer(buf []byte) {
|
||||
b.buf = buf
|
||||
b.r, b.w = 0, 0
|
||||
b.err = nil
|
||||
}
|
||||
|
||||
// Buffered returns the number of bytes that can be read from the current buffer.
|
||||
func (b *ElasticBufReader) Buffered() int {
|
||||
return b.w - b.r
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) Bytes() []byte {
|
||||
return b.buf[b.r:b.w]
|
||||
}
|
||||
|
||||
var errNegativeRead = errors.New("bufio: reader returned negative count from Read")
|
||||
|
||||
// fill reads a new chunk into the buffer.
|
||||
func (b *ElasticBufReader) fill() {
|
||||
// Slide existing data to beginning.
|
||||
if b.r > 0 {
|
||||
copy(b.buf, b.buf[b.r:b.w])
|
||||
b.w -= b.r
|
||||
b.r = 0
|
||||
}
|
||||
|
||||
if b.w >= len(b.buf) {
|
||||
panic("bufio: tried to fill full buffer")
|
||||
}
|
||||
|
||||
// Read new data: try a limited number of times.
|
||||
const maxConsecutiveEmptyReads = 100
|
||||
for i := maxConsecutiveEmptyReads; i > 0; i-- {
|
||||
n, err := b.rd.Read(b.buf[b.w:])
|
||||
if n < 0 {
|
||||
panic(errNegativeRead)
|
||||
}
|
||||
b.w += n
|
||||
if err != nil {
|
||||
b.err = err
|
||||
return
|
||||
}
|
||||
if n > 0 {
|
||||
return
|
||||
}
|
||||
}
|
||||
b.err = io.ErrNoProgress
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) readErr() error {
|
||||
err := b.err
|
||||
b.err = nil
|
||||
return err
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) ReadSlice(delim byte) (line []byte, err error) {
|
||||
for {
|
||||
// Search buffer.
|
||||
if i := bytes.IndexByte(b.buf[b.r:b.w], delim); i >= 0 {
|
||||
line = b.buf[b.r : b.r+i+1]
|
||||
b.r += i + 1
|
||||
break
|
||||
}
|
||||
|
||||
// Pending error?
|
||||
if b.err != nil {
|
||||
line = b.buf[b.r:b.w]
|
||||
b.r = b.w
|
||||
err = b.readErr()
|
||||
break
|
||||
}
|
||||
|
||||
// Buffer full?
|
||||
if b.Buffered() >= len(b.buf) {
|
||||
b.grow(len(b.buf) + defaultBufSize)
|
||||
}
|
||||
|
||||
b.fill() // buffer is not full
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) ReadLine() (line []byte, err error) {
|
||||
line, err = b.ReadSlice('\n')
|
||||
if len(line) == 0 {
|
||||
if err != nil {
|
||||
line = nil
|
||||
}
|
||||
return
|
||||
}
|
||||
err = nil
|
||||
|
||||
if line[len(line)-1] == '\n' {
|
||||
drop := 1
|
||||
if len(line) > 1 && line[len(line)-2] == '\r' {
|
||||
drop = 2
|
||||
}
|
||||
line = line[:len(line)-drop]
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) ReadByte() (byte, error) {
|
||||
for b.r == b.w {
|
||||
if b.err != nil {
|
||||
return 0, b.readErr()
|
||||
}
|
||||
b.fill() // buffer is empty
|
||||
}
|
||||
c := b.buf[b.r]
|
||||
b.r++
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) ReadN(n int) ([]byte, error) {
|
||||
b.grow(n)
|
||||
for b.Buffered() < n {
|
||||
// Pending error?
|
||||
if b.err != nil {
|
||||
buf := b.buf[b.r:b.w]
|
||||
b.r = b.w
|
||||
return buf, b.readErr()
|
||||
}
|
||||
|
||||
b.fill()
|
||||
}
|
||||
|
||||
buf := b.buf[b.r : b.r+n]
|
||||
b.r += n
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
func (b *ElasticBufReader) grow(n int) {
|
||||
if b.w-b.r >= n {
|
||||
return
|
||||
}
|
||||
|
||||
// Slide existing data to beginning.
|
||||
if b.r > 0 {
|
||||
copy(b.buf, b.buf[b.r:b.w])
|
||||
b.w -= b.r
|
||||
b.r = 0
|
||||
}
|
||||
|
||||
// Extend buffer if needed.
|
||||
if d := n - len(b.buf); d > 0 {
|
||||
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()
|
||||
}
|
||||
@ -1,6 +1,7 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
@ -26,39 +27,32 @@ func (e RedisError) Error() string { return string(e) }
|
||||
|
||||
//------------------------------------------------------------------------------
|
||||
|
||||
type MultiBulkParse func(Reader, int64) (interface{}, error)
|
||||
type MultiBulkParse func(*Reader, int64) (interface{}, error)
|
||||
|
||||
type Reader struct {
|
||||
src *ElasticBufReader
|
||||
rd *bufio.Reader
|
||||
_buf []byte
|
||||
}
|
||||
|
||||
func NewReader(src *ElasticBufReader) Reader {
|
||||
return Reader{
|
||||
src: src,
|
||||
func NewReader(rd io.Reader) *Reader {
|
||||
return &Reader{
|
||||
rd: bufio.NewReader(rd),
|
||||
_buf: make([]byte, 64),
|
||||
}
|
||||
}
|
||||
|
||||
func (r Reader) Reset(rd io.Reader) {
|
||||
r.src.Reset(rd)
|
||||
func (r *Reader) Reset(rd io.Reader) {
|
||||
r.rd.Reset(rd)
|
||||
}
|
||||
|
||||
func (r Reader) Buffer() []byte {
|
||||
return r.src.Buffer()
|
||||
}
|
||||
|
||||
func (r Reader) ResetBuffer(buf []byte) {
|
||||
r.src.ResetBuffer(buf)
|
||||
}
|
||||
|
||||
func (r Reader) Bytes() []byte {
|
||||
return r.src.Bytes()
|
||||
}
|
||||
|
||||
func (r Reader) ReadLine() ([]byte, error) {
|
||||
line, err := r.src.ReadLine()
|
||||
func (r *Reader) ReadLine() ([]byte, error) {
|
||||
line, isPrefix, err := r.rd.ReadLine()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if isPrefix {
|
||||
return nil, bufio.ErrBufferFull
|
||||
}
|
||||
if len(line) == 0 {
|
||||
return nil, fmt.Errorf("redis: reply is empty")
|
||||
}
|
||||
@ -68,7 +62,7 @@ func (r Reader) ReadLine() ([]byte, error) {
|
||||
return line, nil
|
||||
}
|
||||
|
||||
func (r Reader) ReadReply(m MultiBulkParse) (interface{}, error) {
|
||||
func (r *Reader) ReadReply(m MultiBulkParse) (interface{}, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@ -93,7 +87,7 @@ func (r Reader) ReadReply(m MultiBulkParse) (interface{}, error) {
|
||||
return nil, fmt.Errorf("redis: can't parse %.100q", line)
|
||||
}
|
||||
|
||||
func (r Reader) ReadIntReply() (int64, error) {
|
||||
func (r *Reader) ReadIntReply() (int64, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
@ -108,7 +102,7 @@ func (r Reader) ReadIntReply() (int64, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (r Reader) ReadString() (string, error) {
|
||||
func (r *Reader) ReadString() (string, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return "", err
|
||||
@ -127,7 +121,7 @@ func (r Reader) ReadString() (string, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (r Reader) readStringReply(line []byte) (string, error) {
|
||||
func (r *Reader) readStringReply(line []byte) (string, error) {
|
||||
if isNilReply(line) {
|
||||
return "", Nil
|
||||
}
|
||||
@ -138,7 +132,7 @@ func (r Reader) readStringReply(line []byte) (string, error) {
|
||||
}
|
||||
|
||||
b := make([]byte, replyLen+2)
|
||||
_, err = io.ReadFull(r.src, b)
|
||||
_, err = io.ReadFull(r.rd, b)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@ -146,7 +140,7 @@ func (r Reader) readStringReply(line []byte) (string, error) {
|
||||
return util.BytesToString(b[:replyLen]), nil
|
||||
}
|
||||
|
||||
func (r Reader) ReadArrayReply(m MultiBulkParse) (interface{}, error) {
|
||||
func (r *Reader) ReadArrayReply(m MultiBulkParse) (interface{}, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@ -165,7 +159,7 @@ func (r Reader) ReadArrayReply(m MultiBulkParse) (interface{}, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (r Reader) ReadArrayLen() (int64, error) {
|
||||
func (r *Reader) ReadArrayLen() (int64, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
@ -180,7 +174,7 @@ func (r Reader) ReadArrayLen() (int64, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (r Reader) ReadScanReply() ([]string, uint64, error) {
|
||||
func (r *Reader) ReadScanReply() ([]string, uint64, error) {
|
||||
n, err := r.ReadArrayLen()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
@ -211,7 +205,7 @@ func (r Reader) ReadScanReply() ([]string, uint64, error) {
|
||||
return keys, cursor, err
|
||||
}
|
||||
|
||||
func (r Reader) ReadInt() (int64, error) {
|
||||
func (r *Reader) ReadInt() (int64, error) {
|
||||
b, err := r.readTmpBytesReply()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
@ -219,7 +213,7 @@ func (r Reader) ReadInt() (int64, error) {
|
||||
return util.ParseInt(b, 10, 64)
|
||||
}
|
||||
|
||||
func (r Reader) ReadUint() (uint64, error) {
|
||||
func (r *Reader) ReadUint() (uint64, error) {
|
||||
b, err := r.readTmpBytesReply()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
@ -227,7 +221,7 @@ func (r Reader) ReadUint() (uint64, error) {
|
||||
return util.ParseUint(b, 10, 64)
|
||||
}
|
||||
|
||||
func (r Reader) ReadFloatReply() (float64, error) {
|
||||
func (r *Reader) ReadFloatReply() (float64, error) {
|
||||
b, err := r.readTmpBytesReply()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
@ -235,7 +229,7 @@ func (r Reader) ReadFloatReply() (float64, error) {
|
||||
return util.ParseFloat(b, 64)
|
||||
}
|
||||
|
||||
func (r Reader) readTmpBytesReply() ([]byte, error) {
|
||||
func (r *Reader) readTmpBytesReply() ([]byte, error) {
|
||||
line, err := r.ReadLine()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@ -252,7 +246,7 @@ func (r Reader) readTmpBytesReply() ([]byte, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (r Reader) _readTmpBytesReply(line []byte) ([]byte, error) {
|
||||
func (r *Reader) _readTmpBytesReply(line []byte) ([]byte, error) {
|
||||
if isNilReply(line) {
|
||||
return nil, Nil
|
||||
}
|
||||
@ -262,11 +256,20 @@ func (r Reader) _readTmpBytesReply(line []byte) ([]byte, error) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
b, err := r.src.ReadN(replyLen + 2)
|
||||
buf := r.buf(replyLen + 2)
|
||||
_, err = io.ReadFull(r.rd, buf)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return b[:replyLen], nil
|
||||
|
||||
return buf[:replyLen], nil
|
||||
}
|
||||
|
||||
func (r *Reader) buf(n int) []byte {
|
||||
if d := n - cap(r._buf); d > 0 {
|
||||
r._buf = append(r._buf, make([]byte, d)...)
|
||||
}
|
||||
return r._buf[:n]
|
||||
}
|
||||
|
||||
func isNilReply(b []byte) bool {
|
||||
|
||||
@ -8,8 +8,8 @@ import (
|
||||
"github.com/go-redis/redis/internal/proto"
|
||||
)
|
||||
|
||||
func newReader(s string) proto.Reader {
|
||||
return proto.NewReader(proto.NewElasticBufReader(strings.NewReader(s)))
|
||||
func newReader(s string) *proto.Reader {
|
||||
return proto.NewReader(strings.NewReader(s))
|
||||
}
|
||||
|
||||
func BenchmarkReader_ParseReply_Status(b *testing.B) {
|
||||
@ -37,7 +37,7 @@ func benchmarkParseReply(b *testing.B, reply string, m proto.MultiBulkParse, wan
|
||||
for i := 0; i < b.N; i++ {
|
||||
buf.WriteString(reply)
|
||||
}
|
||||
p := proto.NewReader(proto.NewElasticBufReader(buf))
|
||||
p := proto.NewReader(buf)
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
@ -48,7 +48,7 @@ func benchmarkParseReply(b *testing.B, reply string, m proto.MultiBulkParse, wan
|
||||
}
|
||||
}
|
||||
|
||||
func multiBulkParse(p proto.Reader, n int64) (interface{}, error) {
|
||||
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)
|
||||
|
||||
@ -1,127 +0,0 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"encoding"
|
||||
"fmt"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
type WriteBuffer struct {
|
||||
buf []byte
|
||||
}
|
||||
|
||||
func NewWriteBuffer() *WriteBuffer {
|
||||
return &WriteBuffer{}
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) Len() int {
|
||||
return len(w.buf)
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) Bytes() []byte {
|
||||
return w.buf
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) Reset() {
|
||||
w.buf = w.buf[:0]
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) Buffer() []byte {
|
||||
return w.buf[:cap(w.buf)]
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) ResetBuffer(buf []byte) {
|
||||
w.buf = buf[:0]
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) Append(args []interface{}) error {
|
||||
w.buf = append(w.buf, ArrayReply)
|
||||
w.buf = strconv.AppendUint(w.buf, uint64(len(args)), 10)
|
||||
w.buf = append(w.buf, '\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")
|
||||
}
|
||||
case encoding.BinaryMarshaler:
|
||||
b, err := v.MarshalBinary()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
w.AppendBytes(b)
|
||||
default:
|
||||
return fmt.Errorf(
|
||||
"redis: can't marshal %T (consider implementing encoding.BinaryMarshaler)", val)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) AppendString(s string) {
|
||||
w.buf = append(w.buf, StringReply)
|
||||
w.buf = strconv.AppendUint(w.buf, uint64(len(s)), 10)
|
||||
w.buf = append(w.buf, '\r', '\n')
|
||||
w.buf = append(w.buf, s...)
|
||||
w.buf = append(w.buf, '\r', '\n')
|
||||
}
|
||||
|
||||
func (w *WriteBuffer) AppendBytes(p []byte) {
|
||||
w.buf = append(w.buf, StringReply)
|
||||
w.buf = strconv.AppendUint(w.buf, uint64(len(p)), 10)
|
||||
w.buf = append(w.buf, '\r', '\n')
|
||||
w.buf = append(w.buf, p...)
|
||||
w.buf = append(w.buf, '\r', '\n')
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
@ -1,6 +1,8 @@
|
||||
package proto_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io/ioutil"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@ -11,21 +13,16 @@ import (
|
||||
)
|
||||
|
||||
var _ = Describe("WriteBuffer", func() {
|
||||
var buf *proto.WriteBuffer
|
||||
var buf *bytes.Buffer
|
||||
var wr *proto.Writer
|
||||
|
||||
BeforeEach(func() {
|
||||
buf = proto.NewWriteBuffer()
|
||||
buf = new(bytes.Buffer)
|
||||
wr = proto.NewWriter(buf)
|
||||
})
|
||||
|
||||
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{}{
|
||||
It("should write args", func() {
|
||||
err := wr.WriteArgs([]interface{}{
|
||||
"string",
|
||||
12,
|
||||
34.56,
|
||||
@ -34,6 +31,10 @@ var _ = Describe("WriteBuffer", func() {
|
||||
nil,
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
err = wr.Flush()
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
Expect(buf.Bytes()).To(Equal([]byte("*6\r\n" +
|
||||
"$6\r\nstring\r\n" +
|
||||
"$2\r\n12\r\n" +
|
||||
@ -45,19 +46,30 @@ var _ = Describe("WriteBuffer", func() {
|
||||
})
|
||||
|
||||
It("should append marshalable args", func() {
|
||||
err := buf.Append([]interface{}{time.Unix(1414141414, 0)})
|
||||
err := wr.WriteArgs([]interface{}{time.Unix(1414141414, 0)})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
err = wr.Flush()
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
|
||||
Expect(buf.Len()).To(Equal(26))
|
||||
})
|
||||
|
||||
})
|
||||
|
||||
func BenchmarkWriteBuffer_Append(b *testing.B) {
|
||||
buf := proto.NewWriteBuffer()
|
||||
buf := proto.NewWriter(ioutil.Discard)
|
||||
args := []interface{}{"hello", "world", "foo", "bar"}
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
buf.Append(args)
|
||||
buf.Reset()
|
||||
err := buf.WriteArgs(args)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
err = buf.Flush()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
159
internal/proto/writer.go
Normal file
159
internal/proto/writer.go
Normal file
@ -0,0 +1,159 @@
|
||||
package proto
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
|
||||
"github.com/go-redis/redis/internal/util"
|
||||
)
|
||||
|
||||
type Writer struct {
|
||||
wr *bufio.Writer
|
||||
|
||||
lenBuf []byte
|
||||
numBuf []byte
|
||||
}
|
||||
|
||||
func NewWriter(wr io.Writer) *Writer {
|
||||
return &Writer{
|
||||
wr: bufio.NewWriter(wr),
|
||||
|
||||
lenBuf: make([]byte, 64),
|
||||
numBuf: make([]byte, 64),
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Writer) WriteArgs(args []interface{}) error {
|
||||
err := w.wr.WriteByte(ArrayReply)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
err = w.writeLen(len(args))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, arg := range args {
|
||||
err := w.writeArg(arg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Writer) writeLen(n int) error {
|
||||
w.lenBuf = strconv.AppendUint(w.lenBuf[:0], uint64(n), 10)
|
||||
w.lenBuf = append(w.lenBuf, '\r', '\n')
|
||||
_, err := w.wr.Write(w.lenBuf)
|
||||
return err
|
||||
}
|
||||
|
||||
func (w *Writer) writeArg(v interface{}) error {
|
||||
switch v := v.(type) {
|
||||
case nil:
|
||||
return w.string("")
|
||||
case string:
|
||||
return w.string(v)
|
||||
case []byte:
|
||||
return w.bytes(v)
|
||||
case int:
|
||||
return w.int(int64(v))
|
||||
case int8:
|
||||
return w.int(int64(v))
|
||||
case int16:
|
||||
return w.int(int64(v))
|
||||
case int32:
|
||||
return w.int(int64(v))
|
||||
case int64:
|
||||
return w.int(v)
|
||||
case uint:
|
||||
return w.uint(uint64(v))
|
||||
case uint8:
|
||||
return w.uint(uint64(v))
|
||||
case uint16:
|
||||
return w.uint(uint64(v))
|
||||
case uint32:
|
||||
return w.uint(uint64(v))
|
||||
case uint64:
|
||||
return w.uint(v)
|
||||
case float32:
|
||||
return w.float(float64(v))
|
||||
case float64:
|
||||
return w.float(v)
|
||||
case bool:
|
||||
if v {
|
||||
return w.int(1)
|
||||
} else {
|
||||
return w.int(0)
|
||||
}
|
||||
case encoding.BinaryMarshaler:
|
||||
b, err := v.MarshalBinary()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return w.bytes(b)
|
||||
default:
|
||||
return fmt.Errorf(
|
||||
"redis: can't marshal %T (implement encoding.BinaryMarshaler)", v)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *Writer) bytes(b []byte) error {
|
||||
err := w.wr.WriteByte(StringReply)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
err = w.writeLen(len(b))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = w.wr.Write(b)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return w.crlf()
|
||||
}
|
||||
|
||||
func (w *Writer) string(s string) error {
|
||||
return w.bytes(util.StringToBytes(s))
|
||||
}
|
||||
|
||||
func (w *Writer) uint(n uint64) error {
|
||||
w.numBuf = strconv.AppendUint(w.numBuf[:0], n, 10)
|
||||
return w.bytes(w.numBuf)
|
||||
}
|
||||
|
||||
func (w *Writer) int(n int64) error {
|
||||
w.numBuf = strconv.AppendInt(w.numBuf[:0], n, 10)
|
||||
return w.bytes(w.numBuf)
|
||||
}
|
||||
|
||||
func (w *Writer) float(f float64) error {
|
||||
w.numBuf = strconv.AppendFloat(w.numBuf[:0], f, 'f', -1, 64)
|
||||
return w.bytes(w.numBuf)
|
||||
}
|
||||
|
||||
func (w *Writer) crlf() error {
|
||||
err := w.wr.WriteByte('\r')
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return w.wr.WriteByte('\n')
|
||||
}
|
||||
|
||||
func (w *Writer) Reset(wr io.Writer) {
|
||||
w.wr.Reset(wr)
|
||||
}
|
||||
|
||||
func (w *Writer) Flush() error {
|
||||
return w.wr.Flush()
|
||||
}
|
||||
Reference in New Issue
Block a user