github.com/pawelgaczynski/gain@v0.4.0-alpha.0.20230821120126-41f1e60a18da/pkg/pool/virtualmem/virtualmem_pool.go (about)

     1  // Copyright (c) 2023 Paweł Gaczyński
     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 virtualmem
    16  
    17  import (
    18  	"math/bits"
    19  
    20  	"github.com/pawelgaczynski/gain/pkg/pool/sync"
    21  )
    22  
    23  const (
    24  	maxVMSize = 64 * 1024 * 1024
    25  )
    26  
    27  var builtinPool = NewPool()
    28  
    29  func Get(size int) *VirtualMem {
    30  	return builtinPool.Get(size)
    31  }
    32  
    33  // Put returns the virtual memory to the built-in pool.
    34  func Put(mem *VirtualMem) {
    35  	mem.Zeroes()
    36  	builtinPool.Put(mem)
    37  }
    38  
    39  // Get retrieves a byte slice of the length requested by the caller from pool or allocates a new one.
    40  func (p *Pool) Get(size int) *VirtualMem {
    41  	if size <= 0 {
    42  		return nil
    43  	}
    44  	size = AdjustBufferSize(size)
    45  
    46  	if size > maxVMSize {
    47  		return NewVirtualMem(size)
    48  	}
    49  	idx := index(uint32(size))
    50  
    51  	ptr := p.pools[idx].Get()
    52  	if ptr == nil {
    53  		return NewVirtualMem(size)
    54  	}
    55  
    56  	return ptr
    57  }
    58  
    59  // Put returns the virtual memory to the pool.
    60  func (p *Pool) Put(mem *VirtualMem) {
    61  	if mem.Size < pageSize || mem.Size > maxVMSize {
    62  		return
    63  	}
    64  
    65  	idx := index(uint32(mem.Size))
    66  	if mem.Size != 1<<idx { // this byte slice is not from Pool.Get(), put it into the previous interval of idx
    67  		idx--
    68  	}
    69  
    70  	p.pools[idx].Put(mem)
    71  }
    72  
    73  func index(n uint32) uint32 {
    74  	return uint32(bits.Len32(n - 1))
    75  }
    76  
    77  type Pool struct {
    78  	pools [32]sync.Pool[*VirtualMem]
    79  }
    80  
    81  func NewPool() Pool {
    82  	var pool Pool
    83  	for i := range pool.pools {
    84  		pool.pools[i] = sync.NewPool[*VirtualMem]()
    85  	}
    86  
    87  	return pool
    88  }