From ae3ab409100708520afc6235f5f9c8f56e8e678c Mon Sep 17 00:00:00 2001
From: Sune Kirkeby <sune.kirkeby@vivino.com>
Date: Wed, 9 Nov 2022 11:08:34 +0100
Subject: [PATCH] Implement memory-profiler command-line flag. (#24)

Add `-memory-profiler` command-line flag to enable pprof memory-dumps at peak usage.
---
 api/rankdb.go                  |   2 +-
 cmd/rankdb/main.go             |  10 ++-
 memprofiler/memory_profiler.go | 112 +++++++++++++++++++++++++++++++++
 3 files changed, 121 insertions(+), 3 deletions(-)
 create mode 100644 memprofiler/memory_profiler.go

diff --git a/api/rankdb.go b/api/rankdb.go
index 08ec89b..4ae7974 100644
--- a/api/rankdb.go
+++ b/api/rankdb.go
@@ -26,7 +26,7 @@ import (
 	"github.com/goadesign/goa"
 	lru "github.com/hashicorp/golang-lru"
 	shutdown "github.com/klauspost/shutdown2"
-	"github.com/mattn/go-colorable"
+	colorable "github.com/mattn/go-colorable"
 	"github.com/sirupsen/logrus"
 )
 
diff --git a/cmd/rankdb/main.go b/cmd/rankdb/main.go
index 967c436..8b3dea4 100644
--- a/cmd/rankdb/main.go
+++ b/cmd/rankdb/main.go
@@ -18,14 +18,16 @@ import (
 	"github.com/Vivino/rankdb/api"
 	"github.com/Vivino/rankdb/log"
 	"github.com/Vivino/rankdb/log/loggoa"
+	"github.com/Vivino/rankdb/memprofiler"
 	goalogrus "github.com/goadesign/goa/logging/logrus"
 	shutdown "github.com/klauspost/shutdown2"
 	"github.com/sirupsen/logrus"
 )
 
 var (
-	configPath  = flag.String("config", "./conf/conf.toml", "Path for config to use.")
-	enableDebug = flag.Bool("restart", false, "Enable rapid restart mode, press ' and <return>.")
+	configPath           = flag.String("config", "./conf/conf.toml", "Path for config to use.")
+	enableDebug          = flag.Bool("restart", false, "Enable rapid restart mode, press ' and <return>.")
+	enableMemoryProfiler = flag.Bool("memory-profiler", false, "Enable memory profiler")
 
 	// SIGUSR2 signal if available.
 	usr2Signal os.Signal
@@ -58,6 +60,10 @@ func main() {
 		})
 	})
 
+	if *enableMemoryProfiler {
+		go memprofiler.Run(ctx, "/var/tmp/rankdb/memory-dumps")
+	}
+
 	if *enableDebug {
 		go func() {
 			for {
diff --git a/memprofiler/memory_profiler.go b/memprofiler/memory_profiler.go
new file mode 100644
index 0000000..c675f64
--- /dev/null
+++ b/memprofiler/memory_profiler.go
@@ -0,0 +1,112 @@
+package memprofiler
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"runtime"
+	"runtime/pprof"
+	"strings"
+	"time"
+
+	"github.com/Vivino/rankdb/log"
+	shutdown "github.com/klauspost/shutdown2"
+)
+
+const dumpTimeLayout = "2006-01-02-150405"
+
+// Run monitors the memory-usage of your application, and dumps profiles at peak usage.
+func Run(ctx context.Context, path string) {
+	// Don't dump in the quiet period while starting up.
+	const (
+		quietPeriod = 3 * time.Minute
+		checkEvery  = 10 * time.Second
+		minRise     = 100 << 20
+	)
+
+	stats := runtime.MemStats{}
+	var high uint64
+	if path == "" {
+		return
+	}
+	ctx, cancel := shutdown.CancelCtx(log.WithValues(ctx, "module", "RAM Monitor"))
+	defer cancel()
+
+	err := os.MkdirAll(path, os.ModePerm)
+	if err != nil {
+		log.Error(ctx, "Unable to create memory dump folder")
+	}
+
+	t := time.NewTicker(checkEvery)
+	defer t.Stop()
+	go cleanOld(ctx, path)
+
+	quietEnds := time.Now().Add(quietPeriod)
+	for {
+		select {
+		case <-t.C:
+			runtime.ReadMemStats(&stats)
+			if stats.Alloc <= high+(minRise) {
+				continue
+			}
+			high = stats.Alloc
+			if time.Now().Before(quietEnds) {
+				log.Info(
+					log.WithValues(ctx, "alloc_mb", stats.Alloc>>20),
+					"In quiet period, skipping dump")
+				continue
+			}
+
+			timeStamp := time.Now().Format(dumpTimeLayout)
+			fn := filepath.Join(path, fmt.Sprintf("%s-%dMB.bin", timeStamp, high>>20))
+
+			log.Info(
+				log.WithValues(ctx, "filename", fn, "alloc_mb", stats.Alloc>>20),
+				"Memory peak, dumping memory profile")
+			f, err := os.Create(fn)
+			if err != nil {
+				log.Error(
+					log.WithValues(ctx, "error", err),
+					"could not create memory profile")
+			}
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				log.Error(
+					log.WithValues(ctx, "error", err),
+					"could not write memory profile")
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+// cleanOld will remove all dumps more than 30 days old.
+func cleanOld(ctx context.Context, path string) {
+	t := time.NewTicker(time.Hour)
+	defer t.Stop()
+	for {
+		select {
+		case <-t.C:
+			timeStamp := time.Now().Add(-31 * 24 * time.Hour).Format(dumpTimeLayout)
+			log.Info(ctx, "Deleting old dumps containing %q", timeStamp)
+
+			err := filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
+				if err == nil || info == nil {
+					return err
+				}
+				if !info.IsDir() && strings.Contains(info.Name(), timeStamp) {
+					return os.Remove(info.Name())
+				}
+				return nil
+			})
+			if err != nil {
+				log.Error(
+					log.WithValues(ctx, "error", err),
+					"error deleting old dumps")
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
+}