aboutsummaryrefslogtreecommitdiff
path: root/vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go
diff options
context:
space:
mode:
Diffstat (limited to 'vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go')
-rw-r--r--vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go90
1 files changed, 90 insertions, 0 deletions
diff --git a/vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go b/vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go
new file mode 100644
index 0000000..fc0fa61
--- /dev/null
+++ b/vendor/github.com/jackc/pgx/v5/pgproto3/chunkreader.go
@@ -0,0 +1,90 @@
+package pgproto3
+
+import (
+ "io"
+
+ "github.com/jackc/pgx/v5/internal/iobufpool"
+)
+
+// chunkReader is a io.Reader wrapper that minimizes IO reads and memory allocations. It allocates memory in chunks and
+// will read as much as will fit in the current buffer in a single call regardless of how large a read is actually
+// requested. The memory returned via Next is only valid until the next call to Next.
+//
+// This is roughly equivalent to a bufio.Reader that only uses Peek and Discard to never copy bytes.
+type chunkReader struct {
+ r io.Reader
+
+ buf *[]byte
+ rp, wp int // buf read position and write position
+
+ minBufSize int
+}
+
+// newChunkReader creates and returns a new chunkReader for r with default configuration. If minBufSize is <= 0 it uses
+// a default value.
+func newChunkReader(r io.Reader, minBufSize int) *chunkReader {
+ if minBufSize <= 0 {
+ // By historical reasons Postgres currently has 8KB send buffer inside,
+ // so here we want to have at least the same size buffer.
+ // @see https://github.com/postgres/postgres/blob/249d64999615802752940e017ee5166e726bc7cd/src/backend/libpq/pqcomm.c#L134
+ // @see https://www.postgresql.org/message-id/0cdc5485-cb3c-5e16-4a46-e3b2f7a41322%40ya.ru
+ //
+ // In addition, testing has found no benefit of any larger buffer.
+ minBufSize = 8192
+ }
+
+ return &chunkReader{
+ r: r,
+ minBufSize: minBufSize,
+ buf: iobufpool.Get(minBufSize),
+ }
+}
+
+// Next returns buf filled with the next n bytes. buf is only valid until next call of Next. If an error occurs, buf
+// will be nil.
+func (r *chunkReader) Next(n int) (buf []byte, err error) {
+ // Reset the buffer if it is empty
+ if r.rp == r.wp {
+ if len(*r.buf) != r.minBufSize {
+ iobufpool.Put(r.buf)
+ r.buf = iobufpool.Get(r.minBufSize)
+ }
+ r.rp = 0
+ r.wp = 0
+ }
+
+ // n bytes already in buf
+ if (r.wp - r.rp) >= n {
+ buf = (*r.buf)[r.rp : r.rp+n : r.rp+n]
+ r.rp += n
+ return buf, err
+ }
+
+ // buf is smaller than requested number of bytes
+ if len(*r.buf) < n {
+ bigBuf := iobufpool.Get(n)
+ r.wp = copy((*bigBuf), (*r.buf)[r.rp:r.wp])
+ r.rp = 0
+ iobufpool.Put(r.buf)
+ r.buf = bigBuf
+ }
+
+ // 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 {
+ r.wp = copy((*r.buf), (*r.buf)[r.rp:r.wp])
+ r.rp = 0
+ }
+
+ // Read at least the required number of bytes from the underlying io.Reader
+ readBytesCount, err := io.ReadAtLeast(r.r, (*r.buf)[r.wp:], minReadCount)
+ r.wp += readBytesCount
+ // fmt.Println("read", n)
+ if err != nil {
+ return nil, err
+ }
+
+ buf = (*r.buf)[r.rp : r.rp+n : r.rp+n]
+ r.rp += n
+ return buf, nil
+}