├── .travis.yml ├── LICENSE ├── README.md ├── chunkreader.go ├── chunkreader_test.go └── go.mod /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | 3 | go: 4 | - 1.x 5 | - tip 6 | 7 | matrix: 8 | allow_failures: 9 | - go: tip 10 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2019 Jack Christensen 2 | 3 | MIT License 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining 6 | a copy of this software and associated documentation files (the 7 | "Software"), to deal in the Software without restriction, including 8 | without limitation the rights to use, copy, modify, merge, publish, 9 | distribute, sublicense, and/or sell copies of the Software, and to 10 | permit persons to whom the Software is furnished to do so, subject to 11 | the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be 14 | included in all copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 19 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 20 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 22 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![](https://godoc.org/github.com/jackc/chunkreader?status.svg)](https://godoc.org/github.com/jackc/chunkreader) 2 | [![Build Status](https://travis-ci.org/jackc/chunkreader.svg)](https://travis-ci.org/jackc/chunkreader) 3 | 4 | # chunkreader 5 | 6 | Package chunkreader provides an io.Reader wrapper that minimizes IO reads and memory allocations. 7 | 8 | Extracted from original implementation in https://github.com/jackc/pgx. 9 | -------------------------------------------------------------------------------- /chunkreader.go: -------------------------------------------------------------------------------- 1 | // Package chunkreader provides an io.Reader wrapper that minimizes IO reads and memory allocations. 2 | package chunkreader 3 | 4 | import ( 5 | "io" 6 | ) 7 | 8 | // ChunkReader is a io.Reader wrapper that minimizes IO reads and memory allocations. It allocates memory in chunks and 9 | // will read as much as will fit in the current buffer in a single call regardless of how large a read is actually 10 | // requested. The memory returned via Next is owned by the caller. This avoids the need for an additional copy. 11 | // 12 | // The downside of this approach is that a large buffer can be pinned in memory even if only a small slice is 13 | // referenced. For example, an entire 4096 byte block could be pinned in memory by even a 1 byte slice. In these rare 14 | // cases it would be advantageous to copy the bytes to another slice. 15 | type ChunkReader struct { 16 | r io.Reader 17 | 18 | buf []byte 19 | rp, wp int // buf read position and write position 20 | 21 | config Config 22 | } 23 | 24 | // Config contains configuration parameters for ChunkReader. 25 | type Config struct { 26 | MinBufLen int // Minimum buffer length 27 | } 28 | 29 | // New creates and returns a new ChunkReader for r with default configuration. 30 | func New(r io.Reader) *ChunkReader { 31 | cr, err := NewConfig(r, Config{}) 32 | if err != nil { 33 | panic("default config can't be bad") 34 | } 35 | 36 | return cr 37 | } 38 | 39 | // NewConfig creates and a new ChunkReader for r configured by config. 40 | func NewConfig(r io.Reader, config Config) (*ChunkReader, error) { 41 | if config.MinBufLen == 0 { 42 | // By historical reasons Postgres currently has 8KB send buffer inside, 43 | // so here we want to have at least the same size buffer. 44 | // @see https://github.com/postgres/postgres/blob/249d64999615802752940e017ee5166e726bc7cd/src/backend/libpq/pqcomm.c#L134 45 | // @see https://www.postgresql.org/message-id/0cdc5485-cb3c-5e16-4a46-e3b2f7a41322%40ya.ru 46 | config.MinBufLen = 8192 47 | } 48 | 49 | return &ChunkReader{ 50 | r: r, 51 | buf: make([]byte, config.MinBufLen), 52 | config: config, 53 | }, nil 54 | } 55 | 56 | // Next returns buf filled with the next n bytes. The caller gains ownership of buf. It is not necessary to make a copy 57 | // of buf. If an error occurs, buf will be nil. 58 | func (r *ChunkReader) Next(n int) (buf []byte, err error) { 59 | // n bytes already in buf 60 | if (r.wp - r.rp) >= n { 61 | buf = r.buf[r.rp : r.rp+n] 62 | r.rp += n 63 | return buf, err 64 | } 65 | 66 | // available space in buf is less than n 67 | if len(r.buf) < n { 68 | r.copyBufContents(r.newBuf(n)) 69 | } 70 | 71 | // buf is large enough, but need to shift filled area to start to make enough contiguous space 72 | minReadCount := n - (r.wp - r.rp) 73 | if (len(r.buf) - r.wp) < minReadCount { 74 | newBuf := r.newBuf(n) 75 | r.copyBufContents(newBuf) 76 | } 77 | 78 | if err := r.appendAtLeast(minReadCount); err != nil { 79 | return nil, err 80 | } 81 | 82 | buf = r.buf[r.rp : r.rp+n] 83 | r.rp += n 84 | return buf, nil 85 | } 86 | 87 | func (r *ChunkReader) appendAtLeast(fillLen int) error { 88 | n, err := io.ReadAtLeast(r.r, r.buf[r.wp:], fillLen) 89 | r.wp += n 90 | return err 91 | } 92 | 93 | func (r *ChunkReader) newBuf(size int) []byte { 94 | if size < r.config.MinBufLen { 95 | size = r.config.MinBufLen 96 | } 97 | return make([]byte, size) 98 | } 99 | 100 | func (r *ChunkReader) copyBufContents(dest []byte) { 101 | r.wp = copy(dest, r.buf[r.rp:r.wp]) 102 | r.rp = 0 103 | r.buf = dest 104 | } 105 | -------------------------------------------------------------------------------- /chunkreader_test.go: -------------------------------------------------------------------------------- 1 | package chunkreader 2 | 3 | import ( 4 | "bytes" 5 | "math/rand" 6 | "testing" 7 | ) 8 | 9 | func TestChunkReaderNextDoesNotReadIfAlreadyBuffered(t *testing.T) { 10 | server := &bytes.Buffer{} 11 | r, err := NewConfig(server, Config{MinBufLen: 4}) 12 | if err != nil { 13 | t.Fatal(err) 14 | } 15 | 16 | src := []byte{1, 2, 3, 4} 17 | server.Write(src) 18 | 19 | n1, err := r.Next(2) 20 | if err != nil { 21 | t.Fatal(err) 22 | } 23 | if bytes.Compare(n1, src[0:2]) != 0 { 24 | t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:2], n1) 25 | } 26 | 27 | n2, err := r.Next(2) 28 | if err != nil { 29 | t.Fatal(err) 30 | } 31 | if bytes.Compare(n2, src[2:4]) != 0 { 32 | t.Fatalf("Expected read bytes to be %v, but they were %v", src[2:4], n2) 33 | } 34 | 35 | if bytes.Compare(r.buf, src) != 0 { 36 | t.Fatalf("Expected r.buf to be %v, but it was %v", src, r.buf) 37 | } 38 | if r.rp != 4 { 39 | t.Fatalf("Expected r.rp to be %v, but it was %v", 4, r.rp) 40 | } 41 | if r.wp != 4 { 42 | t.Fatalf("Expected r.wp to be %v, but it was %v", 4, r.wp) 43 | } 44 | } 45 | 46 | func TestChunkReaderNextExpandsBufAsNeeded(t *testing.T) { 47 | server := &bytes.Buffer{} 48 | r, err := NewConfig(server, Config{MinBufLen: 4}) 49 | if err != nil { 50 | t.Fatal(err) 51 | } 52 | 53 | src := []byte{1, 2, 3, 4, 5, 6, 7, 8} 54 | server.Write(src) 55 | 56 | n1, err := r.Next(5) 57 | if err != nil { 58 | t.Fatal(err) 59 | } 60 | if bytes.Compare(n1, src[0:5]) != 0 { 61 | t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:5], n1) 62 | } 63 | if len(r.buf) != 5 { 64 | t.Fatalf("Expected len(r.buf) to be %v, but it was %v", 5, len(r.buf)) 65 | } 66 | } 67 | 68 | func TestChunkReaderDoesNotReuseBuf(t *testing.T) { 69 | server := &bytes.Buffer{} 70 | r, err := NewConfig(server, Config{MinBufLen: 4}) 71 | if err != nil { 72 | t.Fatal(err) 73 | } 74 | 75 | src := []byte{1, 2, 3, 4, 5, 6, 7, 8} 76 | server.Write(src) 77 | 78 | n1, err := r.Next(4) 79 | if err != nil { 80 | t.Fatal(err) 81 | } 82 | if bytes.Compare(n1, src[0:4]) != 0 { 83 | t.Fatalf("Expected read bytes to be %v, but they were %v", src[0:4], n1) 84 | } 85 | 86 | n2, err := r.Next(4) 87 | if err != nil { 88 | t.Fatal(err) 89 | } 90 | if bytes.Compare(n2, src[4:8]) != 0 { 91 | t.Fatalf("Expected read bytes to be %v, but they were %v", src[4:8], n2) 92 | } 93 | 94 | if bytes.Compare(n1, src[0:4]) != 0 { 95 | t.Fatalf("Expected KeepLast to prevent Next from overwriting buf, expected %v but it was %v", src[0:4], n1) 96 | } 97 | } 98 | 99 | type randomReader struct { 100 | rnd *rand.Rand 101 | } 102 | 103 | // Read reads a random number of random bytes. 104 | func (r *randomReader) Read(p []byte) (n int, err error) { 105 | n = r.rnd.Intn(len(p) + 1) 106 | return r.rnd.Read(p[:n]) 107 | } 108 | 109 | func TestChunkReaderNextFuzz(t *testing.T) { 110 | rr := &randomReader{rnd: rand.New(rand.NewSource(1))} 111 | r, err := NewConfig(rr, Config{MinBufLen: 8192}) 112 | if err != nil { 113 | t.Fatal(err) 114 | } 115 | 116 | randomSizes := rand.New(rand.NewSource(0)) 117 | 118 | for i := 0; i < 100000; i++ { 119 | size := randomSizes.Intn(16384) + 1 120 | buf, err := r.Next(size) 121 | if err != nil { 122 | t.Fatal(err) 123 | } 124 | if len(buf) != size { 125 | t.Fatalf("Expected to get %v bytes but got %v bytes", size, len(buf)) 126 | } 127 | } 128 | } 129 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/jackc/chunkreader/v2 2 | 3 | go 1.12 4 | --------------------------------------------------------------------------------