From 45db8f5863f960e5def39299cce198bb586a62c8 Mon Sep 17 00:00:00 2001
From: Pulkit Kathuria <kevincobain2000@gmail.com>
Date: Sat, 10 Aug 2024 15:46:38 +0900
Subject: [PATCH] (feat) mem limit 100 MB (#22)

---
 main.go       | 17 ++++++++++-------
 pkg/card.go   | 25 +++++++++++++++++++++++++
 pkg/flags.go  |  1 +
 pkg/system.go | 36 ++++++++++++++++++++++++++++++++----
 4 files changed, 68 insertions(+), 11 deletions(-)

diff --git a/main.go b/main.go
index 29ce242..13804de 100644
--- a/main.go
+++ b/main.go
@@ -60,7 +60,7 @@ func main() {
 		watch(filePath)
 	}
 	if f.Every > 0 {
-		if err := gocron.Every(f.Every).Second().Do(pkg.PrintMemUsage); err != nil {
+		if err := gocron.Every(1).Second().Do(pkg.PrintMemUsage, &f); err != nil {
 			slog.Error("Error scheduling memory usage", "error", err.Error())
 			return
 		}
@@ -115,9 +115,6 @@ func syncFilePaths() {
 }
 
 func sendHealthCheck() {
-	if f.MSTeamsHook == "" {
-		return
-	}
 	details := pkg.GetHealthCheckDetails(&f, version)
 	for idx, filePath := range filePaths {
 		details = append(details, gmt.Details{
@@ -125,6 +122,10 @@ func sendHealthCheck() {
 			Message: filePath,
 		})
 	}
+	slog.Info("Sending Health Check Notify", "details", details)
+	if f.MSTeamsHook == "" {
+		return
+	}
 
 	hostname, _ := os.Hostname()
 
@@ -181,13 +182,14 @@ func watch(filePath string) {
 }
 
 func notify(result *pkg.ScanResult) {
+	slog.Info("Sending to MS Teams")
+	details := pkg.GetAlertDetails(&f, result)
+	slog.Info("Sending Alert Notify", "details", details)
+
 	if f.MSTeamsHook == "" {
 		return
 	}
 
-	slog.Info("Sending to MS Teams")
-	details := pkg.GetAlertDetails(&f, result)
-
 	hostname, _ := os.Hostname()
 
 	err := gmt.Send(hostname, details, f.MSTeamsHook, f.Proxy)
@@ -208,6 +210,7 @@ func flags() {
 	flag.Uint64Var(&f.Every, "every", 0, "run every n seconds (0 to run once)")
 	flag.Uint64Var(&f.HealthCheckEvery, "health-check-every", 0, "run health check every n seconds (0 to disable)")
 	flag.IntVar(&f.LogLevel, "log-level", 0, "log level (0=info, 1=debug)")
+	flag.IntVar(&f.MemLimit, "mem-limit", 100, "memory limit in MB (0 to disable)")
 	flag.IntVar(&f.FilePathsCap, "file-paths-cap", 100, "max number of file paths to watch")
 	flag.IntVar(&f.Min, "min", 1, "on minimum num of matches, it should notify")
 	flag.BoolVar(&f.Version, "version", false, "")
diff --git a/pkg/card.go b/pkg/card.go
index cfee5b0..c87d774 100644
--- a/pkg/card.go
+++ b/pkg/card.go
@@ -2,10 +2,35 @@ package pkg
 
 import (
 	"fmt"
+	"runtime"
 
 	gmt "github.com/kevincobain2000/go-msteams/src"
 )
 
+func GetPanicDetails(f *Flags, m *runtime.MemStats) []gmt.Details {
+	return []gmt.Details{
+		{
+			Label:   "File Path",
+			Message: f.FilePath,
+		},
+		{
+			Label:   "Match Pattern",
+			Message: f.Match,
+		},
+		{
+			Label:   "Ignore Pattern",
+			Message: f.Ignore,
+		},
+		{
+			Label:   "Mem Limit (MB) Exceeded",
+			Message: fmt.Sprintf("%d", f.MemLimit),
+		},
+		{
+			Label:   "Alloc (MB)",
+			Message: fmt.Sprintf("%d", BToMb(m.Alloc)),
+		},
+	}
+}
 func GetHealthCheckDetails(f *Flags, version string) []gmt.Details {
 	return []gmt.Details{
 		{
diff --git a/pkg/flags.go b/pkg/flags.go
index 06da2b8..caa6884 100644
--- a/pkg/flags.go
+++ b/pkg/flags.go
@@ -14,6 +14,7 @@ type Flags struct {
 	HealthCheckEvery uint64
 	Proxy            string
 	LogLevel         int
+	MemLimit         int
 	MSTeamsHook      string
 	Version          bool
 }
diff --git a/pkg/system.go b/pkg/system.go
index e840e46..7bf8410 100644
--- a/pkg/system.go
+++ b/pkg/system.go
@@ -5,6 +5,8 @@ import (
 	"os"
 	"os/exec"
 	"runtime"
+
+	gmt "github.com/kevincobain2000/go-msteams/src"
 )
 
 func SystemProxy() string {
@@ -18,16 +20,25 @@ func SystemProxy() string {
 	return ""
 }
 
-func PrintMemUsage() {
+func PrintMemUsage(f *Flags) {
 	var m runtime.MemStats
 	runtime.ReadMemStats(&m)
 	slog.Info("Memory Usage",
-		"Alloc (MB)", bToMb(m.Alloc),
-		"Sys", bToMb(m.Sys),
+		"Alloc (MB)", BToMb(m.Alloc),
+		"Alloc (Bytes)", m.Alloc,
+		"Sys (MB)", BToMb(m.Sys),
+		"Sys (Bytes)", m.Sys,
+		"NumGC", m.NumGC,
+		"HeapAlloc (MB)", BToMb(m.HeapAlloc),
+		"HeapSys (Bytes)", m.HeapSys,
 	)
+	if f.MemLimit > 0 && m.Alloc > uint64(f.MemLimit*1024*1024) {
+		SendHealthCheck(f, &m)
+		panic("Memory Limit Exceeded")
+	}
 }
 
-func bToMb(b uint64) uint64 {
+func BToMb(b uint64) uint64 {
 	return b / 1024 / 1024
 }
 
@@ -36,3 +47,20 @@ func ExecShell(command string) (string, error) {
 	out, err := cmd.CombinedOutput()
 	return string(out), err
 }
+
+func SendHealthCheck(f *Flags, m *runtime.MemStats) {
+	if f.MSTeamsHook == "" {
+		return
+	}
+	details := GetPanicDetails(f, m)
+	slog.Warn("Sending Panic Notify", "details", details)
+
+	hostname, _ := os.Hostname()
+
+	err := gmt.Send(hostname, details, f.MSTeamsHook, f.Proxy)
+	if err != nil {
+		slog.Error("Error sending to Teams", "error", err.Error())
+	} else {
+		slog.Info("Successfully sent to MS Teams")
+	}
+}