diff --git a/cmd/flags.go b/cmd/flags.go index 7ac94e6..df6af4b 100644 --- a/cmd/flags.go +++ b/cmd/flags.go @@ -112,6 +112,12 @@ var ( EnvVars: []string{"PAGES_BRANCHES"}, Value: cli.NewStringSlice("pages"), }, + &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/cmd/main.go b/cmd/main.go index 683e859..f5d0aa9 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -83,7 +83,7 @@ func Serve(ctx *cli.Context) error { // redirectsCache stores redirects in _redirects files redirectsCache := cache.NewKeyValueCache() // clientResponseCache stores responses from the Gitea server - clientResponseCache := cache.NewKeyValueCache() + clientResponseCache := cache.NewKeyValueCacheWithLimit(ctx.Uint64("memory-limit")) giteaClient, err := gitea.NewClient(giteaRoot, giteaAPIToken, clientResponseCache, ctx.Bool("enable-symlink-support"), ctx.Bool("enable-lfs-support")) if err != nil { diff --git a/server/cache/setup.go b/server/cache/setup.go index a5928b0..73ff08f 100644 --- a/server/cache/setup.go +++ b/server/cache/setup.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 +} + +// NewKeyValueCache returns a new mcache that can grow infinitely. func NewKeyValueCache() SetGetKey { return mcache.New() } + +// NewKeyValueCacheWithLimit returns a new mcache with a memory limit. +// If the limit is exceeded, the cache will be cleared. +func NewKeyValueCacheWithLimit(memoryLimit uint64) SetGetKey { + 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 +}