multireader/multireader.go

package multireader

import (
	"errors"
	"io"
	"sort"
	"sync"
)

var (
	ErrMultiReaderClosed = errors.New("multireader closed")
	ErrReaderClosed      = errors.New("reader closed")
	ErrOutOfBound        = errors.New("seek beyond writer")
	ErrUnsupportedWhence = errors.New("unsupported whence")
)

type MultiReader struct {
	mtx          *sync.Mutex
	cond         *sync.Cond
	buffer       []byte
	readers      []*Reader
	pos          []*Reader
	lastWritePos int
	closed       bool
}

type Reader struct {
	multiReader *MultiReader
	mtx         *sync.Mutex
	cond        *sync.Cond
	id          int
	readPos     int
	closed      bool
}

func NewMultiReader(len int) *MultiReader {
	if len < 0 {
		len = 0
	}

	buff := make([]byte, len)
	mtx := &sync.Mutex{}
	return &MultiReader{
		mtx:          mtx,
		cond:         sync.NewCond(mtx),
		buffer:       buff,
		readers:      make([]*Reader, 0),
		pos:          make([]*Reader, 0),
		lastWritePos: 0,
		closed:       false,
	}
}

func (m *MultiReader) Write(data []byte) (int, error) {
	m.mtx.Lock()
	defer m.mtx.Unlock()
	defer m.cond.Broadcast()

	var (
		tmpBuff []byte
	)

	if m.closed {
		return 0, ErrMultiReaderClosed
	}

	if m.lastWritePos+len(data) > len(m.buffer) {
		tmpBuff = make([]byte, m.lastWritePos+len(data)-len(m.buffer))
		m.buffer = append(m.buffer, tmpBuff...)
	}
	copy(m.buffer[m.lastWritePos:], data)
	m.lastWritePos += len(data)

	return len(data), nil
}

func (m *MultiReader) sortLastPos() {
	sort.Slice(m.pos, func(i, j int) bool {
		return m.pos[i].readPos < m.pos[j].readPos
	})
}

func (m *MultiReader) shift() {
	shiftPos := m.pos[0].readPos
	if shiftPos > 0 {
		for _, r := range m.pos {
			r.readPos -= shiftPos
		}
		copy(m.buffer, m.buffer[shiftPos:shiftPos+(m.lastWritePos-shiftPos)])
		m.lastWritePos -= shiftPos
	}
}

func (m *MultiReader) Close() error {
	m.mtx.Lock()
	defer m.mtx.Unlock()
	m.closed = true
	return nil
}

func (m *MultiReader) NewReader() *Reader {
	m.mtx.Lock()
	defer m.mtx.Unlock()

	reader := &Reader{
		multiReader: m,
		mtx:         m.mtx,
		cond:        m.cond,
		id:          len(m.readers),
		readPos:     0,
	}
	m.readers = append(m.readers, reader)
	m.pos = append(m.pos, reader)
	m.sortLastPos()

	return reader
}

func (r *Reader) Read(data []byte) (int, error) {
	r.mtx.Lock()
	defer r.mtx.Unlock()
	defer r.cond.Broadcast()

	var (
		n   int
		err error
	)

	if r.closed {
		return 0, ErrReaderClosed
	}

	for r.readPos >= r.multiReader.lastWritePos {
		if r.multiReader.closed {
			return 0, ErrMultiReaderClosed
		}
		r.cond.Wait()
	}

	n, err = r.readAhead(data)
	_, _ = r.seek(n, io.SeekCurrent)

	return n, err
}

func (r *Reader) seek(offset int, whence int) (int, error) {
	if offset > r.multiReader.lastWritePos || offset < 0 {
		return r.readPos, ErrOutOfBound
	}

	switch whence {
	case io.SeekCurrent:
		r.readPos += offset
	case io.SeekStart:
		r.readPos = offset
	case io.SeekEnd:
		r.readPos = r.multiReader.lastWritePos - offset
	}

	r.multiReader.sortLastPos()
	r.multiReader.shift()

	return r.readPos, nil
}

func (r *Reader) Seek(offset int64, whence int) (int64, error) {
	n, err := r.seek(int(offset), whence)
	return int64(n), err
}

func (r *Reader) readAhead(data []byte) (int, error) {
	var err error
	if r.readPos+len(data) > r.multiReader.lastWritePos {
		copy(data, r.multiReader.buffer[r.readPos:r.multiReader.lastWritePos])
		if r.multiReader.closed {
			err = ErrMultiReaderClosed
		}
		return r.multiReader.lastWritePos - r.readPos, err
	}

	copy(data, r.multiReader.buffer[r.readPos:r.readPos+len(data)])
	return len(data), nil
}

func (r *Reader) ReadAhead(data []byte) (int, error) {
	r.mtx.Lock()
	defer r.mtx.Unlock()
	return r.readAhead(data)
}

func (r *Reader) Close() error {
	r.mtx.Lock()
	defer r.mtx.Unlock()
	r.closed = true
	return nil
}