diff --git a/main.go b/main.go index 13514c1..c4d35b4 100644 --- a/main.go +++ b/main.go @@ -3,11 +3,8 @@ package main import ( "fmt" "os" - "path/filepath" "runtime" - "runtime/pprof" - "runtime/trace" "strconv" "strings" "time" @@ -208,49 +205,18 @@ func main() { } // "-cpuprofile" if args.cpuprofile != "" { - tlog.Info.Printf("Writing CPU profile to %s", args.cpuprofile) - var f *os.File - f, err = os.Create(args.cpuprofile) - if err != nil { - tlog.Fatal.Println(err) - os.Exit(exitcodes.Profiler) - } - err = pprof.StartCPUProfile(f) - if err != nil { - tlog.Fatal.Println(err) - os.Exit(exitcodes.Profiler) - } - defer pprof.StopCPUProfile() + onExitFunc := setupCpuprofile(args.cpuprofile) + defer onExitFunc() } // "-memprofile" if args.memprofile != "" { - tlog.Info.Printf("Writing mem profile to %s", args.memprofile) - var f *os.File - f, err = os.Create(args.memprofile) - if err != nil { - tlog.Fatal.Println(err) - os.Exit(exitcodes.Profiler) - } - defer func() { - pprof.WriteHeapProfile(f) - f.Close() - return - }() + onExitFunc := setupMemprofile(args.memprofile) + defer onExitFunc() } // "-trace" if args.trace != "" { - tlog.Info.Printf("Writing execution trace to %s", args.trace) - f, err := os.Create(args.trace) - if err != nil { - tlog.Fatal.Println(err) - os.Exit(exitcodes.Profiler) - } - err = trace.Start(f) - if err != nil { - tlog.Fatal.Println(err) - os.Exit(exitcodes.Profiler) - } - defer trace.Stop() + onExitFunc := setupTrace(args.trace) + defer onExitFunc() } if args.cpuprofile != "" || args.memprofile != "" || args.trace != "" { tlog.Info.Printf("Note: You must unmount gracefully, otherwise the profile file(s) will stay empty!\n") diff --git a/profiling.go b/profiling.go new file mode 100644 index 0000000..11d6326 --- /dev/null +++ b/profiling.go @@ -0,0 +1,95 @@ +package main + +import ( + "os" + "runtime/pprof" + "runtime/trace" + "time" + + "github.com/rfjakob/gocryptfs/internal/exitcodes" + "github.com/rfjakob/gocryptfs/internal/tlog" +) + +// setupCpuprofile is called to handle a non-empty "-cpuprofile" cli argument +func setupCpuprofile(cpuprofileArg string) func() { + tlog.Info.Printf("Writing CPU profile to %s", cpuprofileArg) + f, err := os.Create(cpuprofileArg) + if err != nil { + tlog.Fatal.Println(err) + os.Exit(exitcodes.Profiler) + } + err = pprof.StartCPUProfile(f) + if err != nil { + tlog.Fatal.Println(err) + os.Exit(exitcodes.Profiler) + } + return func() { + pprof.StopCPUProfile() + } +} + +// setupTrace is called to handle a non-empty "-memprofile" cli argument +func setupMemprofile(memprofileArg string) func() { + tlog.Info.Printf("Will write memory profile to %q", memprofileArg) + f, err := os.Create(memprofileArg) + if err != nil { + tlog.Fatal.Println(err) + os.Exit(exitcodes.Profiler) + } + exiting := false + // Write the memory profile to disk every 60 seconds to get the in-use + // memory stats. + go func() { + for { + time.Sleep(60 * time.Second) + if exiting { + return + } + _, err = f.Seek(0, 0) + if err != nil { + tlog.Warn.Printf("memprofile: Seek failed: %v", err) + return + } + err = f.Truncate(0) + if err != nil { + tlog.Warn.Printf("memprofile: Truncate failed: %v", err) + return + } + err = pprof.WriteHeapProfile(f) + if err == nil { + tlog.Info.Printf("memprofile: periodic write to %q succeeded", + memprofileArg) + } else { + tlog.Warn.Printf("memprofile: periodic WriteHeapProfile failed: %v", err) + return + } + } + }() + // Final write on exit. + return func() { + exiting = true + err = pprof.WriteHeapProfile(f) + if err != nil { + tlog.Warn.Printf("memprofile: on-exit WriteHeapProfile failed: %v", err) + } + f.Close() + } +} + +// setupTrace is called to handle a non-empty "-trace" cli argument +func setupTrace(traceArg string) func() { + tlog.Info.Printf("Writing execution trace to %s", traceArg) + f, err := os.Create(traceArg) + if err != nil { + tlog.Fatal.Println(err) + os.Exit(exitcodes.Profiler) + } + err = trace.Start(f) + if err != nil { + tlog.Fatal.Println(err) + os.Exit(exitcodes.Profiler) + } + return func() { + trace.Stop() + } +}