package pool import (	"context"	"errors"	"fmt"	"io"	"io/fs"	"sync"	"time") // 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//// It can be used as a FIFO to read data from a source and write it out again.type RW struct {	// Written once variables in initialization	pool      *Pool     // pool to get pages from	account   RWAccount // account for a read	accountOn int       // only account on or after this read 	// Shared variables between Read and Write	// Write updates these but Read reads from them	// They must all stay in sync together	mu         sync.Mutex    // protect the shared variables	pages      [][]byte      // backing store	size       int           // size written	lastOffset int           // size in last page	written    chan struct{} // signalled when a write happens	closed     bool          // set after Close and the pages are returned to the pool 	// Read side Variables	out   int // offset we are reading from	reads int // count how many times the data has been read 	reserved [][]byte // reserved buffers} var (	// ErrClosed is returned from all operations on an RW after Close.	// It wraps fs.ErrClosed so errors.Is(err, fs.ErrClosed) is true.	ErrClosed        = fmt.Errorf("pool.RW: %w", fs.ErrClosed)	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 {	rw := &RW{		pool:    pool,		pages:   make([][]byte, 0, 16),		written: make(chan struct{}, 1),	}	return rw} // Reserve bytes of memory.//// This allocates n bytes of memory for later use.//// This is rounded up to the nearest buffer page size.//// Only safe to call once.func (rw *RW) Reserve(n int64) *RW {	rw.mu.Lock()	defer rw.mu.Unlock()	buffers := int((n + int64(rw.pool.bufferSize) - 1) / int64(rw.pool.bufferSize))	rw.reserved = rw.pool.GetN(buffers)	return rw} // 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.//// Not thread safe - call in initialization only.func (rw *RW) SetAccounting(account RWAccount) *RW {	rw.account = account	return rw} // DelayAccountinger enables an accounting delaytype 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.//// Not thread safe - call in initialization only.func (rw *RW) DelayAccounting(i int) {	rw.accountOn = i	rw.reads = 0} // Returns the page and offset of i for reading.//// Returns ErrClosed if the RW has been closed or io.EOF if i is at or// beyond the end of the data written.func (rw *RW) readPage(i int) (page []byte, err error) {	rw.mu.Lock()	defer rw.mu.Unlock()	if rw.closed {		return nil, ErrClosed	}	if i >= rw.size {		return nil, io.EOF	}	// 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:], nil} // account for n bytes being readfunc (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 {		page, err = rw.readPage(rw.out)		if err != nil {			return n, err		}		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 {		page, err = rw.readPage(rw.out)		if err == io.EOF {			return n, nil		} else if err != nil {			return n, err		}		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		}	}} // Get the page we are writing to//// Returns ErrClosed if the RW has been closed.func (rw *RW) writePage() (page []byte, err error) {	rw.mu.Lock()	defer rw.mu.Unlock()	if rw.closed {		return nil, ErrClosed	}	if len(rw.pages) > 0 && rw.lastOffset < rw.pool.bufferSize {		return rw.pages[len(rw.pages)-1][rw.lastOffset:], nil	}	if len(rw.reserved) > 0 {		// Get reserved pages if available		i := len(rw.reserved) - 1		page = rw.reserved[i]		rw.reserved[i] = nil		rw.reserved = rw.reserved[:i]	} else {		page = rw.pool.Get()	}	rw.pages = append(rw.pages, page)	rw.lastOffset = 0	return page, nil} // 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, err = rw.writePage()		if err != nil {			return n, err		}		nn = copy(page, p)		p = p[nn:]		n += nn		rw.mu.Lock()		rw.size += nn		rw.lastOffset += nn		rw.mu.Unlock()		rw.signalWrite() // signal more data available	}	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, err = rw.writePage()		if err != nil {			return n, err		}		nn, err = r.Read(page)		n += int64(nn)		rw.mu.Lock()		rw.size += nn		rw.lastOffset += nn		rw.mu.Unlock()		rw.signalWrite() // signal more data available	}	if err == io.EOF {		err = nil	}	return n, err} // signal that a write has happenedfunc (rw *RW) signalWrite() {	select {	case rw.written <- struct{}{}:	default:	}} // WaitWrite sleeps until a data is written to the RW or Close is// called or the context is cancelled occurs or for a maximum of 1// Second then returns.//// This can be used when calling Read while the buffer is filling up.func (rw *RW) WaitWrite(ctx context.Context) {	timer := time.NewTimer(time.Second)	select {	case <-timer.C:	case <-ctx.Done():	case <-rw.written:	}	timer.Stop()} // 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	rw.mu.Lock()	size := int64(rw.size)	closed := rw.closed	rw.mu.Unlock()	if closed {		return 0, ErrClosed	}	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.//// It is safe to call more than once. All other operations return// ErrClosed afterwards.func (rw *RW) Close() error {	rw.mu.Lock()	defer rw.mu.Unlock()	rw.closed = true	rw.signalWrite() // signal more data available	rw.pool.PutN(rw.pages)	clear(rw.pages)	rw.pages = nil	rw.pool.PutN(rw.reserved)	clear(rw.reserved)	rw.reserved = nil	return nil} // Size returns the number of bytes in the bufferfunc (rw *RW) Size() int64 {	rw.mu.Lock()	defer rw.mu.Unlock()	return int64(rw.size)} // Check interfacesvar (	_ 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))