trpc.group/trpc-go/trpc-go@v1.0.3/internal/allocator/allocator.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package allocator implements byte slice pooling management
    15  // to reduce the pressure of memory allocation.
    16  package allocator
    17  
    18  import (
    19  	"fmt"
    20  	"sync"
    21  )
    22  
    23  const maxPowerToRoundUpInt = 63
    24  
    25  var defaultAllocator = NewClassAllocator()
    26  
    27  // Malloc gets a []byte from pool. The second return param is used to Free.
    28  func Malloc(size int) ([]byte, interface{}) {
    29  	return defaultAllocator.Malloc(size)
    30  }
    31  
    32  // Free releases the bytes to pool.
    33  func Free(bts interface{}) {
    34  	defaultAllocator.Free(bts)
    35  }
    36  
    37  // NewClassAllocator creates a new ClassAllocator.
    38  func NewClassAllocator() *ClassAllocator {
    39  	var pools [maxPowerToRoundUpInt]*sync.Pool
    40  	for i := range pools {
    41  		size := 1 << i
    42  		pools[i] = &sync.Pool{
    43  			New: func() interface{} {
    44  				return make([]byte, size)
    45  			},
    46  		}
    47  	}
    48  	return &ClassAllocator{pools: pools}
    49  }
    50  
    51  // ClassAllocator is a bytes pool. The size of bytes satisfies 1 << n.
    52  type ClassAllocator struct {
    53  	pools [maxPowerToRoundUpInt]*sync.Pool
    54  }
    55  
    56  // Malloc gets a []byte from pool. The second return param is used to Free.
    57  // We may also use first return param to Free bytes, but this causes an additional heap allocation.
    58  // See https://github.com/golang/go/issues/8618 for more details.
    59  func (a *ClassAllocator) Malloc(size int) ([]byte, interface{}) {
    60  	if size <= 0 {
    61  		panic(fmt.Sprintf("invalid alloc size %d", size))
    62  	}
    63  	v := a.pools[powerToRoundUp(size)].Get()
    64  	return v.([]byte)[:size], v
    65  }
    66  
    67  // Free releases the bytes to pool.
    68  func (a *ClassAllocator) Free(bts interface{}) {
    69  	cap := cap(bts.([]byte))
    70  	if cap == 0 {
    71  		panic("free an empty bytes")
    72  	}
    73  	power := powerToRoundUp(cap)
    74  	if 1<<power != cap {
    75  		panic(fmt.Sprintf("cap %d of bts must be power of two", cap))
    76  	}
    77  	a.pools[power].Put(bts)
    78  }
    79  
    80  func powerToRoundUp(n int) int {
    81  	powerOfTwo, power := 1, 0
    82  	for ; n-powerOfTwo > 0; power++ {
    83  		powerOfTwo <<= 1
    84  	}
    85  	return power
    86  }