From 1771c3d14eba6156aa60198fc05eb8775baaa87c Mon Sep 17 00:00:00 2001
From: Gary Rong <garyrong0905@gmail.com>
Date: Mon, 2 Dec 2024 15:28:13 +0800
Subject: [PATCH] triedb/pathdb: parallelize remove layer

---
 triedb/pathdb/lookup.go | 70 ++++++++++++++++++++++++-----------------
 1 file changed, 41 insertions(+), 29 deletions(-)

diff --git a/triedb/pathdb/lookup.go b/triedb/pathdb/lookup.go
index df223dd8894a..990637aa1569 100644
--- a/triedb/pathdb/lookup.go
+++ b/triedb/pathdb/lookup.go
@@ -144,41 +144,53 @@ func (l *lookup) removeLayer(diff *diffLayer) error {
 		lookupRemoveLayerTimer.UpdateSince(now)
 	}(time.Now())
 
-	// TODO(rjl493456442) theoretically the code below could be parallelized,
-	// but it will slow down the other parts of system (e.g., EVM execution)
-	// with unknown reasons.
-	state := diff.rootHash()
+	var (
+		state   = diff.rootHash()
+		lock    sync.RWMutex
+		workers errgroup.Group
+	)
+	workers.SetLimit(runtime.NumCPU() / 2)
+
 	for accountHash, nodes := range diff.nodes.nodes {
-		subset := l.nodes[accountHash]
-		if subset == nil {
-			return fmt.Errorf("unknown node owner %x", accountHash)
-		}
-		for path := range nodes {
+		workers.Go(func() error {
+			lock.RLock()
+			subset := l.nodes[accountHash]
+			if subset == nil {
+				lock.RUnlock()
+				return fmt.Errorf("unknown node owner %x", accountHash)
+			}
+			lock.RUnlock()
+
 			// Traverse the list from oldest to newest to quickly locate the ID
 			// of the stale layer.
-			var found bool
-			for j := 0; j < len(subset[path]); j++ {
-				if subset[path][j] == state {
-					if j == 0 {
-						subset[path] = subset[path][1:] // TODO what if the underlying slice is held forever?
-					} else {
-						subset[path] = append(subset[path][:j], subset[path][j+1:]...)
+			for path := range nodes {
+				var found bool
+				for j := 0; j < len(subset[path]); j++ {
+					if subset[path][j] == state {
+						if j == 0 {
+							subset[path] = subset[path][1:] // TODO what if the underlying slice is held forever?
+						} else {
+							subset[path] = append(subset[path][:j], subset[path][j+1:]...)
+						}
+						found = true
+						break
 					}
-					found = true
-					break
+				}
+				if !found {
+					return fmt.Errorf("failed to delete lookup %x %v", accountHash, []byte(path))
+				}
+				if len(subset[path]) == 0 {
+					returnSlice(subset[path])
+					delete(subset, path)
 				}
 			}
-			if !found {
-				return fmt.Errorf("failed to delete lookup %x %v", accountHash, []byte(path))
-			}
-			if len(subset[path]) == 0 {
-				returnSlice(subset[path])
-				delete(subset, path)
+			if len(subset) == 0 {
+				lock.Lock()
+				delete(l.nodes, accountHash)
+				lock.Unlock()
 			}
-		}
-		if len(subset) == 0 {
-			delete(l.nodes, accountHash)
-		}
+			return nil
+		})
 	}
-	return nil
+	return workers.Wait()
 }