forked from jackc/pgx
-
Notifications
You must be signed in to change notification settings - Fork 1
/
chunkreader.go
89 lines (72 loc) · 1.73 KB
/
chunkreader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
package chunkreader
import (
"io"
)
type ChunkReader struct {
r io.Reader
buf []byte
rp, wp int // buf read position and write position
options Options
}
type Options struct {
MinBufLen int // Minimum buffer length
}
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
}
return &ChunkReader{
r: r,
buf: make([]byte, options.MinBufLen),
options: options,
}, nil
}
// Next returns buf filled with the next n bytes. 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))
}
// 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.newBuf(n)
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
}
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(size int) []byte {
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
}