diff --git a/cmd/crowdsec/main.go b/cmd/crowdsec/main.go index 05a6db01d..bdb04023e 100644 --- a/cmd/crowdsec/main.go +++ b/cmd/crowdsec/main.go @@ -6,6 +6,7 @@ import ( _ "net/http/pprof" "os" "runtime" + "runtime/pprof" "strings" "time" @@ -71,6 +72,7 @@ type Flags struct { DisableCAPI bool Transform string OrderEvent bool + CpuProfile string } type labelsMap map[string]string @@ -179,6 +181,7 @@ func (f *Flags) Parse() { } flag.StringVar(&dumpFolder, "dump-data", "", "dump parsers/buckets raw outputs") + flag.StringVar(&f.CpuProfile, "cpu-profile", "", "write cpu profile to file") flag.Parse() } @@ -352,9 +355,24 @@ func main() { os.Exit(0) } + if flags.CpuProfile != "" { + f, err := os.Create(flags.CpuProfile) + if err != nil { + log.Fatalf("could not create CPU profile: %s", err) + } + log.Infof("CPU profile will be written to %s", flags.CpuProfile) + if err := pprof.StartCPUProfile(f); err != nil { + f.Close() + log.Fatalf("could not start CPU profile: %s", err) + } + defer f.Close() + defer pprof.StopCPUProfile() + } + err := StartRunSvc() if err != nil { - log.Fatal(err) + pprof.StopCPUProfile() + log.Fatal(err) //nolint:gocritic // Disable warning for the defer pprof.StopCPUProfile() call } os.Exit(0) diff --git a/cmd/crowdsec/run_in_svc.go b/cmd/crowdsec/run_in_svc.go index 8b2cb2aee..202053790 100644 --- a/cmd/crowdsec/run_in_svc.go +++ b/cmd/crowdsec/run_in_svc.go @@ -4,6 +4,7 @@ package main import ( "fmt" + "runtime/pprof" log "github.com/sirupsen/logrus" @@ -22,6 +23,10 @@ func StartRunSvc() error { defer trace.CatchPanic("crowdsec/StartRunSvc") + //Always try to stop CPU profiling to avoid passing flags around + //It's a noop if profiling is not enabled + defer pprof.StopCPUProfile() + if cConfig, err = LoadConfig(flags.ConfigFile, flags.DisableAgent, flags.DisableAPI, false); err != nil { return err } diff --git a/cmd/crowdsec/run_in_svc_windows.go b/cmd/crowdsec/run_in_svc_windows.go index 6ce6601b1..991f7ae44 100644 --- a/cmd/crowdsec/run_in_svc_windows.go +++ b/cmd/crowdsec/run_in_svc_windows.go @@ -2,6 +2,7 @@ package main import ( "fmt" + "runtime/pprof" log "github.com/sirupsen/logrus" "golang.org/x/sys/windows/svc" @@ -19,6 +20,10 @@ func StartRunSvc() error { defer trace.CatchPanic("crowdsec/StartRunSvc") + //Always try to stop CPU profiling to avoid passing flags around + //It's a noop if profiling is not enabled + defer pprof.StopCPUProfile() + isRunninginService, err := svc.IsWindowsService() if err != nil { return fmt.Errorf("failed to determine if we are running in windows service mode: %w", err) diff --git a/cmd/crowdsec/serve.go b/cmd/crowdsec/serve.go index 9ec00d9be..a5c8e24cf 100644 --- a/cmd/crowdsec/serve.go +++ b/cmd/crowdsec/serve.go @@ -4,6 +4,7 @@ import ( "fmt" "os" "os/signal" + "runtime/pprof" "syscall" "time" @@ -245,6 +246,10 @@ func HandleSignals(cConfig *csconfig.Config) error { exitChan := make(chan error) + //Always try to stop CPU profiling to avoid passing flags around + //It's a noop if profiling is not enabled + defer pprof.StopCPUProfile() + go func() { defer trace.CatchPanic("crowdsec/HandleSignals") Loop: