Commit 1d08797b authored by Steven Allen's avatar Steven Allen

add a Buffer implementation that uses a buffer pool.

parent deedd3bb
### Applies to buffer.go and buffer_test.go ###
Copyright (c) 2009 The Go Authors. All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:
* Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
* Redistributions in binary form must reproduce the above
copyright notice, this list of conditions and the following disclaimer
in the documentation and/or other materials provided with the
distribution.
* Neither the name of Google Inc. nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
// 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) > 0 && &b.buf[:1][0] != &b.bootstrap[0] {
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
}
// TODO: implement ReadFrom
// 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() {
l := b.Len()
if l == 0 {
b.returnBuf()
b.rOff = 0
} else if l*8 < b.Cap() {
// 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]
}
}
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Modified by stebalien, 2018
package pool
import (
"bytes"
"io"
"math/rand"
"runtime"
"testing"
)
const N = 10000 // make this bigger for a larger (and slower) test
var data string // test data for write tests
var testBytes []byte // test data; same as data but as a slice.
func init() {
testBytes = make([]byte, N)
for i := 0; i < N; i++ {
testBytes[i] = 'a' + byte(i%26)
}
data = string(testBytes)
}
// Verify that contents of buf match the string s.
func check(t *testing.T, testname string, buf *Buffer, s string) {
bytes := buf.Bytes()
str := buf.String()
if buf.Len() != len(bytes) {
t.Errorf("%s: buf.Len() == %d, len(buf.Bytes()) == %d", testname, buf.Len(), len(bytes))
}
if buf.Len() != len(str) {
t.Errorf("%s: buf.Len() == %d, len(buf.String()) == %d", testname, buf.Len(), len(str))
}
if buf.Len() != len(s) {
t.Errorf("%s: buf.Len() == %d, len(s) == %d", testname, buf.Len(), len(s))
}
if string(bytes) != s {
t.Errorf("%s: string(buf.Bytes()) == %q, s == %q", testname, string(bytes), s)
}
}
// Fill buf through n writes of string fus.
// The initial contents of buf corresponds to the string s;
// the result is the final contents of buf returned as a string.
func fillString(t *testing.T, testname string, buf *Buffer, s string, n int, fus string) string {
check(t, testname+" (fill 1)", buf, s)
for ; n > 0; n-- {
m, err := buf.WriteString(fus)
if m != len(fus) {
t.Errorf(testname+" (fill 2): m == %d, expected %d", m, len(fus))
}
if err != nil {
t.Errorf(testname+" (fill 3): err should always be nil, found err == %s", err)
}
s += fus
check(t, testname+" (fill 4)", buf, s)
}
return s
}
// Fill buf through n writes of byte slice fub.
// The initial contents of buf corresponds to the string s;
// the result is the final contents of buf returned as a string.
func fillBytes(t *testing.T, testname string, buf *Buffer, s string, n int, fub []byte) string {
check(t, testname+" (fill 1)", buf, s)
for ; n > 0; n-- {
m, err := buf.Write(fub)
if m != len(fub) {
t.Errorf(testname+" (fill 2): m == %d, expected %d", m, len(fub))
}
if err != nil {
t.Errorf(testname+" (fill 3): err should always be nil, found err == %s", err)
}
s += string(fub)
check(t, testname+" (fill 4)", buf, s)
}
return s
}
func TestNewBuffer(t *testing.T) {
buf := NewBuffer(testBytes)
check(t, "NewBuffer", buf, data)
}
func TestNewBufferString(t *testing.T) {
buf := NewBufferString(data)
check(t, "NewBufferString", buf, data)
}
// Empty buf through repeated reads into fub.
// The initial contents of buf corresponds to the string s.
func empty(t *testing.T, testname string, buf *Buffer, s string, fub []byte) {
check(t, testname+" (empty 1)", buf, s)
for {
n, err := buf.Read(fub)
if n == 0 {
break
}
if err != nil {
t.Errorf(testname+" (empty 2): err should always be nil, found err == %s", err)
}
s = s[n:]
check(t, testname+" (empty 3)", buf, s)
}
check(t, testname+" (empty 4)", buf, "")
}
func TestBasicOperations(t *testing.T) {
var buf Buffer
for i := 0; i < 5; i++ {
check(t, "TestBasicOperations (1)", &buf, "")
buf.Reset()
check(t, "TestBasicOperations (2)", &buf, "")
buf.Truncate(0)
check(t, "TestBasicOperations (3)", &buf, "")
n, err := buf.Write([]byte(data[0:1]))
if n != 1 {
t.Errorf("wrote 1 byte, but n == %d", n)
}
if err != nil {
t.Errorf("err should always be nil, but err == %s", err)
}
check(t, "TestBasicOperations (4)", &buf, "a")
buf.WriteByte(data[1])
check(t, "TestBasicOperations (5)", &buf, "ab")
n, err = buf.Write([]byte(data[2:26]))
if n != 24 {
t.Errorf("wrote 25 bytes, but n == %d", n)
}
check(t, "TestBasicOperations (6)", &buf, string(data[0:26]))
buf.Truncate(26)
check(t, "TestBasicOperations (7)", &buf, string(data[0:26]))
buf.Truncate(20)
check(t, "TestBasicOperations (8)", &buf, string(data[0:20]))
empty(t, "TestBasicOperations (9)", &buf, string(data[0:20]), make([]byte, 5))
empty(t, "TestBasicOperations (10)", &buf, "", make([]byte, 100))
buf.WriteByte(data[1])
c, err := buf.ReadByte()
if err != nil {
t.Error("ReadByte unexpected eof")
}
if c != data[1] {
t.Errorf("ReadByte wrong value c=%v", c)
}
c, err = buf.ReadByte()
if err == nil {
t.Error("ReadByte unexpected not eof")
}
}
}
func TestLargeStringWrites(t *testing.T) {
var buf Buffer
limit := 30
if testing.Short() {
limit = 9
}
for i := 3; i < limit; i += 3 {
s := fillString(t, "TestLargeWrites (1)", &buf, "", 5, data)
empty(t, "TestLargeStringWrites (2)", &buf, s, make([]byte, len(data)/i))
}
check(t, "TestLargeStringWrites (3)", &buf, "")
}
func TestLargeByteWrites(t *testing.T) {
var buf Buffer
limit := 30
if testing.Short() {
limit = 9
}
for i := 3; i < limit; i += 3 {
s := fillBytes(t, "TestLargeWrites (1)", &buf, "", 5, testBytes)
empty(t, "TestLargeByteWrites (2)", &buf, s, make([]byte, len(data)/i))
}
check(t, "TestLargeByteWrites (3)", &buf, "")
}
func TestLargeStringReads(t *testing.T) {
var buf Buffer
for i := 3; i < 30; i += 3 {
s := fillString(t, "TestLargeReads (1)", &buf, "", 5, data[0:len(data)/i])
empty(t, "TestLargeReads (2)", &buf, s, make([]byte, len(data)))
}
check(t, "TestLargeStringReads (3)", &buf, "")
}
func TestLargeByteReads(t *testing.T) {
var buf Buffer
for i := 3; i < 30; i += 3 {
s := fillBytes(t, "TestLargeReads (1)", &buf, "", 5, testBytes[0:len(testBytes)/i])
empty(t, "TestLargeReads (2)", &buf, s, make([]byte, len(data)))
}
check(t, "TestLargeByteReads (3)", &buf, "")
}
func TestMixedReadsAndWrites(t *testing.T) {
var buf Buffer
s := ""
for i := 0; i < 50; i++ {
wlen := rand.Intn(len(data))
if i%2 == 0 {
s = fillString(t, "TestMixedReadsAndWrites (1)", &buf, s, 1, data[0:wlen])
} else {
s = fillBytes(t, "TestMixedReadsAndWrites (1)", &buf, s, 1, testBytes[0:wlen])
}
rlen := rand.Intn(len(data))
fub := make([]byte, rlen)
n, _ := buf.Read(fub)
s = s[n:]
}
empty(t, "TestMixedReadsAndWrites (2)", &buf, s, make([]byte, buf.Len()))
}
func TestNil(t *testing.T) {
var b *Buffer
if b.String() != "<nil>" {
t.Errorf("expected <nil>; got %q", b.String())
}
}
/*
func TestReadFrom(t *testing.T) {
var buf Buffer
for i := 3; i < 30; i += 3 {
s := fillBytes(t, "TestReadFrom (1)", &buf, "", 5, testBytes[0:len(testBytes)/i])
var b Buffer
b.ReadFrom(&buf)
empty(t, "TestReadFrom (2)", &b, s, make([]byte, len(data)))
}
}
*/
func TestWriteTo(t *testing.T) {
var buf Buffer
for i := 3; i < 30; i += 3 {
s := fillBytes(t, "TestWriteTo (1)", &buf, "", 5, testBytes[0:len(testBytes)/i])
var b Buffer
buf.WriteTo(&b)
empty(t, "TestWriteTo (2)", &b, s, make([]byte, len(data)))
}
}
func TestNext(t *testing.T) {
b := []byte{0, 1, 2, 3, 4}
tmp := make([]byte, 5)
for i := 0; i <= 5; i++ {
for j := i; j <= 5; j++ {
for k := 0; k <= 6; k++ {
// 0 <= i <= j <= 5; 0 <= k <= 6
// Check that if we start with a buffer
// of length j at offset i and ask for
// Next(k), we get the right bytes.
buf := NewBuffer(b[0:j])
n, _ := buf.Read(tmp[0:i])
if n != i {
t.Fatalf("Read %d returned %d", i, n)
}
bb := buf.Next(k)
want := k
if want > j-i {
want = j - i
}
if len(bb) != want {
t.Fatalf("in %d,%d: len(Next(%d)) == %d", i, j, k, len(bb))
}
for l, v := range bb {
if v != byte(l+i) {
t.Fatalf("in %d,%d: Next(%d)[%d] = %d, want %d", i, j, k, l, v, l+i)
}
}
}
}
}
}
var readBytesTests = []struct {
buffer string
delim byte
expected []string
err error
}{
{"", 0, []string{""}, io.EOF},
{"a\x00", 0, []string{"a\x00"}, nil},
{"abbbaaaba", 'b', []string{"ab", "b", "b", "aaab"}, nil},
{"hello\x01world", 1, []string{"hello\x01"}, nil},
{"foo\nbar", 0, []string{"foo\nbar"}, io.EOF},
{"alpha\nbeta\ngamma\n", '\n', []string{"alpha\n", "beta\n", "gamma\n"}, nil},
{"alpha\nbeta\ngamma", '\n', []string{"alpha\n", "beta\n", "gamma"}, io.EOF},
}
func TestGrow(t *testing.T) {
x := []byte{'x'}
y := []byte{'y'}
tmp := make([]byte, 72)
for _, startLen := range []int{0, 100, 1000, 10000, 100000} {
xBytes := bytes.Repeat(x, startLen)
for _, growLen := range []int{0, 100, 1000, 10000, 100000} {
buf := NewBuffer(xBytes)
// If we read, this affects buf.off, which is good to test.
readBytes, _ := buf.Read(tmp)
buf.Grow(growLen)
yBytes := bytes.Repeat(y, growLen)
// Check no allocation occurs in write, as long as we're single-threaded.
var m1, m2 runtime.MemStats
runtime.ReadMemStats(&m1)
buf.Write(yBytes)
runtime.ReadMemStats(&m2)
if runtime.GOMAXPROCS(-1) == 1 && m1.Mallocs != m2.Mallocs {
t.Errorf("allocation occurred during write")
}
// Check that buffer has correct data.
if !bytes.Equal(buf.Bytes()[0:startLen-readBytes], xBytes[readBytes:]) {
t.Errorf("bad initial data at %d %d", startLen, growLen)
}
if !bytes.Equal(buf.Bytes()[startLen-readBytes:startLen-readBytes+growLen], yBytes) {
t.Errorf("bad written data at %d %d", startLen, growLen)
}
}
}
}
// Was a bug: used to give EOF reading empty slice at EOF.
func TestReadEmptyAtEOF(t *testing.T) {
b := new(Buffer)
slice := make([]byte, 0)
n, err := b.Read(slice)
if err != nil {
t.Errorf("read error: %v", err)
}
if n != 0 {
t.Errorf("wrong count; got %d want 0", n)
}
}
// Tests that we occasionally compact. Issue 5154.
func TestBufferGrowth(t *testing.T) {
var b Buffer
buf := make([]byte, 1024)
b.Write(buf[0:1])
var cap0 int
for i := 0; i < 5<<10; i++ {
b.Write(buf)
b.Read(buf)
if i == 0 {
cap0 = b.Cap()
}
}
cap1 := b.Cap()
// (*Buffer).grow allows for 2x capacity slop before sliding,
// so set our error threshold at 3x.
if cap1 > cap0*3 {
t.Errorf("buffer cap = %d; too big (grew from %d)", cap1, cap0)
}
}
func BenchmarkWriteByte(b *testing.B) {
const n = 4 << 10
b.SetBytes(n)
buf := NewBuffer(make([]byte, n))
for i := 0; i < b.N; i++ {
buf.Reset()
for i := 0; i < n; i++ {
buf.WriteByte('x')
}
}
}
// From Issue 5154.
func BenchmarkBufferNotEmptyWriteRead(b *testing.B) {
buf := make([]byte, 1024)
for i := 0; i < b.N; i++ {
var b Buffer
b.Write(buf[0:1])
for i := 0; i < 5<<10; i++ {
b.Write(buf)
b.Read(buf)
}
}
}
// Check that we don't compact too often. From Issue 5154.
func BenchmarkBufferFullSmallReads(b *testing.B) {
buf := make([]byte, 1024)
for i := 0; i < b.N; i++ {
var b Buffer
b.Write(buf)
for b.Len()+20 < b.Cap() {
b.Write(buf[:10])
}
for i := 0; i < 5<<10; i++ {
b.Read(buf[:1])
b.Write(buf[:1])
}
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment