memdb: use btree for storage

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Improvements
+
+- [memdb] [\#53](https://github.com/tendermint/tm-db/pull/53) Use a B-tree for storage, which significantly improves range scan performance
+
 ## 0.4.1
 
 **2020-2-26**

backend_test.go

@@ -204,6 +204,7 @@ func TestDBIterator(t *testing.T) {
 	for dbType := range backends {
 		t.Run(fmt.Sprintf("%v", dbType), func(t *testing.T) {
 			testDBIterator(t, dbType)
+			testDBIteratorBlankKey(t, dbType)
 		})
 	}
 }
@@ -311,6 +312,18 @@ func testDBIterator(t *testing.T, backend BackendType) {
 	verifyIterator(t, ritr,
 		[]int64(nil), "reverse iterator from 7 (ex) to 6")
 
+	ritr, err = db.ReverseIterator(int642Bytes(10), nil)
+	require.NoError(t, err)
+	verifyIterator(t, ritr, []int64(nil), "reverse iterator to 10")
+
+	ritr, err = db.ReverseIterator(int642Bytes(6), nil)
+	require.NoError(t, err)
+	verifyIterator(t, ritr, []int64{9, 8, 7}, "reverse iterator to 6")
+
+	ritr, err = db.ReverseIterator(int642Bytes(5), nil)
+	require.NoError(t, err)
+	verifyIterator(t, ritr, []int64{9, 8, 7, 5}, "reverse iterator to 5")
+
 	// verifyIterator(t, db.Iterator(int642Bytes(0), int642Bytes(1)), []int64{0}, "forward iterator from 0 to 1")
 
 	ritr, err = db.ReverseIterator(int642Bytes(8), int642Bytes(9))
@@ -329,7 +342,56 @@ func testDBIterator(t *testing.T, backend BackendType) {
 	require.NoError(t, err)
 	verifyIterator(t, ritr,
 		[]int64(nil), "reverse iterator from 2 (ex) to 4")
+}
+
+func testDBIteratorBlankKey(t *testing.T, backend BackendType) {
+	name := fmt.Sprintf("test_%x", randStr(12))
+	dir := os.TempDir()
+	db := NewDB(name, backend, dir)
+	defer cleanupDBDir(dir, name)
+
+	err := db.Set([]byte(""), []byte{0})
+	require.NoError(t, err)
+	err = db.Set([]byte("a"), []byte{1})
+	require.NoError(t, err)
+	err = db.Set([]byte("b"), []byte{2})
+	require.NoError(t, err)
+
+	value, err := db.Get([]byte(""))
+	require.NoError(t, err)
+	assert.Equal(t, []byte{0}, value)
+
+	i, err := db.Iterator(nil, nil)
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{"", "a", "b"}, "forward")
 
+	i, err = db.Iterator([]byte(""), nil)
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{"", "a", "b"}, "forward from blank")
+
+	i, err = db.Iterator([]byte("a"), nil)
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{"a", "b"}, "forward from a")
+
+	i, err = db.Iterator([]byte(""), []byte("b"))
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{"", "a"}, "forward from blank to b")
+
+	i, err = db.ReverseIterator(nil, nil)
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{"b", "a", ""}, "reverse")
+
+	i, err = db.ReverseIterator([]byte(""), nil)
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{"b", "a", ""}, "reverse to blank")
+
+	i, err = db.ReverseIterator([]byte(""), []byte("a"))
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{""}, "reverse to blank from a")
+
+	i, err = db.ReverseIterator([]byte("a"), nil)
+	require.NoError(t, err)
+	verifyIteratorStrings(t, i, []string{"b", "a"}, "reverse to a")
 }
 
 func verifyIterator(t *testing.T, itr Iterator, expected []int64, msg string) {
@@ -341,3 +403,13 @@ func verifyIterator(t *testing.T, itr Iterator, expected []int64, msg string) {
 	}
 	assert.Equal(t, expected, list, msg)
 }
