A round robin writer

round.go

// (c) 2018 John Lenton
// MIT licensed but I'm lazy

// Package round implements round.Robin, which is an in-memory io.Writer
// (like bytes.Buffer) that uses a fixed amount of memory.
//
// It's a silly toy thing, really; otherwise I'd've made the buffer size
// settable without editing the source. But it's useful to me...
// Enjoy!
package round

import (
	"sync"
)

const bufSize = 16 * 1024 * 1024

// A Robin is a fixed-size circular buffer.
type Robin struct {
	// N is how many bytes have been written
	N   uint64
	mu  sync.Mutex
	end int
	buf [bufSize]byte
}

// Write implements io.Writer
func (w *Robin) Write(p []byte) (n int, err error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.N += uint64(len(p))

	if len(p) > len(w.buf) {
		// the whole thing is too big to fit
		w.end = len(w.buf)
		copy(w.buf[:], p[len(p)-len(w.buf):])
		return len(p), nil
	}

	if w.end == len(w.buf) {
		w.end = 0
	}

	n = copy(w.buf[w.end:], p)
	w.end += n
	if n == len(p) {
		return n, nil
	}
	m := copy(w.buf[:], p[n:])
	w.end = m

	return n + m, nil
}

// Bytes returns a slice of bytes in the buffer. It copies the buffer.
//
// If you want to avoid the copy, add a method that returns two slices
// and use that instead.
func (w *Robin) Bytes() []byte {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.N <= uint64(len(w.buf)) {
		return w.buf[:w.end]
	}
	outbuf := make([]byte, bufSize)
	n := copy(outbuf, w.buf[w.end:])
	copy(outbuf[n:], w.buf[:])

	return outbuf
}

// String of the bytes in the buffer
func (w *Robin) String() string {
	return string(w.Bytes())
}