From 7995342b3b70466bddc9b30f71f146b034b71681 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Fri, 10 May 2024 11:05:26 +0300
Subject: [PATCH] trie: add some attempt at concurrency

---
 trie/trie.go | 68 ++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 50 insertions(+), 18 deletions(-)

diff --git a/trie/trie.go b/trie/trie.go
index 287196346f..1a9c6906bf 100644
--- a/trie/trie.go
+++ b/trie/trie.go
@@ -22,12 +22,14 @@ import (
 	"errors"
 	"fmt"
 	"slices"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/trie/trienode"
 	"github.com/ethereum/go-ethereum/triedb/database"
+	"golang.org/x/sync/errgroup"
 )
 
 // Trie is a Merkle Patricia Trie. Use New to create a trie that sits on
@@ -303,13 +305,9 @@ func (t *Trie) getBatch(origNode node, keys [][]byte, pos int) ([][]byte, node,
 		}
 
 	case *fullNode:
-		var (
-			clone  *fullNode
-			values = make([][]byte, 0, len(keys))
-			fail   error
-			first  int
-		)
-		for first < len(keys) {
+		// Split the keys into subsets groupped by fullnode child nibble
+		offsets := make([][2]int, 0, len(keys))
+		for first := 0; first < len(keys); {
 			last := len(keys)
 			for i, key := range keys[first+1:] {
 				if keys[first][pos] != key[pos] {
@@ -317,19 +315,53 @@ func (t *Trie) getBatch(origNode node, keys [][]byte, pos int) ([][]byte, node,
 					break
 				}
 			}
-			results, newnode, didResolve, err := t.getBatch(n.Children[keys[first][pos]], keys[first:last], pos+1)
-			if err != nil && fail == nil {
-				fail = err
-			}
-			if err == nil && didResolve {
-				if clone == nil {
-					clone = n.copy()
-				}
-				clone.Children[keys[first][pos]] = newnode
-			}
-			values = append(values, results...)
+			offsets = append(offsets, [2]int{first, last})
 			first = last
 		}
+		// Resolve the subkeys sequentially or parallel based on how many there are
+		var (
+			values = make([][]byte, len(keys))
+			clone  *fullNode
+			fail   error
+		)
+		if len(offsets) < 8 {
+			for _, offset := range offsets {
+				results, newnode, didResolve, err := t.getBatch(n.Children[keys[offset[0]][pos]], keys[offset[0]:offset[1]], pos+1)
+				if err == nil && didResolve {
+					if clone == nil {
+						clone = n.copy()
+					}
+					clone.Children[keys[offset[0]][pos]] = newnode
+				}
+				copy(values[offset[0]:], results)
+				if err != nil && fail == nil {
+					fail = err
+				}
+			}
+		} else {
+			var (
+				workers errgroup.Group
+				lock    sync.Mutex
+			)
+			for _, offset := range offsets {
+				offset := offset // Closure for the worker
+				workers.Go(func() error {
+					results, newnode, didResolve, err := t.getBatch(n.Children[keys[offset[0]][pos]], keys[offset[0]:offset[1]], pos+1)
+					if err == nil && didResolve {
+						lock.Lock()
+						if clone == nil {
+							clone = n.copy()
+						}
+						lock.Unlock()
+
+						clone.Children[keys[offset[0]][pos]] = newnode
+					}
+					copy(values[offset[0]:], results)
+					return err
+				})
+				fail = workers.Wait()
+			}
+		}
 		if clone != nil {
 			n = clone
 		}