chunk.go raw

   1  // Copyright 2022 The gVisor Authors.
   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 buffer
  16  
  17  import (
  18  	"fmt"
  19  
  20  	"gvisor.dev/gvisor/pkg/bits"
  21  	"gvisor.dev/gvisor/pkg/sync"
  22  )
  23  
  24  const (
  25  	// This is log2(baseChunkSize). This number is used to calculate which pool
  26  	// to use for a payload size by right shifting the payload size by this
  27  	// number and passing the result to MostSignificantOne64.
  28  	baseChunkSizeLog2 = 6
  29  
  30  	// This is the size of the buffers in the first pool. Each subsequent pool
  31  	// creates payloads 2^(pool index) times larger than the first pool's
  32  	// payloads.
  33  	baseChunkSize = 1 << baseChunkSizeLog2 // 64
  34  
  35  	// MaxChunkSize is largest payload size that we pool. Payloads larger than
  36  	// this will be allocated from the heap and garbage collected as normal.
  37  	MaxChunkSize = baseChunkSize << (numPools - 1) // 64k
  38  
  39  	// The number of chunk pools we have for use.
  40  	numPools = 11
  41  )
  42  
  43  // chunkPools is a collection of pools for payloads of different sizes. The
  44  // size of the payloads doubles in each successive pool.
  45  var chunkPools [numPools]sync.Pool
  46  
  47  func init() {
  48  	for i := 0; i < numPools; i++ {
  49  		chunkSize := baseChunkSize * (1 << i)
  50  		chunkPools[i].New = func() any {
  51  			return &chunk{
  52  				data: make([]byte, chunkSize),
  53  			}
  54  		}
  55  	}
  56  }
  57  
  58  // Precondition: 0 <= size <= maxChunkSize
  59  func getChunkPool(size int) *sync.Pool {
  60  	idx := 0
  61  	if size > baseChunkSize {
  62  		idx = bits.MostSignificantOne64(uint64(size) >> baseChunkSizeLog2)
  63  		if size > 1<<(idx+baseChunkSizeLog2) {
  64  			idx++
  65  		}
  66  	}
  67  	if idx >= numPools {
  68  		panic(fmt.Sprintf("pool for chunk size %d does not exist", size))
  69  	}
  70  	return &chunkPools[idx]
  71  }
  72  
  73  // Chunk represents a slice of pooled memory.
  74  //
  75  // +stateify savable
  76  type chunk struct {
  77  	chunkRefs
  78  	data []byte
  79  }
  80  
  81  func newChunk(size int) *chunk {
  82  	var c *chunk
  83  	if size > MaxChunkSize {
  84  		c = &chunk{
  85  			data: make([]byte, size),
  86  		}
  87  	} else {
  88  		pool := getChunkPool(size)
  89  		c = pool.Get().(*chunk)
  90  		clear(c.data)
  91  	}
  92  	c.InitRefs()
  93  	return c
  94  }
  95  
  96  func (c *chunk) destroy() {
  97  	if len(c.data) > MaxChunkSize {
  98  		c.data = nil
  99  		return
 100  	}
 101  	pool := getChunkPool(len(c.data))
 102  	pool.Put(c)
 103  }
 104  
 105  func (c *chunk) DecRef() {
 106  	c.chunkRefs.DecRef(c.destroy)
 107  }
 108  
 109  func (c *chunk) Clone() *chunk {
 110  	cpy := newChunk(len(c.data))
 111  	copy(cpy.data, c.data)
 112  	return cpy
 113  }
 114