Add environment variables for dumping CPU and memory profiles.

These are intentionally undocumented for now.

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 [<flag>...] [<buildflag>...] [--] <file>...\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)
+		}
+	}
 }