forked from dgraph-io/badger
-
Notifications
You must be signed in to change notification settings - Fork 0
/
arena.go
136 lines (116 loc) · 3.89 KB
/
arena.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
/*
* Copyright 2017 Dgraph Labs, Inc. and Contributors
*
* 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 skl
import (
"sync/atomic"
"unsafe"
"github.com/dgraph-io/badger/y"
)
const (
offsetSize = int(unsafe.Sizeof(uint32(0)))
// Always align nodes on 64-bit boundaries, even on 32-bit architectures,
// so that the node.value field is 64-bit aligned. This is necessary because
// node.getValueOffset uses atomic.LoadUint64, which expects its input
// pointer to be 64-bit aligned.
nodeAlign = int(unsafe.Sizeof(uint64(0))) - 1
)
// Arena should be lock-free.
type Arena struct {
n uint32
buf []byte
}
// newArena returns a new arena.
func newArena(n int64) *Arena {
// Don't store data at position 0 in order to reserve offset=0 as a kind
// of nil pointer.
out := &Arena{
n: 1,
buf: make([]byte, n),
}
return out
}
func (s *Arena) size() int64 {
return int64(atomic.LoadUint32(&s.n))
}
func (s *Arena) reset() {
atomic.StoreUint32(&s.n, 0)
}
// putNode allocates a node in the arena. The node is aligned on a pointer-sized
// boundary. The arena offset of the node is returned.
func (s *Arena) putNode(height int) uint32 {
// Compute the amount of the tower that will never be used, since the height
// is less than maxHeight.
unusedSize := (maxHeight - height) * offsetSize
// Pad the allocation with enough bytes to ensure pointer alignment.
l := uint32(MaxNodeSize - unusedSize + nodeAlign)
n := atomic.AddUint32(&s.n, l)
y.AssertTruef(int(n) <= len(s.buf),
"Arena too small, toWrite:%d newTotal:%d limit:%d",
l, n, len(s.buf))
// Return the aligned offset.
m := (n - l + uint32(nodeAlign)) & ^uint32(nodeAlign)
return m
}
// Put will *copy* val into arena. To make better use of this, reuse your input
// val buffer. Returns an offset into buf. User is responsible for remembering
// size of val. We could also store this size inside arena but the encoding and
// decoding will incur some overhead.
func (s *Arena) putVal(v y.ValueStruct) uint32 {
l := uint32(v.EncodedSize())
n := atomic.AddUint32(&s.n, l)
y.AssertTruef(int(n) <= len(s.buf),
"Arena too small, toWrite:%d newTotal:%d limit:%d",
l, n, len(s.buf))
m := n - l
v.Encode(s.buf[m:])
return m
}
func (s *Arena) putKey(key []byte) uint32 {
l := uint32(len(key))
n := atomic.AddUint32(&s.n, l)
y.AssertTruef(int(n) <= len(s.buf),
"Arena too small, toWrite:%d newTotal:%d limit:%d",
l, n, len(s.buf))
m := n - l
y.AssertTrue(len(key) == copy(s.buf[m:n], key))
return m
}
// getNode returns a pointer to the node located at offset. If the offset is
// zero, then the nil node pointer is returned.
func (s *Arena) getNode(offset uint32) *node {
if offset == 0 {
return nil
}
return (*node)(unsafe.Pointer(&s.buf[offset]))
}
// getKey returns byte slice at offset.
func (s *Arena) getKey(offset uint32, size uint16) []byte {
return s.buf[offset : offset+uint32(size)]
}
// getVal returns byte slice at offset. The given size should be just the value
// size and should NOT include the meta bytes.
func (s *Arena) getVal(offset uint32, size uint16) (ret y.ValueStruct) {
ret.Decode(s.buf[offset : offset+uint32(size)])
return
}
// getNodeOffset returns the offset of node in the arena. If the node pointer is
// nil, then the zero offset is returned.
func (s *Arena) getNodeOffset(nd *node) uint32 {
if nd == nil {
return 0
}
return uint32(uintptr(unsafe.Pointer(nd)) - uintptr(unsafe.Pointer(&s.buf[0])))
}