Merge pull request #185 from multiversx/pprof-and-shutdown

Pprof and shutdown

Author: raduchis

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) {

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)