276 lines
6.9 KiB
Go
276 lines
6.9 KiB
Go
package pool
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
)
|
|
|
|
// RWAccount is a function which will be called after every read
|
|
// from the RW.
|
|
//
|
|
// It may return an error which will be passed back to the user.
|
|
type RWAccount func(n int) error
|
|
|
|
// RW contains the state for the read/writer
|
|
type RW struct {
|
|
pool *Pool // pool to get pages from
|
|
pages [][]byte // backing store
|
|
size int // size written
|
|
out int // offset we are reading from
|
|
lastOffset int // size in last page
|
|
account RWAccount // account for a read
|
|
reads int // count how many times the data has been read
|
|
accountOn int // only account on or after this read
|
|
}
|
|
|
|
var (
|
|
errInvalidWhence = errors.New("pool.RW Seek: invalid whence")
|
|
errNegativeSeek = errors.New("pool.RW Seek: negative position")
|
|
errSeekPastEnd = errors.New("pool.RW Seek: attempt to seek past end of data")
|
|
)
|
|
|
|
// NewRW returns a reader / writer which is backed from pages from the
|
|
// pool passed in.
|
|
//
|
|
// Data can be stored in it by calling Write and read from it by
|
|
// calling Read.
|
|
//
|
|
// When writing it only appends data. Seek only applies to reading.
|
|
func NewRW(pool *Pool) *RW {
|
|
return &RW{
|
|
pool: pool,
|
|
pages: make([][]byte, 0, 16),
|
|
}
|
|
}
|
|
|
|
// SetAccounting should be provided with a function which will be
|
|
// called after every read from the RW.
|
|
//
|
|
// It may return an error which will be passed back to the user.
|
|
func (rw *RW) SetAccounting(account RWAccount) *RW {
|
|
rw.account = account
|
|
return rw
|
|
}
|
|
|
|
// DelayAccountinger enables an accounting delay
|
|
type DelayAccountinger interface {
|
|
// DelayAccounting makes sure the accounting function only
|
|
// gets called on the i-th or later read of the data from this
|
|
// point (counting from 1).
|
|
//
|
|
// This is useful so that we don't account initial reads of
|
|
// the data e.g. when calculating hashes.
|
|
//
|
|
// Set this to 0 to account everything.
|
|
DelayAccounting(i int)
|
|
}
|
|
|
|
// DelayAccounting makes sure the accounting function only gets called
|
|
// on the i-th or later read of the data from this point (counting
|
|
// from 1).
|
|
//
|
|
// This is useful so that we don't account initial reads of the data
|
|
// e.g. when calculating hashes.
|
|
//
|
|
// Set this to 0 to account everything.
|
|
func (rw *RW) DelayAccounting(i int) {
|
|
rw.accountOn = i
|
|
rw.reads = 0
|
|
}
|
|
|
|
// Returns the page and offset of i for reading.
|
|
//
|
|
// Ensure there are pages before calling this.
|
|
func (rw *RW) readPage(i int) (page []byte) {
|
|
// Count a read of the data if we read the first page
|
|
if i == 0 {
|
|
rw.reads++
|
|
}
|
|
pageNumber := i / rw.pool.bufferSize
|
|
offset := i % rw.pool.bufferSize
|
|
page = rw.pages[pageNumber]
|
|
// Clip the last page to the amount written
|
|
if pageNumber == len(rw.pages)-1 {
|
|
page = page[:rw.lastOffset]
|
|
}
|
|
return page[offset:]
|
|
}
|
|
|
|
// account for n bytes being read
|
|
func (rw *RW) accountRead(n int) error {
|
|
if rw.account == nil {
|
|
return nil
|
|
}
|
|
// Don't start accounting until we've reached this many reads
|
|
//
|
|
// rw.reads will be 1 the first time this is called
|
|
// rw.accountOn 2 means start accounting on the 2nd read through
|
|
if rw.reads >= rw.accountOn {
|
|
return rw.account(n)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Read reads up to len(p) bytes into p. It returns the number of
|
|
// bytes read (0 <= n <= len(p)) and any error encountered. If some
|
|
// data is available but not len(p) bytes, Read returns what is
|
|
// available instead of waiting for more.
|
|
func (rw *RW) Read(p []byte) (n int, err error) {
|
|
var (
|
|
nn int
|
|
page []byte
|
|
)
|
|
for len(p) > 0 {
|
|
if rw.out >= rw.size {
|
|
return n, io.EOF
|
|
}
|
|
page = rw.readPage(rw.out)
|
|
nn = copy(p, page)
|
|
p = p[nn:]
|
|
n += nn
|
|
rw.out += nn
|
|
err = rw.accountRead(nn)
|
|
if err != nil {
|
|
return n, err
|
|
}
|
|
}
|
|
return n, nil
|
|
}
|
|
|
|
// WriteTo writes data to w until there's no more data to write or
|
|
// when an error occurs. The return value n is the number of bytes
|
|
// written. Any error encountered during the write is also returned.
|
|
//
|
|
// The Copy function uses WriteTo if available. This avoids an
|
|
// allocation and a copy.
|
|
func (rw *RW) WriteTo(w io.Writer) (n int64, err error) {
|
|
var (
|
|
nn int
|
|
page []byte
|
|
)
|
|
for rw.out < rw.size {
|
|
page = rw.readPage(rw.out)
|
|
nn, err = w.Write(page)
|
|
n += int64(nn)
|
|
rw.out += nn
|
|
if err != nil {
|
|
return n, err
|
|
}
|
|
err = rw.accountRead(nn)
|
|
if err != nil {
|
|
return n, err
|
|
}
|
|
}
|
|
return n, nil
|
|
}
|
|
|
|
// Get the page we are writing to
|
|
func (rw *RW) writePage() (page []byte) {
|
|
if len(rw.pages) > 0 && rw.lastOffset < rw.pool.bufferSize {
|
|
return rw.pages[len(rw.pages)-1][rw.lastOffset:]
|
|
}
|
|
page = rw.pool.Get()
|
|
rw.pages = append(rw.pages, page)
|
|
rw.lastOffset = 0
|
|
return page
|
|
}
|
|
|
|
// Write writes len(p) bytes from p to the underlying data stream. It returns
|
|
// the number of bytes written len(p). It cannot return an error.
|
|
func (rw *RW) Write(p []byte) (n int, err error) {
|
|
var (
|
|
nn int
|
|
page []byte
|
|
)
|
|
for len(p) > 0 {
|
|
page = rw.writePage()
|
|
nn = copy(page, p)
|
|
p = p[nn:]
|
|
n += nn
|
|
rw.size += nn
|
|
rw.lastOffset += nn
|
|
}
|
|
return n, nil
|
|
}
|
|
|
|
// ReadFrom reads data from r until EOF or error. The return value n is the
|
|
// number of bytes read. Any error except EOF encountered during the read is
|
|
// also returned.
|
|
//
|
|
// The Copy function uses ReadFrom if available. This avoids an
|
|
// allocation and a copy.
|
|
func (rw *RW) ReadFrom(r io.Reader) (n int64, err error) {
|
|
var (
|
|
nn int
|
|
page []byte
|
|
)
|
|
for err == nil {
|
|
page = rw.writePage()
|
|
nn, err = r.Read(page)
|
|
n += int64(nn)
|
|
rw.size += nn
|
|
rw.lastOffset += nn
|
|
}
|
|
if err == io.EOF {
|
|
err = nil
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
// Seek sets the offset for the next Read (not Write - this is always
|
|
// appended) to offset, interpreted according to whence: SeekStart
|
|
// means relative to the start of the file, SeekCurrent means relative
|
|
// to the current offset, and SeekEnd means relative to the end (for
|
|
// example, offset = -2 specifies the penultimate byte of the file).
|
|
// Seek returns the new offset relative to the start of the file or an
|
|
// error, if any.
|
|
//
|
|
// Seeking to an offset before the start of the file is an error. Seeking
|
|
// beyond the end of the written data is an error.
|
|
func (rw *RW) Seek(offset int64, whence int) (int64, error) {
|
|
var abs int64
|
|
size := int64(rw.size)
|
|
switch whence {
|
|
case io.SeekStart:
|
|
abs = offset
|
|
case io.SeekCurrent:
|
|
abs = int64(rw.out) + offset
|
|
case io.SeekEnd:
|
|
abs = size + offset
|
|
default:
|
|
return 0, errInvalidWhence
|
|
}
|
|
if abs < 0 {
|
|
return 0, errNegativeSeek
|
|
}
|
|
if abs > size {
|
|
return offset - (abs - size), errSeekPastEnd
|
|
}
|
|
rw.out = int(abs)
|
|
return abs, nil
|
|
}
|
|
|
|
// Close the buffer returning memory to the pool
|
|
func (rw *RW) Close() error {
|
|
for _, page := range rw.pages {
|
|
rw.pool.Put(page)
|
|
}
|
|
rw.pages = nil
|
|
return nil
|
|
}
|
|
|
|
// Size returns the number of bytes in the buffer
|
|
func (rw *RW) Size() int64 {
|
|
return int64(rw.size)
|
|
}
|
|
|
|
// Check interfaces
|
|
var (
|
|
_ io.Reader = (*RW)(nil)
|
|
_ io.ReaderFrom = (*RW)(nil)
|
|
_ io.Writer = (*RW)(nil)
|
|
_ io.WriterTo = (*RW)(nil)
|
|
_ io.Seeker = (*RW)(nil)
|
|
_ io.Closer = (*RW)(nil)
|
|
_ DelayAccountinger = (*RW)(nil)
|
|
)
|