github.com/kelleygo/clashcore@v1.0.2/common/pool/alloc.go (about)

     1  package pool
     2  
     3  // Inspired by https://github.com/xtaci/smux/blob/master/alloc.go
     4  
     5  import (
     6  	"errors"
     7  	"math/bits"
     8  	"sync"
     9  )
    10  
    11  var defaultAllocator = NewAllocator()
    12  
    13  // Allocator for incoming frames, optimized to prevent overwriting after zeroing
    14  type Allocator struct {
    15  	buffers [11]sync.Pool
    16  }
    17  
    18  // NewAllocator initiates a []byte allocator for frames less than 65536 bytes,
    19  // the waste(memory fragmentation) of space allocation is guaranteed to be
    20  // no more than 50%.
    21  func NewAllocator() *Allocator {
    22  	return &Allocator{
    23  		buffers: [...]sync.Pool{ // 64B -> 64K
    24  			{New: func() any { return new([1 << 6]byte) }},
    25  			{New: func() any { return new([1 << 7]byte) }},
    26  			{New: func() any { return new([1 << 8]byte) }},
    27  			{New: func() any { return new([1 << 9]byte) }},
    28  			{New: func() any { return new([1 << 10]byte) }},
    29  			{New: func() any { return new([1 << 11]byte) }},
    30  			{New: func() any { return new([1 << 12]byte) }},
    31  			{New: func() any { return new([1 << 13]byte) }},
    32  			{New: func() any { return new([1 << 14]byte) }},
    33  			{New: func() any { return new([1 << 15]byte) }},
    34  			{New: func() any { return new([1 << 16]byte) }},
    35  		},
    36  	}
    37  }
    38  
    39  // Get a []byte from pool with most appropriate cap
    40  func (alloc *Allocator) Get(size int) []byte {
    41  	switch {
    42  	case size < 0:
    43  		panic("alloc.Get: len out of range")
    44  	case size == 0:
    45  		return nil
    46  	case size > 65536:
    47  		return make([]byte, size)
    48  	default:
    49  		var index uint16
    50  		if size > 64 {
    51  			index = msb(size)
    52  			if size != 1<<index {
    53  				index += 1
    54  			}
    55  			index -= 6
    56  		}
    57  
    58  		buffer := alloc.buffers[index].Get()
    59  		switch index {
    60  		case 0:
    61  			return buffer.(*[1 << 6]byte)[:size]
    62  		case 1:
    63  			return buffer.(*[1 << 7]byte)[:size]
    64  		case 2:
    65  			return buffer.(*[1 << 8]byte)[:size]
    66  		case 3:
    67  			return buffer.(*[1 << 9]byte)[:size]
    68  		case 4:
    69  			return buffer.(*[1 << 10]byte)[:size]
    70  		case 5:
    71  			return buffer.(*[1 << 11]byte)[:size]
    72  		case 6:
    73  			return buffer.(*[1 << 12]byte)[:size]
    74  		case 7:
    75  			return buffer.(*[1 << 13]byte)[:size]
    76  		case 8:
    77  			return buffer.(*[1 << 14]byte)[:size]
    78  		case 9:
    79  			return buffer.(*[1 << 15]byte)[:size]
    80  		case 10:
    81  			return buffer.(*[1 << 16]byte)[:size]
    82  		default:
    83  			panic("invalid pool index")
    84  		}
    85  	}
    86  }
    87  
    88  // Put returns a []byte to pool for future use,
    89  // which the cap must be exactly 2^n
    90  func (alloc *Allocator) Put(buf []byte) error {
    91  	if cap(buf) == 0 || cap(buf) > 65536 {
    92  		return nil
    93  	}
    94  
    95  	bits := msb(cap(buf))
    96  	if cap(buf) != 1<<bits {
    97  		return errors.New("allocator Put() incorrect buffer size")
    98  	}
    99  	if cap(buf) < 1<<6 {
   100  		return nil
   101  	}
   102  	bits -= 6
   103  	buf = buf[:cap(buf)]
   104  
   105  	//nolint
   106  	//lint:ignore SA6002 ignore temporarily
   107  	switch bits {
   108  	case 0:
   109  		alloc.buffers[bits].Put((*[1 << 6]byte)(buf))
   110  	case 1:
   111  		alloc.buffers[bits].Put((*[1 << 7]byte)(buf))
   112  	case 2:
   113  		alloc.buffers[bits].Put((*[1 << 8]byte)(buf))
   114  	case 3:
   115  		alloc.buffers[bits].Put((*[1 << 9]byte)(buf))
   116  	case 4:
   117  		alloc.buffers[bits].Put((*[1 << 10]byte)(buf))
   118  	case 5:
   119  		alloc.buffers[bits].Put((*[1 << 11]byte)(buf))
   120  	case 6:
   121  		alloc.buffers[bits].Put((*[1 << 12]byte)(buf))
   122  	case 7:
   123  		alloc.buffers[bits].Put((*[1 << 13]byte)(buf))
   124  	case 8:
   125  		alloc.buffers[bits].Put((*[1 << 14]byte)(buf))
   126  	case 9:
   127  		alloc.buffers[bits].Put((*[1 << 15]byte)(buf))
   128  	case 10:
   129  		alloc.buffers[bits].Put((*[1 << 16]byte)(buf))
   130  	default:
   131  		panic("invalid pool index")
   132  	}
   133  	return nil
   134  }
   135  
   136  // msb return the pos of most significant bit
   137  func msb(size int) uint16 {
   138  	return uint16(bits.Len32(uint32(size)) - 1)
   139  }