forked from k3s-io/k3s
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bitmap.go
236 lines (200 loc) · 6.24 KB
/
bitmap.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
/*
Copyright 2015 The Kubernetes Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package allocator
import (
"errors"
"math/big"
"math/rand"
"sync"
"time"
)
// AllocationBitmap is a contiguous block of resources that can be allocated atomically.
//
// Each resource has an offset. The internal structure is a bitmap, with a bit for each offset.
//
// If a resource is taken, the bit at that offset is set to one.
// r.count is always equal to the number of set bits and can be recalculated at any time
// by counting the set bits in r.allocated.
//
// TODO: use RLE and compact the allocator to minimize space.
type AllocationBitmap struct {
// strategy carries the details of how to choose the next available item out of the range
strategy bitAllocator
// max is the maximum size of the usable items in the range
max int
// rangeSpec is the range specifier, matching RangeAllocation.Range
rangeSpec string
// lock guards the following members
lock sync.Mutex
// count is the number of currently allocated elements in the range
count int
// allocated is a bit array of the allocated items in the range
allocated *big.Int
}
// AllocationBitmap implements Interface and Snapshottable
var _ Interface = &AllocationBitmap{}
var _ Snapshottable = &AllocationBitmap{}
// bitAllocator represents a search strategy in the allocation map for a valid item.
type bitAllocator interface {
AllocateBit(allocated *big.Int, max, count int) (int, bool)
}
// NewAllocationMap creates an allocation bitmap using the random scan strategy.
func NewAllocationMap(max int, rangeSpec string) *AllocationBitmap {
a := AllocationBitmap{
strategy: randomScanStrategy{
rand: rand.New(rand.NewSource(time.Now().UnixNano())),
},
allocated: big.NewInt(0),
count: 0,
max: max,
rangeSpec: rangeSpec,
}
return &a
}
// NewContiguousAllocationMap creates an allocation bitmap using the contiguous scan strategy.
func NewContiguousAllocationMap(max int, rangeSpec string) *AllocationBitmap {
a := AllocationBitmap{
strategy: contiguousScanStrategy{},
allocated: big.NewInt(0),
count: 0,
max: max,
rangeSpec: rangeSpec,
}
return &a
}
// Allocate attempts to reserve the provided item.
// Returns true if it was allocated, false if it was already in use
func (r *AllocationBitmap) Allocate(offset int) (bool, error) {
r.lock.Lock()
defer r.lock.Unlock()
if r.allocated.Bit(offset) == 1 {
return false, nil
}
r.allocated = r.allocated.SetBit(r.allocated, offset, 1)
r.count++
return true, nil
}
// AllocateNext reserves one of the items from the pool.
// (0, false, nil) may be returned if there are no items left.
func (r *AllocationBitmap) AllocateNext() (int, bool, error) {
r.lock.Lock()
defer r.lock.Unlock()
next, ok := r.strategy.AllocateBit(r.allocated, r.max, r.count)
if !ok {
return 0, false, nil
}
r.count++
r.allocated = r.allocated.SetBit(r.allocated, next, 1)
return next, true, nil
}
// Release releases the item back to the pool. Releasing an
// unallocated item or an item out of the range is a no-op and
// returns no error.
func (r *AllocationBitmap) Release(offset int) error {
r.lock.Lock()
defer r.lock.Unlock()
if r.allocated.Bit(offset) == 0 {
return nil
}
r.allocated = r.allocated.SetBit(r.allocated, offset, 0)
r.count--
return nil
}
const (
// Find the size of a big.Word in bytes.
notZero = uint64(^big.Word(0))
wordPower = (notZero>>8)&1 + (notZero>>16)&1 + (notZero>>32)&1
wordSize = 1 << wordPower
)
// ForEach calls the provided function for each allocated bit. The
// AllocationBitmap may not be modified while this loop is running.
func (r *AllocationBitmap) ForEach(fn func(int)) {
r.lock.Lock()
defer r.lock.Unlock()
words := r.allocated.Bits()
for wordIdx, word := range words {
bit := 0
for word > 0 {
if (word & 1) != 0 {
fn((wordIdx * wordSize * 8) + bit)
word = word &^ 1
}
bit++
word = word >> 1
}
}
}
// Has returns true if the provided item is already allocated and a call
// to Allocate(offset) would fail.
func (r *AllocationBitmap) Has(offset int) bool {
r.lock.Lock()
defer r.lock.Unlock()
return r.allocated.Bit(offset) == 1
}
// Free returns the count of items left in the range.
func (r *AllocationBitmap) Free() int {
r.lock.Lock()
defer r.lock.Unlock()
return r.max - r.count
}
// Snapshot saves the current state of the pool.
func (r *AllocationBitmap) Snapshot() (string, []byte) {
r.lock.Lock()
defer r.lock.Unlock()
return r.rangeSpec, r.allocated.Bytes()
}
// Restore restores the pool to the previously captured state.
func (r *AllocationBitmap) Restore(rangeSpec string, data []byte) error {
r.lock.Lock()
defer r.lock.Unlock()
if r.rangeSpec != rangeSpec {
return errors.New("the provided range does not match the current range")
}
r.allocated = big.NewInt(0).SetBytes(data)
r.count = countBits(r.allocated)
return nil
}
// randomScanStrategy chooses a random address from the provided big.Int, and then
// scans forward looking for the next available address (it will wrap the range if
// necessary).
type randomScanStrategy struct {
rand *rand.Rand
}
func (rss randomScanStrategy) AllocateBit(allocated *big.Int, max, count int) (int, bool) {
if count >= max {
return 0, false
}
offset := rss.rand.Intn(max)
for i := 0; i < max; i++ {
at := (offset + i) % max
if allocated.Bit(at) == 0 {
return at, true
}
}
return 0, false
}
var _ bitAllocator = randomScanStrategy{}
// contiguousScanStrategy tries to allocate starting at 0 and filling in any gaps
type contiguousScanStrategy struct{}
func (contiguousScanStrategy) AllocateBit(allocated *big.Int, max, count int) (int, bool) {
if count >= max {
return 0, false
}
for i := 0; i < max; i++ {
if allocated.Bit(i) == 0 {
return i, true
}
}
return 0, false
}
var _ bitAllocator = contiguousScanStrategy{}