From 764495f402dec4c4b4e2d7a32fb939bb9cda25c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Toni=20Ram=C3=ADrez?=
 <58293609+ToniRamirezM@users.noreply.github.com>
Date: Tue, 18 Jun 2024 09:42:43 +0200
Subject: [PATCH] disable MT in DS Tool (#3714)

---
 tools/datastreamer/config/default.go       | 15 +++++--------
 tools/datastreamer/config/tool.config.toml |  3 +--
 tools/datastreamer/main.go                 | 26 +++++++++++++---------
 3 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/tools/datastreamer/config/default.go b/tools/datastreamer/config/default.go
index 74c6f35b24..3d5e97caad 100644
--- a/tools/datastreamer/config/default.go
+++ b/tools/datastreamer/config/default.go
@@ -3,14 +3,15 @@ package config
 // DefaultValues is the default configuration
 const DefaultValues = `
 [Online]
-URI = "zkevm-sequencer:6900"
+URI = "localhost:6900"
 StreamType = 1
 
 [Offline]
 Port = 6901
-Filename = "datastreamer.bin"
-Version = 1
+Filename = "datastream.bin"
+Version = 3
 ChainID = 1440
+WriteTimeout = "5s"
 UpgradeEtrogBatchNumber = 0
 
 [StateDB]
@@ -22,13 +23,9 @@ Port = "5432"
 EnableLog = false	
 MaxConns = 200
 
-[Executor]
-URI = "zkevm-prover:50071"
-MaxGRPCMessageSize = 100000000
-
 [MerkleTree]
-URI = "zkevm-prover:50061"
-MaxThreads = 20
+URI = "localhost:50061"
+MaxThreads = 0
 CacheFile = ""
 
 [Log]
diff --git a/tools/datastreamer/config/tool.config.toml b/tools/datastreamer/config/tool.config.toml
index f5530b8271..11d4c8bef7 100644
--- a/tools/datastreamer/config/tool.config.toml
+++ b/tools/datastreamer/config/tool.config.toml
@@ -7,7 +7,6 @@ Port = 6901
 Filename = "datastream.bin"
 Version = 3
 ChainID = 1440
-WriteTimeout = "5s"
 UpgradeEtrogBatchNumber = 0
 
 [StateDB]
@@ -21,7 +20,7 @@ MaxConns = 200
 
 [MerkleTree]
 URI = "localhost:50061"
-MaxThreads = 20
+MaxThreads = 0
 CacheFile = "merkle_tree_cache.json"
 
 [Log]
diff --git a/tools/datastreamer/main.go b/tools/datastreamer/main.go
index 4ac7dab2dd..ad73317506 100644
--- a/tools/datastreamer/main.go
+++ b/tools/datastreamer/main.go
@@ -222,15 +222,21 @@ func generate(cliCtx *cli.Context) error {
 	stateDBStorage := pgstatestorage.NewPostgresStorage(state.Config{}, stateSqlDB)
 	log.Debug("Connected to the database")
 
-	mtDBServerConfig := merkletree.Config{URI: c.MerkleTree.URI}
-	var mtDBCancel context.CancelFunc
-	mtDBServiceClient, mtDBClientConn, mtDBCancel := merkletree.NewMTDBServiceClient(cliCtx.Context, mtDBServerConfig)
-	defer func() {
-		mtDBCancel()
-		mtDBClientConn.Close()
-	}()
-	stateTree := merkletree.NewStateTree(mtDBServiceClient)
-	log.Debug("Connected to the merkle tree")
+	var stateTree *merkletree.StateTree
+
+	if c.MerkleTree.MaxThreads > 0 {
+		mtDBServerConfig := merkletree.Config{URI: c.MerkleTree.URI}
+		var mtDBCancel context.CancelFunc
+		mtDBServiceClient, mtDBClientConn, mtDBCancel := merkletree.NewMTDBServiceClient(cliCtx.Context, mtDBServerConfig)
+		defer func() {
+			mtDBCancel()
+			mtDBClientConn.Close()
+		}()
+		stateTree = merkletree.NewStateTree(mtDBServiceClient)
+		log.Debug("Connected to the merkle tree")
+	} else {
+		log.Debug("Merkle tree disabled")
+	}
 
 	stateDB := state.NewState(state.Config{}, stateDBStorage, nil, stateTree, nil, nil, nil)
 
@@ -288,7 +294,7 @@ func generate(cliCtx *cli.Context) error {
 	wg.Wait()
 
 	// Convert imStateRoots to a json and save it to a file
-	if c.MerkleTree.CacheFile != "" {
+	if c.MerkleTree.CacheFile != "" && c.MerkleTree.MaxThreads > 0 {
 		jsonFile, _ := json.Marshal(imStateRoots)
 		err = os.WriteFile(c.MerkleTree.CacheFile, jsonFile, 0644) // nolint:gosec, gomnd
 		if err != nil {