github.com/imannamdari/v2ray-core/v5@v5.0.5/common/bytespool/pool.go (about)

     1  package bytespool
     2  
     3  import "sync"
     4  
     5  func createAllocFunc(size int32) func() interface{} {
     6  	return func() interface{} {
     7  		return make([]byte, size)
     8  	}
     9  }
    10  
    11  // The following parameters controls the size of buffer pools.
    12  // There are numPools pools. Starting from 2k size, the size of each pool is sizeMulti of the previous one.
    13  // Package buf is guaranteed to not use buffers larger than the largest pool.
    14  // Other packets may use larger buffers.
    15  const (
    16  	numPools  = 4
    17  	sizeMulti = 4
    18  )
    19  
    20  var (
    21  	pool     [numPools]sync.Pool
    22  	poolSize [numPools]int32
    23  )
    24  
    25  func init() {
    26  	size := int32(2048)
    27  	for i := 0; i < numPools; i++ {
    28  		pool[i] = sync.Pool{
    29  			New: createAllocFunc(size),
    30  		}
    31  		poolSize[i] = size
    32  		size *= sizeMulti
    33  	}
    34  }
    35  
    36  // GetPool returns a sync.Pool that generates bytes array with at least the given size.
    37  // It may return nil if no such pool exists.
    38  //
    39  // v2ray:api:stable
    40  func GetPool(size int32) *sync.Pool {
    41  	for idx, ps := range poolSize {
    42  		if size <= ps {
    43  			return &pool[idx]
    44  		}
    45  	}
    46  	return nil
    47  }
    48  
    49  // Alloc returns a byte slice with at least the given size. Minimum size of returned slice is 2048.
    50  //
    51  // v2ray:api:stable
    52  func Alloc(size int32) []byte {
    53  	pool := GetPool(size)
    54  	if pool != nil {
    55  		return pool.Get().([]byte)
    56  	}
    57  	return make([]byte, size)
    58  }
    59  
    60  // Free puts a byte slice into the internal pool.
    61  //
    62  // v2ray:api:stable
    63  func Free(b []byte) {
    64  	size := int32(cap(b))
    65  	b = b[0:cap(b)]
    66  	for i := numPools - 1; i >= 0; i-- {
    67  		if size >= poolSize[i] {
    68  			pool[i].Put(b) // nolint: staticcheck
    69  			return
    70  		}
    71  	}
    72  }