From f61b50b1e85ebf4a81f7ba52131858b8c6476bd3 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Mon, 3 Oct 2022 13:37:17 +0200
Subject: [PATCH] eth/protocols/snap: serve snap requests when possible
 (#25644)

This PR makes it so that the snap server responds to trie heal requests when possible, even if the snapshot does not exist. The idea being that it might prolong the lifetime of a state root, so we don't have to pivot quite as often.
---
 cmd/devp2p/internal/ethtest/snap.go | 36 ++++++++++++++++++++++++++---
 eth/protocols/snap/handler.go       | 31 +++++++++++++++----------
 2 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/cmd/devp2p/internal/ethtest/snap.go b/cmd/devp2p/internal/ethtest/snap.go
index 032afeafc..7ecbcef6f 100644
--- a/cmd/devp2p/internal/ethtest/snap.go
+++ b/cmd/devp2p/internal/ethtest/snap.go
@@ -406,8 +406,10 @@ func (s *Suite) TestSnapTrieNodes(t *utesting.T) {
 				{[]byte{0}},
 				{[]byte{1}, []byte{0}},
 			},
-			nBytes:    5000,
-			expHashes: []common.Hash{},
+			nBytes: 5000,
+			expHashes: []common.Hash{
+				common.HexToHash("0x1ee1bb2fbac4d46eab331f3e8551e18a0805d084ed54647883aa552809ca968d"),
+			},
 		},
 		{
 			// The leaf is only a couple of levels down, so the continued trie traversal causes lookup failures.
@@ -437,7 +439,35 @@ func (s *Suite) TestSnapTrieNodes(t *utesting.T) {
 				common.HexToHash("0xbcefee69b37cca1f5bf3a48aebe08b35f2ea1864fa958bb0723d909a0e0d28d8"),
 			},
 		},
-	} {
+		{
+			/*
+				A test against this account, requesting trie nodes for the storage trie
+				{
+				  "balance": "0",
+				  "nonce": 1,
+				  "root": "0xbe3d75a1729be157e79c3b77f00206db4d54e3ea14375a015451c88ec067c790",
+				  "codeHash": "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470",
+				  "storage": {
+				    "0x405787fa12a823e0f2b7631cc41b3ba8828b3321ca811111fa75cd3aa3bb5ace": "02",
+				    "0xb10e2d527612073b26eecdfd717e6a320cf44b4afac2b0732d9fcbe2b7fa0cf6": "01",
+				    "0xc2575a0e9e593c00f959f8c92f12db2869c3395a3b0502d05e2516446f71f85b": "03"
+				  },
+				  "key": "0xf493f79c43bd747129a226ad42529885a4b108aba6046b2d12071695a6627844"
+				}
+			*/
+			root: s.chain.RootAt(999),
+			paths: []snap.TrieNodePathSet{
+				{
+					common.FromHex("0xf493f79c43bd747129a226ad42529885a4b108aba6046b2d12071695a6627844"),
+					[]byte{0},
+				},
+			},
+			nBytes: 5000,
+			expHashes: []common.Hash{
+				common.HexToHash("0xbe3d75a1729be157e79c3b77f00206db4d54e3ea14375a015451c88ec067c790"),
+			},
+		},
+	}[7:] {
 		tc := tc
 		if err := s.snapGetTrieNodes(t, &tc); err != nil {
 			t.Errorf("test %d \n #hashes %x\n root: %#x\n bytes: %d\nfailed: %v", i, len(tc.expHashes), tc.root, tc.nBytes, err)
diff --git a/eth/protocols/snap/handler.go b/eth/protocols/snap/handler.go
index aa245ab7e..e001a3883 100644
--- a/eth/protocols/snap/handler.go
+++ b/eth/protocols/snap/handler.go
@@ -493,14 +493,8 @@ func ServiceGetTrieNodesQuery(chain *core.BlockChain, req *GetTrieNodesPacket, s
 		// We don't have the requested state available, bail out
 		return nil, nil
 	}
+	// The 'snap' might be nil, in which case we cannot serve storage slots.
 	snap := chain.Snapshots().Snapshot(req.Root)
-	if snap == nil {
-		// We don't have the requested state snapshotted yet, bail out.
-		// In reality we could still serve using the account and storage
-		// tries only, but let's protect the node a bit while it's doing
-		// snapshot generation.
-		return nil, nil
-	}
 	// Retrieve trie nodes until the packet size limit is reached
 	var (
 		nodes [][]byte
@@ -524,13 +518,26 @@ func ServiceGetTrieNodesQuery(chain *core.BlockChain, req *GetTrieNodesPacket, s
 			bytes += uint64(len(blob))
 
 		default:
+			var stRoot common.Hash
 			// Storage slots requested, open the storage trie and retrieve from there
-			account, err := snap.Account(common.BytesToHash(pathset[0]))
-			loads++ // always account database reads, even for failures
-			if err != nil || account == nil {
-				break
+			if snap == nil {
+				// We don't have the requested state snapshotted yet (or it is stale),
+				// but can look up the account via the trie instead.
+				account, err := accTrie.TryGetAccountWithPreHashedKey(pathset[0])
+				loads += 8 // We don't know the exact cost of lookup, this is an estimate
+				if err != nil || account == nil {
+					break
+				}
+				stRoot = account.Root
+			} else {
+				account, err := snap.Account(common.BytesToHash(pathset[0]))
+				loads++ // always account database reads, even for failures
+				if err != nil || account == nil {
+					break
+				}
+				stRoot = common.BytesToHash(account.Root)
 			}
-			id := trie.StorageTrieID(req.Root, common.BytesToHash(pathset[0]), common.BytesToHash(account.Root))
+			id := trie.StorageTrieID(req.Root, common.BytesToHash(pathset[0]), stRoot)
 			stTrie, err := trie.NewStateTrie(id, triedb)
 			loads++ // always account database reads, even for failures
 			if err != nil {