1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
|
package bridge
import (
"io"
"net/http/httputil"
)
type bytesBufferPool struct {
channel chan []byte
bufferSize int
}
func (p *bytesBufferPool) Get() (result []byte) {
select {
case result = <-p.channel:
default:
result = make([]byte, p.bufferSize)
}
return result
}
func (p *bytesBufferPool) Put(b []byte) {
select {
case p.channel <- b:
default: //If the pool is full, discard the buffer.
}
}
//NewBytesBufferPool returns new httputil.BufferPool pool.
func NewBytesBufferPool(poolSize, bufferSize int) httputil.BufferPool {
return &bytesBufferPool{
channel: make(chan []byte, poolSize),
bufferSize: bufferSize,
}
}
//CopyBuffer copies bytes from passed in source to destination with provided pool
func CopyWithBufferPool(source io.Reader, destination io.Writer, pool httputil.BufferPool) (int64, error) {
buf := pool.Get()
defer pool.Put(buf)
var written int64
for {
bytesRead, readError := source.Read(buf)
if bytesRead > 0 {
bytesWritten, writeError := destination.Write(buf[:bytesRead])
if bytesWritten > 0 {
written += int64(bytesWritten)
}
if writeError != nil {
return written, writeError
}
if bytesRead != bytesWritten {
return written, io.ErrShortWrite
}
}
if readError != nil {
return written, readError
}
}
}
|