github.com/demonoid81/moby@v0.0.0-20200517203328-62dd8e17c460/pkg/pools/pools.go (about)

     1  // Package pools provides a collection of pools which provide various
     2  // data types with buffers. These can be used to lower the number of
     3  // memory allocations and reuse buffers.
     4  //
     5  // New pools should be added to this package to allow them to be
     6  // shared across packages.
     7  //
     8  // Utility functions which operate on pools should be added to this
     9  // package to allow them to be reused.
    10  package pools // import "github.com/demonoid81/moby/pkg/pools"
    11  
    12  import (
    13  	"bufio"
    14  	"io"
    15  	"sync"
    16  
    17  	"github.com/demonoid81/moby/pkg/ioutils"
    18  )
    19  
    20  const buffer32K = 32 * 1024
    21  
    22  var (
    23  	// BufioReader32KPool is a pool which returns bufio.Reader with a 32K buffer.
    24  	BufioReader32KPool = newBufioReaderPoolWithSize(buffer32K)
    25  	// BufioWriter32KPool is a pool which returns bufio.Writer with a 32K buffer.
    26  	BufioWriter32KPool = newBufioWriterPoolWithSize(buffer32K)
    27  	buffer32KPool      = newBufferPoolWithSize(buffer32K)
    28  )
    29  
    30  // BufioReaderPool is a bufio reader that uses sync.Pool.
    31  type BufioReaderPool struct {
    32  	pool sync.Pool
    33  }
    34  
    35  // newBufioReaderPoolWithSize is unexported because new pools should be
    36  // added here to be shared where required.
    37  func newBufioReaderPoolWithSize(size int) *BufioReaderPool {
    38  	return &BufioReaderPool{
    39  		pool: sync.Pool{
    40  			New: func() interface{} { return bufio.NewReaderSize(nil, size) },
    41  		},
    42  	}
    43  }
    44  
    45  // Get returns a bufio.Reader which reads from r. The buffer size is that of the pool.
    46  func (bufPool *BufioReaderPool) Get(r io.Reader) *bufio.Reader {
    47  	buf := bufPool.pool.Get().(*bufio.Reader)
    48  	buf.Reset(r)
    49  	return buf
    50  }
    51  
    52  // Put puts the bufio.Reader back into the pool.
    53  func (bufPool *BufioReaderPool) Put(b *bufio.Reader) {
    54  	b.Reset(nil)
    55  	bufPool.pool.Put(b)
    56  }
    57  
    58  type bufferPool struct {
    59  	pool sync.Pool
    60  }
    61  
    62  func newBufferPoolWithSize(size int) *bufferPool {
    63  	return &bufferPool{
    64  		pool: sync.Pool{
    65  			New: func() interface{} { s := make([]byte, size); return &s },
    66  		},
    67  	}
    68  }
    69  
    70  func (bp *bufferPool) Get() *[]byte {
    71  	return bp.pool.Get().(*[]byte)
    72  }
    73  
    74  func (bp *bufferPool) Put(b *[]byte) {
    75  	bp.pool.Put(b)
    76  }
    77  
    78  // Copy is a convenience wrapper which uses a buffer to avoid allocation in io.Copy.
    79  func Copy(dst io.Writer, src io.Reader) (written int64, err error) {
    80  	buf := buffer32KPool.Get()
    81  	written, err = io.CopyBuffer(dst, src, *buf)
    82  	buffer32KPool.Put(buf)
    83  	return
    84  }
    85  
    86  // NewReadCloserWrapper returns a wrapper which puts the bufio.Reader back
    87  // into the pool and closes the reader if it's an io.ReadCloser.
    88  func (bufPool *BufioReaderPool) NewReadCloserWrapper(buf *bufio.Reader, r io.Reader) io.ReadCloser {
    89  	return ioutils.NewReadCloserWrapper(r, func() error {
    90  		if readCloser, ok := r.(io.ReadCloser); ok {
    91  			readCloser.Close()
    92  		}
    93  		bufPool.Put(buf)
    94  		return nil
    95  	})
    96  }
    97  
    98  // BufioWriterPool is a bufio writer that uses sync.Pool.
    99  type BufioWriterPool struct {
   100  	pool sync.Pool
   101  }
   102  
   103  // newBufioWriterPoolWithSize is unexported because new pools should be
   104  // added here to be shared where required.
   105  func newBufioWriterPoolWithSize(size int) *BufioWriterPool {
   106  	return &BufioWriterPool{
   107  		pool: sync.Pool{
   108  			New: func() interface{} { return bufio.NewWriterSize(nil, size) },
   109  		},
   110  	}
   111  }
   112  
   113  // Get returns a bufio.Writer which writes to w. The buffer size is that of the pool.
   114  func (bufPool *BufioWriterPool) Get(w io.Writer) *bufio.Writer {
   115  	buf := bufPool.pool.Get().(*bufio.Writer)
   116  	buf.Reset(w)
   117  	return buf
   118  }
   119  
   120  // Put puts the bufio.Writer back into the pool.
   121  func (bufPool *BufioWriterPool) Put(b *bufio.Writer) {
   122  	b.Reset(nil)
   123  	bufPool.pool.Put(b)
   124  }
   125  
   126  // NewWriteCloserWrapper returns a wrapper which puts the bufio.Writer back
   127  // into the pool and closes the writer if it's an io.Writecloser.
   128  func (bufPool *BufioWriterPool) NewWriteCloserWrapper(buf *bufio.Writer, w io.Writer) io.WriteCloser {
   129  	return ioutils.NewWriteCloserWrapper(w, func() error {
   130  		buf.Flush()
   131  		if writeCloser, ok := w.(io.WriteCloser); ok {
   132  			writeCloser.Close()
   133  		}
   134  		bufPool.Put(buf)
   135  		return nil
   136  	})
   137  }