redis/internal/proto/elastic_reader.go

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{
		buf: make([]byte, defaultBufSize),
		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) Read(p []byte) (n int, err error) {
	n = len(p)
	if n == 0 {
		return 0, b.readErr()
	}
	if b.r == b.w {
		if b.err != nil {
			return 0, b.readErr()
		}
		if len(p) >= len(b.buf) {
			// Large read, empty buffer.
			// Read directly into p to avoid copy.
			n, b.err = b.rd.Read(p)
			if n < 0 {
				panic(errNegativeRead)
			}
			return n, b.readErr()
		}
		// One read.
		// Do not use b.fill, which will loop.
		b.r = 0
		b.w = 0
		n, b.err = b.rd.Read(b.buf)
		if n < 0 {
			panic(errNegativeRead)
		}
		if n == 0 {
			return 0, b.readErr()
		}
		b.w += n
	}

	// copy as much as we can
	n = copy(p, b.buf[b.r:b.w])
	b.r += n
	return n, nil
}

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)...)
	}
}