diff --git a/extractor/cli/go-extractor/go-extractor.go b/extractor/cli/go-extractor/go-extractor.go index bc0e7b86c15..483f734b8dc 100644 --- a/extractor/cli/go-extractor/go-extractor.go +++ b/extractor/cli/go-extractor/go-extractor.go @@ -4,6 +4,8 @@ import ( "fmt" "log" "os" + "runtime" + "runtime/pprof" "strings" "github.com/github/codeql-go/extractor/dbscheme" @@ -11,6 +13,8 @@ import ( "github.com/github/codeql-go/extractor" ) +var cpuprofile, memprofile string + func usage() { fmt.Fprintf(os.Stderr, "%s is a program for building a snapshot of a Go code base.\n\n", os.Args[0]) fmt.Fprintf(os.Stderr, "Usage:\n\n %s [...] [...] [--] ...\n\n", os.Args[0]) @@ -25,14 +29,14 @@ func parseFlags(args []string) ([]string, []string, string) { buildFlags := []string{} for i < len(args) && strings.HasPrefix(args[i], "-") { if args[i] == "--" { - i += 1 + i++ break } if strings.HasPrefix(args[i], "--dbscheme=") { dumpDbscheme = strings.TrimPrefix(args[i], "--dbscheme=") } else if args[i] == "--dbscheme" { - i += 1 + i++ dumpDbscheme = args[i] } else if args[i] == "--help" { usage() @@ -41,9 +45,12 @@ func parseFlags(args []string) ([]string, []string, string) { buildFlags = append(buildFlags, args[i]) } - i += 1 + i++ } + cpuprofile = os.Getenv("CODEQL_EXTRACTOR_GO_CPU_PROFILE") + memprofile = os.Getenv("CODEQL_EXTRACTOR_GO_MEM_PROFILE") + return buildFlags, args[i:], dumpDbscheme } @@ -60,6 +67,18 @@ func main() { log.Printf("Dbscheme written to file %s.", dumpDbscheme) } + if cpuprofile != "" { + f, err := os.Create(cpuprofile) + if err != nil { + log.Fatalf("Unable to create CPU profile: %v.", err) + } + defer f.Close() + if err := pprof.StartCPUProfile(f); err != nil { + log.Fatalf("Unable to start CPU profile: %v.", err) + } + defer pprof.StopCPUProfile() + } + if len(patterns) == 0 { log.Println("Nothing to extract.") } else { @@ -68,4 +87,16 @@ func main() { log.Fatal(err) } } + + if memprofile != "" { + f, err := os.Create(memprofile) + if err != nil { + log.Fatalf("Unable to create memory profile: %v", err) + } + defer f.Close() + runtime.GC() // get up-to-date statistics + if err := pprof.WriteHeapProfile(f); err != nil { + log.Fatal("Unable to write memory profile: ", err) + } + } }