arena.go raw
1 /*
2 * SPDX-FileCopyrightText: © Hypermode Inc. <hello@hypermode.com>
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6 package skl
7
8 import (
9 "sync/atomic"
10 "unsafe"
11
12 "github.com/dgraph-io/badger/v4/y"
13 )
14
15 const (
16 offsetSize = int(unsafe.Sizeof(uint32(0)))
17
18 // Always align nodes on 64-bit boundaries, even on 32-bit architectures,
19 // so that the node.value field is 64-bit aligned. This is necessary because
20 // node.getValueOffset uses atomic.LoadUint64, which expects its input
21 // pointer to be 64-bit aligned.
22 nodeAlign = int(unsafe.Sizeof(uint64(0))) - 1
23 )
24
25 // Arena should be lock-free.
26 type Arena struct {
27 n atomic.Uint32
28 buf []byte
29 }
30
31 // newArena returns a new arena.
32 func newArena(n int64) *Arena {
33 // Don't store data at position 0 in order to reserve offset=0 as a kind
34 // of nil pointer.
35 out := &Arena{buf: make([]byte, n)}
36 out.n.Store(1)
37 return out
38 }
39
40 func (s *Arena) size() int64 {
41 return int64(s.n.Load())
42 }
43
44 // putNode allocates a node in the arena. The node is aligned on a pointer-sized
45 // boundary. The arena offset of the node is returned.
46 func (s *Arena) putNode(height int) uint32 {
47 // Compute the amount of the tower that will never be used, since the height
48 // is less than maxHeight.
49 unusedSize := (maxHeight - height) * offsetSize
50
51 // Pad the allocation with enough bytes to ensure pointer alignment.
52 l := uint32(MaxNodeSize - unusedSize + nodeAlign)
53 n := s.n.Add(l)
54 y.AssertTruef(int(n) <= len(s.buf),
55 "Arena too small, toWrite:%d newTotal:%d limit:%d",
56 l, n, len(s.buf))
57
58 // Return the aligned offset.
59 m := (n - l + uint32(nodeAlign)) & ^uint32(nodeAlign)
60 return m
61 }
62
63 // Put will *copy* val into arena. To make better use of this, reuse your input
64 // val buffer. Returns an offset into buf. User is responsible for remembering
65 // size of val. We could also store this size inside arena but the encoding and
66 // decoding will incur some overhead.
67 func (s *Arena) putVal(v y.ValueStruct) uint32 {
68 l := v.EncodedSize()
69 n := s.n.Add(l)
70 y.AssertTruef(int(n) <= len(s.buf),
71 "Arena too small, toWrite:%d newTotal:%d limit:%d",
72 l, n, len(s.buf))
73 m := n - l
74 v.Encode(s.buf[m:])
75 return m
76 }
77
78 func (s *Arena) putKey(key []byte) uint32 {
79 l := uint32(len(key))
80 n := s.n.Add(l)
81 y.AssertTruef(int(n) <= len(s.buf),
82 "Arena too small, toWrite:%d newTotal:%d limit:%d",
83 l, n, len(s.buf))
84 // m is the offset where you should write.
85 // n = new len - key len give you the offset at which you should write.
86 m := n - l
87 // Copy to buffer from m:n
88 y.AssertTrue(len(key) == copy(s.buf[m:n], key))
89 return m
90 }
91
92 // getNode returns a pointer to the node located at offset. If the offset is
93 // zero, then the nil node pointer is returned.
94 func (s *Arena) getNode(offset uint32) *node {
95 if offset == 0 {
96 return nil
97 }
98
99 return (*node)(unsafe.Pointer(&s.buf[offset]))
100 }
101
102 // getKey returns byte slice at offset.
103 func (s *Arena) getKey(offset uint32, size uint16) []byte {
104 return s.buf[offset : offset+uint32(size)]
105 }
106
107 // getVal returns byte slice at offset. The given size should be just the value
108 // size and should NOT include the meta bytes.
109 func (s *Arena) getVal(offset uint32, size uint32) (ret y.ValueStruct) {
110 ret.Decode(s.buf[offset : offset+size])
111 return
112 }
113
114 // getNodeOffset returns the offset of node in the arena. If the node pointer is
115 // nil, then the zero offset is returned.
116 func (s *Arena) getNodeOffset(nd *node) uint32 {
117 if nd == nil {
118 return 0
119 }
120
121 return uint32(uintptr(unsafe.Pointer(nd)) - uintptr(unsafe.Pointer(&s.buf[0])))
122 }
123