Skip to content

Commit

Permalink
part: Use a proper leaf node
Browse files Browse the repository at this point in the history 
Have a proper leaf node instead of storing the leaf inside an empty
node4. This will reduce memory usage and takes out an allocation and
a pointer dereference.

For 1M dummy objects the memory usage is ~15% less and the insert
speed is ~15% faster, with 30% less allocations.

The leaf node is still wasting some space as it shares the header
and the 'leaf' and 'prefix' fields could be removed.

Before:

Benchmark_Insert_RootOnlyWatch-8            4873            213619 ns/op           4681237 objects/sec    148882 B/op       3076 allocs/op
Benchmark_Insert-8                          4107            267379 ns/op           3740016 objects/sec    249387 B/op       4114 allocs/op

Inserting batch 1000/1000 ...
Waiting for reconciliation to finish ...

1000000 objects reconciled in 3.10 seconds (batch size 1000)
Throughput 322570.39 objects per second
Allocated 8012221 objects, 471726kB bytes, 527128kB bytes still in use

After:

Benchmark_Insert_RootOnlyWatch-8            5898            187664 ns/op           5328686 objects/sec    116722 B/op       2071 allocs/op
Benchmark_Insert-8                          4978            257483 ns/op           3883746 objects/sec    217225 B/op       3109 allocs/op

benchmark % go run . -objects 1000000
Inserting batch 1000/1000 ...
Waiting for reconciliation to finish ...

1000000 objects reconciled in 2.73 seconds (batch size 1000)
Throughput 365990.20 objects per second
Allocated 6011516 objects, 409163kB bytes, 484024kB bytes still in use

Signed-off-by: Jussi Maki <[email protected]>
  • Loading branch information
