Problem: memiavl's memory copy behavior is not consistent (#1011)

* Problem: memiavl's memory copy behavior is not consistent Closes: #1009 - make the default behavior safe - provide the unsafe zero-copy option * Update memiavl/db_test.go Signed-off-by: yihuang <[email protected]> * Update memiavl/db_test.go Co-authored-by: mmsqe <[email protected]> Signed-off-by: yihuang <[email protected]> --------- Signed-off-by: yihuang <[email protected]> Co-authored-by: mmsqe <[email protected]> Co-authored-by: mmsqe <[email protected]>
crypto-org-chain · May 3, 2023 · 7a75b41 · 7a75b41
1 parent 0bfc853
commit 7a75b41
Show file tree

Hide file tree

Showing 17 changed files with 114 additions and 35 deletions.
diff --git a/app/app.go b/app/app.go
@@ -152,7 +152,9 @@ const (
 	AddrLen = 20
 
 	FlagMemIAVL  = "store.memiavl"
-	FlagAsyncWAL = "store.async-wal"
+	FlagAsyncWAL = "store.memiavl-async-wal"
+	// don't enable zero-copy together with inter-block cache.
+	FlagZeroCopy = "store.memiavl-zero-copy"
 )
 
 // this line is used by starport scaffolding # stargate/wasm/app/enabledProposals
@@ -352,6 +354,7 @@ func New(
 		// FIXME we are assuming the cms won't be overridden by the other options, but we can't be sure.
 		cms := rootmulti.NewStore(filepath.Join(homePath, "data", "memiavl.db"), logger)
 		cms.SetAsyncWAL(cast.ToBool(appOpts.Get(FlagAsyncWAL)))
+		cms.SetZeroCopy(cast.ToBool(appOpts.Get(FlagZeroCopy)))
 		baseAppOptions = append([]func(*baseapp.BaseApp){setCMS(cms)}, baseAppOptions...)
 	}
 

diff --git a/integration_tests/configs/default.jsonnet b/integration_tests/configs/default.jsonnet
@@ -29,6 +29,7 @@
       'app-config': {
         store: {
           memiavl: true,
+          'memiavl-zero-copy': true,
           streamers: ['versiondb'],
         },
       },

