From 381b94acb55a9318f9e5e3a7e13b49a14028942e Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Wed, 3 Jun 2020 23:11:50 +0200
Subject: [PATCH] Add cpu profiling capability to base module

---
 core/base/module.go    | 12 +++++++++++-
 core/base/profiling.go | 43 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 core/base/profiling.go

diff --git a/core/base/module.go b/core/base/module.go
index f3fdc92d..4c548936 100644
--- a/core/base/module.go
+++ b/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()
+}
diff --git a/core/base/profiling.go b/core/base/profiling.go
new file mode 100644
index 00000000..bf96e9f5
--- /dev/null
+++ b/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
+}