From bb41a4abfa544348ada25875948f915a5d257bc9 Mon Sep 17 00:00:00 2001
From: crapStone <crapstone01@gmail.com>
Date: Fri, 15 Mar 2024 20:35:39 +0100
Subject: [PATCH] use /proc/self/statm to get current usage

---
 cli/flags.go           |  4 ++--
 server/cache/memory.go | 42 ++++++++++++++++++++++++++++++++----------
 server/startup.go      |  2 +-
 3 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/cli/flags.go b/cli/flags.go
index 09a6404..4fc4b64 100644
--- a/cli/flags.go
+++ b/cli/flags.go
@@ -140,8 +140,8 @@ var (
 		},
 		&cli.Uint64Flag{
 			Name:    "memory-limit",
-			Usage:   "maximum size of memory in bytes to use for caching, default: 512MB",
-			Value:   512 * 1024 * 1024,
+			Usage:   "maximum size of memory in KiB to use for caching, default: 512MiB",
+			Value:   512 * 1024,
 			EnvVars: []string{"MAX_MEMORY_SIZE"},
 		},
 
diff --git a/server/cache/memory.go b/server/cache/memory.go
index 743b6d0..2a9fba1 100644
--- a/server/cache/memory.go
+++ b/server/cache/memory.go
@@ -1,7 +1,9 @@
 package cache
 
 import (
-	"runtime"
+	"os"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/OrlovEvgeny/go-mcache"
@@ -10,7 +12,7 @@ import (
 
 type Cache struct {
 	mcache      *mcache.CacheDriver
-	memoryLimit uint64
+	memoryLimit int
 	lastCheck   time.Time
 }
 
@@ -21,7 +23,7 @@ func NewInMemoryCache() ICache {
 
 // NewInMemoryCache returns a new mcache with a memory limit.
 // If the limit is exceeded, the cache will be cleared.
-func NewInMemoryCacheWithLimit(memoryLimit uint64) ICache {
+func NewInMemoryCacheWithLimit(memoryLimit int) ICache {
 	return &Cache{
 		mcache:      mcache.New(),
 		memoryLimit: memoryLimit,
@@ -31,11 +33,10 @@ func NewInMemoryCacheWithLimit(memoryLimit uint64) ICache {
 func (c *Cache) Set(key string, value interface{}, ttl time.Duration) error {
 	now := time.Now()
 
-	// checking memory limit is a "stop the world" operation
-	// so we don't want to do it too often
+	// we don't want to do it too often
 	if now.Sub(c.lastCheck) > (time.Second * 3) {
 		if c.memoryLimitOvershot() {
-			log.Debug().Msg("memory limit exceeded, clearing cache")
+			log.Info().Msg("[cache] memory limit exceeded, clearing cache")
 			c.mcache.Truncate()
 		}
 		c.lastCheck = now
@@ -53,10 +54,31 @@ func (c *Cache) Remove(key string) {
 }
 
 func (c *Cache) memoryLimitOvershot() bool {
-	var stats runtime.MemStats
-	runtime.ReadMemStats(&stats)
+	mem := getCurrentMemory()
 
-	log.Debug().Uint64("bytes", stats.HeapAlloc).Msg("current memory usage")
+	log.Debug().Int("kB", mem).Msg("[cache] current memory usage")
 
-	return stats.HeapAlloc > c.memoryLimit
+	return mem > c.memoryLimit
+}
+
+// getCurrentMemory returns the current memory in KB
+func getCurrentMemory() int {
+	b, err := os.ReadFile("/proc/self/statm")
+	if err != nil {
+		log.Error().Err(err).Msg("[cache] could not read /proc/self/statm")
+		return 0
+	}
+
+	str := string(b)
+	arr := strings.Split(str, " ")
+
+	// convert to pages
+	res, err := strconv.Atoi(arr[1])
+	if err != nil {
+		log.Error().Err(err).Msg("[cache] could not convert string to int")
+		return 0
+	}
+
+	// convert to KB
+	return (res * os.Getpagesize()) / 1024
 }
diff --git a/server/startup.go b/server/startup.go
index 3966938..b7fa9a8 100644
--- a/server/startup.go
+++ b/server/startup.go
@@ -79,7 +79,7 @@ func Serve(ctx *cli.Context) error {
 	// redirectsCache stores redirects in _redirects files
 	redirectsCache := cache.NewInMemoryCache()
 	// clientResponseCache stores responses from the Gitea server
-	clientResponseCache := cache.NewInMemoryCacheWithLimit(ctx.Uint64("memory-limit"))
+	clientResponseCache := cache.NewInMemoryCacheWithLimit(ctx.Int("memory-limit"))
 
 	giteaClient, err := gitea.NewClient(cfg.Gitea, clientResponseCache)
 	if err != nil {