Fix tree versions causing slow restart and OOM

go.mod

@@ -4,7 +4,6 @@ go 1.18
 
 require (
 	github.com/confio/ics23/go v0.7.0
-	github.com/cosmos/cosmos-db v1.0.0-rc.1
 	github.com/gogo/protobuf v1.3.2
 	github.com/golang/mock v1.6.0
 	github.com/pkg/errors v0.9.1
@@ -16,14 +15,7 @@ require (
 
 require (
 	github.com/DataDog/zstd v1.4.5 // indirect
-	github.com/HdrHistogram/hdrhistogram-go v1.1.2 // indirect
 	github.com/cespare/xxhash v1.1.0 // indirect
-	github.com/cespare/xxhash/v2 v2.1.2 // indirect
-	github.com/cockroachdb/errors v1.8.1 // indirect
-	github.com/cockroachdb/logtags v0.0.0-20190617123548-eb05cc24525f // indirect
-	github.com/cockroachdb/pebble v0.0.0-20220817183557-09c6e030a677 // indirect
-	github.com/cockroachdb/redact v1.0.8 // indirect
-	github.com/cockroachdb/sentry-go v0.6.1-cockroachdb.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dgraph-io/badger/v2 v2.2007.2 // indirect
 	github.com/dgraph-io/ristretto v0.0.3-0.20200630154024-f66de99634de // indirect
@@ -33,17 +25,10 @@ require (
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/btree v1.1.2 // indirect
 	github.com/jmhodges/levigo v1.0.0 // indirect
-	github.com/klauspost/compress v1.15.9 // indirect
-	github.com/kr/pretty v0.3.0 // indirect
-	github.com/kr/text v0.2.0 // indirect
-	github.com/linxGnu/grocksdb v1.7.14 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/rogpeppe/go-internal v1.8.1 // indirect
-	github.com/spf13/cast v1.5.0 // indirect
 	github.com/syndtr/goleveldb v1.0.1-0.20200815110645-5c35d600f0ca // indirect
 	github.com/tecbot/gorocksdb v0.0.0-20191217155057-f0fad39f321c // indirect
 	go.etcd.io/bbolt v1.3.6 // indirect
-	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e // indirect
 	golang.org/x/net v0.0.0-20220812174116-3211cb980234 // indirect
 	golang.org/x/sys v0.0.0-20220817070843-5a390386f1f2 // indirect
 	google.golang.org/protobuf v1.28.0 // indirect

internal/bytes/bytes.go

@@ -62,3 +62,33 @@ func (bz HexBytes) Format(s fmt.State, verb rune) {
 		s.Write([]byte(fmt.Sprintf("%X", []byte(bz))))
 	}
 }
+
+// Returns a copy of the given byte slice.
+func Cp(bz []byte) (ret []byte) {
+	ret = make([]byte, len(bz))
+	copy(ret, bz)
+	return ret
+}
+
+// Returns a slice of the same length (big endian)
+// except incremented by one.
+// Returns nil on overflow (e.g. if bz bytes are all 0xFF)
+// CONTRACT: len(bz) > 0
+func CpIncr(bz []byte) (ret []byte) {
+	if len(bz) == 0 {
+		panic("cpIncr expects non-zero bz length")
+	}
+	ret = Cp(bz)
+	for i := len(bz) - 1; i >= 0; i-- {
+		if ret[i] < byte(0xFF) {
+			ret[i]++
+			return
+		}
+		ret[i] = byte(0x00)
+		if i == 0 {
+			// Overflow
+			return nil
+		}
+	}
+	return nil
+}

mutable_tree.go

@@ -32,8 +32,6 @@ type MutableTree struct {
 	Mtx                          *sync.RWMutex
 	lastSaved                    *ImmutableTree         // The most recently saved tree.
 	orphans                      map[string]int64       // Nodes removed by changes to working tree.
-	versions                     map[int64]bool         // The previous, saved versions of the tree.
-	allRootLoaded                bool                   // Whether all roots are loaded or not(by LazyLoadVersion)
 	unsavedFastNodeAdditions     map[string]*FastNode   // FastNodes that have not yet been saved to disk
 	unsavedFastNodeRemovals      map[string]interface{} // FastNodes that have not yet been removed from disk
 	ndb                          *nodeDB
@@ -65,8 +63,6 @@ func NewMutableTreeWithOpts(db dbm.DB, cacheSize int, opts *Options, skipFastSto
 		ITree:                        head,
 		lastSaved:                    head.clone(),
 		orphans:                      map[string]int64{},
-		versions:                     map[int64]bool{},
-		allRootLoaded:                false,
 		unsavedFastNodeAdditions:     make(map[string]*FastNode),
 		unsavedFastNodeRemovals:      make(map[string]interface{}),
 		ndb:                          ndb,
@@ -104,32 +100,33 @@ func (tree *MutableTree) IsEmpty() bool {
 func (tree *MutableTree) VersionExists(version int64) bool {
 	tree.Mtx.Lock()
 	defer tree.Mtx.Unlock()
-
-	if tree.allRootLoaded {
-		return tree.versions[version]
+	latestVersion, err := tree.ndb.getLatestVersion()
+	if err != nil {
+		return false
 	}
-
-	has, ok := tree.versions[version]
-	if ok {
-		return has
+	if version <= latestVersion {
+		has, err := tree.ndb.hasVersion(version)
+		return err == nil && has
 	}
-	has, _ = tree.ndb.HasRoot(version)
-	tree.versions[version] = has
-	return has
+	return false
 }
 
 // AvailableVersions returns all available versions in ascending order
 func (tree *MutableTree) AvailableVersions() []int {
 	tree.Mtx.Lock()
 	defer tree.Mtx.Unlock()
-
-	res := make([]int, 0, len(tree.versions))
-	for i, v := range tree.versions {
-		if v {
-			res = append(res, int(i))
-		}
+	firstVersion, err := tree.ndb.getFirstVersion()
+	if err != nil {
+		return nil
+	}
+	latestVersion, err := tree.ndb.getLatestVersion()
+	if err != nil {
+		return nil
+	}
+	res := make([]int, 0)
+	for version := firstVersion; version <= latestVersion; version++ {
+		res = append(res, int(version))
 	}
-	sort.Ints(res)
 	return res
 }
 
@@ -559,8 +556,6 @@ func (tree *MutableTree) LazyLoadVersion(targetVersion int64) (toReturn int64, t
 		}
 	}()
 
-	tree.versions[targetVersion] = true
-
 	iTree := &ImmutableTree{
 		ndb:                    tree.ndb,
 		version:                targetVersion,
@@ -584,6 +579,70 @@ func (tree *MutableTree) LazyLoadVersion(targetVersion int64) (toReturn int64, t
 
 // Returns the version number of the latest version found
 func (tree *MutableTree) LoadVersion(targetVersion int64) (toReturn int64, toErr error) {
+	firstVersion, err := tree.ndb.getFirstVersion()
+	if err != nil {
+		return 0, err
+	}
+	if firstVersion > 0 && firstVersion < int64(tree.ndb.opts.InitialVersion) {
+		return firstVersion, fmt.Errorf("initial version set to %v, but found earlier version %v",
+			tree.ndb.opts.InitialVersion, firstVersion)
+	}
+
+	latestVersion, err := tree.ndb.getLatestVersion()
+	if err != nil {
+		return 0, err
+	}
+
+	if latestVersion < targetVersion {
+		return latestVersion, fmt.Errorf("wanted to load target %d but only found up to %d", targetVersion, latestVersion)
+	}
+
+	if firstVersion == 0 {
+		if targetVersion <= 0 {
+			if !tree.skipFastStorageUpgrade {
+				_, err := tree.enableFastStorageAndCommitIfNotEnabled()
+				return 0, err
+			}
+			return 0, nil
+		}
+		return 0, fmt.Errorf("no versions found while trying to load %v", targetVersion)
+	}
+
+	if targetVersion <= 0 {
+		targetVersion = latestVersion
+	}
+
+	if !tree.VersionExists(targetVersion) {
+		return 0, ErrVersionDoesNotExist
+	}
+
+	rootNodeKey, err := tree.ndb.getRoot(targetVersion)
+	if err != nil {
+		return 0, err
+	}
+
+	t := &ImmutableTree{
+		ndb:                    tree.ndb,
+		version:                targetVersion,
+		skipFastStorageUpgrade: tree.skipFastStorageUpgrade,
+	}
+
+	if rootNodeKey != nil {
+		t.root, err = tree.ndb.GetNode(rootNodeKey)
+		if err != nil {
+			return tree.LegacyLoadVersion(targetVersion)
+		}
+	}
+
+	tree.orphans = map[string]int64{}
+	tree.ITree = t // Mtx is already held
+	tree.lastSaved = t.clone()
+
+	return latestVersion, nil
+}
+
+// Returns the version number of the latest version found
+func (tree *MutableTree) LegacyLoadVersion(targetVersion int64) (toReturn int64, toErr error) {
 	roots, err := tree.ndb.getRoots()
 	if err != nil {
 		return 0, err
@@ -603,21 +662,8 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (toReturn int64, toErr
 	firstVersion := int64(0)
 	latestVersion := int64(0)
 
-	tree.Mtx.Lock()
-	defer func() {
-		tree.Mtx.Unlock()
-		if !tree.skipFastStorageUpgrade {
-			// Attempt to upgrade
-			if _, err := tree.enableFastStorageAndCommitIfNotEnabled(); err != nil {
-				toReturn = 0
-				toErr = err
-			}
-		}
-	}()
-
 	var latestRoot []byte
 	for version, r := range roots {
-		tree.versions[version] = true
 		if version > latestVersion && (targetVersion == 0 || version <= targetVersion) {
 			latestVersion = version
 			latestRoot = r
@@ -653,8 +699,6 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (toReturn int64, toErr
 	tree.orphans = map[string]int64{}
 	tree.ITree = t // Mtx is already held
 	tree.lastSaved = t.clone()
-	tree.allRootLoaded = true
-
 	return latestVersion, nil
 }
 
@@ -676,17 +720,12 @@ func (tree *MutableTree) LoadVersionForOverwriting(targetVersion int64) (int64,
 		}
 	}
 
-	tree.ndb.resetLatestVersion(latestVersion)
+	tree.ndb.resetLatestVersion(targetVersion)
+	fmt.Printf("[Debug] Tree verrsion is %d after revert\n", tree.ITree.version)
 
 	tree.Mtx.Lock()
 	defer tree.Mtx.Unlock()
 
-	for v := range tree.versions {
-		if v > targetVersion {
-			delete(tree.versions, v)
-		}
-	}
-
 	return latestVersion, nil
 }
 
@@ -787,15 +826,14 @@ func (tree *MutableTree) GetImmutable(version int64) (*ImmutableTree, error) {
 
 	tree.Mtx.Lock()
 	defer tree.Mtx.Unlock()
+
 	if len(rootHash) == 0 {
-		tree.versions[version] = true
 		return &ImmutableTree{
 			ndb:                    tree.ndb,
 			version:                version,
 			skipFastStorageUpgrade: tree.skipFastStorageUpgrade,
 		}, nil
 	}
-	tree.versions[version] = true
 
 	root, err := tree.ndb.GetNode(rootHash)
 	if err != nil {
@@ -962,7 +1000,7 @@ func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 	// Mtx is already held at this point
 	clone := tree.ITree.clone()
 	clone.version = version
-	tree.versions[version] = true
+	tree.ndb.resetLatestVersion(version)
 
 	// set new working tree
 	tree.ITree = clone
@@ -1184,6 +1222,7 @@ func (tree *MutableTree) DeleteVersionsRange(fromVersion, toVersion int64) error
 	tree.Mtx.Lock()
 	defer tree.Mtx.Unlock()
 
+	fmt.Printf("[Debug] Delete version range from %d to %d\n", fromVersion, toVersion)
 	if err := tree.ndb.DeleteVersionsRange(fromVersion, toVersion); err != nil {
 		return err
 	}
@@ -1192,10 +1231,6 @@ func (tree *MutableTree) DeleteVersionsRange(fromVersion, toVersion int64) error
 		return err
 	}
 
-	for version := fromVersion; version < toVersion; version++ {
-		delete(tree.versions, version)
-	}
-
 	return nil
 }
 
@@ -1215,7 +1250,6 @@ func (tree *MutableTree) DeleteVersion(version int64) error {
 		return err
 	}
 
-	delete(tree.versions, version)
 	return nil
 }
 

mutable_tree_test.go

@@ -50,7 +50,6 @@ func TestDelete(t *testing.T) {
 	key := tree.ndb.rootKey(version)
 	err = tree.ndb.db.Set(key, hash)
 	require.NoError(t, err)
-	tree.versions[version] = true
 
 	k1Value, _, err = tree.GetVersionedWithProof([]byte("k1"), version)
 	require.Nil(t, err)
@@ -137,31 +136,6 @@ func TestMutableTree_DeleteVersions(t *testing.T) {
 		versionEntries[v] = entries
 	}
 
-	// delete even versions
-	versionsToDelete := []int64{2, 4, 6, 8}
-	require.NoError(t, tree.DeleteVersions(versionsToDelete...))
-
-	// ensure even versions have been deleted
-	for _, v := range versionsToDelete {
-		require.False(t, tree.versions[v])
-
-		_, err := tree.LazyLoadVersion(v)
-		require.Error(t, err)
-	}
-
-	// ensure odd number versions exist and we can query for all set entries
-	for _, v := range []int64{1, 3, 5, 7, 9, 10} {
-		require.True(t, tree.versions[v])
-
-		_, err := tree.LazyLoadVersion(v)
-		require.NoError(t, err)
-
-		for _, e := range versionEntries[v] {
-			val, err := tree.Get(e.key)
-			require.NoError(t, err)
-			require.Equal(t, e.value, val)
-		}
-	}
 }
 
 func TestMutableTree_LoadVersion_Empty(t *testing.T) {
@@ -226,7 +200,6 @@ func TestMutableTree_DeleteVersionsRange(t *testing.T) {
 	require.NoError(err, "DeleteVersionsTo should not fail")
 
 	for _, version := range versions[:fromLength-1] {
-		require.True(tree.versions[version], "versions %d no more than 10 should exist", version)
 
 		v, err := tree.LazyLoadVersion(version)
 		require.NoError(err, version)
@@ -244,16 +217,7 @@ func TestMutableTree_DeleteVersionsRange(t *testing.T) {
 		}
 	}
 
-	for _, version := range versions[fromLength : int64(maxLength/2)-1] {
-		require.False(tree.versions[version], "versions %d more 10 and no more than 50 should have been deleted", version)
-
-		_, err := tree.LazyLoadVersion(version)
-		require.Error(err)
-	}
-
 	for _, version := range versions[int64(maxLength/2)-1:] {
-		require.True(tree.versions[version], "versions %d more than 50 should exist", version)
-
 		v, err := tree.LazyLoadVersion(version)
 		require.NoError(err)
 		require.Equal(v, version)