diff --git a/cmd/chainsimulator/flags.go b/cmd/chainsimulator/flags.go index 3b955b5..e3efd18 100644 --- a/cmd/chainsimulator/flags.go +++ b/cmd/chainsimulator/flags.go @@ -147,6 +147,10 @@ var ( Name: "fetch-configs-and-close", Usage: "This flag is used to specify to fetch all configs and close the chain simulator after", } + enableProfiling = cli.BoolFlag{ + Name: "enable-profiling", + Usage: "Boolean option for enabling CPU profiling. If set, CPU profile will be saved to a file.", + } ) func applyFlags(ctx *cli.Context, cfg *config.Config) { diff --git a/cmd/chainsimulator/main.go b/cmd/chainsimulator/main.go index e4205b2..12b9cec 100644 --- a/cmd/chainsimulator/main.go +++ b/cmd/chainsimulator/main.go @@ -3,14 +3,17 @@ package main import ( "errors" "fmt" + "net/http" "os" "os/signal" "runtime/debug" + "runtime/pprof" "strconv" "strings" "syscall" "time" + "github.com/gin-gonic/gin" "github.com/multiversx/mx-chain-core-go/core" "github.com/multiversx/mx-chain-core-go/core/check" "github.com/multiversx/mx-chain-core-go/core/closing" @@ -86,6 +89,7 @@ func main() { skipConfigsDownload, fetchConfigsAndClose, pathWhereToSaveLogs, + enableProfiling, } app.Authors = []cli.Author{ @@ -177,6 +181,23 @@ func startChainSimulator(ctx *cli.Context) error { return err } + // CPU profiling setup - only if enable-profiling flag is set + var profileFile *os.File + profilingEnabled := ctx.GlobalBool(enableProfiling.Name) + if profilingEnabled { + pathLogsSave := ctx.GlobalString(pathWhereToSaveLogs.Name) + profileFile, err = startCPUProfiling(pathLogsSave, startTimeUnix) + if err != nil { + return fmt.Errorf("%w while starting CPU profiling", err) + } + + // Ensure pprof is stopped and file is synced/closed even on early exits + defer func() { + log.Info("stopping CPU profile (defer)") + stopCPUProfiling(profileFile) + }() + } + var alterConfigsError error argsChainSimulator := chainSimulator.ArgsChainSimulator{ BypassTxSignatureCheck: bypassTxsSignature, @@ -277,7 +298,28 @@ func startChainSimulator(ctx *cli.Context) error { return err } - err = endpointsProc.ExtendProxyServer(proxyInstance.GetHttpServer()) + // Create a channel for programmatic shutdown + shutdownChan := make(chan struct{}) + + // Add a shutdown endpoint before extending the proxy server + httpServer := proxyInstance.GetHttpServer() + ginEngine, ok := httpServer.Handler.(*gin.Engine) + if !ok { + return fmt.Errorf("cannot cast httpServer.Handler to gin.Engine") + } + + ginEngine.POST("/simulator/shutdown", func(c *gin.Context) { + log.Info("shutdown requested via HTTP endpoint") + c.JSON(http.StatusOK, gin.H{"message": "shutdown initiated"}) + + // Trigger shutdown in a goroutine to allow the response to be sent + go func() { + time.Sleep(100 * time.Millisecond) + close(shutdownChan) + }() + }) + + err = endpointsProc.ExtendProxyServer(httpServer) if err != nil { return err } @@ -289,9 +331,19 @@ func startChainSimulator(ctx *cli.Context) error { interrupt := make(chan os.Signal, 1) signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM) - <-interrupt - log.Info("close") + // Wait for either signal or programmatic shutdown + select { + case sig := <-interrupt: + log.Info("close", "signal", sig) + case <-shutdownChan: + log.Info("close", "trigger", "HTTP shutdown endpoint") + } + + // Stop CPU profiling FIRST and flush to disk (only if profiling is enabled) + if profilingEnabled { + stopCPUProfiling(profileFile) + } generator.Close() @@ -306,6 +358,43 @@ func startChainSimulator(ctx *cli.Context) error { return nil } +func startCPUProfiling(pathLogsSave string, startTimeUnix int64) (*os.File, error) { + timestampMilisecond := time.Unix(startTimeUnix, 0).UnixNano() / 1000000 + cpuProfilePath := fmt.Sprintf("%s/cpu-%d.pprof", pathLogsSave, timestampMilisecond) + + profileFile, err := os.Create(cpuProfilePath) + if err != nil { + return nil, fmt.Errorf("could not create CPU profile: %w", err) + } + + if err := pprof.StartCPUProfile(profileFile); err != nil { + _ = profileFile.Close() + return nil, fmt.Errorf("could not start CPU profile: %w", err) + } + + log.Info("CPU profiling started", "path", cpuProfilePath) + return profileFile, nil +} + +func stopCPUProfiling(profileFile *os.File) { + if profileFile == nil { + return + } + + log.Info("stopping CPU profile") + pprof.StopCPUProfile() + + if err := profileFile.Sync(); err != nil { + log.Error("error syncing CPU profile file", "err", err) + } + + if err := profileFile.Close(); err != nil { + log.Error("error closing CPU profile file", "err", err) + } else { + log.Info("CPU profile file closed successfully") + } +} + func initializeLogger(ctx *cli.Context, cfg config.Config) (closing.Closer, error) { logLevelFlagValue := ctx.GlobalString(logLevel.Name) err := logger.SetLogLevel(logLevelFlagValue)