feat: add the tree Close api (backport #905) (#912)

Co-authored-by: cool-developer <[email protected]>

mutable_tree.go

@@ -1105,3 +1105,13 @@ func (tree *MutableTree) SaveChangeSet(cs *ChangeSet) (int64, error) {
 	_, version, err := tree.SaveVersion()
 	return version, err
 }
+
+// Close closes the tree.
+func (tree *MutableTree) Close() error {
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
+
+	tree.ImmutableTree = nil
+	tree.lastSaved = nil
+	return tree.ndb.Close()
+}

mutable_tree_test.go

@@ -1445,3 +1445,16 @@ func TestMutableTree_InitialVersion_FirstVersion(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, initialVersion+1, node.nodeKey.version)
 }
+
+func TestMutableTreeClose(t *testing.T) {
+	db := dbm.NewMemDB()
+	tree := NewMutableTree(db, 0, true, log.NewNopLogger())
+
+	_, err := tree.Set([]byte("hello"), []byte("world"))
+	require.NoError(t, err)
+
+	_, _, err = tree.SaveVersion()
+	require.NoError(t, err)
+
+	require.NoError(t, tree.Close())
+}

nodedb.go

@@ -1002,6 +1002,22 @@ func (ndb *nodeDB) traverseOrphans(prevVersion, curVersion int64, fn func(*Node)
 	return nil
 }
 
+// Close the nodeDB.
+func (ndb *nodeDB) Close() error {
+	ndb.mtx.Lock()
+	defer ndb.mtx.Unlock()
+
+	if ndb.batch != nil {
+		if err := ndb.batch.Close(); err != nil {
+			return err
+		}
+		ndb.batch = nil
+	}
+
+	// skip the db.Close() since it can be used by other trees
+	return nil
+}
+
 // Utility and test functions
 
 func (ndb *nodeDB) leafNodes() ([]*Node, error) {