Skip to content

Commit

Permalink
fix: close CPU profile
Browse files Browse the repository at this point in the history
Signed-off-by: guoguangwu <[email protected]>
  • Loading branch information
testwill authored and developStorm committed May 22, 2024
1 parent 990c9c3 commit 6106c1a
Showing 1 changed file with 10 additions and 4 deletions.
14 changes: 10 additions & 4 deletions bin/bin.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ func dumpHeapProfile() {
// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), start tracking
// CPU profiling in the configured file. Caller is responsible for invoking
// stopCPUProfile() when finished.
func startCPUProfile() {
func startCPUProfile() *os.File {
if file := getCPUProfileFile(); file != "" {
now := time.Now()
fullFile := getFormattedFile(file, now)
Expand All @@ -72,24 +72,30 @@ func startCPUProfile() {
if err := pprof.StartCPUProfile(f); err != nil {
log.Fatal("could not start CPU profile: ", err)
}
return f
}

return nil
}

// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), stop profiling
// CPU usage.
func stopCPUProfile() {
func stopCPUProfile(f *os.File) {
if getCPUProfileFile() != "" {
pprof.StopCPUProfile()
}
if f != nil {
f.Close()
}
}

// ZGrab2Main should be called by func main() in a binary. The caller is
// responsible for importing any modules in use. This allows clients to easily
// include custom sets of scan modules by creating new main packages with custom
// sets of ZGrab modules imported with side-effects.
func ZGrab2Main() {
startCPUProfile()
defer stopCPUProfile()
f := startCPUProfile()
defer stopCPUProfile(f)
defer dumpHeapProfile()
_, moduleType, flag, err := zgrab2.ParseCommandLine(os.Args[1:])

Expand Down

0 comments on commit 6106c1a

Please sign in to comment.