Add sync.Pool to async reader

This commit is contained in:
Nick Craig-Wood 2017-02-14 22:28:18 +00:00
parent 6fc114d681
commit 3f778d70f7
3 changed files with 68 additions and 43 deletions

View file

@ -351,13 +351,13 @@ func NewAccountSizeNameWithBuffer(in io.ReadCloser, size int64, name string) *Ac
const bufSize = 1024 * 1024 const bufSize = 1024 * 1024
var buffers int var buffers int
if size >= int64(Config.BufferSize) { if size >= int64(Config.BufferSize) {
buffers = int(int64(Config.BufferSize) / bufSize) buffers = int(int64(Config.BufferSize) / asyncBufferSize)
} else { } else {
buffers = int(size / bufSize) buffers = int(size / asyncBufferSize)
} }
// On big files add a buffer // On big files add a buffer
if buffers > 0 { if buffers > 0 {
newIn, err := newAsyncReader(in, buffers, bufSize) newIn, err := newAsyncReader(in, buffers)
if err != nil { if err != nil {
Errorf(name, "Failed to make buffer: %v", err) Errorf(name, "Failed to make buffer: %v", err)
} else { } else {

View file

@ -2,10 +2,17 @@ package fs
import ( import (
"io" "io"
"sync"
"github.com/pkg/errors" "github.com/pkg/errors"
) )
const asyncBufferSize = 128 * 1024
var asyncBufferPool = sync.Pool{
New: func() interface{} { return newBuffer() },
}
// asyncReader will do async read-ahead from the input reader // asyncReader will do async read-ahead from the input reader
// and make the data available as an io.Reader. // and make the data available as an io.Reader.
// This should be fully transparent, except that once an error // This should be fully transparent, except that once an error
@ -13,25 +20,21 @@ import (
type asyncReader struct { type asyncReader struct {
in io.ReadCloser // Input reader in io.ReadCloser // Input reader
ready chan *buffer // Buffers ready to be handed to the reader ready chan *buffer // Buffers ready to be handed to the reader
reuse chan *buffer // Buffers to reuse for input reading token chan struct{} // Tokens which allow a buffer to be taken
exit chan struct{} // Closes when finished exit chan struct{} // Closes when finished
buffers int // Number of buffers buffers int // Number of buffers
err error // If an error has occurred it is here err error // If an error has occurred it is here
cur *buffer // Current buffer being served cur *buffer // Current buffer being served
exited chan struct{} // Channel is closed been the async reader shuts down exited chan struct{} // Channel is closed been the async reader shuts down
closed bool // Has the parent reader been closed?
} }
// newAsyncReader returns a reader that will asynchronously read from // newAsyncReader returns a reader that will asynchronously read from
// the supplied Reader into a number of buffers each with a given size. // the supplied Reader into a number of buffers each of size asyncBufferSize
// It will start reading from the input at once, maybe even before this // It will start reading from the input at once, maybe even before this
// function has returned. // function has returned.
// The input can be read from the returned reader. // The input can be read from the returned reader.
// When done use Close to release the buffers and close the supplied input. // When done use Close to release the buffers and close the supplied input.
func newAsyncReader(rd io.ReadCloser, buffers, size int) (io.ReadCloser, error) { func newAsyncReader(rd io.ReadCloser, buffers int) (io.ReadCloser, error) {
if size <= 0 {
return nil, errors.New("buffer size too small")
}
if buffers <= 0 { if buffers <= 0 {
return nil, errors.New("number of buffers too small") return nil, errors.New("number of buffers too small")
} }
@ -39,35 +42,36 @@ func newAsyncReader(rd io.ReadCloser, buffers, size int) (io.ReadCloser, error)
return nil, errors.New("nil reader supplied") return nil, errors.New("nil reader supplied")
} }
a := &asyncReader{} a := &asyncReader{}
a.init(rd, buffers, size) a.init(rd, buffers)
return a, nil return a, nil
} }
func (a *asyncReader) init(rd io.ReadCloser, buffers, size int) { func (a *asyncReader) init(rd io.ReadCloser, buffers int) {
a.in = rd a.in = rd
a.ready = make(chan *buffer, buffers) a.ready = make(chan *buffer, buffers)
a.reuse = make(chan *buffer, buffers) a.token = make(chan struct{}, buffers)
a.exit = make(chan struct{}, 0) a.exit = make(chan struct{}, 0)
a.exited = make(chan struct{}, 0) a.exited = make(chan struct{}, 0)
a.buffers = buffers a.buffers = buffers
a.cur = nil a.cur = nil
// Create buffers // Create tokens
for i := 0; i < buffers; i++ { for i := 0; i < buffers; i++ {
a.reuse <- newBuffer(size) a.token <- struct{}{}
} }
// Start async reader // Start async reader
go func() { go func() {
// Ensure that when we exit this is signalled. // Ensure that when we exit this is signalled.
defer close(a.exited) defer close(a.exited)
defer close(a.ready)
for { for {
select { select {
case b := <-a.reuse: case <-a.token:
b := a.getBuffer()
err := b.read(a.in) err := b.read(a.in)
a.ready <- b a.ready <- b
if err != nil { if err != nil {
close(a.ready)
return return
} }
case <-a.exit: case <-a.exit:
@ -77,11 +81,23 @@ func (a *asyncReader) init(rd io.ReadCloser, buffers, size int) {
}() }()
} }
// return the buffer to the pool (clearing it)
func (a *asyncReader) putBuffer(b *buffer) {
b.clear()
asyncBufferPool.Put(b)
}
// get a buffer from the pool
func (a *asyncReader) getBuffer() *buffer {
return asyncBufferPool.Get().(*buffer)
}
// Read will return the next available data. // Read will return the next available data.
func (a *asyncReader) fill() (err error) { func (a *asyncReader) fill() (err error) {
if a.cur.isEmpty() { if a.cur.isEmpty() {
if a.cur != nil { if a.cur != nil {
a.reuse <- a.cur a.putBuffer(a.cur)
a.token <- struct{}{}
a.cur = nil a.cur = nil
} }
b, ok := <-a.ready b, ok := <-a.ready
@ -139,18 +155,24 @@ func (a *asyncReader) WriteTo(w io.Writer) (n int64, err error) {
// Close will ensure that the underlying async reader is shut down. // Close will ensure that the underlying async reader is shut down.
// It will also close the input supplied on newAsyncReader. // It will also close the input supplied on newAsyncReader.
func (a *asyncReader) Close() (err error) { func (a *asyncReader) Close() (err error) {
// Return if already closed
select { select {
case <-a.exited: case <-a.exit:
return
default: default:
}
// Close and wait for go routine
close(a.exit) close(a.exit)
<-a.exited <-a.exited
// Return any outstanding buffers to the Pool
if a.cur != nil {
a.putBuffer(a.cur)
}
for b := range a.ready {
a.putBuffer(b)
} }
if !a.closed {
a.closed = true
return a.in.Close() return a.in.Close()
} }
return nil
}
// Internal buffer // Internal buffer
// If an error is present, it must be returned // If an error is present, it must be returned
@ -159,11 +181,20 @@ type buffer struct {
buf []byte buf []byte
err error err error
offset int offset int
size int
} }
func newBuffer(size int) *buffer { func newBuffer() *buffer {
return &buffer{buf: make([]byte, size), err: nil, size: size} return &buffer{
buf: make([]byte, asyncBufferSize),
err: nil,
}
}
// clear returns the buffer to its full size and clears the members
func (b *buffer) clear() {
b.buf = b.buf[:cap(b.buf)]
b.err = nil
b.offset = 0
} }
// isEmpty returns true is offset is at end of // isEmpty returns true is offset is at end of
@ -183,7 +214,7 @@ func (b *buffer) isEmpty() bool {
// Any error encountered during the read is returned. // Any error encountered during the read is returned.
func (b *buffer) read(rd io.Reader) error { func (b *buffer) read(rd io.Reader) error {
var n int var n int
n, b.err = ReadFill(rd, b.buf[0:b.size]) n, b.err = ReadFill(rd, b.buf)
b.buf = b.buf[0:n] b.buf = b.buf[0:n]
b.offset = 0 b.offset = 0
return b.err return b.err

View file

@ -15,7 +15,7 @@ import (
func TestAsyncReader(t *testing.T) { func TestAsyncReader(t *testing.T) {
buf := ioutil.NopCloser(bytes.NewBufferString("Testbuffer")) buf := ioutil.NopCloser(bytes.NewBufferString("Testbuffer"))
ar, err := newAsyncReader(buf, 4, 10000) ar, err := newAsyncReader(buf, 4)
require.NoError(t, err) require.NoError(t, err)
var dst = make([]byte, 100) var dst = make([]byte, 100)
@ -40,7 +40,7 @@ func TestAsyncReader(t *testing.T) {
// Test Close without reading everything // Test Close without reading everything
buf = ioutil.NopCloser(bytes.NewBuffer(make([]byte, 50000))) buf = ioutil.NopCloser(bytes.NewBuffer(make([]byte, 50000)))
ar, err = newAsyncReader(buf, 4, 100) ar, err = newAsyncReader(buf, 4)
require.NoError(t, err) require.NoError(t, err)
err = ar.Close() err = ar.Close()
require.NoError(t, err) require.NoError(t, err)
@ -49,7 +49,7 @@ func TestAsyncReader(t *testing.T) {
func TestAsyncWriteTo(t *testing.T) { func TestAsyncWriteTo(t *testing.T) {
buf := ioutil.NopCloser(bytes.NewBufferString("Testbuffer")) buf := ioutil.NopCloser(bytes.NewBufferString("Testbuffer"))
ar, err := newAsyncReader(buf, 4, 10000) ar, err := newAsyncReader(buf, 4)
require.NoError(t, err) require.NoError(t, err)
var dst = &bytes.Buffer{} var dst = &bytes.Buffer{}
@ -68,20 +68,14 @@ func TestAsyncWriteTo(t *testing.T) {
func TestAsyncReaderErrors(t *testing.T) { func TestAsyncReaderErrors(t *testing.T) {
// test nil reader // test nil reader
_, err := newAsyncReader(nil, 4, 10000) _, err := newAsyncReader(nil, 4)
require.Error(t, err) require.Error(t, err)
// invalid buffer number // invalid buffer number
buf := ioutil.NopCloser(bytes.NewBufferString("Testbuffer")) buf := ioutil.NopCloser(bytes.NewBufferString("Testbuffer"))
_, err = newAsyncReader(buf, 0, 10000) _, err = newAsyncReader(buf, 0)
require.Error(t, err) require.Error(t, err)
_, err = newAsyncReader(buf, -1, 10000) _, err = newAsyncReader(buf, -1)
require.Error(t, err)
// invalid buffer size
_, err = newAsyncReader(buf, 4, 0)
require.Error(t, err)
_, err = newAsyncReader(buf, 4, -1)
require.Error(t, err) require.Error(t, err)
} }
@ -161,7 +155,7 @@ func TestAsyncReaderSizes(t *testing.T) {
bufsize := bufsizes[k] bufsize := bufsizes[k]
read := readmaker.fn(strings.NewReader(text)) read := readmaker.fn(strings.NewReader(text))
buf := bufio.NewReaderSize(read, bufsize) buf := bufio.NewReaderSize(read, bufsize)
ar, _ := newAsyncReader(ioutil.NopCloser(buf), l, 100) ar, _ := newAsyncReader(ioutil.NopCloser(buf), l)
s := bufreader.fn(ar) s := bufreader.fn(ar)
// "timeout" expects the Reader to recover, asyncReader does not. // "timeout" expects the Reader to recover, asyncReader does not.
if s != text && readmaker.name != "timeout" { if s != text && readmaker.name != "timeout" {
@ -200,7 +194,7 @@ func TestAsyncReaderWriteTo(t *testing.T) {
bufsize := bufsizes[k] bufsize := bufsizes[k]
read := readmaker.fn(strings.NewReader(text)) read := readmaker.fn(strings.NewReader(text))
buf := bufio.NewReaderSize(read, bufsize) buf := bufio.NewReaderSize(read, bufsize)
ar, _ := newAsyncReader(ioutil.NopCloser(buf), l, 100) ar, _ := newAsyncReader(ioutil.NopCloser(buf), l)
dst := &bytes.Buffer{} dst := &bytes.Buffer{}
wt := ar.(io.WriterTo) wt := ar.(io.WriterTo)
_, err := wt.WriteTo(dst) _, err := wt.WriteTo(dst)