github.com/sohaha/zlsgo@v1.7.13-0.20240501141223-10dd1a906f76/zutil/buffpool.go (about)

     1  package zutil
     2  
     3  import (
     4  	"bytes"
     5  	"sync"
     6  )
     7  
     8  var bufTypes = [...]int{
     9  	0, 16, 32, 64, 128, 256, 512, 1024, 2048, 5120, 1 << 20, (1 << 19) * 10, (1 << 20) * 10, (1 << 19) * 100, (1 << 20) * 100,
    10  }
    11  
    12  const bufTypeNum = len(bufTypes)
    13  
    14  var (
    15  	bufPools [bufTypeNum]sync.Pool
    16  	BuffSize = bufTypes[3]
    17  )
    18  
    19  func init() {
    20  	for i := 0; i < bufTypeNum; i++ {
    21  		l := bufTypes[i]
    22  		bufPools[i].New = func() interface{} {
    23  			return bytes.NewBuffer(make([]byte, 0, l))
    24  		}
    25  	}
    26  }
    27  
    28  func GetBuff(ss ...int) *bytes.Buffer {
    29  	size := BuffSize
    30  	if len(ss) > 0 {
    31  		size = ss[0]
    32  	}
    33  	if size > 0 {
    34  		if size <= bufTypes[bufTypeNum-1] {
    35  			for i := 0; i < bufTypeNum; i++ {
    36  				if size <= bufTypes[i] {
    37  					return bufPools[i].Get().(*bytes.Buffer)
    38  				}
    39  			}
    40  		}
    41  		return bytes.NewBuffer(make([]byte, 0, size))
    42  	}
    43  
    44  	return bufPools[0].Get().(*bytes.Buffer)
    45  }
    46  
    47  func PutBuff(buffer *bytes.Buffer) {
    48  	size := buffer.Cap()
    49  	buffer.Reset()
    50  	if size > bufTypes[bufTypeNum-1] {
    51  		bufPools[0].Put(buffer)
    52  		return
    53  	}
    54  	for i := 1; i < bufTypeNum; i++ {
    55  		if size <= bufTypes[i] {
    56  			if size == bufTypes[i] {
    57  				bufPools[i].Put(buffer)
    58  			} else {
    59  				bufPools[i-1].Put(buffer)
    60  			}
    61  			return
    62  		}
    63  	}
    64  }