diff --git a/mutable_tree.go b/mutable_tree.go index 1649675de..c5965681e 100644 --- a/mutable_tree.go +++ b/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() +} diff --git a/mutable_tree_test.go b/mutable_tree_test.go index 7475c0882..8fb21aeac 100644 --- a/mutable_tree_test.go +++ b/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()) +} diff --git a/nodedb.go b/nodedb.go index 17f279425..3d3cd0370 100644 --- a/nodedb.go +++ b/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) {