From 421ef3bf9cbb5303eb11c0115576d4a746119823 Mon Sep 17 00:00:00 2001
From: blotus <sebastien@crowdsec.net>
Date: Tue, 16 Jan 2024 11:40:29 +0100
Subject: [PATCH] add cpu-profile flag (#2723)

---
 cmd/crowdsec/main.go               | 20 +++++++++++++++++++-
 cmd/crowdsec/run_in_svc.go         |  5 +++++
 cmd/crowdsec/run_in_svc_windows.go |  5 +++++
 cmd/crowdsec/serve.go              |  5 +++++
 4 files changed, 34 insertions(+), 1 deletion(-)

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: