From 4681643241b849475aa7f50c90e845202549d0ee Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 20:51:52 -0700
Subject: [PATCH 01/18] find trie root

---
 cmd/util/cmd/find-trie-root/cmd.go | 151 +++++++++++++++++++++++++++++
 cmd/util/cmd/root.go               |   2 +
 2 files changed, 153 insertions(+)
 create mode 100644 cmd/util/cmd/find-trie-root/cmd.go

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
new file mode 100644
index 00000000000..a02530069af
--- /dev/null
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -0,0 +1,151 @@
+package find_trie_root
+
+import (
+	"encoding/hex"
+	"fmt"
+	"math"
+
+	prometheusWAL "github.com/onflow/wal/wal"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+	"github.com/spf13/cobra"
+
+	"github.com/onflow/flow-go/ledger"
+	"github.com/onflow/flow-go/ledger/common/hash"
+	"github.com/onflow/flow-go/ledger/complete/wal"
+)
+
+var (
+	flagExecutionStateDir string
+	flagRootHash          string
+	flagFrom              int
+	flagTo                int
+)
+
+var Cmd = &cobra.Command{
+	Use:   "find-trie-root",
+	Short: "find trie root",
+	Run:   run,
+}
+
+func init() {
+	Cmd.Flags().StringVarP(&flagExecutionStateDir, "execution-state-dir", "e", "/var/flow/data/execution", "directory to the execution state")
+	_ = Cmd.MarkFlagRequired("execution-state-dir")
+
+	Cmd.Flags().StringVar(&flagRootHash, "root-hash", "",
+		"ledger root hash (hex-encoded, 64 characters)")
+	_ = Cmd.MarkFlagRequired("root-hash")
+
+	Cmd.Flags().IntVar(&flagFrom, "from", 0, "from segment")
+	Cmd.Flags().IntVar(&flagTo, "to", math.MaxInt32, "to segment")
+}
+
+func run(*cobra.Command, []string) {
+	rootHash, err := parseInput(flagRootHash)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot parse input")
+	}
+
+	segment, offset, err := searchRootHashInSegments(rootHash, flagExecutionStateDir, flagFrom, flagTo)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot find root hash in segments")
+	}
+	log.Info().Msgf("found root hash in segment %d at offset %d", segment, offset)
+}
+
+func parseInput(rootHashStr string) (ledger.RootHash, error) {
+	rootHashBytes, err := hex.DecodeString(rootHashStr)
+	if err != nil {
+		return ledger.RootHash(hash.DummyHash), fmt.Errorf("cannot decode root hash: %w", err)
+	}
+	rootHash, err := ledger.ToRootHash(rootHashBytes)
+	if err != nil {
+		return ledger.RootHash(hash.DummyHash), fmt.Errorf("invalid root hash: %w", err)
+	}
+	return rootHash, nil
+}
+
+func searchRootHashInSegments(
+	expectedHash ledger.RootHash,
+	dir string,
+	wantFrom, wantTo int,
+) (int, int64, error) {
+	log := zerolog.Logger{}
+	w, err := prometheusWAL.NewSize(log, prometheus.DefaultRegisterer, dir, wal.SegmentSize, false)
+	if err != nil {
+		return 0, 0, fmt.Errorf("cannot create WAL: %w", err)
+	}
+
+	from, to, err := prometheusWAL.Segments(dir)
+	if err != nil {
+		return 0, 0, fmt.Errorf("cannot get segments: %w", err)
+	}
+
+	if wantFrom > to {
+		return 0, 0, fmt.Errorf("from segment %d is greater than the last segment %d", wantFrom, to)
+	}
+
+	if wantTo < from {
+		return 0, 0, fmt.Errorf("to segment %d is less than the first segment %d", wantTo, from)
+	}
+
+	if wantFrom > from {
+		from = wantFrom
+	}
+
+	if wantTo < to {
+		to = wantTo
+	}
+
+	log.Info().
+		Str("dir", dir).
+		Int("from", from).
+		Int("to", to).
+		Int("want-from", wantFrom).
+		Int("want-to", wantTo).
+		Msgf("searching for trie root hash %x in segments [%d,%d]", expectedHash, wantFrom, wantTo)
+
+	sr, err := prometheusWAL.NewSegmentsRangeReader(log, prometheusWAL.SegmentRange{
+		Dir:   w.Dir(),
+		First: from,
+		Last:  to,
+	})
+
+	if err != nil {
+		return 0, 0, fmt.Errorf("cannot create WAL segments reader: %w", err)
+	}
+
+	defer sr.Close()
+
+	reader := prometheusWAL.NewReader(sr)
+
+	for reader.Next() {
+		record := reader.Record()
+		operation, rootHash, _, err := wal.Decode(record)
+		if err != nil {
+			return 0, 0, fmt.Errorf("cannot decode LedgerWAL record: %w", err)
+		}
+
+		log.Debug().
+			Uint8("operation", uint8(operation)).
+			Str("root-hash", rootHash.String()).
+			Msgf("read LedgerWAL record")
+
+		switch operation {
+		case wal.WALUpdate:
+			if rootHash == expectedHash {
+				log.Info().Msgf("found expected trie root hash %x", rootHash)
+				return reader.Segment(), reader.Offset(), nil
+			}
+		default:
+		}
+
+		err = reader.Err()
+		if err != nil {
+			return 0, 0, fmt.Errorf("cannot read LedgerWAL: %w", err)
+		}
+	}
+
+	return 0, 0, fmt.Errorf("finish reading all segment files from %d to %d, but not found", from, to)
+}
diff --git a/cmd/util/cmd/root.go b/cmd/util/cmd/root.go
index ba18f16d1f2..4ed05b17e18 100644
--- a/cmd/util/cmd/root.go
+++ b/cmd/util/cmd/root.go
@@ -18,6 +18,7 @@ import (
 	ledger_json_exporter "github.com/onflow/flow-go/cmd/util/cmd/export-json-execution-state"
 	export_json_transactions "github.com/onflow/flow-go/cmd/util/cmd/export-json-transactions"
 	find_inconsistent_result "github.com/onflow/flow-go/cmd/util/cmd/find-inconsistent-result"
+	find_trie_root "github.com/onflow/flow-go/cmd/util/cmd/find-trie-root"
 	read_badger "github.com/onflow/flow-go/cmd/util/cmd/read-badger/cmd"
 	read_execution_state "github.com/onflow/flow-go/cmd/util/cmd/read-execution-state"
 	read_hotstuff "github.com/onflow/flow-go/cmd/util/cmd/read-hotstuff/cmd"
@@ -81,6 +82,7 @@ func addCommands() {
 	rootCmd.AddCommand(export_json_transactions.Cmd)
 	rootCmd.AddCommand(read_hotstuff.RootCmd)
 	rootCmd.AddCommand(find_inconsistent_result.Cmd)
+	rootCmd.AddCommand(find_trie_root.Cmd)
 	rootCmd.AddCommand(update_commitment.Cmd)
 }
 

From 7bbceca8946cb7f30213a287622e95420f6d9d39 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 21:18:53 -0700
Subject: [PATCH 02/18] add output-dir

---
 cmd/util/cmd/find-trie-root/cmd.go | 73 +++++++++++++++++++++++++++++-
 1 file changed, 72 insertions(+), 1 deletion(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index a02530069af..22be1d18fcf 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -21,6 +21,7 @@ var (
 	flagRootHash          string
 	flagFrom              int
 	flagTo                int
+	flagOutputDir         string
 )
 
 var Cmd = &cobra.Command{
@@ -30,7 +31,7 @@ var Cmd = &cobra.Command{
 }
 
 func init() {
-	Cmd.Flags().StringVarP(&flagExecutionStateDir, "execution-state-dir", "e", "/var/flow/data/execution", "directory to the execution state")
+	Cmd.Flags().StringVar(&flagExecutionStateDir, "execution-state-dir", "/var/flow/data/execution", "directory to the execution state")
 	_ = Cmd.MarkFlagRequired("execution-state-dir")
 
 	Cmd.Flags().StringVar(&flagRootHash, "root-hash", "",
@@ -39,6 +40,8 @@ func init() {
 
 	Cmd.Flags().IntVar(&flagFrom, "from", 0, "from segment")
 	Cmd.Flags().IntVar(&flagTo, "to", math.MaxInt32, "to segment")
+
+	Cmd.Flags().StringVar(&flagOutputDir, "output-dir", "", "output directory")
 }
 
 func run(*cobra.Command, []string) {
@@ -52,6 +55,17 @@ func run(*cobra.Command, []string) {
 		log.Fatal().Err(err).Msg("cannot find root hash in segments")
 	}
 	log.Info().Msgf("found root hash in segment %d at offset %d", segment, offset)
+
+	if len(flagOutputDir) == 0 {
+		return
+	}
+
+	err = copyWAL(flagExecutionStateDir, flagOutputDir, segment, rootHash)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot copy WAL")
+	}
+
+	log.Info().Msgf("copied WAL to %s", flagOutputDir)
 }
 
 func parseInput(rootHashStr string) (ledger.RootHash, error) {
@@ -149,3 +163,60 @@ func searchRootHashInSegments(
 
 	return 0, 0, fmt.Errorf("finish reading all segment files from %d to %d, but not found", from, to)
 }
+
+func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) error {
+	writer, err := prometheusWAL.NewSize(log.Logger, prometheus.DefaultRegisterer, outputDir, wal.SegmentSize, false)
+	if err != nil {
+		return fmt.Errorf("cannot create writer WAL: %w", err)
+	}
+
+	defer writer.Close()
+
+	w, err := prometheusWAL.NewSize(log.Logger, prometheus.DefaultRegisterer, dir, wal.SegmentSize, false)
+	if err != nil {
+		return fmt.Errorf("cannot create WAL: %w", err)
+	}
+
+	sr, err := prometheusWAL.NewSegmentsRangeReader(log.Logger, prometheusWAL.SegmentRange{
+		Dir:   w.Dir(),
+		First: segment,
+		Last:  segment,
+	})
+	if err != nil {
+		return fmt.Errorf("cannot create WAL segments reader: %w", err)
+	}
+
+	defer sr.Close()
+
+	reader := prometheusWAL.NewReader(sr)
+
+	for reader.Next() {
+		record := reader.Record()
+		operation, rootHash, update, err := wal.Decode(record)
+		if err != nil {
+			return fmt.Errorf("cannot decode LedgerWAL record: %w", err)
+		}
+
+		bytes := wal.EncodeUpdate(update)
+		_, err = writer.Log(bytes)
+		if err != nil {
+			return fmt.Errorf("cannot write LedgerWAL record: %w", err)
+		}
+
+		switch operation {
+		case wal.WALUpdate:
+			if rootHash == expectedRoot {
+				log.Info().Msgf("found expected trie root hash %x, finish writing", rootHash)
+				return nil
+			}
+		default:
+		}
+
+		err = reader.Err()
+		if err != nil {
+			return fmt.Errorf("cannot read LedgerWAL: %w", err)
+		}
+	}
+
+	return fmt.Errorf("finish reading all segment files from %d to %d, but not found", segment, segment)
+}

From 4606afcf4a988293866485f5589f969c1797724a Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 21:24:12 -0700
Subject: [PATCH 03/18] use equals

---
 cmd/util/cmd/find-trie-root/cmd.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 22be1d18fcf..f186114407b 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -148,7 +148,7 @@ func searchRootHashInSegments(
 
 		switch operation {
 		case wal.WALUpdate:
-			if rootHash == expectedHash {
+			if rootHash.Equals(expectedHash) {
 				log.Info().Msgf("found expected trie root hash %x", rootHash)
 				return reader.Segment(), reader.Offset(), nil
 			}
@@ -205,7 +205,7 @@ func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) e
 
 		switch operation {
 		case wal.WALUpdate:
-			if rootHash == expectedRoot {
+			if rootHash.Equals(expectedRoot) {
 				log.Info().Msgf("found expected trie root hash %x, finish writing", rootHash)
 				return nil
 			}

From c1024c24e5938f0f259bbaa4dfa3137b52c70a45 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 21:45:37 -0700
Subject: [PATCH 04/18] update log

---
 cmd/util/cmd/find-trie-root/cmd.go | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index f186114407b..fc69fa6ddcf 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -4,6 +4,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"math"
+	"os"
 
 	prometheusWAL "github.com/onflow/wal/wal"
 	"github.com/prometheus/client_golang/prometheus"
@@ -85,8 +86,8 @@ func searchRootHashInSegments(
 	dir string,
 	wantFrom, wantTo int,
 ) (int, int64, error) {
-	log := zerolog.Logger{}
-	w, err := prometheusWAL.NewSize(log, prometheus.DefaultRegisterer, dir, wal.SegmentSize, false)
+	lg := zerolog.New(os.Stderr).With().Timestamp().Logger()
+	w, err := prometheusWAL.NewSize(lg, prometheus.DefaultRegisterer, dir, wal.SegmentSize, false)
 	if err != nil {
 		return 0, 0, fmt.Errorf("cannot create WAL: %w", err)
 	}
@@ -120,7 +121,7 @@ func searchRootHashInSegments(
 		Int("want-to", wantTo).
 		Msgf("searching for trie root hash %x in segments [%d,%d]", expectedHash, wantFrom, wantTo)
 
-	sr, err := prometheusWAL.NewSegmentsRangeReader(log, prometheusWAL.SegmentRange{
+	sr, err := prometheusWAL.NewSegmentsRangeReader(lg, prometheusWAL.SegmentRange{
 		Dir:   w.Dir(),
 		First: from,
 		Last:  to,

From d5ef271df2e14b9071ab7cbf866d7ad679e6a3eb Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 22:06:43 -0700
Subject: [PATCH 05/18] fix root hash

---
 cmd/util/cmd/find-trie-root/cmd.go | 2 +-
 ledger/complete/wal/encoding.go    | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index fc69fa6ddcf..52c109658f4 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -119,7 +119,7 @@ func searchRootHashInSegments(
 		Int("to", to).
 		Int("want-from", wantFrom).
 		Int("want-to", wantTo).
-		Msgf("searching for trie root hash %x in segments [%d,%d]", expectedHash, wantFrom, wantTo)
+		Msgf("searching for trie root hash %v in segments [%d,%d]", expectedHash, wantFrom, wantTo)
 
 	sr, err := prometheusWAL.NewSegmentsRangeReader(lg, prometheusWAL.SegmentRange{
 		Dir:   w.Dir(),
diff --git a/ledger/complete/wal/encoding.go b/ledger/complete/wal/encoding.go
index 8bc5f8d6d13..2839b90b194 100644
--- a/ledger/complete/wal/encoding.go
+++ b/ledger/complete/wal/encoding.go
@@ -61,6 +61,7 @@ func Decode(data []byte) (operation WALOperation, rootHash ledger.RootHash, upda
 	switch operation {
 	case WALUpdate:
 		update, err = ledger.DecodeTrieUpdate(data[1:])
+		rootHash = update.RootHash
 		return
 	case WALDelete:
 		var rootHashBytes []byte

From 1018f74c9aa376d28d2d5d831305dedeba470503 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 22:18:34 -0700
Subject: [PATCH 06/18] fix outputer

---
 cmd/util/cmd/find-trie-root/cmd.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 52c109658f4..1eb759b6d91 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -166,14 +166,14 @@ func searchRootHashInSegments(
 }
 
 func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) error {
-	writer, err := prometheusWAL.NewSize(log.Logger, prometheus.DefaultRegisterer, outputDir, wal.SegmentSize, false)
+	writer, err := prometheusWAL.NewSize(log.Logger, nil, outputDir, wal.SegmentSize, false)
 	if err != nil {
 		return fmt.Errorf("cannot create writer WAL: %w", err)
 	}
 
 	defer writer.Close()
 
-	w, err := prometheusWAL.NewSize(log.Logger, prometheus.DefaultRegisterer, dir, wal.SegmentSize, false)
+	w, err := prometheusWAL.NewSize(log.Logger, nil, dir, wal.SegmentSize, false)
 	if err != nil {
 		return fmt.Errorf("cannot create WAL: %w", err)
 	}

From 540a27dc3a739b1b8ccca6af0e0a10b3f4d2352b Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 7 Jun 2024 22:28:49 -0700
Subject: [PATCH 07/18] fix log

---
 cmd/util/cmd/find-trie-root/cmd.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 1eb759b6d91..bc227eb7e37 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -150,7 +150,7 @@ func searchRootHashInSegments(
 		switch operation {
 		case wal.WALUpdate:
 			if rootHash.Equals(expectedHash) {
-				log.Info().Msgf("found expected trie root hash %x", rootHash)
+				log.Info().Msgf("found expected trie root hash %v", rootHash)
 				return reader.Segment(), reader.Offset(), nil
 			}
 		default:
@@ -207,7 +207,7 @@ func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) e
 		switch operation {
 		case wal.WALUpdate:
 			if rootHash.Equals(expectedRoot) {
-				log.Info().Msgf("found expected trie root hash %x, finish writing", rootHash)
+				log.Info().Msgf("found expected trie root hash %v, finish writing", rootHash)
 				return nil
 			}
 		default:

From 7fddffef4cd32da25807ba5194f5032fc5c904c6 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 14 Jun 2024 14:41:09 -0700
Subject: [PATCH 08/18] add comment

---
 cmd/util/cmd/find-trie-root/cmd.go | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index bc227eb7e37..a87e3902991 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -25,9 +25,18 @@ var (
 	flagOutputDir         string
 )
 
+// find trie root hash from the wal files.
+// useful for state extraction and rolling back executed height.
+// for instance, when extracting state for a target height, it requires the wal files
+// has the trie root hash of the target block as the latest few records. If not the case,
+// then it is necessary to trim the wal files to the last record with the target trie root hash.
+// in order to do that, this command can be used to find the trie root hash in the wal files,
+// and copy the wal that contains the trie root hash to a new directory and trim it to
+// have the target trie root hash as the last record.
+// after that, the new wal file can be used to extract the state for the target height.
 var Cmd = &cobra.Command{
 	Use:   "find-trie-root",
-	Short: "find trie root",
+	Short: "find trie root hash from the wal files",
 	Run:   run,
 }
 

From 10d6aee33f9e0eed975dc490cc8796c3c908a22a Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 14 Jun 2024 15:13:52 -0700
Subject: [PATCH 09/18] fix test case

---
 ledger/complete/wal/encoding_test.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/ledger/complete/wal/encoding_test.go b/ledger/complete/wal/encoding_test.go
index 6580d0a683f..500ae5a8f60 100644
--- a/ledger/complete/wal/encoding_test.go
+++ b/ledger/complete/wal/encoding_test.go
@@ -7,7 +7,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/onflow/flow-go/ledger"
-	"github.com/onflow/flow-go/ledger/common/hash"
 	"github.com/onflow/flow-go/ledger/common/testutils"
 	realWAL "github.com/onflow/flow-go/ledger/complete/wal"
 )
@@ -44,7 +43,7 @@ func TestUpdate(t *testing.T) {
 		operation, stateCommitment, up, err := realWAL.Decode(data)
 		require.NoError(t, err)
 		assert.Equal(t, realWAL.WALUpdate, operation)
-		assert.Equal(t, stateCommitment, ledger.RootHash(hash.DummyHash))
+		assert.Equal(t, stateCommitment, ledger.RootHash(rootHash))
 		assert.Equal(t, update, up)
 	})
 }

From d16b8ea07b84003570e67c6d9945b24efb44d697 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 14 Jun 2024 17:18:24 -0700
Subject: [PATCH 10/18] address review comments

---
 cmd/util/cmd/find-trie-root/cmd.go | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index a87e3902991..38736acc433 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -7,7 +7,6 @@ import (
 	"os"
 
 	prometheusWAL "github.com/onflow/wal/wal"
-	"github.com/prometheus/client_golang/prometheus"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
@@ -60,6 +59,10 @@ func run(*cobra.Command, []string) {
 		log.Fatal().Err(err).Msg("cannot parse input")
 	}
 
+	if flagExecutionStateDir == flagOutputDir {
+		log.Fatal().Msg("output directory cannot be the same as the execution state directory")
+	}
+
 	segment, offset, err := searchRootHashInSegments(rootHash, flagExecutionStateDir, flagFrom, flagTo)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot find root hash in segments")
@@ -96,16 +99,15 @@ func searchRootHashInSegments(
 	wantFrom, wantTo int,
 ) (int, int64, error) {
 	lg := zerolog.New(os.Stderr).With().Timestamp().Logger()
-	w, err := prometheusWAL.NewSize(lg, prometheus.DefaultRegisterer, dir, wal.SegmentSize, false)
-	if err != nil {
-		return 0, 0, fmt.Errorf("cannot create WAL: %w", err)
-	}
-
 	from, to, err := prometheusWAL.Segments(dir)
 	if err != nil {
 		return 0, 0, fmt.Errorf("cannot get segments: %w", err)
 	}
 
+	if from < 0 {
+		return 0, 0, fmt.Errorf("no segments found in %s", dir)
+	}
+
 	if wantFrom > to {
 		return 0, 0, fmt.Errorf("from segment %d is greater than the last segment %d", wantFrom, to)
 	}
@@ -122,7 +124,7 @@ func searchRootHashInSegments(
 		to = wantTo
 	}
 
-	log.Info().
+	lg.Info().
 		Str("dir", dir).
 		Int("from", from).
 		Int("to", to).
@@ -131,7 +133,7 @@ func searchRootHashInSegments(
 		Msgf("searching for trie root hash %v in segments [%d,%d]", expectedHash, wantFrom, wantTo)
 
 	sr, err := prometheusWAL.NewSegmentsRangeReader(lg, prometheusWAL.SegmentRange{
-		Dir:   w.Dir(),
+		Dir:   dir,
 		First: from,
 		Last:  to,
 	})

From fa6f02d8dc5350b85b1b87f5e74e36634b35e6ee Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Mon, 17 Jun 2024 08:26:13 -0700
Subject: [PATCH 11/18] revert decoding changes

---
 cmd/util/cmd/find-trie-root/cmd.go   | 30 ++++++++++++++++++----------
 ledger/complete/wal/encoding.go      |  5 ++++-
 ledger/complete/wal/encoding_test.go |  3 ++-
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 38736acc433..37b9f047373 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -148,18 +148,23 @@ func searchRootHashInSegments(
 
 	for reader.Next() {
 		record := reader.Record()
-		operation, rootHash, _, err := wal.Decode(record)
+		operation, _, update, err := wal.Decode(record)
 		if err != nil {
 			return 0, 0, fmt.Errorf("cannot decode LedgerWAL record: %w", err)
 		}
 
-		log.Debug().
+		lg = lg.With().
 			Uint8("operation", uint8(operation)).
-			Str("root-hash", rootHash.String()).
-			Msgf("read LedgerWAL record")
+			Logger()
 
 		switch operation {
 		case wal.WALUpdate:
+			rootHash := update.RootHash
+
+			lg.Debug().
+				Str("root-hash", rootHash.String()).
+				Msg("found WALUpdate")
+
 			if rootHash.Equals(expectedHash) {
 				log.Info().Msgf("found expected trie root hash %v", rootHash)
 				return reader.Segment(), reader.Offset(), nil
@@ -204,19 +209,22 @@ func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) e
 
 	for reader.Next() {
 		record := reader.Record()
-		operation, rootHash, update, err := wal.Decode(record)
+		operation, _, update, err := wal.Decode(record)
 		if err != nil {
 			return fmt.Errorf("cannot decode LedgerWAL record: %w", err)
 		}
 
-		bytes := wal.EncodeUpdate(update)
-		_, err = writer.Log(bytes)
-		if err != nil {
-			return fmt.Errorf("cannot write LedgerWAL record: %w", err)
-		}
-
 		switch operation {
 		case wal.WALUpdate:
+
+			bytes := wal.EncodeUpdate(update)
+			_, err = writer.Log(bytes)
+			if err != nil {
+				return fmt.Errorf("cannot write LedgerWAL record: %w", err)
+			}
+
+			rootHash := update.RootHash
+
 			if rootHash.Equals(expectedRoot) {
 				log.Info().Msgf("found expected trie root hash %v, finish writing", rootHash)
 				return nil
diff --git a/ledger/complete/wal/encoding.go b/ledger/complete/wal/encoding.go
index 2839b90b194..dce4f084f93 100644
--- a/ledger/complete/wal/encoding.go
+++ b/ledger/complete/wal/encoding.go
@@ -51,6 +51,10 @@ func EncodeDelete(rootHash ledger.RootHash) []byte {
 	return buf
 }
 
+// Decode decodes the given data into a WAL operation, root hash and trie update.
+// It returns (WALDelete, rootHash, nil, nil) if the operation is WALDelete.
+// It returns (WALUpdate, hash.DummyHash, update, nil) if the operation is WALUpdate.
+// To read the root hash of the trie update, use update.RootHash.
 func Decode(data []byte) (operation WALOperation, rootHash ledger.RootHash, update *ledger.TrieUpdate, err error) {
 	if len(data) < 4 { // 1 byte op + 2 size + actual data = 4 minimum
 		err = fmt.Errorf("data corrupted, too short to represent operation - hexencoded data: %x", data)
@@ -61,7 +65,6 @@ func Decode(data []byte) (operation WALOperation, rootHash ledger.RootHash, upda
 	switch operation {
 	case WALUpdate:
 		update, err = ledger.DecodeTrieUpdate(data[1:])
-		rootHash = update.RootHash
 		return
 	case WALDelete:
 		var rootHashBytes []byte
diff --git a/ledger/complete/wal/encoding_test.go b/ledger/complete/wal/encoding_test.go
index 500ae5a8f60..6580d0a683f 100644
--- a/ledger/complete/wal/encoding_test.go
+++ b/ledger/complete/wal/encoding_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/onflow/flow-go/ledger"
+	"github.com/onflow/flow-go/ledger/common/hash"
 	"github.com/onflow/flow-go/ledger/common/testutils"
 	realWAL "github.com/onflow/flow-go/ledger/complete/wal"
 )
@@ -43,7 +44,7 @@ func TestUpdate(t *testing.T) {
 		operation, stateCommitment, up, err := realWAL.Decode(data)
 		require.NoError(t, err)
 		assert.Equal(t, realWAL.WALUpdate, operation)
-		assert.Equal(t, stateCommitment, ledger.RootHash(rootHash))
+		assert.Equal(t, stateCommitment, ledger.RootHash(hash.DummyHash))
 		assert.Equal(t, update, up)
 	})
 }

From 3153ec2a7cc6456d6e7eae99ac1205fd5e99a6dd Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Mon, 17 Jun 2024 16:31:34 -0700
Subject: [PATCH 12/18] implement backup wal files

---
 cmd/util/cmd/find-trie-root/cmd.go | 150 +++++++++++++++++++++++++----
 1 file changed, 129 insertions(+), 21 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 37b9f047373..1d8825e531e 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -21,7 +21,8 @@ var (
 	flagRootHash          string
 	flagFrom              int
 	flagTo                int
-	flagOutputDir         string
+	flagBackupDir         string
+	flagTrimAsLatestWAL   bool
 )
 
 // find trie root hash from the wal files.
@@ -50,7 +51,9 @@ func init() {
 	Cmd.Flags().IntVar(&flagFrom, "from", 0, "from segment")
 	Cmd.Flags().IntVar(&flagTo, "to", math.MaxInt32, "to segment")
 
-	Cmd.Flags().StringVar(&flagOutputDir, "output-dir", "", "output directory")
+	Cmd.Flags().StringVar(&flagBackupDir, "backup-dir", "", "directory for backup wal files. must be not exist or empty folder. required when --trim-as-latest-wal flag is set to true.")
+
+	Cmd.Flags().BoolVar(&flagTrimAsLatestWAL, "trim-as-latest-wal", false, "trim the wal file to the last record with the target trie root hash")
 }
 
 func run(*cobra.Command, []string) {
@@ -59,26 +62,46 @@ func run(*cobra.Command, []string) {
 		log.Fatal().Err(err).Msg("cannot parse input")
 	}
 
-	if flagExecutionStateDir == flagOutputDir {
-		log.Fatal().Msg("output directory cannot be the same as the execution state directory")
+	if flagExecutionStateDir == flagBackupDir {
+		log.Fatal().Msg("backup directory cannot be the same as the execution state directory")
 	}
 
 	segment, offset, err := searchRootHashInSegments(rootHash, flagExecutionStateDir, flagFrom, flagTo)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot find root hash in segments")
 	}
-	log.Info().Msgf("found root hash in segment %d at offset %d", segment, offset)
 
-	if len(flagOutputDir) == 0 {
+	segmentFile := prometheusWAL.SegmentName(flagExecutionStateDir, segment)
+
+	log.Info().Msgf("found root hash in segment %d at offset %d, segment file: %v", segment, offset, segmentFile)
+
+	if !flagTrimAsLatestWAL {
+		log.Info().Msg("not trimming WAL. Exiting. to trim the WAL, use --trim-as-latest-wal flag")
 		return
 	}
 
-	err = copyWAL(flagExecutionStateDir, flagOutputDir, segment, rootHash)
+	if len(flagBackupDir) == 0 {
+		log.Error().Msgf("backup directory is not provided")
+		return
+	}
+
+	// genereate a segment file to the temporary folder with the root hash as its last record
+	newSegmentFile, err := findRootHashAndCreateTrimmed(flagExecutionStateDir, segment, rootHash)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot copy WAL")
 	}
 
-	log.Info().Msgf("copied WAL to %s", flagOutputDir)
+	log.Info().Msgf("successfully copied WAL to the temporary folder %v", newSegmentFile)
+
+	// before replacing the last wal file with the newly generated one, backup the rollbacked wals
+	// then move the last segment file to the execution state directory
+	err = backupRollbackedWALsAndMoveLastSegmentFile(segment, flagExecutionStateDir, flagBackupDir, newSegmentFile)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot backup rollbacked WALs")
+	}
+
+	log.Info().Msgf("successfully trimmed WAL %v the trie root hash %v as its last record, original wal files are moved to %v",
+		segment, rootHash, flagBackupDir)
 }
 
 func parseInput(rootHashStr string) (ledger.RootHash, error) {
@@ -181,26 +204,33 @@ func searchRootHashInSegments(
 	return 0, 0, fmt.Errorf("finish reading all segment files from %d to %d, but not found", from, to)
 }
 
-func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) error {
-	writer, err := prometheusWAL.NewSize(log.Logger, nil, outputDir, wal.SegmentSize, false)
+// findRootHashAndCreateTrimmed finds the root hash in the segment file from the given dir folder
+// and creates a new segment file with the expected root hash as the last record in a temporary folder.
+// it return the path to the new segment file.
+func findRootHashAndCreateTrimmed(dir string, segment int, expectedRoot ledger.RootHash) (string, error) {
+	tmpFolder, err := os.MkdirTemp("", "flow-last-segment-file")
 	if err != nil {
-		return fmt.Errorf("cannot create writer WAL: %w", err)
+		return "", fmt.Errorf("cannot create temporary folder: %w", err)
 	}
 
-	defer writer.Close()
+	newSegmentFile := prometheusWAL.SegmentName(tmpFolder, segment)
+
+	log.Info().Msgf("writing new segment file to %v", newSegmentFile)
 
-	w, err := prometheusWAL.NewSize(log.Logger, nil, dir, wal.SegmentSize, false)
+	writer, err := prometheusWAL.NewSize(log.Logger, nil, tmpFolder, wal.SegmentSize, false)
 	if err != nil {
-		return fmt.Errorf("cannot create WAL: %w", err)
+		return "", fmt.Errorf("cannot create writer WAL: %w", err)
 	}
 
+	defer writer.Close()
+
 	sr, err := prometheusWAL.NewSegmentsRangeReader(log.Logger, prometheusWAL.SegmentRange{
-		Dir:   w.Dir(),
+		Dir:   dir,
 		First: segment,
 		Last:  segment,
 	})
 	if err != nil {
-		return fmt.Errorf("cannot create WAL segments reader: %w", err)
+		return "", fmt.Errorf("cannot create WAL segments reader: %w", err)
 	}
 
 	defer sr.Close()
@@ -211,7 +241,7 @@ func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) e
 		record := reader.Record()
 		operation, _, update, err := wal.Decode(record)
 		if err != nil {
-			return fmt.Errorf("cannot decode LedgerWAL record: %w", err)
+			return "", fmt.Errorf("cannot decode LedgerWAL record: %w", err)
 		}
 
 		switch operation {
@@ -220,23 +250,101 @@ func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) e
 			bytes := wal.EncodeUpdate(update)
 			_, err = writer.Log(bytes)
 			if err != nil {
-				return fmt.Errorf("cannot write LedgerWAL record: %w", err)
+				return "", fmt.Errorf("cannot write LedgerWAL record: %w", err)
 			}
 
 			rootHash := update.RootHash
 
 			if rootHash.Equals(expectedRoot) {
 				log.Info().Msgf("found expected trie root hash %v, finish writing", rootHash)
-				return nil
+				return newSegmentFile, nil
 			}
 		default:
 		}
 
 		err = reader.Err()
 		if err != nil {
-			return fmt.Errorf("cannot read LedgerWAL: %w", err)
+			return "", fmt.Errorf("cannot read LedgerWAL: %w", err)
 		}
 	}
 
-	return fmt.Errorf("finish reading all segment files from %d to %d, but not found", segment, segment)
+	return "", fmt.Errorf("finish reading all segment files from %d to %d, but not found", segment, segment)
+}
+
+func checkFolderNotExistOrEmpty(folderPath string) (bool, error) {
+	// Check if the folder exists
+	info, err := os.Stat(folderPath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return true, nil
+		}
+		return false, nil
+	}
+
+	// Check if the path is a directory
+	if !info.IsDir() {
+		return false, fmt.Errorf("The path is not a directory.")
+	}
+
+	// Check if the folder is empty
+	files, err := os.ReadDir(folderPath)
+	if err != nil {
+		return false, fmt.Errorf("Cannot read the folder.")
+	}
+
+	return len(files) == 0, nil
+}
+
+// backup new wals before replacing
+func backupRollbackedWALsAndMoveLastSegmentFile(
+	segment int, walDir, backupDir string, newSegmentFile string) error {
+	// making sure the backup dir is empty
+	empty, err := checkFolderNotExistOrEmpty(backupDir)
+	if err != nil {
+		return fmt.Errorf("cannot check backup directory: %w", err)
+	}
+
+	if !empty {
+		return fmt.Errorf("backup directory %s is not empty", backupDir)
+	}
+
+	// Create the backup directory
+	err = os.MkdirAll(backupDir, os.ModePerm)
+	if err != nil {
+		return fmt.Errorf("cannot create backup directory: %w", err)
+	}
+
+	first, last, err := prometheusWAL.Segments(walDir)
+	if err != nil {
+		return fmt.Errorf("cannot get segments: %w", err)
+	}
+
+	if segment < first {
+		return fmt.Errorf("segment %d is less than the first segment %d", segment, first)
+	}
+
+	// backup all the segment files that have higher number than the given segment, including
+	// the segment file itself, since it will be replaced.
+	for i := segment; i <= last; i++ {
+		segmentFile := prometheusWAL.SegmentName(walDir, i)
+		backupFile := prometheusWAL.SegmentName(backupDir, i)
+
+		log.Info().Msgf("backup segment file %s to %s, %v/%v", segmentFile, backupFile, i, last)
+		err := os.Rename(segmentFile, backupFile)
+		if err != nil {
+			return fmt.Errorf("cannot move segment file %s to %s: %w", segmentFile, backupFile, err)
+		}
+	}
+
+	// after backup the segment files, replace the last segment file
+	segmentToBeReplaced := prometheusWAL.SegmentName(walDir, segment)
+
+	log.Info().Msgf("moving segment file %s to %s", newSegmentFile, segmentToBeReplaced)
+
+	err = os.Rename(newSegmentFile, segmentToBeReplaced)
+	if err != nil {
+		return fmt.Errorf("cannot move segment file %s to %s: %w", newSegmentFile, segmentToBeReplaced, err)
+	}
+
+	return nil
 }

From d3dfed9227ed88d46229b591ea09c53bda112b63 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Mon, 17 Jun 2024 16:55:20 -0700
Subject: [PATCH 13/18] copy file

---
 cmd/util/cmd/find-trie-root/cmd.go | 33 +++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 1d8825e531e..7cad7535af5 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -3,6 +3,7 @@ package find_trie_root
 import (
 	"encoding/hex"
 	"fmt"
+	"io"
 	"math"
 	"os"
 
@@ -63,7 +64,7 @@ func run(*cobra.Command, []string) {
 	}
 
 	if flagExecutionStateDir == flagBackupDir {
-		log.Fatal().Msg("backup directory cannot be the same as the execution state directory")
+		log.Fatal().Msg("--backup-dir directory cannot be the same as the execution state directory")
 	}
 
 	segment, offset, err := searchRootHashInSegments(rootHash, flagExecutionStateDir, flagFrom, flagTo)
@@ -81,7 +82,7 @@ func run(*cobra.Command, []string) {
 	}
 
 	if len(flagBackupDir) == 0 {
-		log.Error().Msgf("backup directory is not provided")
+		log.Error().Msgf("--backup-dir directory is not provided")
 		return
 	}
 
@@ -341,10 +342,36 @@ func backupRollbackedWALsAndMoveLastSegmentFile(
 
 	log.Info().Msgf("moving segment file %s to %s", newSegmentFile, segmentToBeReplaced)
 
-	err = os.Rename(newSegmentFile, segmentToBeReplaced)
+	_, err = copyFile(newSegmentFile, segmentToBeReplaced)
 	if err != nil {
 		return fmt.Errorf("cannot move segment file %s to %s: %w", newSegmentFile, segmentToBeReplaced, err)
 	}
 
 	return nil
 }
+
+func copyFile(src, dst string) (int64, error) {
+	sourceFile, err := os.Open(src)
+	if err != nil {
+		return 0, err
+	}
+	defer sourceFile.Close()
+
+	destinationFile, err := os.Create(dst)
+	if err != nil {
+		return 0, err
+	}
+	defer destinationFile.Close()
+
+	bytesCopied, err := io.Copy(destinationFile, sourceFile)
+	if err != nil {
+		return 0, err
+	}
+
+	err = destinationFile.Sync()
+	if err != nil {
+		return 0, err
+	}
+
+	return bytesCopied, nil
+}

From 3d2e8f91cb244c4b398cbe63ec06cc3cf3b4f841 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Mon, 17 Jun 2024 17:06:00 -0700
Subject: [PATCH 14/18] fix copy file

---
 cmd/util/cmd/find-trie-root/cmd.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 7cad7535af5..8aa41f0b820 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -214,7 +214,9 @@ func findRootHashAndCreateTrimmed(dir string, segment int, expectedRoot ledger.R
 		return "", fmt.Errorf("cannot create temporary folder: %w", err)
 	}
 
-	newSegmentFile := prometheusWAL.SegmentName(tmpFolder, segment)
+	// the new segment file will be created in the temporary folder
+	// and it's always 00000000
+	newSegmentFile := prometheusWAL.SegmentName(tmpFolder, 0)
 
 	log.Info().Msgf("writing new segment file to %v", newSegmentFile)
 

From 00a37c352b703b3c8a449ed4fa66e1db08d5e5b7 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Mon, 17 Jun 2024 21:40:00 -0700
Subject: [PATCH 15/18] address review comments

---
 cmd/util/cmd/find-trie-root/cmd.go | 59 ++++++++++++------------------
 1 file changed, 23 insertions(+), 36 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 8aa41f0b820..b07225c16fd 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -3,9 +3,9 @@ package find_trie_root
 import (
 	"encoding/hex"
 	"fmt"
-	"io"
 	"math"
 	"os"
+	"path/filepath"
 
 	prometheusWAL "github.com/onflow/wal/wal"
 	"github.com/rs/zerolog"
@@ -86,8 +86,26 @@ func run(*cobra.Command, []string) {
 		return
 	}
 
+	// create a temporary folder in the backup folder to store the new segment file
+	tmpFolder := filepath.Join(flagBackupDir, "flow-last-segment-file")
+
+	log.Info().Msgf("creating temporary folder %v", tmpFolder)
+
+	err = os.Mkdir(tmpFolder, os.ModePerm)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot create temporary folder")
+	}
+
+	defer func() {
+		log.Info().Msgf("removing temporary folder %v", tmpFolder)
+		err := os.RemoveAll(tmpFolder)
+		if err != nil {
+			log.Error().Err(err).Msg("cannot remove temporary folder")
+		}
+	}()
+
 	// genereate a segment file to the temporary folder with the root hash as its last record
-	newSegmentFile, err := findRootHashAndCreateTrimmed(flagExecutionStateDir, segment, rootHash)
+	newSegmentFile, err := findRootHashAndCreateTrimmed(flagExecutionStateDir, segment, rootHash, tmpFolder)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot copy WAL")
 	}
@@ -208,12 +226,7 @@ func searchRootHashInSegments(
 // findRootHashAndCreateTrimmed finds the root hash in the segment file from the given dir folder
 // and creates a new segment file with the expected root hash as the last record in a temporary folder.
 // it return the path to the new segment file.
-func findRootHashAndCreateTrimmed(dir string, segment int, expectedRoot ledger.RootHash) (string, error) {
-	tmpFolder, err := os.MkdirTemp("", "flow-last-segment-file")
-	if err != nil {
-		return "", fmt.Errorf("cannot create temporary folder: %w", err)
-	}
-
+func findRootHashAndCreateTrimmed(dir string, segment int, expectedRoot ledger.RootHash, tmpFolder string) (string, error) {
 	// the new segment file will be created in the temporary folder
 	// and it's always 00000000
 	newSegmentFile := prometheusWAL.SegmentName(tmpFolder, 0)
@@ -281,7 +294,7 @@ func checkFolderNotExistOrEmpty(folderPath string) (bool, error) {
 		if os.IsNotExist(err) {
 			return true, nil
 		}
-		return false, nil
+		return false, err
 	}
 
 	// Check if the path is a directory
@@ -344,36 +357,10 @@ func backupRollbackedWALsAndMoveLastSegmentFile(
 
 	log.Info().Msgf("moving segment file %s to %s", newSegmentFile, segmentToBeReplaced)
 
-	_, err = copyFile(newSegmentFile, segmentToBeReplaced)
+	err = os.Rename(newSegmentFile, segmentToBeReplaced)
 	if err != nil {
 		return fmt.Errorf("cannot move segment file %s to %s: %w", newSegmentFile, segmentToBeReplaced, err)
 	}
 
 	return nil
 }
-
-func copyFile(src, dst string) (int64, error) {
-	sourceFile, err := os.Open(src)
-	if err != nil {
-		return 0, err
-	}
-	defer sourceFile.Close()
-
-	destinationFile, err := os.Create(dst)
-	if err != nil {
-		return 0, err
-	}
-	defer destinationFile.Close()
-
-	bytesCopied, err := io.Copy(destinationFile, sourceFile)
-	if err != nil {
-		return 0, err
-	}
-
-	err = destinationFile.Sync()
-	if err != nil {
-		return 0, err
-	}
-
-	return bytesCopied, nil
-}

From 3bf99dfc8049a997d57fa01e4ce271bf051c13ea Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Tue, 18 Jun 2024 10:44:15 -0700
Subject: [PATCH 16/18] fix logging

---
 cmd/util/cmd/find-trie-root/cmd.go | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index b07225c16fd..50f843c0026 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -195,15 +195,12 @@ func searchRootHashInSegments(
 			return 0, 0, fmt.Errorf("cannot decode LedgerWAL record: %w", err)
 		}
 
-		lg = lg.With().
-			Uint8("operation", uint8(operation)).
-			Logger()
-
 		switch operation {
 		case wal.WALUpdate:
 			rootHash := update.RootHash
 
-			lg.Debug().
+			log.Debug().
+				Uint8("operation", uint8(operation)).
 				Str("root-hash", rootHash.String()).
 				Msg("found WALUpdate")
 

From e60aa2d8fb5b231b9e438a7f043cbdb1ae748fbe Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Tue, 18 Jun 2024 11:17:47 -0700
Subject: [PATCH 17/18] adjust folder creation

---
 cmd/util/cmd/find-trie-root/cmd.go | 49 ++++++++++++++++--------------
 1 file changed, 27 insertions(+), 22 deletions(-)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 50f843c0026..0fa558dc39a 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -42,7 +42,8 @@ var Cmd = &cobra.Command{
 }
 
 func init() {
-	Cmd.Flags().StringVar(&flagExecutionStateDir, "execution-state-dir", "/var/flow/data/execution", "directory to the execution state")
+	Cmd.Flags().StringVar(&flagExecutionStateDir, "execution-state-dir", "/var/flow/data/execution",
+		"directory to the execution state")
 	_ = Cmd.MarkFlagRequired("execution-state-dir")
 
 	Cmd.Flags().StringVar(&flagRootHash, "root-hash", "",
@@ -52,9 +53,11 @@ func init() {
 	Cmd.Flags().IntVar(&flagFrom, "from", 0, "from segment")
 	Cmd.Flags().IntVar(&flagTo, "to", math.MaxInt32, "to segment")
 
-	Cmd.Flags().StringVar(&flagBackupDir, "backup-dir", "", "directory for backup wal files. must be not exist or empty folder. required when --trim-as-latest-wal flag is set to true.")
+	Cmd.Flags().StringVar(&flagBackupDir, "backup-dir", "",
+		"directory for backup wal files. must be not exist or empty folder. required when --trim-as-latest-wal flag is set to true.")
 
-	Cmd.Flags().BoolVar(&flagTrimAsLatestWAL, "trim-as-latest-wal", false, "trim the wal file to the last record with the target trie root hash")
+	Cmd.Flags().BoolVar(&flagTrimAsLatestWAL, "trim-as-latest-wal", false,
+		"trim the wal file to the last record with the target trie root hash")
 }
 
 func run(*cobra.Command, []string) {
@@ -67,6 +70,16 @@ func run(*cobra.Command, []string) {
 		log.Fatal().Msg("--backup-dir directory cannot be the same as the execution state directory")
 	}
 
+	// making sure the backup dir is empty
+	empty, err := checkFolderIsEmpty(flagBackupDir)
+	if err != nil {
+		log.Fatal().Msgf("--backup-dir directory %v must exist and empty", flagBackupDir)
+	}
+
+	if !empty {
+		log.Fatal().Msgf("--backup-dir directory %v must be empty", flagBackupDir)
+	}
+
 	segment, offset, err := searchRootHashInSegments(rootHash, flagExecutionStateDir, flagFrom, flagTo)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot find root hash in segments")
@@ -114,7 +127,8 @@ func run(*cobra.Command, []string) {
 
 	// before replacing the last wal file with the newly generated one, backup the rollbacked wals
 	// then move the last segment file to the execution state directory
-	err = backupRollbackedWALsAndMoveLastSegmentFile(segment, flagExecutionStateDir, flagBackupDir, newSegmentFile)
+	err = backupRollbackedWALsAndMoveLastSegmentFile(
+		segment, flagExecutionStateDir, flagBackupDir, newSegmentFile)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot backup rollbacked WALs")
 	}
@@ -223,7 +237,8 @@ func searchRootHashInSegments(
 // findRootHashAndCreateTrimmed finds the root hash in the segment file from the given dir folder
 // and creates a new segment file with the expected root hash as the last record in a temporary folder.
 // it return the path to the new segment file.
-func findRootHashAndCreateTrimmed(dir string, segment int, expectedRoot ledger.RootHash, tmpFolder string) (string, error) {
+func findRootHashAndCreateTrimmed(
+	dir string, segment int, expectedRoot ledger.RootHash, tmpFolder string) (string, error) {
 	// the new segment file will be created in the temporary folder
 	// and it's always 00000000
 	newSegmentFile := prometheusWAL.SegmentName(tmpFolder, 0)
@@ -284,11 +299,17 @@ func findRootHashAndCreateTrimmed(dir string, segment int, expectedRoot ledger.R
 	return "", fmt.Errorf("finish reading all segment files from %d to %d, but not found", segment, segment)
 }
 
-func checkFolderNotExistOrEmpty(folderPath string) (bool, error) {
+func checkFolderIsEmpty(folderPath string) (bool, error) {
 	// Check if the folder exists
 	info, err := os.Stat(folderPath)
 	if err != nil {
 		if os.IsNotExist(err) {
+			// create the folder if not exist
+			err = os.MkdirAll(folderPath, os.ModePerm)
+			if err != nil {
+				return false, fmt.Errorf("Cannot create the folder.")
+			}
+
 			return true, nil
 		}
 		return false, err
@@ -311,22 +332,6 @@ func checkFolderNotExistOrEmpty(folderPath string) (bool, error) {
 // backup new wals before replacing
 func backupRollbackedWALsAndMoveLastSegmentFile(
 	segment int, walDir, backupDir string, newSegmentFile string) error {
-	// making sure the backup dir is empty
-	empty, err := checkFolderNotExistOrEmpty(backupDir)
-	if err != nil {
-		return fmt.Errorf("cannot check backup directory: %w", err)
-	}
-
-	if !empty {
-		return fmt.Errorf("backup directory %s is not empty", backupDir)
-	}
-
-	// Create the backup directory
-	err = os.MkdirAll(backupDir, os.ModePerm)
-	if err != nil {
-		return fmt.Errorf("cannot create backup directory: %w", err)
-	}
-
 	first, last, err := prometheusWAL.Segments(walDir)
 	if err != nil {
 		return fmt.Errorf("cannot get segments: %w", err)

From daaaa5cc99420626c8a9da5cba2ab2dcffd4b86b Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Tue, 18 Jun 2024 11:27:13 -0700
Subject: [PATCH 18/18] log folder creation

---
 cmd/util/cmd/find-trie-root/cmd.go | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cmd/util/cmd/find-trie-root/cmd.go b/cmd/util/cmd/find-trie-root/cmd.go
index 0fa558dc39a..d50b4d9b3dd 100644
--- a/cmd/util/cmd/find-trie-root/cmd.go
+++ b/cmd/util/cmd/find-trie-root/cmd.go
@@ -304,6 +304,8 @@ func checkFolderIsEmpty(folderPath string) (bool, error) {
 	info, err := os.Stat(folderPath)
 	if err != nil {
 		if os.IsNotExist(err) {
+			log.Info().Msgf("folder %v does not exist, creating the folder", folderPath)
+
 			// create the folder if not exist
 			err = os.MkdirAll(folderPath, os.ModePerm)
 			if err != nil {