buffer/writer.go

package buffer

import (
	"errors"
	"io"
)

type writer struct {
	out         io.Writer
	options     Options
	buffer      []byte
	offset, len int
	err         error
}

var errClosed = errors.New("buffer closed")

func (w *writer) write(p []byte) (int, error) {
	var n int
	for {
		if w.err != nil {
			return n, w.err
		}

		if len(p) == 0 {
			return n, nil
		}

		if len(w.buffer) == 0 {
			w.buffer, w.err = w.options.Pool.Get()
			if len(w.buffer) == 0 && w.err == nil {
				w.err = ErrZeroAllocation
			}

			continue
		}

		if w.offset+w.len == len(w.buffer) {
			w.flush()
			continue
		}

		ni := copy(w.buffer[w.offset+w.len:], p)
		w.len += ni
		p = p[ni:]
		n += ni
	}
}

func (w *writer) flush() error {
	var zeroWrite bool
	for {
		if w.err != nil {
			return w.err
		}

		if w.len == 0 {
			w.offset = 0
			return nil
		}

		var n int
		n, w.err = w.out.Write(w.buffer[w.offset : w.offset+w.len])
		if n == 0 && w.err == nil && zeroWrite {
			w.err = io.ErrShortWrite
		}

		zeroWrite = n == 0 && w.err == nil
		w.offset += n
		w.len -= n
		if w.err != nil {
			w.options.Pool.Put(w.buffer)
			w.buffer = nil
		}
	}
}

func (w *writer) close() error {
	if w.err != nil {
		return w.err
	}

	w.flush()
	if w.err != nil {
		return w.err
	}

	w.err = errClosed
	w.options.Pool.Put(w.buffer)
	w.buffer = nil
	return nil
}