fix: Better lock handling on close (backport #1024)

CHANGELOG.md

@@ -1,11 +1,20 @@
 # Changelog
 
+## v1.2.2 November 26, 2024
+
+### Bug Fixes
+
+- [#1007](https://github.com/cosmos/iavl/pull/1007) Add the extra check for the reformatted root node in `GetNode`
+
+## v1.2.1 July 31, 2024
 
 ### Improvements
 
 - [#961](https://github.com/cosmos/iavl/pull/961) Add new `GetLatestVersion` API to get the latest version.
 - [#965](https://github.com/cosmos/iavl/pull/965) Use expected interface for expected IAVL `Logger`.
 - [#970](https://github.com/cosmos/iavl/pull/970) Close the pruning process when the nodeDB is closed.
+- [#980](https://github.com/cosmos/iavl/pull/980) Use the `sdk/core/store.KVStoreWithBatch` interface instead of `iavl/db.DB` interface
+- [#1018](https://github.com/cosmos/iavl/pull/1018) Cache first version for legacy versions, fix performance regression after upgrade.
 
 ## v1.2.0 May 13, 2024
 

mutable_tree_test.go

@@ -1453,3 +1453,27 @@ func TestMutableTreeClose(t *testing.T) {
 
 	require.NoError(t, tree.Close())
 }
+
+func TestReferenceRootPruning(t *testing.T) {
+	memDB := dbm.NewMemDB()
+	tree := NewMutableTree(memDB, 0, true, NewNopLogger())
+
+	_, err := tree.Set([]byte("foo"), []byte("bar"))
+	require.NoError(t, err)
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	_, err = tree.Set([]byte("foo1"), []byte("bar"))
+	require.NoError(t, err)
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	err = tree.DeleteVersionsTo(1)
+	require.NoError(t, err)
+
+	_, err = tree.Set([]byte("foo"), []byte("bar*"))
+	require.NoError(t, err)
+}

nodedb.go

@@ -159,6 +159,20 @@ func (ndb *nodeDB) GetNode(nk []byte) (*Node, error) {
 	if err != nil {
 		return nil, fmt.Errorf("can't get node %v: %v", nk, err)
 	}
+	if buf == nil && !isLegcyNode {
+		// if the node is reformatted by pruning, check against (version, 0)
+		nKey := GetNodeKey(nk)
+		if nKey.nonce == 1 {
+			nodeKey = ndb.nodeKey((&NodeKey{
+				version: nKey.version,
+				nonce:   0,
+			}).GetKey())
+			buf, err = ndb.db.Get(nodeKey)
+			if err != nil {
+				return nil, fmt.Errorf("can't get the reformatted node %v: %v", nk, err)
+			}
+		}
+	}
 	if buf == nil {
 		return nil, fmt.Errorf("Value missing for key %v corresponding to nodeKey %x", nk, nodeKey)
 	}
@@ -588,7 +602,7 @@ func (ndb *nodeDB) startPruning() {
 	for {
 		select {
 		case <-ndb.ctx.Done():
-			ndb.done <- struct{}{}
+			close(ndb.done)
 			return
 		default:
 			ndb.mtx.Lock()
@@ -726,6 +740,7 @@ func (ndb *nodeDB) getFirstVersion() (int64, error) {
 	if itr.Valid() {
 		var version int64
 		legacyRootKeyFormat.Scan(itr.Key(), &version)
+		ndb.resetFirstVersion(version)
 		return version, nil
 	}
 	// Find the first version
@@ -1106,14 +1121,15 @@ func (ndb *nodeDB) traverseOrphans(prevVersion, curVersion int64, fn func(*Node)
 
 // Close the nodeDB.
 func (ndb *nodeDB) Close() error {
-	ndb.mtx.Lock()
-	defer ndb.mtx.Unlock()
-
 	ndb.cancel()
+
 	if ndb.opts.AsyncPruning {
 		<-ndb.done // wait for the pruning process to finish
 	}
 
+	ndb.mtx.Lock()
+	defer ndb.mtx.Unlock()
+
 	if ndb.batch != nil {
 		if err := ndb.batch.Close(); err != nil {
 			return err

nodedb_test.go

@@ -444,4 +444,5 @@ func TestCloseNodeDB(t *testing.T) {
 	opts.AsyncPruning = true
 	ndb := newNodeDB(db, 0, opts, NewNopLogger())
 	require.NoError(t, ndb.Close())
+	require.NoError(t, ndb.Close()) // must not block or fail on second call
 }