Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit e1eda90

Browse files
jackcjackc
committed
Add ChunkReader
0 parents  commit e1eda90

File tree

2 files changed

+234
-0
lines changed

2 files changed

+234
-0
lines changed

‎chunkreader.go‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
package chunkreader
2+
3+
import (
4+
"io"
5+
)
6+
7+
type ChunkReader struct {
8+
r io.Reader
9+
10+
buf []byte
11+
rp, wp int // buf read position and write position
12+
taken bool
13+
14+
options Options
15+
}
16+
17+
type Options struct {
18+
MinBufLen int // Minimum buffer length
19+
BlockLen int // Increments to expand buffer (e.g. a 8000 byte request with a BlockLen of 1024 would yield a buffer len of 8192)
20+
}
21+
22+
func NewChunkReader(r io.Reader) *ChunkReader {
23+
cr, err := NewChunkReaderEx(r, Options{})
24+
if err != nil {
25+
panic("default options can't be bad")
26+
}
27+
28+
return cr
29+
}
30+
31+
func NewChunkReaderEx(r io.Reader, options Options) (*ChunkReader, error) {
32+
if options.MinBufLen == 0 {
33+
options.MinBufLen = 4096
34+
}
35+
if options.BlockLen == 0 {
36+
options.BlockLen = 512
37+
}
38+
39+
return &ChunkReader{
40+
r: r,
41+
buf: make([]byte, options.MinBufLen),
42+
options: options,
43+
}, nil
44+
}
45+
46+
// Next returns buf filled with the next n bytes. buf is only valid until the
47+
// next call to Next. If an error occurs, buf will be nil.
48+
func (r *ChunkReader) Next(n int) (buf []byte, err error) {
49+
// n bytes already in buf
50+
if (r.wp - r.rp) >= n {
51+
buf = r.buf[r.rp : r.rp+n]
52+
r.rp += n
53+
return buf, err
54+
}
55+
56+
// available space in buf is less than n
57+
if len(r.buf) < n {
58+
r.copyBufContents(r.newBuf(n))
59+
r.taken = false
60+
}
61+
62+
// buf is large enough, but need to shift filled area to start to make enough contiguous space
63+
minReadCount := n - (r.wp - r.rp)
64+
if (len(r.buf) - r.wp) < minReadCount {
65+
newBuf := r.buf
66+
if r.taken {
67+
newBuf = r.newBuf(n)
68+
r.taken = false
69+
}
70+
r.copyBufContents(newBuf)
71+
}
72+
73+
if err := r.appendAtLeast(minReadCount); err != nil {
74+
return nil, err
75+
}
76+
77+
buf = r.buf[r.rp : r.rp+n]
78+
r.rp += n
79+
return buf, nil
80+
}
81+
82+
// KeepLast prevents the last data retrieved by Next from being reused by the
83+
// ChunkReader.
84+
func (r *ChunkReader) KeepLast() {
85+
r.taken = true
86+
}
87+
88+
func (r *ChunkReader) appendAtLeast(fillLen int) error {
89+
n, err := io.ReadAtLeast(r.r, r.buf[r.wp:], fillLen)
90+
r.wp += n
91+
return err
92+
}
93+
94+
func (r *ChunkReader) newBuf(min int) []byte {
95+
size := ((min / r.options.BlockLen) + 1) * r.options.BlockLen
96+
if size < r.options.MinBufLen {
97+
size = r.options.MinBufLen
98+
}
99+
return make([]byte, size)
100+
}
101+
102+
func (r *ChunkReader) copyBufContents(dest []byte) {
103+
r.wp = copy(dest, r.buf[r.rp:r.wp])
104+
r.rp = 0
105+
r.buf = dest
106+
}

