From e1eda90e29497aecea3ea0bb7a09e62536233498 Mon Sep 17 00:00:00 2001 From: Jack Christensen Date: Sun, 12 Feb 2017 21:46:15 -0600 Subject: [PATCH] Add ChunkReader --- chunkreader.go | 106 ++++++++++++++++++++++++++++++++++++ chunkreader_test.go | 128 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 234 insertions(+) create mode 100644 chunkreader.go create mode 100644 chunkreader_test.go diff --git a/chunkreader.go b/chunkreader.go new file mode 100644 index 00000000..f9d6555c --- /dev/null +++ b/chunkreader.go @@ -0,0 +1,106 @@ +package chunkreader + +import ( + "io" +) + +type ChunkReader struct { + r io.Reader + + buf []byte + rp, wp int // buf read position and write position + taken bool + + options Options +} + +type Options struct { + MinBufLen int // Minimum buffer length + BlockLen int // Increments to expand buffer (e.g. a 8000 byte request with a BlockLen of 1024 would yield a buffer len of 8192) +} + +func NewChunkReader(r io.Reader) *ChunkReader { + cr, err := NewChunkReaderEx(r, Options{}) + if err != nil { + panic("default options can't be bad") + } + + return cr +} + +func NewChunkReaderEx(r io.Reader, options Options) (*ChunkReader, error) { + if options.MinBufLen == 0 { + options.MinBufLen = 4096 + } + if options.BlockLen == 0 { + options.BlockLen = 512 + } + + return &ChunkReader{ + r: r, + buf: make([]byte, options.MinBufLen), + options: options, + }, nil +} + +// Next returns buf filled with the next n bytes. buf is only valid until the +// next call to Next. If an error occurs, buf will be nil. +func (r *ChunkReader) Next(n int) (buf []byte, err error) { + // n bytes already in buf + if (r.wp - r.rp) >= n { + buf = r.buf[r.rp : r.rp+n] + r.rp += n + return buf, err + } + + // available space in buf is less than n + if len(r.buf) < n { + r.copyBufContents(r.newBuf(n)) + r.taken = false + } + + // buf is large enough, but need to shift filled area to start to make enough contiguous space + minReadCount := n - (r.wp - r.rp) + if (len(r.buf) - r.wp) < minReadCount { + newBuf := r.buf + if r.taken { + newBuf = r.newBuf(n) + r.taken = false + } + r.copyBufContents(newBuf) + } + + if err := r.appendAtLeast(minReadCount); err != nil { + return nil, err + } + + buf = r.buf[r.rp : r.rp+n] + r.rp += n + return buf, nil +} + +// KeepLast prevents the last data retrieved by Next from being reused by the +// ChunkReader. +func (r *ChunkReader) KeepLast() { + r.taken = true +} + +func (r *ChunkReader) appendAtLeast(fillLen int) error { + n, err := io.ReadAtLeast(r.r, r.buf[r.wp:], fillLen) + r.wp += n + return err +} + +func (r *ChunkReader) newBuf(min int) []byte { + size := ((min / r.options.BlockLen) + 1) * r.options.BlockLen + if size < r.options.MinBufLen { + size = r.options.MinBufLen + } + return make([]byte, size) +} + +func (r *ChunkReader) copyBufContents(dest []byte) { + r.wp = copy(dest, r.buf[r.rp:r.wp]) + r.rp = 0 + r.buf = dest +} diff --git a/chunkreader_test.go b/chunkreader_test.go new file mode 100644 index 00000000..9c19ff4a --- /dev/null +++ b/chunkreader_test.go @@ -0,0 +1,128 @@ +package chunkreader + +import ( + "bytes" + "testing" +) + +func TestChunkReaderNextDoesNotReadIfAlreadyBuffered(t *testing.T) { + server := &bytes.Buffer{} + r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 2}) + if err != nil { + t.Fatal(err) + } + + src := []byte{1, 2, 3, 4} + server.Write(src) + + n1, err := r.Next(2) + if err != nil { + t.Fatal(err) + } + if bytes.Compare(n1, src[0:2]) != 0 { + t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:2], n1) + } + + n2, err := r.Next(2) + if err != nil { + t.Fatal(err) + } + if bytes.Compare(n2, src[2:4]) != 0 { + t.Fatalf("Expected read bytes to be %v, but they were %v", src[2:4], n2) + } + + if bytes.Compare(r.buf, src) != 0 { + t.Fatalf("Expected r.buf to be %v, but it was %v", src, r.buf) + } + if r.rp != 4 { + t.Fatalf("Expected r.rp to be %v, but it was %v", 4, r.rp) + } + if r.wp != 4 { + t.Fatalf("Expected r.wp to be %v, but it was %v", 4, r.wp) + } +} + +func TestChunkReaderNextExpandsBufAsNeeded(t *testing.T) { + server := &bytes.Buffer{} + r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 2}) + if err != nil { + t.Fatal(err) + } + + src := []byte{1, 2, 3, 4, 5, 6, 7, 8} + server.Write(src) + + n1, err := r.Next(5) + if err != nil { + t.Fatal(err) + } + if bytes.Compare(n1, src[0:5]) != 0 { + t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:5], n1) + } + if len(r.buf) != 6 { + t.Fatalf("Expected len(r.buf) to be %v, but it was %v", 6, len(r.buf)) + } +} + +func TestChunkReaderNextReusesBuf(t *testing.T) { + server := &bytes.Buffer{} + r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 1}) + if err != nil { + t.Fatal(err) + } + + src := []byte{1, 2, 3, 4, 5, 6, 7, 8} + server.Write(src) + + n1, err := r.Next(4) + if err != nil { + t.Fatal(err) + } + if bytes.Compare(n1, src[0:4]) != 0 { + t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:4], n1) + } + + n2, err := r.Next(4) + if err != nil { + t.Fatal(err) + } + if bytes.Compare(n2, src[4:8]) != 0 { + t.Fatalf("Expected read bytes to be %v, but they were %v", src[4:8], n2) + } + + if bytes.Compare(n1, src[4:8]) != 0 { + t.Fatalf("Expected Next to have reused buf, %v found instead of %v", src[4:8], n1) + } +} + +func TestChunkReaderKeepLastPreventsBufReuse(t *testing.T) { + server := &bytes.Buffer{} + r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 1}) + if err != nil { + t.Fatal(err) + } + + src := []byte{1, 2, 3, 4, 5, 6, 7, 8} + server.Write(src) + + n1, err := r.Next(4) + if err != nil { + t.Fatal(err) + } + if bytes.Compare(n1, src[0:4]) != 0 { + t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:4], n1) + } + r.KeepLast() + + n2, err := r.Next(4) + if err != nil { + t.Fatal(err) + } + if bytes.Compare(n2, src[4:8]) != 0 { + t.Fatalf("Expected read bytes to be %v, but they were %v", src[4:8], n2) + } + + if bytes.Compare(n1, src[0:4]) != 0 { + t.Fatalf("Expected KeepLast to prevent Next from overwriting buf, expected %v but it was %v", src[0:4], n1) + } +}