Merge pull request #6121 from onflow/bastian/profile-whole-state-extr…

…action Profile whole state extraction, not just migration
onflow · Jun 19, 2024 · 0c0c0d2 · 0c0c0d2
2 parents 3a86c0a + 0ad28cc
commit 0c0c0d2
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/cmd/util/cmd/execution-state-extract/cmd.go b/cmd/util/cmd/execution-state-extract/cmd.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"runtime/pprof"
 	"strings"
 
 	"github.com/rs/zerolog/log"
@@ -178,6 +179,20 @@ func init() {
 }
 
 func run(*cobra.Command, []string) {
+	if flagCPUProfile != "" {
+		f, err := os.Create(flagCPUProfile)
+		if err != nil {
+			log.Fatal().Err(err).Msg("could not create CPU profile")
+		}
+
+		err = pprof.StartCPUProfile(f)
+		if err != nil {
+			log.Fatal().Err(err).Msg("could not start CPU profile")
+		}
+
+		defer pprof.StopCPUProfile()
+	}
+
 	var stateCommitment flow.StateCommitment
 
 	if len(flagBlockHash) > 0 && len(flagStateCommitment) > 0 {

diff --git a/cmd/util/cmd/execution-state-extract/execution_state_extract.go b/cmd/util/cmd/execution-state-extract/execution_state_extract.go
@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"math"
 	"os"
-	"runtime/pprof"
 	syncAtomic "sync/atomic"
 	"time"
 
@@ -325,19 +324,6 @@ func newMigration(
 	nWorker int,
 ) ledger.Migration {
 	return func(payloads []*ledger.Payload) ([]*ledger.Payload, error) {
-		if flagCPUProfile != "" {
-			f, err := os.Create(flagCPUProfile)
-			if err != nil {
-				logger.Fatal().Err(err).Msg("could not create CPU profile")
-			}
-
-			err = pprof.StartCPUProfile(f)
-			if err != nil {
-				logger.Fatal().Err(err).Msg("could not start CPU profile")
-			}
-
-			defer pprof.StopCPUProfile()
-		}
 
 		if len(migrations) == 0 {
 			return payloads, nil