303 lines
6.7 KiB
Go
303 lines
6.7 KiB
Go
// This is a derivitive work of Go's bytes.Buffer implementation.
|
|
//
|
|
// Originally copyright 2009 The Go Authors. All rights reserved.
|
|
//
|
|
// Modifications copyright 2018 Steven Allen. All rights reserved.
|
|
//
|
|
// Use of this source code is governed by both a BSD-style and an MIT-style
|
|
// license that can be found in the LICENSE_BSD and LICENSE files.
|
|
|
|
package pool
|
|
|
|
import (
|
|
"io"
|
|
)
|
|
|
|
// Buffer is a buffer like bytes.Buffer that:
|
|
//
|
|
// 1. Uses a buffer pool.
|
|
// 2. Frees memory on read.
|
|
//
|
|
// If you only have a few buffers and read/write at a steady rate, *don't* use
|
|
// this package, it'll be slower.
|
|
//
|
|
// However:
|
|
//
|
|
// 1. If you frequently create/destroy buffers, this implementation will be
|
|
// significantly nicer to the allocator.
|
|
// 2. If you have many buffers with bursty traffic, this implementation will use
|
|
// significantly less memory.
|
|
type Buffer struct {
|
|
// Pool is the buffer pool to use. If nil, this Buffer will use the
|
|
// global buffer pool.
|
|
Pool *BufferPool
|
|
|
|
buf []byte
|
|
rOff int
|
|
|
|
// Preallocated slice for samll reads/writes.
|
|
// This is *really* important for performance and only costs 8 words.
|
|
bootstrap [64]byte
|
|
}
|
|
|
|
// NewBuffer constructs a new buffer initialized to `buf`.
|
|
// Unlike `bytes.Buffer`, we *copy* the buffer but don't reuse it (to ensure
|
|
// that we *only* use buffers from the pool).
|
|
func NewBuffer(buf []byte) *Buffer {
|
|
b := new(Buffer)
|
|
if len(buf) > 0 {
|
|
b.buf = b.getBuf(len(buf))
|
|
copy(b.buf, buf)
|
|
}
|
|
return b
|
|
}
|
|
|
|
// NewBufferString is identical to NewBuffer *except* that it allows one to
|
|
// initialize the buffer from a string (without having to allocate an
|
|
// intermediate bytes slice).
|
|
func NewBufferString(buf string) *Buffer {
|
|
b := new(Buffer)
|
|
if len(buf) > 0 {
|
|
b.buf = b.getBuf(len(buf))
|
|
copy(b.buf, buf)
|
|
}
|
|
return b
|
|
}
|
|
|
|
func (b *Buffer) grow(n int) int {
|
|
wOff := len(b.buf)
|
|
bCap := cap(b.buf)
|
|
|
|
if bCap >= wOff+n {
|
|
b.buf = b.buf[:wOff+n]
|
|
return wOff
|
|
}
|
|
|
|
bSize := b.Len()
|
|
|
|
minCap := 2*bSize + n
|
|
|
|
// Slide if cap >= minCap.
|
|
// Reallocate otherwise.
|
|
if bCap >= minCap {
|
|
copy(b.buf, b.buf[b.rOff:])
|
|
} else {
|
|
// Needs new buffer.
|
|
newBuf := b.getBuf(minCap)
|
|
copy(newBuf, b.buf[b.rOff:])
|
|
b.returnBuf()
|
|
b.buf = newBuf
|
|
}
|
|
|
|
b.rOff = 0
|
|
b.buf = b.buf[:bSize+n]
|
|
return bSize
|
|
}
|
|
|
|
func (b *Buffer) getPool() *BufferPool {
|
|
if b.Pool == nil {
|
|
return GlobalPool
|
|
}
|
|
return b.Pool
|
|
}
|
|
|
|
func (b *Buffer) returnBuf() {
|
|
if cap(b.buf) > len(b.bootstrap) {
|
|
b.getPool().Put(b.buf)
|
|
}
|
|
b.buf = nil
|
|
}
|
|
|
|
func (b *Buffer) getBuf(n int) []byte {
|
|
if n <= len(b.bootstrap) {
|
|
return b.bootstrap[:n]
|
|
}
|
|
return b.getPool().Get(n)
|
|
}
|
|
|
|
// Len returns the number of bytes that can be read from this buffer.
|
|
func (b *Buffer) Len() int {
|
|
return len(b.buf) - b.rOff
|
|
}
|
|
|
|
// Cap returns the current capacity of the buffer.
|
|
//
|
|
// Note: Buffer *may* re-allocate when writing (or growing by) `n` bytes even if
|
|
// `Cap() < Len() + n` to avoid excessive copying.
|
|
func (b *Buffer) Cap() int {
|
|
return cap(b.buf)
|
|
}
|
|
|
|
// Bytes returns the slice of bytes currently buffered in the Buffer.
|
|
//
|
|
// The buffer returned by Bytes is valid until the next call grow, truncate,
|
|
// read, or write. Really, just don't touch the Buffer until you're done with
|
|
// the return value of this function.
|
|
func (b *Buffer) Bytes() []byte {
|
|
return b.buf[b.rOff:]
|
|
}
|
|
|
|
// String returns the string representation of the buffer.
|
|
//
|
|
// It returns `<nil>` the buffer is a nil pointer.
|
|
func (b *Buffer) String() string {
|
|
if b == nil {
|
|
return "<nil>"
|
|
}
|
|
return string(b.buf[b.rOff:])
|
|
}
|
|
|
|
// WriteString writes a string to the buffer.
|
|
//
|
|
// This function is identical to Write except that it allows one to write a
|
|
// string directly without allocating an intermediate byte slice.
|
|
func (b *Buffer) WriteString(buf string) (int, error) {
|
|
wOff := b.grow(len(buf))
|
|
return copy(b.buf[wOff:], buf), nil
|
|
}
|
|
|
|
// Truncate truncates the Buffer.
|
|
//
|
|
// Panics if `n > b.Len()`.
|
|
//
|
|
// This function may free memory by shrinking the internal buffer.
|
|
func (b *Buffer) Truncate(n int) {
|
|
if n < 0 || n > b.Len() {
|
|
panic("truncation out of range")
|
|
}
|
|
b.buf = b.buf[:b.rOff+n]
|
|
b.shrink()
|
|
}
|
|
|
|
// Reset is equivalent to Truncate(0).
|
|
func (b *Buffer) Reset() {
|
|
b.returnBuf()
|
|
b.rOff = 0
|
|
}
|
|
|
|
// ReadByte reads a single byte from the Buffer.
|
|
func (b *Buffer) ReadByte() (byte, error) {
|
|
if b.rOff >= len(b.buf) {
|
|
return 0, io.EOF
|
|
}
|
|
c := b.buf[b.rOff]
|
|
b.rOff++
|
|
return c, nil
|
|
}
|
|
|
|
// WriteByte writes a single byte to the Buffer.
|
|
func (b *Buffer) WriteByte(c byte) error {
|
|
wOff := b.grow(1)
|
|
b.buf[wOff] = c
|
|
return nil
|
|
}
|
|
|
|
// Grow grows the internal buffer such that `n` bytes can be written without
|
|
// reallocating.
|
|
func (b *Buffer) Grow(n int) {
|
|
wOff := b.grow(n)
|
|
b.buf = b.buf[:wOff]
|
|
}
|
|
|
|
// Next is an alternative to `Read` that returns a byte slice instead of taking
|
|
// one.
|
|
//
|
|
// The returned byte slice is valid until the next read, write, grow, or
|
|
// truncate.
|
|
func (b *Buffer) Next(n int) []byte {
|
|
m := b.Len()
|
|
if m < n {
|
|
n = m
|
|
}
|
|
data := b.buf[b.rOff : b.rOff+n]
|
|
b.rOff += n
|
|
return data
|
|
}
|
|
|
|
// Write writes the byte slice to the buffer.
|
|
func (b *Buffer) Write(buf []byte) (int, error) {
|
|
wOff := b.grow(len(buf))
|
|
return copy(b.buf[wOff:], buf), nil
|
|
}
|
|
|
|
// WriteTo copies from the buffer into the given writer until the buffer is
|
|
// empty.
|
|
func (b *Buffer) WriteTo(w io.Writer) (int64, error) {
|
|
if b.rOff < len(b.buf) {
|
|
n, err := w.Write(b.buf[b.rOff:])
|
|
b.rOff += n
|
|
if b.rOff > len(b.buf) {
|
|
panic("invalid write count")
|
|
}
|
|
b.shrink()
|
|
return int64(n), err
|
|
}
|
|
return 0, nil
|
|
}
|
|
|
|
// MinRead is the minimum slice size passed to a Read call by
|
|
// Buffer.ReadFrom. As long as the Buffer has at least MinRead bytes beyond
|
|
// what is required to hold the contents of r, ReadFrom will not grow the
|
|
// underlying buffer.
|
|
const MinRead = 512
|
|
|
|
// ReadFrom reads from the given reader into the buffer.
|
|
func (b *Buffer) ReadFrom(r io.Reader) (int64, error) {
|
|
n := int64(0)
|
|
for {
|
|
wOff := b.grow(MinRead)
|
|
// Use *entire* buffer.
|
|
b.buf = b.buf[:cap(b.buf)]
|
|
|
|
read, err := r.Read(b.buf[wOff:])
|
|
b.buf = b.buf[:wOff+read]
|
|
n += int64(read)
|
|
switch err {
|
|
case nil:
|
|
case io.EOF:
|
|
err = nil
|
|
fallthrough
|
|
default:
|
|
b.shrink()
|
|
return n, err
|
|
}
|
|
}
|
|
}
|
|
|
|
// Read reads at most `len(buf)` bytes from the internal buffer into the given
|
|
// buffer.
|
|
func (b *Buffer) Read(buf []byte) (int, error) {
|
|
if len(buf) == 0 {
|
|
return 0, nil
|
|
}
|
|
if b.rOff >= len(b.buf) {
|
|
return 0, io.EOF
|
|
}
|
|
n := copy(buf, b.buf[b.rOff:])
|
|
b.rOff += n
|
|
b.shrink()
|
|
return n, nil
|
|
}
|
|
|
|
func (b *Buffer) shrink() {
|
|
c := b.Cap()
|
|
// Either nil or bootstrap.
|
|
if c <= len(b.bootstrap) {
|
|
return
|
|
}
|
|
|
|
l := b.Len()
|
|
if l == 0 {
|
|
// Shortcut if empty.
|
|
b.returnBuf()
|
|
b.rOff = 0
|
|
} else if l*8 < c {
|
|
// Only shrink when capacity > 8x length. Avoids shrinking too aggressively.
|
|
newBuf := b.getBuf(l)
|
|
copy(newBuf, b.buf[b.rOff:])
|
|
b.returnBuf()
|
|
b.rOff = 0
|
|
b.buf = newBuf[:l]
|
|
}
|
|
}
|