github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/buf/alloc.go (about)

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