@joamaki
joamaki committed Apr 16, 2024
1 parent c61760f commit 9749ce1
Show file tree
Hide file tree
Showing 5 changed files with 102 additions and 55 deletions.
8 changes: 4 additions & 4 deletions part/iterator.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ func lowerbound[T any](start *header[T], key []byte) *Iterator[T] {

var traverseToMin func(n *header[T])
traverseToMin = func(n *header[T]) {
if n.leaf != nil {
if leaf := n.getLeaf(); leaf != nil {
edges = append(edges, []*header[T]{n})
return
}
Expand Down Expand Up @@ -163,9 +163,9 @@ func (it *Iterator[T]) Next() (key []byte, value T, ok bool) {
if node.size() > 0 {
it.next = append(it.next, node.children())
}
if node.leaf != nil {
key = node.leaf.key
value = node.leaf.value
if leaf := node.getLeaf(); leaf != nil {
key = leaf.key
value = leaf.value
ok = true
return
}
Expand Down
69 changes: 53 additions & 16 deletions part/node.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,17 +15,13 @@ type nodeKind uint8

const (
nodeKindUnknown = iota
nodeKindLeaf
nodeKind4
nodeKind16
nodeKind48
nodeKind256
)

type leaf[T any] struct {
key []byte
value T
}

type header[T any] struct {
flags uint16 // kind(4b) | unused(3b) | size(9b)
prefix []byte // the compressed prefix, [0] is the key
Expand All @@ -47,6 +43,8 @@ const sizeMask = uint16(0b0000_000_1111_1111_1)

func (n *header[T]) cap() int {
switch n.kind() {
case nodeKindLeaf:
return 0
case nodeKind4:
return 4
case nodeKind16:
Expand All @@ -60,6 +58,17 @@ func (n *header[T]) cap() int {
}
}

func (n *header[T]) isLeaf() bool {
return n.kind() == nodeKindLeaf
}

func (n *header[T]) getLeaf() *leaf[T] {
if n.isLeaf() {
return (*leaf[T])(unsafe.Pointer(n))
}
return n.leaf
}

func (n *header[T]) size() int {
return int(n.flags & sizeMask)
}
Expand Down Expand Up @@ -94,6 +103,9 @@ func (n *header[T]) node256() *node256[T] {
func (n *header[T]) clone(watch bool) *header[T] {
var nCopy *header[T]
switch n.kind() {
case nodeKindLeaf:
l := *n.getLeaf()
nCopy = (&l).self()
case nodeKind4:
n4 := *n.node4()
nCopy = (&n4).self()
Expand All @@ -114,17 +126,20 @@ func (n *header[T]) clone(watch bool) *header[T] {
} else {
nCopy.watch = nil
}
if n.leaf != nil {
nCopy.leaf = &leaf[T]{
key: n.leaf.key,
value: n.leaf.value,
}
}
return nCopy
}

func (n *header[T]) promote(watch bool) *header[T] {
switch n.kind() {
case nodeKindLeaf:
node4 := &node4[T]{}
node4.prefix = n.prefix
node4.leaf = n.getLeaf()
node4.setKind(nodeKind4)
if watch {
node4.watch = make(chan struct{})
}
return node4.self()
case nodeKind4:
node4 := n.node4()
node16 := &node16[T]{header: *n}
Expand Down Expand Up @@ -169,6 +184,8 @@ func (n *header[T]) printTree(level int) {

var children []*header[T]
switch n.kind() {
case nodeKindLeaf:
fmt.Printf("leaf[%v]:", n.prefix)
case nodeKind4:
fmt.Printf("node4[%v]:", n.prefix)
children = n.node4().children[:n.size()]
Expand All @@ -184,8 +201,8 @@ func (n *header[T]) printTree(level int) {
default:
panic("unknown node kind")
}
if n.leaf != nil {
fmt.Printf(" %v -> %v", n.leaf.key, n.leaf.value)
if leaf := n.getLeaf(); leaf != nil {
fmt.Printf(" %v -> %v", leaf.key, leaf.value)
}
fmt.Printf("(%p)\n", n)

Expand All @@ -198,6 +215,8 @@ func (n *header[T]) printTree(level int) {

func (n *header[T]) children() []*header[T] {
switch n.kind() {
case nodeKindLeaf:
return nil
case nodeKind4:
return n.node4().children[0:n.size():4]
case nodeKind16:
Expand Down Expand Up @@ -257,6 +276,24 @@ func (n *header[T]) remove(idx int) {
n.setSize(n.size() - 1)
}

type leaf[T any] struct {
header[T]
key []byte
value T
}

func newLeaf[T any](o *options, prefix, key []byte, value T) *leaf[T] {
leaf := &leaf[T]{key: key, value: value}
leaf.prefix = prefix
leaf.setKind(nodeKindLeaf)

if !o.rootOnlyWatch {
leaf.watch = make(chan struct{})
}

return leaf
}

type node4[T any] struct {
header[T]
children [4]*header[T]
Expand All @@ -277,7 +314,7 @@ type node256[T any] struct {
children [256]*header[T]
}

func newRoot[T any]() *header[T] {
func newNode4[T any]() *header[T] {
n := &node4[T]{header: header[T]{watch: make(chan struct{})}}
n.setKind(nodeKind4)
return n.self()
Expand All @@ -293,8 +330,8 @@ func search[T any](root *header[T], key []byte) (value T, watch <-chan struct{},
key = key[len(this.prefix):]

if len(key) == 0 {
if this.leaf != nil {
value = this.leaf.value
if leaf := this.getLeaf(); leaf != nil {
value = leaf.value
watch = this.watch
ok = true
}
Expand Down
22 changes: 19 additions & 3 deletions part/part_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,10 +114,12 @@ func Test_delete(t *testing.T) {

txn = tree.Txn()
for _, i := range keys {
_, _, ok := txn.Get(intKey(i))
v, _, ok := txn.Get(intKey(i))
assert.True(t, ok)
_, hadOld = txn.Delete(intKey(i))
assert.EqualValues(t, v, i)
v, hadOld = txn.Delete(intKey(i))
assert.True(t, hadOld)
assert.EqualValues(t, v, i)
_, _, ok = txn.Get(intKey(i))
assert.False(t, ok)
}
Expand Down Expand Up @@ -255,7 +257,7 @@ func Test_watch(t *testing.T) {
t.Fatal("expected to find 'a'")
}
if string(v) != "b" {
t.Fatal("expected value 'b'")
t.Fatalf("expected value 'b', got '%s'", v)
}
}

Expand Down Expand Up @@ -354,6 +356,20 @@ func Test_deleteNonExistantCommonPrefix(t *testing.T) {
}
}

func Test_replace(t *testing.T) {
tree := New[int]()
key := binary.BigEndian.AppendUint32(nil, uint32(0))

var v int
var hadOld bool
_, hadOld, tree = tree.Insert(key, 1)
require.False(t, hadOld)

v, hadOld, tree = tree.Insert(key, 2)
require.True(t, hadOld)
require.EqualValues(t, 1, v)
}

func Test_prefix(t *testing.T) {
tree := New[[]byte]()
ins := func(s string) { _, _, tree = tree.Insert([]byte(s), []byte(s)) }
Expand Down
2 changes: 1 addition & 1 deletion part/tree.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func New[T any](opts ...Option) *Tree[T] {
opt(&o)
}
return &Tree[T]{
root: newRoot[T](),
root: newNode4[T](),
size: 0,
opts: &o,
}
Expand Down
56 changes: 25 additions & 31 deletions part/txn.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,20 +128,10 @@ func (txn *Txn[T]) cloneNode(n *header[T]) *header[T] {

func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, hadOld bool, newRoot *header[T]) {
fullKey := key
mkLeafNode := func(prefix []byte) *header[T] {
newLeaf := &node4[T]{}
newLeaf.leaf = &leaf[T]{key: fullKey, value: value}
newLeaf.setKind(nodeKind4)
if !txn.opts.rootOnlyWatch {
newLeaf.header.watch = make(chan struct{})
}
newLeaf.prefix = prefix
return newLeaf.self()
}

if root.size() == 0 && root.leaf == nil {
if root.size() == 0 && root.getLeaf() == nil {
txn.watches[root.watch] = struct{}{}
newRoot = mkLeafNode(key)
newRoot = newLeaf(txn.opts, key, key, value).self()
if newRoot.watch == nil {
newRoot.watch = make(chan struct{})
}
Expand All @@ -155,14 +145,15 @@ func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, had
if bytes.HasPrefix(key, this.prefix) {
key = key[len(this.prefix):]
if len(key) == 0 {
if this.leaf != nil {
oldValue = this.leaf.value
if this.isLeaf() {
// This is a leaf node and we just cloned it. Update the value.
leaf := this.getLeaf()
oldValue = leaf.value
leaf.value = value
hadOld = true

this.leaf.key = fullKey
this.leaf.value = value
} else {
this.leaf = &leaf[T]{key: fullKey, value: value}
// This is a non-leaf node, create/replace the existing leaf.
this.leaf = newLeaf(txn.opts, key, fullKey, value)
}
return
}
Expand All @@ -181,7 +172,7 @@ func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, had
// it fits, just need a clone.
this = txn.cloneNode(this)
}
this.insert(idx, mkLeafNode(key).self())
this.insert(idx, newLeaf(txn.opts, key, fullKey, value).self())
*thisp = this
return
}
Expand All @@ -202,7 +193,7 @@ func (txn *Txn[T]) insert(root *header[T], key []byte, value T) (oldValue T, had
this.prefix = this.prefix[len(newPrefix):]
key = key[len(newPrefix):]

newLeaf := mkLeafNode(key)
newLeaf := newLeaf(txn.opts, key, fullKey, value).self()
newNode := &node4[T]{
header: header[T]{prefix: newPrefix},
}
Expand Down Expand Up @@ -242,12 +233,13 @@ func (txn *Txn[T]) delete(root *header[T], key []byte) (oldValue T, hadOld bool,
parents := txn.deleteParents[:1] // Placeholder for root

// Find the target node and record the path to it.
var leaf *leaf[T]
for {
if bytes.HasPrefix(key, this.prefix) {
key = key[len(this.prefix):]
if len(key) == 0 {
if this.leaf == nil {
// Not found.
leaf = this.getLeaf()
if leaf == nil {
return
}
// Target node found!
Expand All @@ -265,23 +257,25 @@ func (txn *Txn[T]) delete(root *header[T], key []byte) (oldValue T, hadOld bool,
}
}

oldValue = this.leaf.value
oldValue = leaf.value
hadOld = true

// The target was found, rebuild the tree from the root upwards.
newRoot = txn.cloneNode(root)
parents[0].node = newRoot

if this == root {
// Target is the root, clear it.
newRoot.leaf = nil
if newRoot.size() == 0 {
// No children so we can clear the prefix.
newRoot.prefix = nil
if root.isLeaf() || newRoot.size() == 0 {
// Replace leaf or empty root with a node4
newRoot = newNode4[T]()
} else {
newRoot = txn.cloneNode(root)
newRoot.leaf = nil
}
return
}

// The target was found, rebuild the tree from the root upwards.
newRoot = txn.cloneNode(root)
parents[0].node = newRoot

for i := len(parents) - 1; i > 0; i-- {
parent := &parents[i-1]
target := &parents[i]
Expand Down

0 comments on commit 9749ce1

Please sign in to comment.