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 {