From 6106c1ad2186502a73e5b4e9acae5ffb03cf546d Mon Sep 17 00:00:00 2001 From: guoguangwu Date: Wed, 22 May 2024 14:04:04 +0800 Subject: [PATCH] fix: close CPU profile Signed-off-by: guoguangwu --- bin/bin.go | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/bin/bin.go b/bin/bin.go index dc4904612..50f7a055f 100644 --- a/bin/bin.go +++ b/bin/bin.go @@ -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) @@ -72,15 +72,21 @@ 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 @@ -88,8 +94,8 @@ func stopCPUProfile() { // 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:])