sectionreader.go

package sstream

import (
	"bytes"
	"io"
	"sync"

	"golang.org/x/net/context"
)

// SectionReader reads a section of bytes from a Reader into an internal buffer.
// The read operation stops if an error occurs or the context is canceled. The
// internal buffer is reused to prevent memory allocations.
type SectionReader struct {
	buf   []byte
	piece *bytes.Buffer
	mu    sync.Mutex
}

// NewSectionReader returns an instance of SectionReader. SectionReader will use
// the supplied buf as its internal buffer. The supplied buffer will be reused
// in future requests to prevent memory allocations.
// buf must have a capacity greater than 0.
func NewSectionReader(buf *bytes.Buffer) *SectionReader {
	if buf == nil {
		panic("buf was nil in NewSectionReader")
	}
	if buf.Cap() < 1 {
		panic("buf with capacity less than 1 in NewSectionReader")
	}
	return &SectionReader{piece: buf}
}

// Read implements io.Reader. It reads from the internal buffer which was filled
// from a call to RetrieveSection.
func (s *SectionReader) Read(p []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.piece.Read(p)
}

// RetrieveSection reads n bytes from the supplied reader. When the call
// returns, the retrieved bytes can be read by calling Read. It returns an error
// if a read fails, EOF is returned early, or the operation is canceled. n must
// be less than or equal to the capacity of the buf passed to NewSectionReader.
//
// Adapted code from:
// https://golang.org/src/io/io.go
func (s *SectionReader) RetrieveSection(ctx context.Context, r io.Reader, n int) (err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if ctx == nil {
		panic("ctx was nil in SectionReader.RetrieveSection")
	}
	if r == nil {
		panic("r was nil in SectionReader.RetrieveSection")
	}
	if n <= 0 {
		panic("n was less than or equal to 0 in SectionReader.RetrieveSection")
	}
	if s.piece == nil {
		s.piece = bytes.NewBuffer(make([]byte, 0, 30*1024))
	}
	s.piece.Reset()
	if s.buf == nil {
		s.buf = make([]byte, 30*1024)
	} else {
		if cap(s.buf) > len(s.buf) {
			s.buf = s.buf[0:cap(s.buf)]
		}
	}
	var read int
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		if read >= n {
			break
		}
		if len(s.buf) > n-read {
			s.buf = s.buf[0 : n-read]
		}
		nr, er := r.Read(s.buf)
		if nr > 0 {
			nw, ew := s.piece.Write(s.buf[0:nr])
			if nw > 0 {
				read += nw
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er == io.EOF {
			if read == n {
				break
			}
			err = er
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	return
}