package bytebufferpool import ( "errors" "io" "time" ) // ByteBuffer provides byte buffer, which can be used for minimizing // memory allocations. // // ByteBuffer may be used with functions appending data to the given []byte // slice. See example code for details. // // Use Get for obtaining an empty byte buffer. type ByteBuffer struct { // B is a byte buffer to use in append-like workloads. // See example code for details. B []byte } // Len returns the size of the byte buffer. func (b *ByteBuffer) Len() int { return len(b.B) } // Copy copies from src to dst until either EOF is reached // on src or an error occurs. It returns the number of bytes // copied and the first error encountered while copying, if any. func (b *ByteBuffer) Copy(dst io.Writer, src io.Reader) (written int64, err error) { b.Reset() b.B = make([]byte, 32*1024) for { nr, er := src.Read(b.B) if nr > 0 { nw, ew := dst.Write(b.B[0:nr]) if nw > 0 { written += int64(nw) } if ew != nil { err = ew break } if nr != nw { err = io.ErrShortWrite break } } if er != nil { if er != io.EOF { err = er } break } } return written, err } // CopyWithIdleDuration copies from src to dst until either EOF is reached // on src, or an error occurs, or idle time out. It returns the number of bytes // copied and the first error encountered while copying, if any. func (b *ByteBuffer) CopyWithIdleDuration(dst io.Writer, src io.Reader, idle time.Duration) (written int64, err error) { if idle == 0 { return b.Copy(dst, src) } b.Reset() b.B = make([]byte, 32*1024) for { var nr int var er error idleChan := make(chan struct{}, 1) go func() { nr, er = src.Read(b.B) idleChan <- struct{}{} }() select { case <-idleChan: case <-time.After(idle): return written, errors.New("idle time out") } if nr > 0 { nw, ew := dst.Write(b.B[0:nr]) if nw > 0 { written += int64(nw) } if ew != nil { err = ew break } if nr != nw { err = io.ErrShortWrite break } } if er != nil { if er != io.EOF { err = er } break } } return written, err } // ReadFrom implements io.ReaderFrom. // // The function appends all the data read from r to b. func (b *ByteBuffer) ReadFrom(r io.Reader) (int64, error) { p := b.B nStart := int64(len(p)) nMax := int64(cap(p)) n := nStart if nMax == 0 { nMax = 64 p = make([]byte, nMax) } else { p = p[:nMax] } for { if n == nMax { nMax *= 2 bNew := make([]byte, nMax) copy(bNew, p) p = bNew } nn, err := r.Read(p[n:]) n += int64(nn) if err != nil { b.B = p[:n] n -= nStart if err == io.EOF { return n, nil } return n, err } } } // WriteTo implements io.WriterTo. func (b *ByteBuffer) WriteTo(w io.Writer) (int64, error) { n, err := w.Write(b.B) return int64(n), err } // Bytes returns b.B, i.e. all the bytes accumulated in the buffer. // // The purpose of this function is bytes.Buffer compatibility. func (b *ByteBuffer) Bytes() []byte { return b.B } // Write implements io.Writer - it appends p to ByteBuffer.B func (b *ByteBuffer) Write(p []byte) (int, error) { b.B = append(b.B, p...) return len(p), nil } // WriteByte appends the byte c to the buffer. // // The purpose of this function is bytes.Buffer compatibility. // // The function always returns nil. func (b *ByteBuffer) WriteByte(c byte) error { b.B = append(b.B, c) return nil } // WriteString appends s to ByteBuffer.B. func (b *ByteBuffer) WriteString(s string) (int, error) { b.B = append(b.B, s...) return len(s), nil } // Set sets ByteBuffer.B to p. func (b *ByteBuffer) Set(p []byte) { b.B = append(b.B[:0], p...) } // SetString sets ByteBuffer.B to s. func (b *ByteBuffer) SetString(s string) { b.B = append(b.B[:0], s...) } // String returns string representation of ByteBuffer.B. func (b *ByteBuffer) String() string { return string(b.B) } // Reset makes ByteBuffer.B empty. func (b *ByteBuffer) Reset() { b.B = b.B[:0] }