From 14b012486d5c3c7a601b9e8d200ee71042601659 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Fri, 6 Oct 2023 14:50:14 +0200
Subject: [PATCH] Formalize Memory Monitoring

by extracting the interval and threshold into configuration options.

Related to f670b07e.
---
 cmd/poseidon/main.go       | 29 +++++++++++++++++------------
 configuration.example.yaml |  9 +++++++--
 internal/config/config.go  |  9 +++++++--
 3 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/cmd/poseidon/main.go b/cmd/poseidon/main.go
index 59f5e9c..6456686 100644
--- a/cmd/poseidon/main.go
+++ b/cmd/poseidon/main.go
@@ -85,8 +85,8 @@ func shutdownSentry() {
 }
 
 func initProfiling(options config.Profiling) (cancel func()) {
-	if options.Enabled {
-		profile, err := os.Create(options.File)
+	if options.CPUEnabled {
+		profile, err := os.Create(options.CPUFile)
 		if err != nil {
 			log.WithError(err).Error("Error while opening the profile file")
 		}
@@ -97,7 +97,7 @@ func initProfiling(options config.Profiling) (cancel func()) {
 		}
 
 		cancel = func() {
-			if options.Enabled {
+			if options.CPUEnabled {
 				log.Debug("Stopping CPU profiler")
 				pprof.StopCPUProfile()
 				if err := profile.Close(); err != nil {
@@ -112,18 +112,20 @@ func initProfiling(options config.Profiling) (cancel func()) {
 }
 
 // watchMemoryAndAlert monitors the memory usage of Poseidon and sends an alert if it exceeds a threshold.
-func watchMemoryAndAlert() {
-	// We assume that Poseidon usually takes about 50-300 MB of memory. Therefore, we specify the threshold of 1 GB.
-	// Improve: Make this value dynamic or relative.
-	const threshold = 1 * 1000 * 1000 * 1000
-	const interval = 5 * time.Second
+func watchMemoryAndAlert(options config.Profiling) {
+	if options.MemoryInterval == 0 {
+		return
+	}
 
+	var exceeded bool
 	for {
 		var stats runtime.MemStats
 		runtime.ReadMemStats(&stats)
 		log.WithField("heap", stats.HeapAlloc).Trace("Current Memory Usage")
 
-		if stats.HeapAlloc >= threshold {
+		const megabytesToBytes = 1000 * 1000
+		if !exceeded && stats.HeapAlloc >= uint64(options.MemoryThreshold)*megabytesToBytes {
+			exceeded = true
 			log.WithField("heap", stats.HeapAlloc).Warn("Memory Threshold exceeded")
 
 			err := pprof.Lookup("heap").WriteTo(os.Stderr, 1)
@@ -135,10 +137,13 @@ func watchMemoryAndAlert() {
 			if err != nil {
 				log.WithError(err).Warn("Failed to log the goroutines")
 			}
+		} else if exceeded {
+			exceeded = false
+			log.WithField("heap", stats.HeapAlloc).Info("Memory Threshold no longer exceeded")
 		}
 
 		select {
-		case <-time.After(interval):
+		case <-time.After(time.Duration(options.MemoryInterval) * time.Millisecond):
 			continue
 		case <-context.Background().Done():
 			return
@@ -270,13 +275,13 @@ func main() {
 		log.WithError(err).Warn("Could not initialize configuration")
 	}
 	logging.InitializeLogging(config.Config.Logger.Level, config.Config.Logger.Formatter)
-	initSentry(&config.Config.Sentry, config.Config.Profiling.Enabled)
+	initSentry(&config.Config.Sentry, config.Config.Profiling.CPUEnabled)
 
 	cancelInflux := monitoring.InitializeInfluxDB(&config.Config.InfluxDB)
 	defer cancelInflux()
 
 	stopProfiling := initProfiling(config.Config.Profiling)
-	go watchMemoryAndAlert()
+	go watchMemoryAndAlert(config.Config.Profiling)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	server := initServer(ctx)
diff --git a/configuration.example.yaml b/configuration.example.yaml
index dae5473..a66191a 100644
--- a/configuration.example.yaml
+++ b/configuration.example.yaml
@@ -64,10 +64,15 @@ logger:
 # Configuration of the embedded profiler
 profiling:
   # Enables the runtime profiler
-  enabled: false
+  cpuenabled: false
   # The file to which the profile is written to.
   # The default location `cmd/poseidon/default.pgo` will be picked up during the build process to create a profile-guided build.
-  file: cmd/poseidon/default.pgo
+  cpufile: cmd/poseidon/default.pgo
+  # If set, a memory watchdog will be started that monitors the memory usage of Poseidon and alerts if the threshold is exceeded.
+  # The value defines the interval in milliseconds in which the memory usage is checked.
+  memoryinterval: 30_000
+  # The Threshold in MB of memory usage at which Poseidon will start alerting.
+  memorythreshold: 1_000
 
 # Configuration of the sentry logging
 sentry:
diff --git a/internal/config/config.go b/internal/config/config.go
index 14ef176..eb9abbd 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -55,6 +55,9 @@ var (
 			Level:     "INFO",
 			Formatter: dto.FormatterText,
 		},
+		Profiling: Profiling{
+			MemoryThreshold: 1_000,
+		},
 		Sentry: sentry.ClientOptions{
 			AttachStacktrace: true,
 		},
@@ -130,8 +133,10 @@ type Logger struct {
 
 // Profiling configures the usage of a runtime profiler to create optimized binaries.
 type Profiling struct {
-	Enabled bool
-	File    string
+	CPUEnabled      bool
+	CPUFile         string
+	MemoryInterval  uint
+	MemoryThreshold uint
 }
 
 // InfluxDB configures the usage of an Influx db monitoring.