From a9b6615286996d8f920172b25b4187b7038d7e16 Mon Sep 17 00:00:00 2001
From: crapStone <me@crapstone.dev>
Date: Mon, 5 Feb 2024 19:16:35 +0100
Subject: [PATCH] create a workaround to fix the oom problem

---
 cli/flags.go           |  6 +++++
 server/cache/memory.go | 57 +++++++++++++++++++++++++++++++++++++++++-
 server/startup.go      |  2 +-
 3 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/cli/flags.go b/cli/flags.go
index 097cf4f..09a6404 100644
--- a/cli/flags.go
+++ b/cli/flags.go
@@ -138,6 +138,12 @@ var (
 			Aliases: []string{"config"},
 			EnvVars: []string{"CONFIG_FILE"},
 		},
+		&cli.Uint64Flag{
+			Name:    "memory-limit",
+			Usage:   "maximum size of memory in bytes to use for caching, default: 512MB",
+			Value:   512 * 1024 * 1024,
+			EnvVars: []string{"MAX_MEMORY_SIZE"},
+		},
 
 		// ############################
 		// ### ACME Client Settings ###
diff --git a/server/cache/memory.go b/server/cache/memory.go
index 093696f..743b6d0 100644
--- a/server/cache/memory.go
+++ b/server/cache/memory.go
@@ -1,7 +1,62 @@
 package cache
 
-import "github.com/OrlovEvgeny/go-mcache"
+import (
+	"runtime"
+	"time"
 
+	"github.com/OrlovEvgeny/go-mcache"
+	"github.com/rs/zerolog/log"
+)
+
+type Cache struct {
+	mcache      *mcache.CacheDriver
+	memoryLimit uint64
+	lastCheck   time.Time
+}
+
+// NewInMemoryCache returns a new mcache that can grow infinitely.
 func NewInMemoryCache() ICache {
 	return mcache.New()
 }
+
+// NewInMemoryCache returns a new mcache with a memory limit.
+// If the limit is exceeded, the cache will be cleared.
+func NewInMemoryCacheWithLimit(memoryLimit uint64) ICache {
+	return &Cache{
+		mcache:      mcache.New(),
+		memoryLimit: memoryLimit,
+	}
+}
+
+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
+	if now.Sub(c.lastCheck) > (time.Second * 3) {
+		if c.memoryLimitOvershot() {
+			log.Debug().Msg("memory limit exceeded, clearing cache")
+			c.mcache.Truncate()
+		}
+		c.lastCheck = now
+	}
+
+	return c.mcache.Set(key, value, ttl)
+}
+
+func (c *Cache) Get(key string) (interface{}, bool) {
+	return c.mcache.Get(key)
+}
+
+func (c *Cache) Remove(key string) {
+	c.mcache.Remove(key)
+}
+
+func (c *Cache) memoryLimitOvershot() bool {
+	var stats runtime.MemStats
+	runtime.ReadMemStats(&stats)
+
+	log.Debug().Uint64("bytes", stats.HeapAlloc).Msg("current memory usage")
+
+	return stats.HeapAlloc > c.memoryLimit
+}
diff --git a/server/startup.go b/server/startup.go
index ffdabb7..3966938 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.NewInMemoryCache()
+	clientResponseCache := cache.NewInMemoryCacheWithLimit(ctx.Uint64("memory-limit"))
 
 	giteaClient, err := gitea.NewClient(cfg.Gitea, clientResponseCache)
 	if err != nil {