Add command line switches for profiling

flags.go

@@ -18,6 +18,9 @@ var connPortVar string
 var modeVar string
 var gatewayVar string
 
+var blockProfilingVar bool
+var cpuProfilingVar bool
+
 func printUsage() {
 	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
 	fmt.Fprintf(os.Stderr, "%s <server address>\n", os.Args[0])
@@ -33,6 +36,8 @@ func parseFlags() {
 	flag.StringVar(&modeVar, "mode", "client", "Whether the process starts a server or as a client")
 	flag.StringVar(&networkAddrVar, "network", "192.168.69.1/24", "Address for this interface with netmask")
 	flag.StringVar(&gatewayVar, "gw", "", "(Client only) Set the default gateway to this value")
+	flag.BoolVar(&blockProfilingVar, "blockProfile", false, "Enable block profiling")
+	flag.BoolVar(&cpuProfilingVar, "cpuProfile", false, "Enable CPU profiling")
 
 	flag.Usage = printUsage
 	flag.Parse()
@@ -50,5 +55,10 @@ func parseFlags() {
 		}
 	}
 
+	if cpuProfilingVar && blockProfilingVar {
+		fmt.Fprintf(os.Stderr, "Err: Cannot enable both block and CPU profiling at once.\n")
+		os.Exit(2)
+	}
+
 	serverAddressVar = flag.Arg(0)
 }

main.go

@@ -5,6 +5,8 @@ import (
 	"log"
 	"os"
 	"os/signal"
+	"runtime"
+	"runtime/pprof"
 	"subnet"
 	"syscall"
 )
@@ -13,6 +15,18 @@ func main() {
 	parseFlags()
 	fatalErrChan := make(chan error)
 
+	if cpuProfilingVar {
+		cpuFile := startCPUProfile()
+		defer cpuFile.Close()
+		defer pprof.StopCPUProfile()
+	}
+
+	if blockProfilingVar {
+		blockFile := startBlockProfile()
+		defer blockFile.Close()
+		defer pprof.Lookup("block").WriteTo(blockFile, 1)
+	}
+
 	switch modeVar {
 	case "client":
 		c, err := subnet.NewClient(serverAddressVar, connPortVar, networkAddrVar, interfaceNameVar, gatewayVar, ourCertPathVar, ourKeyPathVar, caCertPathVar)
@@ -57,3 +71,23 @@ func waitInterrupt(fatalErrChan chan error) {
 		log.Println("Fatal internal error: ", err)
 	}
 }
+
+func startCPUProfile() *os.File {
+	f, err := os.Create("cpu.prof")
+	if err != nil {
+		log.Fatal("could not create CPU profile: ", err)
+	}
+	if err := pprof.StartCPUProfile(f); err != nil {
+		log.Fatal("could not start CPU profile: ", err)
+	}
+	return f
+}
+
+func startBlockProfile() *os.File {
+	f, err := os.Create("block.prof")
+	if err != nil {
+		log.Fatal("could not create block profile: ", err)
+	}
+	runtime.SetBlockProfileRate(1)
+	return f
+}