Skip to content

Commit

Permalink
add cpu-profile flag (#2723)
Browse files Browse the repository at this point in the history
  • Loading branch information
blotus authored Jan 16, 2024
1 parent 08794c5 commit 421ef3b
Show file tree
Hide file tree
Showing 4 changed files with 34 additions and 1 deletion.
20 changes: 19 additions & 1 deletion cmd/crowdsec/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
_ "net/http/pprof"
"os"
"runtime"
"runtime/pprof"
"strings"
"time"

Expand Down Expand Up @@ -71,6 +72,7 @@ type Flags struct {
DisableCAPI bool
Transform string
OrderEvent bool
CpuProfile string
}

type labelsMap map[string]string
Expand Down Expand Up @@ -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()
}

Expand Down Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions cmd/crowdsec/run_in_svc.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ package main

import (
"fmt"
"runtime/pprof"

log "github.com/sirupsen/logrus"

Expand All @@ -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
}
Expand Down
5 changes: 5 additions & 0 deletions cmd/crowdsec/run_in_svc_windows.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package main

import (
"fmt"
"runtime/pprof"

log "github.com/sirupsen/logrus"
"golang.org/x/sys/windows/svc"
Expand All @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions cmd/crowdsec/serve.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"
"os"
"os/signal"
"runtime/pprof"
"syscall"
"time"

Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 421ef3b

Please sign in to comment.