|
| 1 | +package adaptivepool |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "errors" |
| 6 | + "fmt" |
| 7 | + "io" |
| 8 | + "sync" |
| 9 | +) |
| 10 | + |
| 11 | +// ReaderBufferer buffers data from [io.Reader]s and [io.ReadCloser]s into |
| 12 | +// [BufferedReader]s that, upon calling their `Close` method, will put the data |
| 13 | +// back into an [AdaptivePool] for reuse. |
| 14 | +type ReaderBufferer struct { |
| 15 | + bufPool AdaptivePool[[]byte] |
| 16 | + rdPool sync.Pool |
| 17 | +} |
| 18 | + |
| 19 | +// NewReaderBufferer returns a new ReaderBufferer. The `minCap` and `thresh` |
| 20 | +// arguments will be the values of the internal [NormalSlice.MinCap] and |
| 21 | +// [NormalSlice.Threshold], respectively. Example: |
| 22 | +// |
| 23 | +// rb := NewReaderBufferer(512, 2, 500) |
| 24 | +func NewReaderBufferer(minCap int, thresh, maxN float64) *ReaderBufferer { |
| 25 | + return new(ReaderBufferer).init(minCap, thresh, maxN) |
| 26 | +} |
| 27 | + |
| 28 | +func (p *ReaderBufferer) init(minCap int, thresh, |
| 29 | + maxN float64) *ReaderBufferer { |
| 30 | + p.rdPool.New = newBytesReader |
| 31 | + p.bufPool.init(NormalSlice[byte]{ |
| 32 | + MinCap: minCap, |
| 33 | + Threshold: thresh, |
| 34 | + }, maxN) |
| 35 | + return p |
| 36 | +} |
| 37 | + |
| 38 | +func newBytesReader() any { |
| 39 | + return bytes.NewReader(nil) |
| 40 | +} |
| 41 | + |
| 42 | +// Stats returns the statistics from the internal AdaptivePool. |
| 43 | +func (p *ReaderBufferer) Stats() Stats { |
| 44 | + return p.bufPool.Stats() |
| 45 | +} |
| 46 | + |
| 47 | +// Reader buffers the contents of the given io.Reader in a BufferedReader. |
| 48 | +func (p *ReaderBufferer) Reader(r io.Reader) (*BufferedReader, error) { |
| 49 | + return p.buf(r, nil) |
| 50 | +} |
| 51 | + |
| 52 | +// ReadCloser buffers the contents of the given io.ReadCloser in a |
| 53 | +// BufferedReader. It always calls Close, and it fails if it returns an error. |
| 54 | +func (p *ReaderBufferer) ReadCloser(rc io.ReadCloser) (*BufferedReader, error) { |
| 55 | + return p.buf(rc, rc) |
| 56 | +} |
| 57 | + |
| 58 | +func (p *ReaderBufferer) buf(r io.Reader, |
| 59 | + c io.Closer) (*BufferedReader, error) { |
| 60 | + buf := p.bufPool.Get() |
| 61 | + bytesBuf := bytes.NewBuffer(buf) |
| 62 | + n, readErr := bytesBuf.ReadFrom(r) |
| 63 | + if readErr != nil && c == nil { |
| 64 | + p.put(buf) |
| 65 | + return nil, fmt.Errorf("read io.Reader: %w; bytes read: %v", readErr, n) |
| 66 | + } |
| 67 | + buf = bytesBuf.Bytes() |
| 68 | + |
| 69 | + var closeErr error |
| 70 | + if c != nil { |
| 71 | + closeErr = c.Close() |
| 72 | + if readErr == nil && closeErr != nil { |
| 73 | + p.put(buf) |
| 74 | + return nil, fmt.Errorf("close io.ReadCloser: %w; bytes read: %v", |
| 75 | + closeErr, n) |
| 76 | + } |
| 77 | + } |
| 78 | + |
| 79 | + if readErr != nil || closeErr != nil { |
| 80 | + p.put(buf) |
| 81 | + return nil, fmt.Errorf("buffer io.ReadCloser: read error: %w; close"+ |
| 82 | + " error: %w; bytes read: %v", readErr, closeErr, n) |
| 83 | + } |
| 84 | + |
| 85 | + rd := p.rdPool.Get().(*bytes.Reader) |
| 86 | + rd.Reset(buf) |
| 87 | + |
| 88 | + return &BufferedReader{ |
| 89 | + reader: rd, |
| 90 | + buf: buf, |
| 91 | + release: p.release, |
| 92 | + }, nil |
| 93 | +} |
| 94 | + |
| 95 | +func (p *ReaderBufferer) release(buf []byte, rd *bytes.Reader) { |
| 96 | + rd.Reset(nil) |
| 97 | + p.rdPool.Put(rd) |
| 98 | + p.put(buf) |
| 99 | +} |
| 100 | + |
| 101 | +func (p *ReaderBufferer) put(buf []byte) { |
| 102 | + if cap(buf) > 0 { |
| 103 | + clear(buf[:cap(buf)]) |
| 104 | + p.bufPool.Put(buf[:0]) |
| 105 | + } |
| 106 | +} |
| 107 | + |
| 108 | +// NOTE: we explicitly do not want to offer io.ReaderAt in BufferedReader |
| 109 | +// because, as per its docs, "Clients of ReadAt can execute parallel ReadAt |
| 110 | +// calls on the same input source". This means that we should add a sync.RWMutex |
| 111 | +// to protect the underlying implementation and make it more heavyweight in |
| 112 | +// order to guard the parallel ReadAt operations from potential Close |
| 113 | +// operations. Clients can still use the Seek method and then Read as a |
| 114 | +// sequential workaround. |
| 115 | + |
| 116 | +// BufferedReader holds a read-only buffer of the contents extracted from an |
| 117 | +// [io.Reader] or [io.ReadCloser]. Its `Close` method releases internal buffers |
| 118 | +// for reuse, and after that it will be empty. It is not safe for concurrent |
| 119 | +// use. |
| 120 | +type BufferedReader struct { |
| 121 | + reader *bytes.Reader |
| 122 | + buf []byte |
| 123 | + release func([]byte, *bytes.Reader) |
| 124 | +} |
| 125 | + |
| 126 | +// Bytes returns the internal buffered []byte, transferring their ownership to |
| 127 | +// the caller. The data will not be later put back into a pool by the |
| 128 | +// implementation, and subsequent calls to any method will behave as if `Close` |
| 129 | +// had been called. Subsequent calls to this method return nil, the same as if |
| 130 | +// `Close` had been called before. |
| 131 | +func (bb *BufferedReader) Bytes() []byte { |
| 132 | + if bb.reader != nil { |
| 133 | + bb.release(nil, bb.reader) |
| 134 | + buf := bb.buf |
| 135 | + *bb = BufferedReader{} |
| 136 | + return buf |
| 137 | + } |
| 138 | + return nil |
| 139 | +} |
| 140 | + |
| 141 | +// Len returns the number of unread bytes. |
| 142 | +func (bb *BufferedReader) Len() int { |
| 143 | + if bb.reader != nil { |
| 144 | + return bb.reader.Len() |
| 145 | + } |
| 146 | + return 0 |
| 147 | +} |
| 148 | + |
| 149 | +// Read is part of the implementation of the io.Reader interface. |
| 150 | +func (bb *BufferedReader) Read(p []byte) (int, error) { |
| 151 | + if bb.reader != nil { |
| 152 | + return bb.reader.Read(p) |
| 153 | + } |
| 154 | + return 0, io.EOF |
| 155 | +} |
| 156 | + |
| 157 | +// Close is part of the implementation of the io.Closer interface. This method |
| 158 | +// releases the internal buffer for reuse. After this, the *BufferedReader will |
| 159 | +// be empty. This method is idempotent and always returns a nil error. |
| 160 | +func (bb *BufferedReader) Close() error { |
| 161 | + if bb.reader != nil { |
| 162 | + bb.release(bb.buf, bb.reader) |
| 163 | + *bb = BufferedReader{} |
| 164 | + } |
| 165 | + return nil |
| 166 | +} |
| 167 | + |
| 168 | +// Seek is part of the implementation of the io.Seeker interface. |
| 169 | +func (bb *BufferedReader) Seek(offset int64, whence int) (int64, error) { |
| 170 | + if bb.reader != nil { |
| 171 | + return bb.reader.Seek(offset, whence) |
| 172 | + } |
| 173 | + |
| 174 | + switch whence { |
| 175 | + case io.SeekStart, io.SeekCurrent, io.SeekEnd: |
| 176 | + default: |
| 177 | + return 0, errors.New("BufferedReader.Seek: invalid whence") |
| 178 | + } |
| 179 | + if offset < 0 { |
| 180 | + return 0, errors.New("BufferedReader.Seek: negative position") |
| 181 | + } |
| 182 | + |
| 183 | + return 0, nil |
| 184 | +} |
| 185 | + |
| 186 | +// ReadByte is part of the implementation of the io.ByteReader interface. |
| 187 | +func (bb *BufferedReader) ReadByte() (byte, error) { |
| 188 | + if bb.reader != nil { |
| 189 | + return bb.reader.ReadByte() |
| 190 | + } |
| 191 | + return 0, io.EOF |
| 192 | +} |
| 193 | + |
| 194 | +// UnreadByte is part of the implementation of the io.ByteScanner interface. |
| 195 | +func (bb *BufferedReader) UnreadByte() error { |
| 196 | + if bb.reader != nil { |
| 197 | + return bb.reader.UnreadByte() |
| 198 | + } |
| 199 | + return errors.New("BufferedReader.UnreadByte: resource closed") |
| 200 | +} |
| 201 | + |
| 202 | +// ReadRune is part of the implementation of the io.RuneReader interface. |
| 203 | +func (bb *BufferedReader) ReadRune() (r rune, size int, err error) { |
| 204 | + if bb.reader != nil { |
| 205 | + return bb.reader.ReadRune() |
| 206 | + } |
| 207 | + return 0, 0, io.EOF |
| 208 | +} |
| 209 | + |
| 210 | +// UnreadRune is part of the implementation of the io.RuneScanner interface. |
| 211 | +func (bb *BufferedReader) UnreadRune() error { |
| 212 | + if bb.reader != nil { |
| 213 | + return bb.reader.UnreadRune() |
| 214 | + } |
| 215 | + return errors.New("BufferedReader.UnreadRune: resource closed") |
| 216 | +} |
| 217 | + |
| 218 | +// WriteTo is part of the implementation of the io.WriterTo interface. |
| 219 | +func (bb *BufferedReader) WriteTo(w io.Writer) (n int64, err error) { |
| 220 | + if bb.reader != nil { |
| 221 | + return bb.reader.WriteTo(w) |
| 222 | + } |
| 223 | + return 0, nil |
| 224 | +} |
0 commit comments