github.com/GoWebProd/gip@v0.0.0-20230623090727-b60d41d5d320/slab/slab.go (about)

     1  package slab
     2  
     3  import (
     4  	"math"
     5  	"reflect"
     6  	"sort"
     7  	"unsafe"
     8  
     9  	"github.com/GoWebProd/gip/allocator"
    10  	"github.com/GoWebProd/gip/pool"
    11  )
    12  
    13  const (
    14  	addrBits = 48
    15  	cntBits  = 64 - addrBits
    16  
    17  	addrMask = (1 << addrBits) - 1
    18  	cntMask  = math.MaxUint64 ^ addrMask
    19  )
    20  
    21  type Slab struct {
    22  	minSize int
    23  	maxSize int
    24  	pools   []pool.Pool[[]byte]
    25  	sizes   []int
    26  }
    27  
    28  func New(min int, max int, growFactor float64) *Slab {
    29  	s := &Slab{
    30  		minSize: min,
    31  		maxSize: max,
    32  		pools:   make([]pool.Pool[[]byte], 0, 16),
    33  		sizes:   make([]int, 0, 16),
    34  	}
    35  
    36  	last := 0.0
    37  
    38  	for i := float64(min); ; i *= growFactor {
    39  		if i-last < 1.0 {
    40  			i = math.Trunc(last + 1)
    41  		}
    42  
    43  		s.pools = append(s.pools, pool.Pool[[]byte]{})
    44  		s.sizes = append(s.sizes, int(i))
    45  		last = i
    46  
    47  		if i >= float64(max) {
    48  			break
    49  		}
    50  	}
    51  
    52  	return s
    53  }
    54  
    55  func (s *Slab) Get(size int) *[]byte {
    56  	pool, idx := s.findPool(size)
    57  	if pool == nil {
    58  		d := allocator.Alloc(size)
    59  
    60  		pack(&d, math.MaxUint16)
    61  
    62  		return &d
    63  	}
    64  
    65  	data := pool.Get()
    66  	if data == nil {
    67  		data = new([]byte)
    68  		*data = allocator.Alloc(s.sizes[idx])
    69  	}
    70  
    71  	*data = (*data)[:size]
    72  
    73  	pack(data, uintptr(idx))
    74  
    75  	return data
    76  }
    77  
    78  func (s *Slab) Put(data *[]byte) {
    79  	if data == nil {
    80  		return
    81  	}
    82  
    83  	(*data) = (*data)[:cap(*data)]
    84  	idx := unpack(data)
    85  
    86  	if int(idx) > len(s.sizes) {
    87  		allocator.Free(*data)
    88  
    89  		return
    90  	}
    91  
    92  	s.pools[idx].Put(data)
    93  }
    94  
    95  func (s *Slab) findPool(size int) (*pool.Pool[[]byte], int) {
    96  	idx := sort.Search(len(s.pools), func(i int) bool {
    97  		return s.sizes[i] >= size
    98  	})
    99  
   100  	if len(s.pools) == idx {
   101  		if idx == 0 || s.sizes[idx-1] < size {
   102  			return nil, 0
   103  		}
   104  
   105  		idx--
   106  	}
   107  
   108  	return &s.pools[idx], idx
   109  }
   110  
   111  func pack(data *[]byte, cnt uintptr) {
   112  	h := (*reflect.SliceHeader)(unsafe.Pointer(data))
   113  
   114  	h.Data = h.Data | cnt<<addrBits
   115  }
   116  
   117  func unpack(data *[]byte) uintptr {
   118  	h := (*reflect.SliceHeader)(unsafe.Pointer(data))
   119  
   120  	cnt := h.Data & cntMask
   121  	h.Data = h.Data & addrMask
   122  
   123  	return cnt >> addrBits
   124  }