‎chunkreader_test.go‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
package chunkreader
2+
3+
import (
4+
"bytes"
5+
"testing"
6+
)
7+
8+
func TestChunkReaderNextDoesNotReadIfAlreadyBuffered(t *testing.T) {
9+
server := &bytes.Buffer{}
10+
r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 2})
11+
if err != nil {
12+
t.Fatal(err)
13+
}
14+
15+
src := []byte{1, 2, 3, 4}
16+
server.Write(src)
17+
18+
n1, err := r.Next(2)
19+
if err != nil {
20+
t.Fatal(err)
21+
}
22+
if bytes.Compare(n1, src[0:2]) != 0 {
23+
t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:2], n1)
24+
}
25+
26+
n2, err := r.Next(2)
27+
if err != nil {
28+
t.Fatal(err)
29+
}
30+
if bytes.Compare(n2, src[2:4]) != 0 {
31+
t.Fatalf("Expected read bytes to be %v, but they were %v", src[2:4], n2)
32+
}
33+
34+
if bytes.Compare(r.buf, src) != 0 {
35+
t.Fatalf("Expected r.buf to be %v, but it was %v", src, r.buf)
36+
}
37+
if r.rp != 4 {
38+
t.Fatalf("Expected r.rp to be %v, but it was %v", 4, r.rp)
39+
}
40+
if r.wp != 4 {
41+
t.Fatalf("Expected r.wp to be %v, but it was %v", 4, r.wp)
42+
}
43+
}
44+
45+
func TestChunkReaderNextExpandsBufAsNeeded(t *testing.T) {
46+
server := &bytes.Buffer{}
47+
r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 2})
48+
if err != nil {
49+
t.Fatal(err)
50+
}
51+
52+
src := []byte{1, 2, 3, 4, 5, 6, 7, 8}
53+
server.Write(src)
54+
55+
n1, err := r.Next(5)
56+
if err != nil {
57+
t.Fatal(err)
58+
}
59+
if bytes.Compare(n1, src[0:5]) != 0 {
60+
t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:5], n1)
61+
}
62+
if len(r.buf) != 6 {
63+
t.Fatalf("Expected len(r.buf) to be %v, but it was %v", 6, len(r.buf))
64+
}
65+
}
66+
67+
func TestChunkReaderNextReusesBuf(t *testing.T) {
68+
server := &bytes.Buffer{}
69+
r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 1})
70+
if err != nil {
71+
t.Fatal(err)
72+
}
73+
74+
src := []byte{1, 2, 3, 4, 5, 6, 7, 8}
75+
server.Write(src)
76+
77+
n1, err := r.Next(4)
78+
if err != nil {
79+
t.Fatal(err)
80+
}
81+
if bytes.Compare(n1, src[0:4]) != 0 {
82+
t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:4], n1)
83+
}
84+
85+
n2, err := r.Next(4)
86+
if err != nil {
87+
t.Fatal(err)
88+
}
89+
if bytes.Compare(n2, src[4:8]) != 0 {
90+
t.Fatalf("Expected read bytes to be %v, but they were %v", src[4:8], n2)
91+
}
92+
93+
if bytes.Compare(n1, src[4:8]) != 0 {
94+
t.Fatalf("Expected Next to have reused buf, %v found instead of %v", src[4:8], n1)
95+
}
96+
}
97+
98+
func TestChunkReaderKeepLastPreventsBufReuse(t *testing.T) {
99+
server := &bytes.Buffer{}
100+
r, err := NewChunkReaderEx(server, Options{MinBufLen: 4, BlockLen: 1})
101+
if err != nil {
102+
t.Fatal(err)
103+
}
104+
105+
src := []byte{1, 2, 3, 4, 5, 6, 7, 8}
106+
server.Write(src)
107+
108+
n1, err := r.Next(4)
109+
if err != nil {
110+
t.Fatal(err)
111+
}
112+
if bytes.Compare(n1, src[0:4]) != 0 {
113+
t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:4], n1)
114+
}
115+
r.KeepLast()
116+
117+
n2, err := r.Next(4)
118+
if err != nil {
119+
t.Fatal(err)
120+
}
121+
if bytes.Compare(n2, src[4:8]) != 0 {
122+
t.Fatalf("Expected read bytes to be %v, but they were %v", src[4:8], n2)
123+
}
124+
125+
if bytes.Compare(n1, src[0:4]) != 0 {
126+
t.Fatalf("Expected KeepLast to prevent Next from overwriting buf, expected %v but it was %v", src[0:4], n1)
127+
}
128+
}

0 commit comments

Comments
 (0)