Merge pull request #69 from safing/feature/cpu-profiling

Add cpu profiling capability to base module

core/base/module.go

@@ -8,8 +8,12 @@ import (
 	_ "github.com/safing/portbase/rng"
 )
 
+var (
+	module *modules.Module
+)
+
 func init() {
-	modules.Register("base", nil, registerDatabases, nil, "database", "config", "rng")
+	module = modules.Register("base", nil, start, nil, "database", "config", "rng")
 
 	// For prettier subsystem graph, printed with --print-subsystem-graph
 	/*
@@ -23,3 +27,9 @@ func init() {
 		)
 	*/
 }
+
+func start() error {
+	startProfiling()
+
+	return registerDatabases()
+}

core/base/profiling.go

@@ -0,0 +1,43 @@
+package base
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"os"
+	"runtime/pprof"
+)
+
+var (
+	cpuProfile string
+)
+
+func init() {
+	flag.StringVar(&cpuProfile, "cpuprofile", "", "write cpu profile to `file`")
+}
+
+func startProfiling() {
+	if cpuProfile != "" {
+		module.StartWorker("cpu profiler", cpuProfiler)
+	}
+}
+
+func cpuProfiler(ctx context.Context) error {
+	f, err := os.Create(cpuProfile)
+	if err != nil {
+		return fmt.Errorf("could not create CPU profile: %s", err)
+	}
+	if err := pprof.StartCPUProfile(f); err != nil {
+		return fmt.Errorf("could not start CPU profile: %s", err)
+	}
+
+	// wait for shutdown
+	<-ctx.Done()
+
+	pprof.StopCPUProfile()
+	err = f.Close()
+	if err != nil {
+		return fmt.Errorf("failed to close CPU profile file: %s", err)
+	}
+	return nil
+}