github.com/viant/toolbox@v0.34.5/bridge/byte_buffer_pool.go (about)

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