diff --git a/memiavl/benchmark_test.go b/memiavl/benchmark_test.go
@@ -38,7 +38,7 @@ func BenchmarkRandomGet(b *testing.B) {
 	snapshot, err := OpenSnapshot(snapshotDir)
 	require.NoError(b, err)
 	defer snapshot.Close()
-	diskTree := NewFromSnapshot(snapshot)
+	diskTree := NewFromSnapshot(snapshot, true)
 
 	require.Equal(b, targetValue, tree.Get(targetKey))
 	require.Equal(b, targetValue, diskTree.Get(targetKey))

diff --git a/memiavl/db.go b/memiavl/db.go
@@ -65,6 +65,8 @@ type Options struct {
 	TargetVersion uint32
 	// Write WAL asynchronously, it's ok in blockchain case because we can always replay the raw blocks.
 	AsyncWAL bool
+	// ZeroCopy if true, the get and iterator methods could return a slice pointing to mmaped blob files.
+	ZeroCopy bool
 }
 
 const (
@@ -73,13 +75,13 @@ const (
 
 func Load(dir string, opts Options) (*DB, error) {
 	currentDir := currentPath(dir)
-	mtree, err := LoadMultiTree(currentDir)
+	mtree, err := LoadMultiTree(currentDir, opts.ZeroCopy)
 	if err != nil {
 		if opts.CreateIfMissing && os.IsNotExist(err) {
 			if err := initEmptyDB(dir, opts.InitialVersion); err != nil {
 				return nil, err
 			}
-			mtree, err = LoadMultiTree(currentDir)
+			mtree, err = LoadMultiTree(currentDir, opts.ZeroCopy)
 		}
 		if err != nil {
 			return nil, err
@@ -326,7 +328,7 @@ func (db *DB) Reload() error {
 }
 
 func (db *DB) reload() error {
-	mtree, err := LoadMultiTree(currentPath(db.dir))
+	mtree, err := LoadMultiTree(currentPath(db.dir), db.zeroCopy)
 	if err != nil {
 		return err
 	}
@@ -376,7 +378,7 @@ func (db *DB) RewriteSnapshotBackground() error {
 			ch <- snapshotResult{err: err}
 			return
 		}
-		mtree, err := LoadMultiTree(currentPath(db.dir))
+		mtree, err := LoadMultiTree(currentPath(cloned.dir), cloned.zeroCopy)
 		if err != nil {
 			ch <- snapshotResult{err: err}
 			return

diff --git a/memiavl/db_test.go b/memiavl/db_test.go
@@ -2,7 +2,9 @@ package memiavl
 
 import (
 	"encoding/hex"
+	"errors"
 	"os"
+	"runtime/debug"
 	"strconv"
 	"testing"
 	"time"
@@ -233,3 +235,33 @@ func TestLoadVersion(t *testing.T) {
 		require.Equal(t, expItems, collectIter(tmp.TreeByName("test").Iterator(nil, nil, true)))
 	}
 }
+
+func TestZeroCopy(t *testing.T) {
+	db, err := Load(t.TempDir(), Options{InitialStores: []string{"test"}, CreateIfMissing: true, ZeroCopy: true})
+	require.NoError(t, err)
+	db.Commit([]*NamedChangeSet{
+		{Name: "test", Changeset: ChangeSets[0]},
+	})
+	require.NoError(t, errors.Join(
+		db.RewriteSnapshot(),
+		db.Reload(),
+	))
+
+	value := db.TreeByName("test").Get([]byte("hello"))
+	require.Equal(t, []byte("world"), value)
+
+	db.SetZeroCopy(false)
+	valueCloned := db.TreeByName("test").Get([]byte("hello"))
+	require.Equal(t, []byte("world"), valueCloned)
+
+	require.NoError(t, db.Close())
+
+	require.Equal(t, []byte("world"), valueCloned)
+
+	// accessing the zero-copy value after the db is closed triggers segment fault.
+	// reset global panic on fault setting after function finished
+	defer debug.SetPanicOnFault(debug.SetPanicOnFault(true))
+	require.Panics(t, func() {
+		require.Equal(t, []byte("world"), value)
+	})
+}
diff --git a/memiavl/export.go b/memiavl/export.go
@@ -15,7 +15,7 @@ func (db *DB) Snapshot(height uint64, protoWriter protoio.Writer) error {
 		return fmt.Errorf("height overflows uint32: %d", height)
 	}
 
-	mtree, err := LoadMultiTree(snapshotPath(db.dir, uint32(height)))
+	mtree, err := LoadMultiTree(snapshotPath(db.dir, uint32(height)), true)
 	if err != nil {
 		return errors.Wrapf(err, "invalid snapshot height: %d", height)
 	}

diff --git a/memiavl/iterator.go b/memiavl/iterator.go
@@ -10,6 +10,7 @@ type Iterator struct {
 	// domain of iteration, end is exclusive
 	start, end []byte
 	ascending  bool
+	zeroCopy   bool
 
 	// cache the next key-value pair
 	key, value []byte
@@ -21,12 +22,13 @@ type Iterator struct {
 
 var _ dbm.Iterator = (*Iterator)(nil)
 
-func NewIterator(start, end []byte, ascending bool, root Node) dbm.Iterator {
+func NewIterator(start, end []byte, ascending bool, root Node, zeroCopy bool) dbm.Iterator {
 	iter := &Iterator{
 		start:     start,
 		end:       end,
 		ascending: ascending,
 		valid:     true,
+		zeroCopy:  zeroCopy,
 	}
 
 	if root != nil {
@@ -54,11 +56,17 @@ func (iter *Iterator) Error() error {
 
 // Key implements dbm.Iterator
 func (iter *Iterator) Key() []byte {
+	if !iter.zeroCopy {
+		return bytes.Clone(iter.key)
+	}
 	return iter.key
 }
 
 // Value implements dbm.Iterator
 func (iter *Iterator) Value() []byte {
+	if !iter.zeroCopy {
+		return bytes.Clone(iter.value)
+	}
 	return iter.value
 }
 

diff --git a/memiavl/mmap.go b/memiavl/mmap.go
@@ -38,10 +38,18 @@ func NewMmap(path string) (*MmapFile, error) {
 
 // Close closes the file and mmap handles
 func (m *MmapFile) Close() error {
-	return errors.Join(m.file.Close(), mmap.Munmap(m.data, m.handle))
+	return errors.Join(mmap.Munmap(m.data, m.handle), m.file.Close())
 }
 
 // Data returns the mmap-ed buffer
 func (m *MmapFile) Data() []byte {
 	return m.data
 }
+
+func Mmap(f *os.File) ([]byte, *[mmap.MaxMapSize]byte, error) {
+	fi, err := f.Stat()
+	if err != nil {
+		return nil, nil, err
+	}
+	return mmap.Mmap(f, int(fi.Size()))
+}
diff --git a/memiavl/multitree.go b/memiavl/multitree.go
@@ -45,6 +45,8 @@ type MultiTree struct {
 	// it always corresponds to the wal entry with index 1.
 	initialVersion uint32
 
+	zeroCopy bool
+
 	trees          []namedTree
 	treesByName    map[string]int // reversed index of the trees
 	lastCommitInfo storetypes.CommitInfo
@@ -57,10 +59,11 @@ func NewEmptyMultiTree(initialVersion uint32) *MultiTree {
 	return &MultiTree{
 		initialVersion: initialVersion,
 		treesByName:    make(map[string]int),
+		zeroCopy:       true,
 	}
 }
 
-func LoadMultiTree(dir string) (*MultiTree, error) {
+func LoadMultiTree(dir string, zeroCopy bool) (*MultiTree, error) {
 	// load commit info
 	bz, err := os.ReadFile(filepath.Join(dir, MetadataFileName))
 	if err != nil {
@@ -94,7 +97,7 @@ func LoadMultiTree(dir string) (*MultiTree, error) {
 		if err != nil {
 			return nil, err
 		}
-		treeMap[name] = NewFromSnapshot(snapshot)
+		treeMap[name] = NewFromSnapshot(snapshot, zeroCopy)
 	}
 
 	sort.Strings(treeNames)
@@ -112,6 +115,7 @@ func LoadMultiTree(dir string) (*MultiTree, error) {
 		treesByName:    treesByName,
 		lastCommitInfo: *metadata.CommitInfo,
 		metadata:       metadata,
+		zeroCopy:       zeroCopy,
 	}
 	// initial version is nesserary for wal index conversion
 	mtree.setInitialVersion(metadata.InitialVersion)
@@ -152,6 +156,13 @@ func (t *MultiTree) setInitialVersion(initialVersion int64) {
 	}
 }
 
+func (t *MultiTree) SetZeroCopy(zeroCopy bool) {
+	t.zeroCopy = zeroCopy
+	for _, entry := range t.trees {
+		entry.tree.SetZeroCopy(zeroCopy)
+	}
+}
+
 // Copy returns a snapshot of the tree which won't be corrupted by further modifications on the main tree.
 func (t *MultiTree) Copy() *MultiTree {
 	trees := make([]namedTree, len(t.trees))

diff --git a/memiavl/proof_test.go b/memiavl/proof_test.go
@@ -43,7 +43,7 @@ func TestProofs(t *testing.T) {
 			require.NoError(t, tree.WriteSnapshot(tmpDir, false))
 			snapshot, err := OpenSnapshot(tmpDir)
 			require.NoError(t, err)
-			ptree := NewFromSnapshot(snapshot)
+			ptree := NewFromSnapshot(snapshot, true)
 			defer ptree.Close()
 
 			proof, err = ptree.GetMembershipProof(tc.existKey)

diff --git a/memiavl/snapshot.go b/memiavl/snapshot.go
@@ -12,7 +12,6 @@ import (
 	"path/filepath"
 
 	"github.com/ledgerwatch/erigon-lib/etl"
-	"github.com/ledgerwatch/erigon-lib/mmap"
 	"github.com/ledgerwatch/erigon-lib/recsplit"
 )
 
@@ -539,14 +538,6 @@ func buildIndex(input *os.File, idxPath, tmpDir string, count int) error {
 	return nil
 }
 
-func Mmap(f *os.File) ([]byte, *[mmap.MaxMapSize]byte, error) {
-	fi, err := f.Stat()
-	if err != nil {
-		return nil, nil, err
-	}
-	return mmap.Mmap(f, int(fi.Size()))
-}
-
 func createFile(name string) (*os.File, error) {
 	return os.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
 }
diff --git a/memiavl/snapshot_test.go b/memiavl/snapshot_test.go
@@ -24,7 +24,7 @@ func TestSnapshotEncodingRoundTrip(t *testing.T) {
 	snapshot, err := OpenSnapshot(snapshotDir)
 	require.NoError(t, err)
 
-	tree2 := NewFromSnapshot(snapshot)
+	tree2 := NewFromSnapshot(snapshot, true)
 
 	require.Equal(t, tree.Version(), tree2.Version())
 	require.Equal(t, tree.RootHash(), tree2.RootHash())
@@ -40,7 +40,7 @@ func TestSnapshotEncodingRoundTrip(t *testing.T) {
 	// test modify tree loaded from snapshot
 	snapshot, err = OpenSnapshot(snapshotDir)
 	require.NoError(t, err)
-	tree3 := NewFromSnapshot(snapshot)
+	tree3 := NewFromSnapshot(snapshot, true)
 	hash, v, err := tree3.ApplyChangeSet(ChangeSets[len(ChangeSets)-1], true)
 	require.NoError(t, err)
 	require.Equal(t, RefHashes[len(ChangeSets)-1], hash)

diff --git a/memiavl/tree.go b/memiavl/tree.go
@@ -1,6 +1,7 @@
 package memiavl
 
 import (
+	"bytes"
 	"crypto/sha256"
 	"errors"
 	"fmt"
@@ -20,6 +21,9 @@ type Tree struct {
 	snapshot *Snapshot
 
 	initialVersion, cowVersion uint32
+
+	// when true, the get and iterator methods could return a slice pointing to mmaped blob files.
+	zeroCopy bool
 }
 
 // NewEmptyTree creates an empty tree at an arbitrary version.
@@ -28,7 +32,11 @@ func NewEmptyTree(version uint64) *Tree {
 		panic("version overflows uint32")
 	}
 
-	return &Tree{version: uint32(version)}
+	return &Tree{
+		version: uint32(version),
+		// no need to copy if the tree is not backed by snapshot
+		zeroCopy: true,
+	}
 }
 
 // New creates an empty tree at genesis version
@@ -45,10 +53,11 @@ func NewWithInitialVersion(initialVersion uint32) *Tree {
 }
 
 // NewFromSnapshot mmap the blob files and create the root node.
-func NewFromSnapshot(snapshot *Snapshot) *Tree {
+func NewFromSnapshot(snapshot *Snapshot, zeroCopy bool) *Tree {
 	tree := &Tree{
 		version:  snapshot.Version(),
 		snapshot: snapshot,
+		zeroCopy: zeroCopy,
 	}
 
 	if !snapshot.IsEmpty() {
@@ -58,6 +67,10 @@ func NewFromSnapshot(snapshot *Snapshot) *Tree {
 	return tree
 }
 
+func (t *Tree) SetZeroCopy(zeroCopy bool) {
+	t.zeroCopy = zeroCopy
+}
+
 func (t *Tree) IsEmpty() bool {
 	return t.root == nil
 }
@@ -141,6 +154,9 @@ func (t *Tree) GetWithIndex(key []byte) (int64, []byte) {
 	}
 
 	value, index := t.root.Get(key)
+	if !t.zeroCopy {
+		value = bytes.Clone(value)
+	}
 	return int64(index), value
 }
 
@@ -152,15 +168,16 @@ func (t *Tree) GetByIndex(index int64) ([]byte, []byte) {
 		return nil, nil
 	}
 
-	return t.root.GetByIndex(uint32(index))
+	key, value := t.root.GetByIndex(uint32(index))
+	if !t.zeroCopy {
+		key = bytes.Clone(key)
+		value = bytes.Clone(value)
+	}
+	return key, value
 }
 
 func (t *Tree) Get(key []byte) []byte {
-	if t.root == nil {
-		return nil
-	}
-
-	value, _ := t.root.Get(key)
+	_, value := t.GetWithIndex(key)
 	return value
 }
 
@@ -169,7 +186,7 @@ func (t *Tree) Has(key []byte) bool {
 }
 
 func (t *Tree) Iterator(start, end []byte, ascending bool) dbm.Iterator {
-	return NewIterator(start, end, ascending, t.root)
+	return NewIterator(start, end, ascending, t.root, t.zeroCopy)
 }
 
 func (t *Tree) Close() error {