From 8f1f1e4973fe853edbb7a8fa2fc203d13440fc1b Mon Sep 17 00:00:00 2001
From: Bui Quang Minh <minh.bui@skymavis.com>
Date: Thu, 10 Oct 2024 15:43:23 +0700
Subject: [PATCH] core: add many storage updates benchmark

This benchmark is intended to be used with mainnet data. The benchmark includes
many storage updates to Axie contract's storage.
---
 core/blockchain_test.go | 83 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 83 insertions(+)

diff --git a/core/blockchain_test.go b/core/blockchain_test.go
index 6359bb687..47b660c1b 100644
--- a/core/blockchain_test.go
+++ b/core/blockchain_test.go
@@ -24,6 +24,7 @@ import (
 	"math/big"
 	"math/rand"
 	"os"
+	"os/exec"
 	"sync"
 	"testing"
 	"time"
@@ -41,6 +42,7 @@ import (
 	"github.com/ethereum/go-ethereum/consensus/ethash"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/state"
+	"github.com/ethereum/go-ethereum/core/state/snapshot"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/core/vm"
 	"github.com/ethereum/go-ethereum/crypto"
@@ -4399,3 +4401,84 @@ func TestBlockChain_2000StorageUpdate(t *testing.T) {
 		t.Fatalf("Block hash mismatches, exp %s got %s", common.Hash{}, blockHash)
 	}
 }
+
+// This benchmark is intended to be used with mainnet data, so mainnet chaindata's directory
+// is needed to run this benchmark
+func BenchmarkManyStorageUpdate(b *testing.B) {
+	const (
+		// Fill the chaindata's parent directory
+		datadir   = ""
+		numInsert = state.ParallelInsertThreshold + 1
+	)
+
+	var (
+		diskdb       ethdb.Database
+		err          error
+		axieContract = common.HexToAddress("0x32950db2a7164ae833121501c797d79e7b79d74c")
+		value        = common.HexToHash("0x11")
+	)
+	defer func() {
+		if diskdb != nil {
+			diskdb.Close()
+			cmd := exec.Command("../script/overlay_chaindata.sh", "-d", datadir, "-c")
+			if err := cmd.Run(); err != nil {
+				b.Fatal(err)
+			}
+		}
+	}()
+
+	keys := make([]common.Hash, 0, numInsert)
+	for i := 0; i < numInsert; i++ {
+		hash := crypto.Keccak256Hash(big.NewInt(int64(i)).Bytes())
+		keys = append(keys, hash)
+	}
+
+	b.StopTimer()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		cmd := exec.Command("../script/overlay_chaindata.sh", "-d", datadir)
+		if err := cmd.Run(); err != nil {
+			b.Fatal(err)
+		}
+
+		diskdb, err = rawdb.NewPebbleDBDatabase(datadir+"/chaindata", 1024, 500000, "", false, false)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		engine := ethash.NewFaker()
+		chain, err := NewBlockChain(diskdb, nil, params.TestChainConfig, engine, vm.Config{}, nil, nil)
+		if err != nil {
+			b.Fatalf("failed to create tester chain: %v", err)
+		}
+		headBlock := chain.CurrentBlock()
+
+		database := state.NewDatabase(diskdb)
+		snapshot, err := snapshot.New(diskdb, database.TrieDB(), 256, headBlock.Root(), true, true, false)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		statedb, err := state.New(headBlock.Root(), database, snapshot)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		b.StartTimer()
+		for i := 0; i < numInsert; i++ {
+			statedb.SetState(axieContract, keys[i], value)
+		}
+		_, err = statedb.Commit(true)
+		if err != nil {
+			b.Fatal(err)
+		}
+		b.StopTimer()
+
+		diskdb.Close()
+		cmd = exec.Command("../script/overlay_chaindata.sh", "-d", datadir, "-c")
+		if err := cmd.Run(); err != nil {
+			b.Fatal(err)
+		}
+		diskdb = nil
+	}
+}