github.com/zuoyebang/bitalosdb@v1.1.1-0.20240516111551-79a8c4d8ce20/internal/bytepools/bytepools.go (about)

     1  // Copyright 2021 The Bitalosdb author(hustxrb@163.com) and other contributors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package bytepools
    16  
    17  import (
    18  	"sync"
    19  )
    20  
    21  const (
    22  	poolsNum       = 13
    23  	maxPoolBufSize = 36 << 10
    24  )
    25  
    26  var ReaderBytePools = NewBytePools()
    27  
    28  type BytePools struct {
    29  	pools [poolsNum]sync.Pool
    30  }
    31  
    32  func NewBytePools() *BytePools {
    33  	p := new(BytePools)
    34  	for i := 0; i < poolsNum; i++ {
    35  		size := 16 * (1 << i)
    36  		if i == poolsNum-1 {
    37  			size = maxPoolBufSize
    38  		}
    39  		p.pools[i] = sync.Pool{
    40  			New: func() interface{} {
    41  				return make([]byte, size)
    42  			},
    43  		}
    44  	}
    45  	return p
    46  }
    47  
    48  func (p *BytePools) getIndex(size int) int {
    49  	if size <= 128 {
    50  		switch {
    51  		case size <= 16:
    52  			return 0
    53  		case size <= 32:
    54  			return 1
    55  		case size <= 64:
    56  			return 2
    57  		default:
    58  			return 3
    59  		}
    60  	} else if size <= 2048 {
    61  		switch {
    62  		case size <= 256:
    63  			return 4
    64  		case size <= 512:
    65  			return 5
    66  		case size <= 1024:
    67  			return 6
    68  		default:
    69  			return 7
    70  		}
    71  	} else {
    72  		switch {
    73  		case size <= 4096:
    74  			return 8
    75  		case size <= 8192:
    76  			return 9
    77  		case size <= 16384:
    78  			return 10
    79  		case size <= 32768:
    80  			return 11
    81  		case size <= maxPoolBufSize:
    82  			return 12
    83  		default:
    84  			return -1
    85  		}
    86  	}
    87  }
    88  
    89  func (p *BytePools) Get(size int) interface{} {
    90  	index := p.getIndex(size)
    91  	if index == -1 {
    92  		return make([]byte, size)
    93  	}
    94  	return p.pools[index].Get()
    95  }
    96  
    97  func (p *BytePools) Put(x []byte) {
    98  	index := p.getIndex(len(x))
    99  	if index >= 0 {
   100  		p.pools[index].Put(x)
   101  	}
   102  }
   103  
   104  func (p *BytePools) GetBytePool(size int) ([]byte, func()) {
   105  	v := p.Get(size).([]byte)
   106  	return v, func() {
   107  		p.PutBytePool(v)
   108  	}
   109  }
   110  
   111  func (p *BytePools) GetMaxBytePool() ([]byte, func()) {
   112  	v := p.pools[poolsNum-1].Get().([]byte)
   113  	return v, func() {
   114  		p.pools[poolsNum-1].Put(v)
   115  	}
   116  }
   117  
   118  func (p *BytePools) MakeValue(v []byte) ([]byte, func()) {
   119  	size := len(v)
   120  	pool := p.Get(size).([]byte)
   121  	copy(pool[:size], v)
   122  	return pool[:size], func() {
   123  		p.PutBytePool(pool)
   124  	}
   125  }
   126  
   127  func (p *BytePools) PutBytePool(v []byte) {
   128  	if len(v) > maxPoolBufSize {
   129  		return
   130  	}
   131  	p.Put(v)
   132  }