From 4b24853187aba4878f90a36b56eee0e792ce4491 Mon Sep 17 00:00:00 2001
From: codchen <codchen03@gmail.com>
Date: Wed, 1 Mar 2023 19:16:38 +0800
Subject: [PATCH] fix

---
 immutable_tree.go |  4 ++--
 iterator.go       | 18 +++++++++++-------
 mutable_tree.go   |  3 ++-
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/immutable_tree.go b/immutable_tree.go
index 63e2f8f..ef16ea8 100644
--- a/immutable_tree.go
+++ b/immutable_tree.go
@@ -19,7 +19,7 @@ type ImmutableTree struct {
 	ndb                    *nodeDB
 	version                int64
 	skipFastStorageUpgrade bool
-	mtx                    sync.RWMutex
+	mtx                    *sync.RWMutex
 }
 
 // NewImmutableTree creates both in-memory and persistent instances
@@ -259,7 +259,7 @@ func (t *ImmutableTree) Iterator(start, end []byte, ascending bool) (dbm.Iterato
 		}
 	}
 	t.mtx.Lock()
-	return NewIterator(start, end, ascending, t, &t.mtx), nil
+	return NewIterator(start, end, ascending, t, t.mtx), nil
 }
 
 // IterateRange makes a callback for all nodes with key between start and end non-inclusive.
diff --git a/iterator.go b/iterator.go
index f8b0e74..b23c7c5 100644
--- a/iterator.go
+++ b/iterator.go
@@ -181,8 +181,9 @@ type Iterator struct {
 
 	t *traversal
 
-	mtx    *sync.RWMutex
-	locked bool
+	mtx       *sync.RWMutex
+	locked    bool
+	unlockMtx *sync.Mutex
 }
 
 var _ dbm.Iterator = (*Iterator)(nil)
@@ -190,10 +191,11 @@ var _ dbm.Iterator = (*Iterator)(nil)
 // Returns a new iterator over the immutable tree. If the tree is nil, the iterator will be invalid.
 func NewIterator(start, end []byte, ascending bool, tree *ImmutableTree, mtx *sync.RWMutex) dbm.Iterator {
 	iter := &Iterator{
-		start:  start,
-		end:    end,
-		mtx:    mtx,
-		locked: true,
+		start:     start,
+		end:       end,
+		mtx:       mtx,
+		locked:    true,
+		unlockMtx: &sync.Mutex{},
 	}
 
 	if tree == nil {
@@ -270,12 +272,14 @@ func (iter *Iterator) IsFast() bool {
 }
 
 func (iter *Iterator) unlock() {
+	iter.unlockMtx.Lock()
+	defer iter.unlockMtx.Unlock()
 	if iter.mtx == nil {
 		return
 	}
 	if !iter.locked {
 		return
 	}
-	defer iter.mtx.Unlock()
+	iter.mtx.Unlock()
 	iter.locked = false
 }
diff --git a/mutable_tree.go b/mutable_tree.go
index 9bffdb8..8fe214b 100644
--- a/mutable_tree.go
+++ b/mutable_tree.go
@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"fmt"
 	"sort"
+	"sync"
 
 	"github.com/pkg/errors"
 	dbm "github.com/tendermint/tm-db"
@@ -46,7 +47,7 @@ func NewMutableTree(db dbm.DB, cacheSize int, skipFastStorageUpgrade bool) (*Mut
 // NewMutableTreeWithOpts returns a new tree with the specified options.
 func NewMutableTreeWithOpts(db dbm.DB, cacheSize int, opts *Options, skipFastStorageUpgrade bool) (*MutableTree, error) {
 	ndb := newNodeDB(db, cacheSize, opts)
-	head := &ImmutableTree{ndb: ndb, skipFastStorageUpgrade: skipFastStorageUpgrade}
+	head := &ImmutableTree{ndb: ndb, skipFastStorageUpgrade: skipFastStorageUpgrade, mtx: &sync.RWMutex{}}
 
 	return &MutableTree{
 		ImmutableTree:            head,