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