Problem: memiavl leaf node separation optimization is not implemented

memiavl/benchmark_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"math/rand"
+	"sort"
 	"testing"
 
 	iavlcache "github.com/cosmos/iavl/cache"
@@ -146,6 +147,23 @@ func BenchmarkRandomGet(b *testing.B) {
 			_ = m[string(targetKey)]
 		}
 	})
+
+	b.Run("binary-search", func(b *testing.B) {
+		// the last benchmark sort the items in place
+		sort.Slice(items, func(i, j int) bool {
+			return bytes.Compare(items[i].key, items[j].key) < 0
+		})
+		i := sort.Search(len(items), func(i int) bool { return bytes.Compare(items[i].key, targetKey) != -1 })
+		require.Equal(b, targetValue, items[i].value)
+
+		cmp := func(i int) bool { return bytes.Compare(items[i].key, targetKey) != -1 }
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			n := sort.Search(len(items), cmp)
+			_ = items[n].value
+		}
+	})
 }
 
 func BenchmarkRandomSet(b *testing.B) {

memiavl/export.go

@@ -1,6 +1,7 @@
 package memiavl
 
 import (
+	"context"
 	"fmt"
 	"math"
 
@@ -10,6 +11,11 @@ import (
 	protoio "github.com/gogo/protobuf/io"
 )
 
+// exportBufferSize is the number of nodes to buffer in the exporter. It improves throughput by
+// processing multiple nodes per context switch, but take care to avoid excessive memory usage,
+// especially since callers may export several IAVL stores in parallel (e.g. the Cosmos SDK).
+const exportBufferSize = 32
+
 func (db *DB) Snapshot(height uint64, protoWriter protoio.Writer) error {
 	if height > math.MaxUint32 {
 		return fmt.Errorf("height overflows uint32: %d", height)
@@ -59,26 +65,89 @@ func (db *DB) Snapshot(height uint64, protoWriter protoio.Writer) error {
 
 type Exporter struct {
 	snapshot *Snapshot
-	i        uint32
-	count    int
+	ch       chan *iavl.ExportNode
+	cancel   context.CancelFunc
+}
+
+func newExporter(snapshot *Snapshot) *Exporter {
+	ctx, cancel := context.WithCancel(context.Background())
+	exporter := &Exporter{
+		snapshot: snapshot,
+		ch:       make(chan *iavl.ExportNode, exportBufferSize),
+		cancel:   cancel,
+	}
+	go exporter.export(ctx)
+	return exporter
+}
+
+func (e *Exporter) export(ctx context.Context) {
+	defer close(e.ch)
+
+	if e.snapshot.leavesLen() == 0 {
+		return
+	}
+
+	if e.snapshot.leavesLen() == 1 {
+		leaf := e.snapshot.Leaf(0)
+		e.ch <- &iavl.ExportNode{
+			Height:  0,
+			Version: int64(leaf.Version()),
+			Key:     leaf.Key(),
+			Value:   leaf.Value(),
+		}
+		return
+	}
+
+	var pendingTrees int
+	var i, j uint32
+	for ; i < uint32(e.snapshot.nodesLen()); i++ {
+		// pending branch node
+		node := e.snapshot.nodesLayout.Node(i)
+		for pendingTrees < int(node.PreTrees())+2 {
+			// add more leaf nodes
+			leaf := e.snapshot.leavesLayout.Leaf(j)
+			key, value := e.snapshot.KeyValue(leaf.KeyOffset())
+			enode := &iavl.ExportNode{
+				Height:  0,
+				Version: int64(leaf.Version()),
+				Key:     key,
+				Value:   value,
+			}
+			j++
+			pendingTrees++
+
+			select {
+			case e.ch <- enode:
+			case <-ctx.Done():
+				return
+			}
+		}
+		enode := &iavl.ExportNode{
+			Height:  int8(node.Height()),
+			Version: int64(node.Version()),
+			Key:     e.snapshot.LeafKey(node.KeyLeaf()),
+		}
+		pendingTrees--
+
+		select {
+		case e.ch <- enode:
+		case <-ctx.Done():
+			return
+		}
+	}
 }
 
 func (e *Exporter) Next() (*iavl.ExportNode, error) {
-	if int(e.i) >= e.count {
-		return nil, iavl.ExportDone
+	if exportNode, ok := <-e.ch; ok {
+		return exportNode, nil
 	}
-	node := e.snapshot.Node(e.i)
-	e.i++
+	return nil, iavl.ExportDone
+}
 
-	height := node.Height()
-	var value []byte
-	if height == 0 {
-		value = node.Value()
+// Close closes the exporter. It is safe to call multiple times.
+func (e *Exporter) Close() {
+	e.cancel()
+	for range e.ch { // drain channel
 	}
-	return &iavl.ExportNode{
-		Height:  int8(height),
-		Version: int64(node.Version()),
-		Key:     node.Key(),
-		Value:   value,
-	}, nil
+	e.snapshot = nil
 }

memiavl/import.go

@@ -144,22 +144,24 @@ func doImport(dir string, version int64, nodes <-chan *iavl.ExportNode, writeHas
 			}
 		}
 
-		switch len(i.indexStack) {
+		switch len(i.leavesStack) {
 		case 0:
-			return EmptyRootNodeIndex, nil
+			return 0, nil
 		case 1:
-			return i.indexStack[0], nil
+			return i.leafCounter, nil
 		default:
-			return 0, fmt.Errorf("invalid node structure, found stack size %v after imported", len(i.indexStack))
+			return 0, fmt.Errorf("invalid node structure, found stack size %v after imported", len(i.leavesStack))
 		}
 	})
 }
 
 type importer struct {
 	snapshotWriter
 
-	indexStack []uint32
-	nodeStack  []*MemNode
+	// keep track of how many leaves has been written before the pending nodes
+	leavesStack []uint32
+	// keep track of the pending nodes
+	nodeStack []*MemNode
 }
 
 func (i *importer) Add(n *iavl.ExportNode) error {
@@ -169,24 +171,23 @@ func (i *importer) Add(n *iavl.ExportNode) error {
 
 	if n.Height == 0 {
 		node := &MemNode{
-			height:  uint8(n.Height),
+			height:  0,
 			size:    1,
 			version: uint32(n.Version),
 			key:     n.Key,
 			value:   n.Value,
 		}
 		nodeHash := node.Hash()
-		idx, err := i.writeLeaf(node.version, node.key, node.value, nodeHash)
-		if err != nil {
+		if err := i.writeLeaf(node.version, node.key, node.value, nodeHash); err != nil {
 			return err
 		}
-		i.indexStack = append(i.indexStack, idx)
+		i.leavesStack = append(i.leavesStack, i.leafCounter)
 		i.nodeStack = append(i.nodeStack, node)
 		return nil
 	}
 
 	// branch node
-	leftIndex := i.indexStack[len(i.indexStack)-2]
+	keyLeaf := i.leavesStack[len(i.leavesStack)-2]
 	leftNode := i.nodeStack[len(i.nodeStack)-2]
 	rightNode := i.nodeStack[len(i.nodeStack)-1]
 
@@ -199,13 +200,13 @@ func (i *importer) Add(n *iavl.ExportNode) error {
 		right:   rightNode,
 	}
 	nodeHash := node.Hash()
-	idx, err := i.writeBranch(node.version, uint32(node.size), node.height, leftIndex+1, nodeHash)
-	if err != nil {
+	preTrees := uint8(len(i.nodeStack) - 2)
+	if err := i.writeBranch(node.version, uint32(node.size), node.height, preTrees, keyLeaf, nodeHash); err != nil {
 		return err
 	}
 
-	i.indexStack = i.indexStack[:len(i.indexStack)-2]
-	i.indexStack = append(i.indexStack, idx)
+	i.leavesStack = i.leavesStack[:len(i.leavesStack)-2]
+	i.leavesStack = append(i.leavesStack, i.leafCounter)
 
 	i.nodeStack = i.nodeStack[:len(i.nodeStack)-2]
 	i.nodeStack = append(i.nodeStack, node)

memiavl/iterator.go

@@ -82,7 +82,7 @@ func (iter *Iterator) Next() {
 		afterStart := iter.start == nil || startCmp < 0
 		beforeEnd := iter.end == nil || bytes.Compare(key, iter.end) < 0
 
-		if isLeaf(node) {
+		if node.IsLeaf() {
 			startOrAfter := afterStart || startCmp == 0
 			if startOrAfter && beforeEnd {
 				iter.key = key

memiavl/layout_little_endian.go

@@ -30,6 +30,10 @@ func (node NodeLayout) Height() uint8 {
 	return node.data[OffsetHeight]
 }
 
+func (node NodeLayout) PreTrees() uint8 {
+	return node.data[OffsetPreTrees]
+}
+
 func (node NodeLayout) Version() uint32 {
 	return binary.LittleEndian.Uint32(node.data[OffsetVersion : OffsetVersion+4])
 }
@@ -38,14 +42,44 @@ func (node NodeLayout) Size() uint32 {
 	return binary.LittleEndian.Uint32(node.data[OffsetSize : OffsetSize+4])
 }
 
-func (node NodeLayout) KeyOffset() uint64 {
-	return binary.LittleEndian.Uint64(node.data[OffsetKeyOffset : OffsetKeyOffset+8])
-}
-
-func (node NodeLayout) KeyNode() uint32 {
-	return binary.LittleEndian.Uint32(node.data[OffsetKeyNode : OffsetKeyNode+4])
+func (node NodeLayout) KeyLeaf() uint32 {
+	return binary.LittleEndian.Uint32(node.data[OffsetKeyLeaf : OffsetKeyLeaf+4])
 }
 
 func (node NodeLayout) Hash() []byte {
 	return node.data[OffsetHash : OffsetHash+SizeHash]
 }
+
+// Leaves is a continuously stored IAVL nodes
+type Leaves struct {
+	data []byte
+}
+
+func NewLeaves(data []byte) (Leaves, error) {
+	return Leaves{data}, nil
+}
+
+func (leaves Leaves) Leaf(i uint32) LeafLayout {
+	offset := int(i) * SizeLeaf
+	return LeafLayout{data: (*[SizeLeaf]byte)(leaves.data[offset : offset+SizeLeaf])}
+}
+
+type LeafLayout struct {
+	data *[SizeLeaf]byte
+}
+
+func (leaf LeafLayout) Version() uint32 {
+	return binary.LittleEndian.Uint32(leaf.data[OffsetLeafVersion : OffsetLeafVersion+4])
+}
+
+func (leaf LeafLayout) KeyLength() uint32 {
+	return binary.LittleEndian.Uint32(leaf.data[OffsetLeafKeyLen : OffsetLeafKeyLen+4])
+}
+
+func (leaf LeafLayout) KeyOffset() uint64 {
+	return binary.LittleEndian.Uint64(leaf.data[OffsetLeafKeyOffset : OffsetLeafKeyOffset+8])
+}
+
+func (leaf LeafLayout) Hash() []byte {
+	return leaf.data[OffsetLeafHash : OffsetLeafHash+32]
+}

memiavl/layout_native.go

@@ -43,6 +43,10 @@ func (node *nodeLayout) Height() uint8 {
 	return uint8(node.data[0])
 }
 
+func (node NodeLayout) PreTrees() uint8 {
+	return uint8(node.data[0] >> 8)
+}
+
 func (node *nodeLayout) Version() uint32 {
 	return node.data[1]
 }
@@ -51,7 +55,7 @@ func (node *nodeLayout) Size() uint32 {
 	return node.data[2]
 }
 
-func (node *nodeLayout) KeyNode() uint32 {
+func (node *nodeLayout) KeyLeaf() uint32 {
 	return node.data[3]
 }
 
@@ -62,3 +66,51 @@ func (node *nodeLayout) KeyOffset() uint64 {
 func (node *nodeLayout) Hash() []byte {
 	return node.hash[:]
 }
+
+type LeafLayout = *leafLayout
+
+// Nodes is a continuously stored IAVL nodes
+type Leaves struct {
+	leaves []leafLayout
+}
+
+func NewLeaves(buf []byte) (Leaves, error) {
+	// check alignment and size of the buffer
+	p := unsafe.Pointer(unsafe.SliceData(buf))
+	if uintptr(p)%unsafe.Alignof(leafLayout{}) != 0 {
+		return Leaves{}, errors.New("input buffer is not aligned")
+	}
+	size := int(unsafe.Sizeof(leafLayout{}))
+	if len(buf)%size != 0 {
+		return Leaves{}, errors.New("input buffer length is not correct")
+	}
+	leaves := unsafe.Slice((*leafLayout)(p), len(buf)/size)
+	return Leaves{leaves}, nil
+}
+
+func (leaves Leaves) Leaf(i uint32) LeafLayout {
+	return &leaves.leaves[i]
+}
+
+type leafLayout struct {
+	version   uint32
+	keyLen    uint32
+	keyOffset uint64
+	hash      [32]byte
+}
+
+func (leaf *leafLayout) Version() uint32 {
+	return leaf.version
+}
+
+func (leaf *leafLayout) KeyLength() uint32 {
+	return leaf.keyLen
+}
+
+func (leaf *leafLayout) KeyOffset() uint64 {
+	return leaf.keyOffset
+}
+
+func (leaf *leafLayout) Hash() []byte {
+	return leaf.hash[:]
+}

memiavl/mem_node.go

@@ -34,6 +34,10 @@ func (node *MemNode) Height() uint8 {
 	return node.height
 }
 
+func (node *MemNode) IsLeaf() bool {
+	return node.height == 0
+}
+
 func (node *MemNode) Size() int64 {
 	return node.size
 }