buffer/lib.go

// Package buffer provides pooled Buffer IO for Go programs.
//
// It implements a reader similar to bufio.Reader. The underlying memory buffers can be used from a synchronized
// pool. It implements a writer that can be used to write data to an underlying writer in larger, buffered,
// chunks.
package buffer

import (
	"errors"
	"io"
)

// BufferPool defines the interface for the used buffer pool. The buffered reader or writer can be used either
// with the built-in default pool, noop pool, or with a custom pool implementation.
type BufferPool interface {

	// Get should return a non-zero length byte slice. In case it returns a zero length byte slice, or an
	// explicit error, the read operations will fail.
	//
	// It is OK but not recommended to return varying sizes of byte slices.
	Get() ([]byte, error)

	// The reader puts back the byte slices taken by Get, using Put, when it does not use them anymore.
	Put([]byte)
}

// Options provides options for the Reader.
type Options struct {

	// BufferPool defines the buffer pool to be used. It defaults to the pool created by DefaultPool(). It
	// is expected to explicitly set the Pool instance, otherwise, not defining any globals, each Reader
	// instance will create its own pool.
	BufferPool BufferPool

	// InitialSegment allows to initialize the buffer with a memory space already available to the user code
	// without requesting it from the used pool.
	InitialSegment []byte
}

// ContentFunc wraps a function to implement the io.WriterTo interface. Function implementations should be
// ready to be executed in goroutines other than what they were created in.
type ContentFunc func(io.Writer) (int64, error)

// Reader wraps an underlying io.Reader or io.WriterTo, and provides buffered IO via its methods. Initialize it
// via BufferedReader or BufferedContent.
//
// It reads from the underlying source until the first error, but only returns an error when the buffer is
// empty. Once the underlying reader returned an error, it does not attempt to read from it anymore.
//
// The reader does not support concurrent access.
type Reader struct {
	reader *reader
}

// Writer wraps an underlying io.Writer and provides buffered IO via its methods. Initialize it via
// BufferedWriter.
//
// It writes the input bytes into an internal buffer, and flushes them to the underlying writer only when the
// buffer is full or when the writer is closed. If necessary, the buffer can be explicitly flushed.
//
// The writer does not support concurrent access.
type Writer struct {
	writer *writer
}

var (
	// ErrZeroAllocation is returned when the used pool returned a zero length byte slice.
	ErrZeroAllocation = errors.New("zero allocation")

	// ErrAbort is returned to the writer process when using buffered content, when the reader
	// experienced an error. ErrAbort is returned to the reader process, if Close() was called and no read
	// error was received before it.
	ErrAbort = errors.New("read aborted")
)

// DefultPool initializes a synchronized pool that stores and returns byte slices of allocSize length. It can be
// used with multiple readers concurrently.
func DefaultPool(allocSize int) BufferPool {
	if allocSize <= 0 {
		allocSize = 1 << 12
	}

	return newPool(allocSize)
}

// NoPool returns a noop pool that allocates memory on every call to Get() and does nothing on Put().
func NoPool(allocSize int) BufferPool {
	if allocSize <= 0 {
		allocSize = 1 << 12
	}

	return noPool{allocSize: allocSize}
}

// WriteTo implements the WriterTo interface.
func (f ContentFunc) WriteTo(w io.Writer) (int64, error) {
	return f(w)
}

// BufferedReader creates a buffered reader using the input reader as the underlying source.
func BufferedReader(in io.Reader, o Options) Reader {
	if in == nil {
		return Reader{}
	}

	if o.BufferPool == nil {
		o.BufferPool = DefaultPool(1 << 12)
	}

	r := Reader{reader: &reader{options: o, in: in}}
	if len(o.InitialSegment) > 0 {
		r.reader.segments = [][]byte{o.InitialSegment}
	}

	return r
}

// BufferedContent creates a buffered reader using the input content (io.WriterTo) as the underlying source.
//
// It is similar to an io.Pipe, but with dynamic and pooled buffering internally. The individual Write calls are
// blocked until the reading side requests more data.
//
// The provided WriterTo instances need to be safe to call in goroutines other than they were created in. If the
// writer function returns with nil error, it will be interpreted as EOF on the reader side. When the reader
// side experiences an error, and the writer still has content to be written, the passed in io.Writer will
// return an ErrAbort error.
func BufferedContent(c io.WriterTo, o Options) Reader {
	if c == nil {
		return Reader{}
	}

	if o.BufferPool == nil {
		o.BufferPool = DefaultPool(1 << 12)
	}

	return Reader{reader: &reader{options: o, in: mkcontent(c)}}
}

// Read reads max len(p) copied to p and returns how many bytes were read.
//
// It only returns an error when the buffer is empty. It only returns an error when the BufferPool.Get fails or
// the underlying reader returns an error.
//
// It may return zero read length and nil error, but only if the underlying reader did so.
//
// It does not block when the buffer is not empty, even if the underlying reader would block on the next read.
//
// It only allocates a single segment of memory, or none if one was provided during innitialization.
func (r Reader) Read(p []byte) (int, error) {
	if r.reader == nil {
		return 0, io.EOF
	}

	return r.reader.read(p)
}

