123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277 |
- 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)
- )
|