tree: performance improvements for asset tree cache based on review

common/prove/witness_helper.go

@@ -203,11 +203,11 @@ func (w *WitnessHelper) constructAccountWitness(
 		}
 		// it means this is a registerZNS tx
 		if proverAccounts == nil {
-			if accountKey != w.assetTrees.Size() {
+			if accountKey != w.assetTrees.GetNextAccountIndex() {
 				return accountRootBefore, accountsInfoBefore, merkleProofsAccountAssetsBefore, merkleProofsAccountBefore,
 					fmt.Errorf("invalid key")
 			}
-			w.assetTrees.Add(tree.NewEmptyAccountAssetTreeFunc(w.treeCtx, accountKey, finalityBlockNr))
+			w.assetTrees.UpdateCache(accountKey, tree.NewEmptyAccountAssetTreeFunc(w.treeCtx, finalityBlockNr))
 			cryptoAccount = cryptoTypes.EmptyAccount(accountKey, tree.NilAccountAssetRoot)
 			// update account info
 			accountInfo, err := w.accountModel.GetConfirmedAccountByIndex(accountKey)

core/executor/register_zns_executor.go

@@ -96,7 +96,7 @@ func (e *RegisterZnsExecutor) ApplyTransaction() error {
 		return err
 	}
 
-	bc.StateDB().AccountAssetTrees.Add(tree.NewEmptyAccountAssetTreeFunc(bc.StateDB().TreeCtx, txInfo.AccountIndex, uint64(bc.CurrentBlock().BlockHeight)))
+	bc.StateDB().AccountAssetTrees.UpdateCache(txInfo.AccountIndex, tree.NewEmptyAccountAssetTreeFunc(bc.StateDB().TreeCtx, uint64(bc.CurrentBlock().BlockHeight)))
 
 	stateCache := e.bc.StateDB()
 	stateCache.SetPendingNewAccount(txInfo.AccountIndex, formatAccount)

core/statedb/statedb.go

@@ -717,7 +717,7 @@ func (s *StateDB) GetPendingNonce(accountIndex int64) (int64, error) {
 }
 
 func (s *StateDB) GetNextAccountIndex() int64 {
-	return s.AccountAssetTrees.Size()
+	return s.AccountAssetTrees.GetNextAccountIndex()
 }
 
 func (s *StateDB) GetNextNftIndex() int64 {

tree/account_tree.go

@@ -52,20 +52,16 @@ func InitAccountTree(
 	opts := ctx.Options(blockHeight)
 
 	// init account state trees
-	accountAssetTrees = NewLazyTreeCache(types.AssetTreeCacheSize)
-	for index := int64(0); index < accountNums; index++ {
-		// create account assets tree
-		accountAssetTrees.AddToIndex(index, func() bsmt.SparseMerkleTree {
-			tree, err := bsmt.NewBASSparseMerkleTree(bsmt.NewHasher(mimc.NewMiMC()),
-				SetNamespace(ctx, accountAssetNamespace(index)), AssetTreeHeight, NilAccountAssetNodeHash,
-				opts...)
-			if err != nil {
-				logx.Errorf("unable to create new tree by assets: %s", err.Error())
-				panic(err.Error())
-			}
-			return tree
-		})
-	}
+	accountAssetTrees = NewLazyTreeCache(types.AssetTreeCacheSize, accountNums-1, func(index int64) bsmt.SparseMerkleTree {
+		tree, err := bsmt.NewBASSparseMerkleTree(bsmt.NewHasher(mimc.NewMiMC()),
+			SetNamespace(ctx, accountAssetNamespace(index)), AssetTreeHeight, NilAccountAssetNodeHash,
+			opts...)
+		if err != nil {
+			logx.Errorf("unable to create new tree by assets: %s", err.Error())
+			panic(err.Error())
+		}
+		return tree
+	})
 	accountTree, err = bsmt.NewBASSparseMerkleTree(bsmt.NewHasher(mimc.NewMiMC()),
 		SetNamespace(ctx, AccountPrefix), AccountTreeHeight, NilAccountNodeHash,
 		opts...)
@@ -89,7 +85,7 @@ func InitAccountTree(
 			}
 		}
 
-		for i := int64(0); i < accountAssetTrees.Size(); i++ {
+		for i := int64(0); i < accountNums; i++ {
 			_, err := accountAssetTrees.Get(i).Commit(nil)
 			if err != nil {
 				logx.Errorf("unable to set asset to tree: %s", err.Error())
@@ -115,11 +111,11 @@ func InitAccountTree(
 		}
 	}
 
-	for i := int64(0); i < accountAssetTrees.Size(); i++ {
-		account := accountAssetTrees.Get(i)
-		if account.LatestVersion() > bsmt.Version(blockHeight) && !account.IsEmpty() {
-			logx.Infof("asset tree %d version [%d] is higher than block, rollback to %d", i, account.LatestVersion(), blockHeight)
-			err := account.Rollback(bsmt.Version(blockHeight))
+	for i := int64(0); i < accountNums; i++ {
+		asset := accountAssetTrees.Get(i)
+		if asset.LatestVersion() > bsmt.Version(blockHeight) && !asset.IsEmpty() {
+			logx.Infof("asset tree %d version [%d] is higher than block, rollback to %d", i, asset.LatestVersion(), blockHeight)
+			err := asset.Rollback(bsmt.Version(blockHeight))
 			if err != nil {
 				logx.Errorf("unable to rollback asset [%d] tree: %s, version: %d", i, err.Error(), blockHeight)
 				return nil, nil, err
@@ -261,15 +257,18 @@ func AccountToNode(
 
 func NewEmptyAccountAssetTreeFunc(
 	ctx *Context,
-	index int64,
 	blockHeight uint64,
-) func() bsmt.SparseMerkleTree {
-	return func() bsmt.SparseMerkleTree {
-		tree, _ := bsmt.NewBASSparseMerkleTree(
+) func(index int64) bsmt.SparseMerkleTree {
+	return func(index int64) bsmt.SparseMerkleTree {
+		tree, err := bsmt.NewBASSparseMerkleTree(
 			bsmt.NewHasher(mimc.NewMiMC()),
 			SetNamespace(ctx, accountAssetNamespace(index)),
 			AssetTreeHeight, NilAccountAssetNodeHash,
 			ctx.Options(int64(blockHeight))...)
+		if err != nil {
+			logx.Errorf("unable to create new tree by assets: %s", err.Error())
+			panic(err.Error())
+		}
 		return tree
 	}
 }

tree/asset_tree_cache.go

@@ -7,61 +7,80 @@ import (
 	lru "github.com/hashicorp/golang-lru"
 )
 
+// Lazy init cache for asset trees
 type AssetTreeCache struct {
-	funcMap    map[int64]func() bsmt.SparseMerkleTree
-	funcLock   sync.RWMutex
-	commitMap  map[int64]bool
-	CommitLock sync.RWMutex
-	treeCache  *lru.Cache
+	initFunction      func(index int64) bsmt.SparseMerkleTree
+	nextAccountNumber int64
+	mainLock          sync.RWMutex
+	changes           map[int64]bool
+	changesLock       sync.RWMutex
+	treeCache         *lru.Cache
 }
 
-func NewLazyTreeCache(maxSize int) *AssetTreeCache {
-	cache := AssetTreeCache{funcMap: make(map[int64]func() bsmt.SparseMerkleTree), commitMap: make(map[int64]bool)}
-	cache.treeCache, _ = lru.NewWithEvict(maxSize, cache.OnDelete)
+// Creates new AssetTreeCache
+// maxSize defines the maximum size of currently initialized trees
+// accountNumber defines the number of accounts to create/or next index for new account
+func NewLazyTreeCache(maxSize int, accountNumber int64, f func(index int64) bsmt.SparseMerkleTree) *AssetTreeCache {
+	cache := AssetTreeCache{initFunction: f, nextAccountNumber: accountNumber, changes: make(map[int64]bool, maxSize*10)}
+	cache.treeCache, _ = lru.NewWithEvict(maxSize, cache.onDelete)
 	return &cache
 }
 
-func (c *AssetTreeCache) AddToIndex(i int64, f func() bsmt.SparseMerkleTree) {
-	c.funcLock.Lock()
-	c.funcMap[i] = f
-	c.funcLock.Unlock()
+// Updates current cache with new(updated) init function and with latest account index
+func (c *AssetTreeCache) UpdateCache(accountNumber int64, f func(index int64) bsmt.SparseMerkleTree) {
+	c.mainLock.Lock()
+	c.initFunction = f
+	if c.nextAccountNumber < accountNumber {
+		c.nextAccountNumber = accountNumber
+	}
+	c.mainLock.Unlock()
 }
 
-func (c *AssetTreeCache) Add(f func() bsmt.SparseMerkleTree) {
-	c.funcLock.Lock()
-	c.funcMap[int64(len(c.funcMap))] = f
-	c.funcLock.Unlock()
+// Returns index of next account
+func (c *AssetTreeCache) GetNextAccountIndex() int64 {
+	c.mainLock.RLock()
+	defer c.mainLock.RUnlock()
+	return c.nextAccountNumber + 1
 }
 
+// Returns asset tree based on account index
 func (c *AssetTreeCache) Get(i int64) (tree bsmt.SparseMerkleTree) {
-	c.funcLock.RLock()
-	c.treeCache.ContainsOrAdd(i, c.funcMap[i]())
-	c.funcLock.RUnlock()
+	c.mainLock.RLock()
+	c.treeCache.ContainsOrAdd(i, c.initFunction(i))
+	c.mainLock.RUnlock()
 	if tmpTree, ok := c.treeCache.Get(i); ok {
 		tree = tmpTree.(bsmt.SparseMerkleTree)
 	}
 	return
 }
 
-func (c *AssetTreeCache) NeedsCommit(i int64) bool {
-	if c.treeCache.Contains(i) {
-		if tree, ok := c.treeCache.Peek(i); ok {
-			return (tree.(bsmt.SparseMerkleTree).LatestVersion()-tree.(bsmt.SparseMerkleTree).RecentVersion() > 1)
+// Returns slice of indexes of asset trees that were changned
+// and resets cache of the marked changes of asset trees.
+func (c *AssetTreeCache) GetChanges() []int64 {
+	c.mainLock.Lock()
+	c.changesLock.Lock()
+	defer c.mainLock.Unlock()
+	defer c.changesLock.Unlock()
+	for _, key := range c.treeCache.Keys() {
+		tree, _ := c.treeCache.Peek(key)
+		if tree.(bsmt.SparseMerkleTree).LatestVersion()-tree.(bsmt.SparseMerkleTree).RecentVersion() > 1 {
+			c.changes[key.(int64)] = true
 		}
 	}
-	c.CommitLock.RLock()
-	defer c.funcLock.RUnlock()
-	return c.commitMap[i]
-}
-
-func (c *AssetTreeCache) OnDelete(k, v interface{}) {
-	c.CommitLock.Lock()
-	c.commitMap[k.(int64)] = (v.(bsmt.SparseMerkleTree).LatestVersion()-v.(bsmt.SparseMerkleTree).RecentVersion() > 1)
-	c.CommitLock.Unlock()
+	ret := make([]int64, 0, len(c.changes))
+	for key := range c.changes {
+		ret = append(ret, key)
+	}
+	// Cleans map
+	c.changes = make(map[int64]bool, len(c.changes))
+	return ret
 }
 
-func (c *AssetTreeCache) Size() int64 {
-	c.funcLock.RLock()
-	defer c.funcLock.RUnlock()
-	return int64(len(c.funcMap))
+// Internal method to that marks if changes happend to tree eviced from LRU
+func (c *AssetTreeCache) onDelete(k, v interface{}) {
+	c.changesLock.Lock()
+	if v.(bsmt.SparseMerkleTree).LatestVersion()-v.(bsmt.SparseMerkleTree).RecentVersion() > 1 {
+		c.changes[k.(int64)] = true
+	}
+	c.changesLock.Unlock()
 }

tree/util.go

@@ -122,7 +122,8 @@ func CommitTrees(
 	liquidityTree bsmt.SparseMerkleTree,
 	nftTree bsmt.SparseMerkleTree) error {
 
-	totalTask := assetTrees.Size() + 3
+	assetTreeChanges := assetTrees.GetChanges()
+	totalTask := len(assetTreeChanges) + 3
 	errChan := make(chan error, totalTask)
 	defer close(errChan)
 
@@ -142,17 +143,15 @@ func CommitTrees(
 		return err
 	}
 
-	for idx := int64(0); idx < assetTrees.Size(); idx++ {
+	for _, idx := range assetTreeChanges {
 		err := func(i int64) error {
 			return pool.Submit(func() {
-				if assetTrees.NeedsCommit(i) {
-					asset := assetTrees.Get(i)
-					version := asset.LatestVersion()
-					ver, err := asset.Commit(&version)
-					if err != nil {
-						errChan <- errors.Wrapf(err, "unable to commit asset tree [%d], tree ver: %d, prune ver: %d", i, ver, version)
-						return
-					}
+				asset := assetTrees.Get(i)
+				version := asset.LatestVersion()
+				ver, err := asset.Commit(&version)
+				if err != nil {
+					errChan <- errors.Wrapf(err, "unable to commit asset tree [%d], tree ver: %d, prune ver: %d", i, ver, version)
+					return
 				}
 				errChan <- nil
 			})
@@ -194,7 +193,7 @@ func CommitTrees(
 		return err
 	}
 
-	for i := int64(0); i < totalTask; i++ {
+	for i := 0; i < totalTask; i++ {
 		err := <-errChan
 		if err != nil {
 			return err
@@ -212,7 +211,8 @@ func RollBackTrees(
 	liquidityTree bsmt.SparseMerkleTree,
 	nftTree bsmt.SparseMerkleTree) error {
 
-	totalTask := assetTrees.Size() + 3
+	assetTreeChanges := assetTrees.GetChanges()
+	totalTask := len(assetTreeChanges) + 3
 	errChan := make(chan error, totalTask)
 	defer close(errChan)
 
@@ -231,17 +231,15 @@ func RollBackTrees(
 		return err
 	}
 
-	for idx := int64(0); idx < assetTrees.Size(); idx++ {
+	for _, idx := range assetTreeChanges {
 		err := func(i int64) error {
 			return pool.Submit(func() {
-				if assetTrees.NeedsCommit(i) {
-					asset := assetTrees.Get(i)
-					version := asset.RecentVersion()
-					err := asset.Rollback(version)
-					if err != nil {
-						errChan <- errors.Wrapf(err, "unable to rollback asset tree [%d], ver: %d", i, version)
-						return
-					}
+				asset := assetTrees.Get(i)
+				version := asset.RecentVersion()
+				err := asset.Rollback(version)
+				if err != nil {
+					errChan <- errors.Wrapf(err, "unable to rollback asset tree [%d], ver: %d", i, version)
+					return
 				}
 				errChan <- nil
 			})
@@ -279,7 +277,7 @@ func RollBackTrees(
 		return err
 	}
 
-	for i := int64(0); i < totalTask; i++ {
+	for i := 0; i < totalTask; i++ {
 		err := <-errChan
 		if err != nil {
 			return err