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 }