github.com/songzhibin97/gkit@v1.2.13/cache/buffer/byte_pool.go (about)

     1  package buffer
     2  
     3  import "sync"
     4  
     5  const (
     6  	minShift = 6
     7  	maxShift = 18
     8  
     9  	// errSlot: 如果没有找到对应的 slot 返回 -1
    10  	errSlot = -1
    11  )
    12  
    13  var localBytePool = newBytePool()
    14  
    15  // byteSlot 槽区
    16  type byteSlot struct {
    17  	// defaultSize 默认 slot 存储 []byte 大小
    18  	defaultSize int
    19  
    20  	// pool 实际上 pool池
    21  	pool sync.Pool
    22  }
    23  
    24  // bytePool []byte pool
    25  type bytePool struct {
    26  	minShift int
    27  	minSize  int
    28  	maxSize  int
    29  
    30  	// 维护 byteSlot
    31  	pool []*byteSlot
    32  }
    33  
    34  // slot 根据 size 获取到 对应的 byteSlot 下标
    35  func (b *bytePool) slot(size int) int {
    36  	// 超过阈值
    37  	if size > b.maxSize {
    38  		return errSlot
    39  	}
    40  	slot := 0
    41  	shift := 0
    42  	if size > b.minSize {
    43  		size--
    44  		for size > 0 {
    45  			size = size >> 1
    46  			shift++
    47  		}
    48  		slot = shift - b.minShift
    49  	}
    50  	return slot
    51  }
    52  
    53  // get 根据 size 从 bytePool 中获取到 *[]byte
    54  func (b *bytePool) get(size int) *[]byte {
    55  	slot := b.slot(size)
    56  	if slot == errSlot {
    57  		// 如果需要的 []byte 大于 设置的 pool 返回 errSlot
    58  		// 触发 errSlot 会手动创建一个 []byte 返回
    59  		ret := newBytes(size)
    60  		return &ret
    61  	}
    62  	v := b.pool[slot].pool.Get()
    63  	if v == nil {
    64  		// 如果返回的 v == nil
    65  		// 手动创建...
    66  		ret := newBytes(b.pool[slot].defaultSize)
    67  		ret = ret[0:size]
    68  		return &ret
    69  	}
    70  	ret := v.(*[]byte)
    71  	*ret = (*ret)[0:size]
    72  	return ret
    73  }
    74  
    75  // put 将 *[]byte 归还给 bytePool
    76  func (b *bytePool) put(buf *[]byte) {
    77  	if buf == nil {
    78  		return
    79  	}
    80  	// 获取到 size 大小
    81  	size := cap(*buf)
    82  	slot := b.slot(size)
    83  	if slot == errSlot || size != b.pool[slot].defaultSize {
    84  		// 说明是特殊创建的 []byte 直接释放
    85  		// != defaultSize 有可能在执行过程中扩容了
    86  		return
    87  	}
    88  	// 归还
    89  	b.pool[slot].pool.Put(buf)
    90  }
    91  
    92  // newBytes 手动创建 []byte
    93  func newBytes(size int) []byte {
    94  	return make([]byte, size)
    95  }
    96  
    97  // newBytePool 实例化
    98  func newBytePool() *bytePool {
    99  	b := &bytePool{
   100  		minShift: minShift,
   101  		minSize:  1 << minShift,
   102  		maxSize:  1 << maxShift,
   103  	}
   104  	for i := (uint)(0); i < maxShift-minShift; i++ {
   105  		slot := &byteSlot{defaultSize: 1 << (i + minShift)}
   106  		b.pool = append(b.pool, slot)
   107  	}
   108  	return b
   109  }
   110  
   111  // BytePoolContainer 暴露给外部使用的容器对象
   112  type BytePoolContainer struct {
   113  	bytes []*[]byte
   114  	*bytePool
   115  }
   116  
   117  // Reset 将 bytes 中缓存的buffer全部归还给 pool中
   118  func (B *BytePoolContainer) Reset() {
   119  	for _, buf := range B.bytes {
   120  		B.put(buf)
   121  	}
   122  	B.bytes = B.bytes[:0]
   123  }
   124  
   125  func (B *BytePoolContainer) Get(size int) *[]byte {
   126  	buf := B.get(size)
   127  	B.bytes = append(B.bytes, buf)
   128  	return buf
   129  }
   130  
   131  // NewBytePoolContainer 实例化外部容器
   132  func NewBytePoolContainer() *BytePoolContainer {
   133  	return &BytePoolContainer{
   134  		bytePool: localBytePool,
   135  	}
   136  }
   137  
   138  // GetBytes 提供外部接口 获取 size 大小的 buffer
   139  func GetBytes(size int) *[]byte {
   140  	return localBytePool.get(size)
   141  }
   142  
   143  // PutBytes 提供外部接口 将buffer 放回 pool中
   144  func PutBytes(buf *[]byte) {
   145  	localBytePool.put(buf)
   146  }