+
+func verifyIteratorStrings(t *testing.T, itr Iterator, expected []string, msg string) {
+	var list []string
+	for itr.Valid() {
+		key := itr.Key()
+		list = append(list, string(key))
+		itr.Next()
+	}
+	assert.Equal(t, expected, list, msg)
+}

common_test.go

@@ -201,6 +201,38 @@ func (mockIterator) Error() error {
 func (mockIterator) Close() {
 }
 
+func benchmarkRangeScans(b *testing.B, db DB, dbSize int64) {
+	b.StopTimer()
+
+	rangeSize := int64(10000)
+	if dbSize < rangeSize {
+		b.Errorf("db size %v cannot be less than range size %v", dbSize, rangeSize)
+	}
+
+	for i := int64(0); i < dbSize; i++ {
+		bytes := int642Bytes(i)
+		err := db.Set(bytes, bytes)
+		if err != nil {
+			// require.NoError() is very expensive (according to profiler), so check manually
+			b.Fatal(b, err)
+		}
+	}
+	b.StartTimer()
+
+	for i := 0; i < b.N; i++ {
+		start := rand.Int63n(dbSize - rangeSize)
+		end := start + rangeSize
+		iter, err := db.Iterator(int642Bytes(start), int642Bytes(end))
+		require.NoError(b, err)
+		count := 0
+		for ; iter.Valid(); iter.Next() {
+			count++
+		}
+		iter.Close()
+		require.EqualValues(b, rangeSize, count)
+	}
+}
+
 func benchmarkRandomReadsWrites(b *testing.B, db DB) {
 	b.StopTimer()
 
@@ -217,23 +249,29 @@ func benchmarkRandomReadsWrites(b *testing.B, db DB) {
 	for i := 0; i < b.N; i++ {
 		// Write something
 		{
-			idx := int64(rand.Int()) % numItems // nolint:gosec testing file, so accepting weak random number generator
+			idx := rand.Int63n(numItems) // nolint:gosec testing file, so accepting weak random number generator
 			internal[idx]++
 			val := internal[idx]
 			idxBytes := int642Bytes(idx)
 			valBytes := int642Bytes(val)
 			//fmt.Printf("Set %X -> %X\n", idxBytes, valBytes)
 			err := db.Set(idxBytes, valBytes)
-			b.Error(err)
+			if err != nil {
+				// require.NoError() is very expensive (according to profiler), so check manually
+				b.Fatal(b, err)
+			}
 		}
 
 		// Read something
 		{
-			idx := int64(rand.Int()) % numItems // nolint:gosec testing file, so accepting weak random number generator
+			idx := rand.Int63n(numItems) // nolint:gosec testing file, so accepting weak random number generator
 			valExp := internal[idx]
 			idxBytes := int642Bytes(idx)
 			valBytes, err := db.Get(idxBytes)
-			b.Error(err)
+			if err != nil {
+				// require.NoError() is very expensive (according to profiler), so check manually
+				b.Fatal(b, err)
+			}
 			//fmt.Printf("Get %X -> %X\n", idxBytes, valBytes)
 			if valExp == 0 {
 				if !bytes.Equal(valBytes, nil) {

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/facebookgo/stack v0.0.0-20160209184415-751773369052 // indirect
 	github.com/facebookgo/subset v0.0.0-20150612182917-8dac2c3c4870 // indirect
 	github.com/gogo/protobuf v1.3.1
+	github.com/google/btree v1.0.0
 	github.com/jmhodges/levigo v1.0.0
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.5.1

go.sum

@@ -28,6 +28,8 @@ github.com/golang/protobuf v1.3.2 h1:6nsPYzhq5kReh6QImI3k5qWzO4PEbvbIW2cwSfR/6xs
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
 github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/google/btree v1.0.0 h1:0udJVsspx3VBr5FwtLhQQtuAsVc79tTq0ocGIPAU6qo=
+github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/go-cmp v0.2.0 h1:+dTQ8DZQJz0Mb/HjFlkptS1FeQ4cWSnN941F8aEG4SQ=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI=