github.com/pawelgaczynski/gain@v0.4.0-alpha.0.20230821120126-41f1e60a18da/pkg/pool/byteslice/byteslice_pool.go (about)

     1  // Copyright (c) 2023 Paweł Gaczyński
     2  // Copyright (c) 2021 Andy Pan
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package byteslice
    17  
    18  import (
    19  	"math"
    20  	"math/bits"
    21  	"reflect"
    22  	"runtime"
    23  	"unsafe"
    24  
    25  	"github.com/pawelgaczynski/gain/pkg/pool/sync"
    26  )
    27  
    28  var builtinPool = NewByteSlicePool()
    29  
    30  // Pool consists of 32 sync.Pool, representing byte slices of length from 0 to 32 in powers of 2.
    31  type Pool struct {
    32  	pools [32]sync.Pool[unsafe.Pointer]
    33  }
    34  
    35  // Get returns a byte slice with given length from the built-in pool.
    36  func Get(size int) []byte {
    37  	return builtinPool.Get(size)
    38  }
    39  
    40  // Put returns the byte slice to the built-in pool.
    41  func Put(buf []byte) {
    42  	builtinPool.Put(buf)
    43  }
    44  
    45  // Get retrieves a byte slice of the length requested by the caller from pool or allocates a new one.
    46  func (p *Pool) Get(size int) []byte {
    47  	if size <= 0 {
    48  		return nil
    49  	}
    50  
    51  	if size > math.MaxInt32 {
    52  		return make([]byte, size)
    53  	}
    54  	idx := index(uint32(size))
    55  
    56  	ptr := p.pools[idx].Get()
    57  	if ptr == nil {
    58  		return make([]byte, 1<<idx)[:size]
    59  	}
    60  
    61  	var buf []byte
    62  	sh := (*reflect.SliceHeader)(unsafe.Pointer(&buf))
    63  	sh.Data = uintptr(ptr)
    64  	sh.Len = size
    65  	sh.Cap = 1 << idx
    66  
    67  	runtime.KeepAlive(ptr)
    68  
    69  	return buf
    70  }
    71  
    72  // Put returns the byte slice to the pool.
    73  func (p *Pool) Put(buf []byte) {
    74  	size := cap(buf)
    75  	if size == 0 || size > math.MaxInt32 {
    76  		return
    77  	}
    78  
    79  	idx := index(uint32(size))
    80  	if size != 1<<idx { // this byte slice is not from Pool.Get(), put it into the previous interval of idx
    81  		idx--
    82  	}
    83  	// array pointer
    84  	p.pools[idx].Put(unsafe.Pointer(&buf[:1][0]))
    85  }
    86  
    87  func index(n uint32) uint32 {
    88  	return uint32(bits.Len32(n - 1))
    89  }
    90  
    91  func NewByteSlicePool() Pool {
    92  	var byteSlicePool Pool
    93  	for i := range byteSlicePool.pools {
    94  		byteSlicePool.pools[i] = sync.NewPool[unsafe.Pointer]()
    95  	}
    96  
    97  	return byteSlicePool
    98  }