From a155a06dffc66d16a43d484f464ae2c411ade6e9 Mon Sep 17 00:00:00 2001
From: Gary Rong <garyrong0905@gmail.com>
Date: Thu, 19 Oct 2023 16:47:04 +0800
Subject: [PATCH] eth/protocols/snap: skip healing if the storage trie is
 already complete

---
 eth/protocols/snap/metrics.go   |  8 ++++++++
 eth/protocols/snap/sync.go      | 18 ++++++++++++++++--
 eth/protocols/snap/sync_test.go |  2 +-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/eth/protocols/snap/metrics.go b/eth/protocols/snap/metrics.go
index 8123b6c57f..83f553de19 100644
--- a/eth/protocols/snap/metrics.go
+++ b/eth/protocols/snap/metrics.go
@@ -42,4 +42,12 @@ var (
 	// boundaryAccountNodesGauge is the metric to track how many boundary trie
 	// nodes in storage tries are met.
 	boundaryStorageNodesGauge = metrics.NewRegisteredGauge("eth/protocols/snap/sync/boundary/storage", nil)
+
+	// smallStorageGauge is the metric to track how many storages are retrieved
+	// in a single request.
+	smallStorageGauge = metrics.NewRegisteredGauge("eth/protocols/snap/sync/storage/small", nil)
+
+	// largeStorageGauge is the metric to track how many storages are large enough
+	// to retrieved concurrently.
+	largeStorageGauge = metrics.NewRegisteredGauge("eth/protocols/snap/sync/storage/large", nil)
 )
diff --git a/eth/protocols/snap/sync.go b/eth/protocols/snap/sync.go
index a8dc8a5c78..e1faf8b8cd 100644
--- a/eth/protocols/snap/sync.go
+++ b/eth/protocols/snap/sync.go
@@ -2014,6 +2014,7 @@ func (s *Syncer) processStorageResponse(res *storageResponse) {
 			if res.subTask == nil && res.mainTask.needState[j] && (i < len(res.hashes)-1 || !res.cont) {
 				res.mainTask.needState[j] = false
 				res.mainTask.pend--
+				smallStorageGauge.Inc(1)
 			}
 			// If the last contract was chunked, mark it as needing healing
 			// to avoid writing it out to disk prematurely.
@@ -2049,7 +2050,11 @@ func (s *Syncer) processStorageResponse(res *storageResponse) {
 						log.Debug("Chunked large contract", "initiators", len(keys), "tail", lastKey, "chunks", chunks)
 					}
 					r := newHashRange(lastKey, chunks)
-
+					if chunks == 1 {
+						smallStorageGauge.Inc(1)
+					} else {
+						largeStorageGauge.Inc(1)
+					}
 					// Our first task is the one that was just filled by this response.
 					batch := ethdb.HookedBatch{
 						Batch: s.db.NewBatch(),
@@ -2187,7 +2192,16 @@ func (s *Syncer) processStorageResponse(res *storageResponse) {
 	// Large contracts could have generated new trie nodes, flush them to disk
 	if res.subTask != nil {
 		if res.subTask.done {
-			res.subTask.genTrie.Commit()
+			root := res.subTask.genTrie.Commit()
+			accountHash := res.accounts[len(res.accounts)-1]
+			if root == res.subTask.root && rawdb.HasStorageTrieNode(s.db, accountHash, nil, root) {
+				// If the chunk's root is an overflown but full delivery, clear the heal request
+				for i, account := range res.mainTask.res.hashes {
+					if account == accountHash {
+						res.mainTask.needHeal[i] = false
+					}
+				}
+			}
 		}
 		if res.subTask.genBatch.ValueSize() > ethdb.IdealBatchSize || res.subTask.done {
 			if err := res.subTask.genBatch.Write(); err != nil {
diff --git a/eth/protocols/snap/sync_test.go b/eth/protocols/snap/sync_test.go
index d4a47635d9..5d4099a814 100644
--- a/eth/protocols/snap/sync_test.go
+++ b/eth/protocols/snap/sync_test.go
@@ -1871,7 +1871,7 @@ func verifyTrie(scheme string, db ethdb.KeyValueStore, root common.Hash, t *test
 // TestSyncAccountPerformance tests how efficient the snap algo is at minimizing
 // state healing
 func TestSyncAccountPerformance(t *testing.T) {
-	t.SkipNow()
+	t.Parallel()
 
 	testSyncAccountPerformance(t, rawdb.HashScheme)
 	testSyncAccountPerformance(t, rawdb.PathScheme)