// ReadBytes reads until the first occurence of delimiter in the input, within max length.
//
// It returns the bytes and nil error, when the delimiter was found within max. In this case it consumes
// the buffer until and including the delimiter. If the underlying reader returned meanwhile a non-nil error,
// including EOF, it will be returned on subsequent reads, but only after the internal buffer was consumed.
//
// It returns zero length bytes and nil error, when the delimiter was not found within max.
//
// It blocks only when less data is buffered than max, the delimiter was not found within the range defined by
// max, and the underlying reader blocks.
func (r Reader) ReadBytes(delimiter []byte, max int) ([]byte, error) {
	if r.reader == nil {
		return nil, io.EOF
	}

	return r.reader.readBytes(delimiter, max)
}

// ReadUTF8 returns at most max UTF8 characters read from the underlying reader.
//
// When valid UTF8 characters were found, it returns the characters, the number of bytes consumed, and nil
// error.
//
// It only returns an error, including EOF, when the underlying buffer is empty.
//
// It may return zero length characters and nil error, if the underlying reader returned zero length read with
// nil error.
//
// It supports recovery of UTF8 streams by skipping the invalid characters. In such cases, first it returns the
// valid characters with the number of bytes consumed and nil error, then in a subsequent call, it returns
// zero characters, 1 and nil error.
//
// It blocks only when there is not enough data buffered to determine the next UTF8 character and the underlying
// reader blocks.
func (r Reader) ReadUTF8(max int) ([]rune, int, error) {
	if r.reader == nil {
		return nil, 0, io.EOF
	}

	return r.reader.readUTF8(max)
}

// Peek returns at most max read bytes and keeps them buffered.
//
// If the underlying reader returned an error, it returns an error when the underlying buffer was fully consumed
// by Read, ReadBytes, ReadUTF8 or WriteTo.
//
// The returned byte slice is not a copy of the buffered bytes, and therefore should not be modified.
//
// It blocks when there is less data buffered than the range defined by max and the underlying reader blocks.
func (r Reader) Peek(max int) ([]byte, error) {
	if r.reader == nil {
		return nil, io.EOF
	}

	return r.reader.peek(max)
}

// Buffered returns the buffered bytes.
//
// The returned byte slice is not a copy of the buffered bytes, and therefore should not be modified.
func (r Reader) Buffered() []byte {
	if r.reader == nil {
		return nil
	}

	return r.reader.buffered()
}

// WriteTo implements the io.WriterTo interface. It copies all the data from the underlying reader to the
// provided writer using a buffer from the pool.
//
// It is important that the provided writer must not modify the slice data, as defined in the io.Writer
// interface documentation.
//
// It blocks when the underlying reader blocks.
func (r Reader) WriteTo(w io.Writer) (int64, error) {
	if r.reader == nil {
		return 0, nil
	}

	return r.reader.writeTo(w)
}

// Close releases the resource held by the Reader, and puts the underlying byte buffers, that were
// requested from the pool, back into the used pool. The reader cannot be used for read operations after Close()
// was called.
//
// It does not close the underlying reader.
func (r Reader) Close() {
	if r.reader == nil {
		return
	}

	r.reader.free()
}

// BufferedWriter initializes a Writer.
func BufferedWriter(out io.Writer, o Options) Writer {
	if out == nil {
		return Writer{}
	}

	if o.BufferPool == nil {
		o.BufferPool = DefaultPool(1 << 12)
	}

	w := Writer{writer: &writer{out: out, options: o}}
	if len(o.InitialSegment) > 0 {
		w.writer.buffer = o.InitialSegment
	}

	return w
}

// Write writes to the writer's buffer, and if the buffer is full, it causes writing out the buffer's contents
// to the underlying writer.
//
// It blocks when there is not enough buffer space for the data being written and the underlying writer blocks.
func (w Writer) Write(p []byte) (int, error) {
	if w.writer == nil {
		return 0, errors.New("unitialized writer")
	}

	return w.writer.write(p)
}

// ReadFrom implements the io.ReaderFrom interface. It copies all the data from the provided reader to the
// underlying writer using a buffer from the pool.
//
// It flushes the buffer at the end of the input stream.
func (w Writer) ReadFrom(r io.Reader) (int64, error) {
	if w.writer == nil {
		return 0, errors.New("unitialized writer")
	}

	return w.writer.readFrom(r)
}

// Flush forces the writer to flush the buffered content to the underlying writer. After flushed, the writer
// still accepts further writes.
//
// It blocks when there is buffered data to write out and the underlying writer blocks.
func (w Writer) Flush() error {
	if w.writer == nil {
		return nil
	}

	return w.writer.flush()
}

// Close flushes the buffered content if any and closes the writer. After closed, the writer does not accept
// further writes.
//
// It blocks when there is buffered data to write out and the underlying writer blocks.
//
// It does not close the underlying writer.
func (w Writer) Close() error {
	if w.writer == nil {
		return nil
	}

	return w.writer.close()
}