github.com/dubbogo/gost@v1.14.0/bytes/bytes_pool.go (about)

     1  /*
     2   * Licensed to the Apache Software Foundation (ASF) under one or more
     3   * contributor license agreements.  See the NOTICE file distributed with
     4   * this work for additional information regarding copyright ownership.
     5   * The ASF licenses this file to You under the Apache License, Version 2.0
     6   * (the "License"); you may not use this file except in compliance with
     7   * the License.  You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   */
    17  
    18  package gxbytes
    19  
    20  import (
    21  	"sync"
    22  )
    23  
    24  // BytesPool hold specific size []byte
    25  type BytesPool struct {
    26  	sizes  []int // sizes declare the cap of each slot
    27  	slots  []sync.Pool
    28  	length int
    29  }
    30  
    31  var defaultBytesPool = NewBytesPool([]int{512, 1 << 10, 4 << 10, 16 << 10, 64 << 10})
    32  
    33  // NewBytesPool creates a memory pool.
    34  func NewBytesPool(slotSize []int) *BytesPool {
    35  	bp := &BytesPool{}
    36  	bp.sizes = slotSize
    37  	bp.length = len(bp.sizes)
    38  
    39  	bp.slots = make([]sync.Pool, bp.length)
    40  	for i, size := range bp.sizes {
    41  		size := size
    42  		bp.slots[i] = sync.Pool{New: func() interface{} {
    43  			buf := make([]byte, 0, size)
    44  			return &buf
    45  		}}
    46  	}
    47  	return bp
    48  }
    49  
    50  // SetDefaultBytesPool change default pool options
    51  func SetDefaultBytesPool(bp *BytesPool) {
    52  	defaultBytesPool = bp
    53  }
    54  
    55  func (bp *BytesPool) findIndex(size int) int {
    56  	for i := 0; i < bp.length; i++ {
    57  		if bp.sizes[i] >= size {
    58  			return i
    59  		}
    60  	}
    61  	return bp.length
    62  }
    63  
    64  // AcquireBytes get specific length []byte
    65  func (bp *BytesPool) AcquireBytes(size int) *[]byte {
    66  	idx := bp.findIndex(size)
    67  	if idx >= bp.length {
    68  		buf := make([]byte, size, size)
    69  		return &buf
    70  	}
    71  
    72  	bufp := bp.slots[idx].Get().(*[]byte)
    73  	buf := (*bufp)[:size]
    74  	return &buf
    75  }
    76  
    77  // ReleaseBytes ...
    78  func (bp *BytesPool) ReleaseBytes(bufp *[]byte) {
    79  	bufCap := cap(*bufp)
    80  	idx := bp.findIndex(bufCap)
    81  	if idx >= bp.length || bp.sizes[idx] != bufCap {
    82  		return
    83  	}
    84  
    85  	bp.slots[idx].Put(bufp)
    86  }
    87  
    88  // AcquireBytes called by defaultBytesPool
    89  func AcquireBytes(size int) *[]byte { return defaultBytesPool.AcquireBytes(size) }
    90  
    91  // ReleaseBytes called by defaultBytesPool
    92  func ReleaseBytes(bufp *[]byte) { defaultBytesPool.ReleaseBytes(bufp) }