iobufpool.go

Documentation: github.com/jackc/pgx/v5/internal/iobufpool

     1  // Package iobufpool implements a global segregated-fit pool of buffers for IO.
     2  //
     3  // It uses *[]byte instead of []byte to avoid the sync.Pool allocation with Put. Unfortunately, using a pointer to avoid
     4  // an allocation is purposely not documented. https://github.com/golang/go/issues/16323
     5  package iobufpool
     6  
     7  import "sync"
     8  
     9  const minPoolExpOf2 = 8
    10  
    11  var pools [18]*sync.Pool
    12  
    13  func init() {
    14  	for i := range pools {
    15  		bufLen := 1 << (minPoolExpOf2 + i)
    16  		pools[i] = &sync.Pool{
    17  			New: func() any {
    18  				buf := make([]byte, bufLen)
    19  				return &buf
    20  			},
    21  		}
    22  	}
    23  }
    24  
    25  // Get gets a []byte of len size with cap <= size*2.
    26  func Get(size int) *[]byte {
    27  	i := getPoolIdx(size)
    28  	if i >= len(pools) {
    29  		buf := make([]byte, size)
    30  		return &buf
    31  	}
    32  
    33  	ptrBuf := (pools[i].Get().(*[]byte))
    34  	*ptrBuf = (*ptrBuf)[:size]
    35  
    36  	return ptrBuf
    37  }
    38  
    39  func getPoolIdx(size int) int {
    40  	size--
    41  	size >>= minPoolExpOf2
    42  	i := 0
    43  	for size > 0 {
    44  		size >>= 1
    45  		i++
    46  	}
    47  
    48  	return i
    49  }
    50  
    51  // Put returns buf to the pool.
    52  func Put(buf *[]byte) {
    53  	i := putPoolIdx(cap(*buf))
    54  	if i < 0 {
    55  		return
    56  	}
    57  
    58  	pools[i].Put(buf)
    59  }
    60  
    61  func putPoolIdx(size int) int {
    62  	minPoolSize := 1 << minPoolExpOf2
    63  	for i := range pools {
    64  		if size == minPoolSize<<i {
    65  			return i
    66  		}
    67  	}
    68  
    69  	return -1
    70  }
    71
View as plain text