133 lines
1.9 KiB
Go
133 lines
1.9 KiB
Go
package buffer
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
)
|
|
|
|
type writer struct {
|
|
out io.Writer
|
|
options Options
|
|
buffer []byte
|
|
offset, len int
|
|
err error
|
|
}
|
|
|
|
var errClosed = errors.New("buffer closed")
|
|
|
|
func (w *writer) write(p []byte) (int, error) {
|
|
var n int
|
|
for {
|
|
if w.err != nil {
|
|
return n, w.err
|
|
}
|
|
|
|
if len(p) == 0 {
|
|
return n, nil
|
|
}
|
|
|
|
if len(w.buffer) == 0 {
|
|
w.buffer, w.err = w.options.Pool.Get()
|
|
if len(w.buffer) == 0 && w.err == nil {
|
|
w.err = ErrZeroAllocation
|
|
}
|
|
|
|
continue
|
|
}
|
|
|
|
if w.offset+w.len == len(w.buffer) {
|
|
w.flush()
|
|
continue
|
|
}
|
|
|
|
ni := copy(w.buffer[w.offset+w.len:], p)
|
|
w.len += ni
|
|
p = p[ni:]
|
|
n += ni
|
|
}
|
|
}
|
|
|
|
func (w *writer) readFrom(r io.Reader) (int64, error) {
|
|
var n int64
|
|
for {
|
|
if errors.Is(w.err, io.EOF) {
|
|
err := w.err
|
|
w.err = nil
|
|
w.flush()
|
|
if w.err != nil {
|
|
return n, w.err
|
|
}
|
|
|
|
w.err = err
|
|
return n, nil
|
|
}
|
|
|
|
if w.err != nil {
|
|
return n, w.err
|
|
}
|
|
|
|
if len(w.buffer) == 0 {
|
|
w.buffer, w.err = w.options.Pool.Get()
|
|
if len(w.buffer) == 0 && w.err == nil {
|
|
w.err = ErrZeroAllocation
|
|
}
|
|
|
|
continue
|
|
}
|
|
|
|
if w.offset+w.len == len(w.buffer) {
|
|
w.flush()
|
|
continue
|
|
}
|
|
|
|
var ni int
|
|
ni, w.err = r.Read(w.buffer[w.offset+w.len:])
|
|
w.len += ni
|
|
n += int64(ni)
|
|
}
|
|
}
|
|
|
|
func (w *writer) flush() error {
|
|
var zeroWrite bool
|
|
for {
|
|
if w.err != nil {
|
|
return w.err
|
|
}
|
|
|
|
if w.len == 0 {
|
|
w.offset = 0
|
|
return nil
|
|
}
|
|
|
|
var n int
|
|
n, w.err = w.out.Write(w.buffer[w.offset : w.offset+w.len])
|
|
if n == 0 && w.err == nil && zeroWrite {
|
|
w.err = io.ErrShortWrite
|
|
}
|
|
|
|
zeroWrite = n == 0 && w.err == nil
|
|
w.offset += n
|
|
w.len -= n
|
|
if w.err != nil {
|
|
w.options.Pool.Put(w.buffer)
|
|
w.buffer = nil
|
|
}
|
|
}
|
|
}
|
|
|
|
func (w *writer) close() error {
|
|
if w.err != nil {
|
|
return w.err
|
|
}
|
|
|
|
w.flush()
|
|
if w.err != nil {
|
|
return w.err
|
|
}
|
|
|
|
w.err = errClosed
|
|
w.options.Pool.Put(w.buffer)
|
|
w.buffer = nil
|
|
return nil
|
|
}
|