github.com/blend/go-sdk@v1.20220411.3/vault/buffer_pool.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package vault
     9  
    10  import (
    11  	"bytes"
    12  	"sync"
    13  )
    14  
    15  // NewBufferPool returns a new BufferPool.
    16  // bufferSize is the size of the returned buffers pre-allocated size in bytes.
    17  // Typically this is something between 256 bytes and 1kb.
    18  func NewBufferPool(bufferSize int) *BufferPool {
    19  	bp := &BufferPool{}
    20  	bp.Pool = sync.Pool{
    21  		New: func() interface{} {
    22  			b := &Buffer{
    23  				Buffer: bytes.NewBuffer(make([]byte, bufferSize)),
    24  				pool:   bp,
    25  			}
    26  			return b
    27  		},
    28  	}
    29  	return bp
    30  }
    31  
    32  // BufferPool is a sync.Pool of bytes.Buffer.
    33  type BufferPool struct {
    34  	sync.Pool
    35  }
    36  
    37  // Get returns a pooled bytes.Buffer instance.
    38  func (bp *BufferPool) Get() *Buffer {
    39  	buf := bp.Pool.Get().(*Buffer)
    40  	buf.Reset()
    41  	return buf
    42  }
    43  
    44  // Put returns the pooled instance.
    45  func (bp *BufferPool) Put(b *Buffer) {
    46  	bp.Pool.Put(b)
    47  }
    48  
    49  // Buffer is a bytes.Buffer with a reference back to the buffer pool.
    50  // It returns itself to the pool on close.
    51  type Buffer struct {
    52  	*bytes.Buffer
    53  	pool *BufferPool
    54  }
    55  
    56  // Close returns the buffer to the pool.
    57  func (b *Buffer) Close() error {
    58  	b.Reset()
    59  	b.pool.Put(b)
    60  	return nil
    61  }