av/ring/ring.go

/*
NAME
  ring.go - a structure that encapsulates a Buffer datastructure with conccurency
  functionality

DESCRIPTION
  See Readme.md

AUTHOR
  Dan Kortschak <dan@ausocean.org>

LICENSE
  ring.go is Copyright (C) 2017 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
 for more details.

  You should have received a copy of the GNU General Public License
  along with revid in gpl.txt.  If not, see http://www.gnu.org/licenses.
*/

// Package ring provides a ring buffer of io.ReadWriters.
package ring

import (
	"bytes"
	"errors"
	"io"
	"time"
)

var (
	ErrTimeout = errors.New("ring: buffer cycle timeout")
	ErrDropped = errors.New("ring: dropped old write")
	ErrStall   = errors.New("ring: unable to dump old write")
	ErrTooLong = errors.New("ring: write to long for buffer element")
)

// Buffer implements a ring buffer.
//
// The buffer has a writable head and a readable tail with a queue from the head
// to the tail. Concurrent read a write operations are safe.
type Buffer struct {
	head, tail  *bytes.Buffer
	full, empty chan *bytes.Buffer
	timeout     time.Duration
}

// NewBuffer returns a Buffer with len elements of the given size. The timeout
// parameter specifies how long a write operation will wait before failing with
// a temporary timeout error.
func NewBuffer(len, size int, timeout time.Duration) *Buffer {
	if len <= 0 || size <= 0 {
		return nil
	}
	b := Buffer{
		full:    make(chan *bytes.Buffer, len),
		empty:   make(chan *bytes.Buffer, len),
		timeout: timeout,
	}
	for i := 0; i < len; i++ {
		b.empty <- bytes.NewBuffer(make([]byte, 0, size))
	}
	return &b
}

// Len returns the number of full buffer elements.
func (b *Buffer) Len() int {
	return len(b.full)
}

// Write writes the bytes in b to the next current or next available element of the ring buffer
// it returns the number of bytes written and any error.
// If no element can be gained within the timeout or stolen from the queue, ErrStall is returned
// and if the len(p) is greater than the buffer's element size, ErrTooLong is returned. If a
// write was successful but a previous write was dropped, ErrDropped is returned.
//
// Write is safe to use concurrently with Read, but may not be used concurrently with another
// write operation.
func (b *Buffer) Write(p []byte) (int, error) {
	var dropped bool
	if b.head == nil {
		timer := time.NewTimer(b.timeout)
		select {
		case <-timer.C:
			select {
			case b.head = <-b.full:
				b.head.Reset()
				dropped = true
			default:
				// This should never happen.
				return 0, ErrStall
			}
		case b.head = <-b.empty:
			timer.Stop()
		}
	}
	if len(p) > b.head.Cap() {
		return 0, ErrTooLong
	}
	if len(p) > b.head.Cap()-b.head.Len() {
		b.full <- b.head
		b.head = nil
		return b.Write(p)
	}
	n, err := b.head.Write(p)
	if b.head.Cap()-b.head.Len() == 0 {
		b.full <- b.head
		b.head = nil
	}
	if dropped && err == nil {
		err = ErrDropped
	}
	return n, err
}

// Flush puts the currently writing element of the buffer into the queue for reading. Flush
// is idempotent.
//
// Flush is safe to use concurrently with Read, but may not be used concurrently with another
// another write operation.
func (b *Buffer) Flush() {
	if b.head == nil {
		return
	}
	b.full <- b.head
	b.head = nil
}

// Close closes the buffer. The buffer may not be written to after a call to close, but can
// be drained by calls to Read.
//
// Flush is safe to use concurrently with Read, but may not be used concurrently with another
// another write operation.
func (b *Buffer) Close() error {
	b.Flush()
	close(b.full)
	return nil
}

// Next gets the next element from the queue ready for reading, returning ErrTimeout if no
// element is available within the timeout. If the Buffer has been closed Next returns io.EOF.
//
// Next is safe to use concurrently with write operations, but may not be used concurrently with
// another Read call or Next call. A goroutine calling Next must not call Flush or Close.
func (b *Buffer) Next(timeout time.Duration) error {
	if b.tail == nil {
		timer := time.NewTimer(timeout)
		var ok bool
		select {
		case <-timer.C:
			return ErrTimeout
		case b.tail, ok = <-b.full:
			timer.Stop()
			if !ok {
				return io.EOF
			}
		}
	}
	return nil
}

// Read reads bytes from the current tail of the ring buffer into p and returns the number of
// bytes read and any error.
//
// Read is safe to use concurrently with write operations, but may not be used concurrently with
// another Read call or Next call. A goroutine calling Read must not call Flush or Close.
func (b *Buffer) Read(p []byte) (int, error) {
	if b.tail == nil {
		return 0, io.EOF
	}
	n, err := b.tail.Read(p)
	if b.tail.Len() == 0 {
		b.tail.Reset()
		b.empty <- b.tail
		b.tail = nil
	}
	return n, err
}