From a57047ff22804d05d16a551baf638c8ff9242ed4 Mon Sep 17 00:00:00 2001
From: Will Norris <will@willnorris.com>
Date: Sun, 17 Sep 2017 11:00:49 +0000
Subject: [PATCH] replace in-memory cache with size-limited lrucache

Fixes #4
---
 README.md              |  7 +++++--
 cmd/imageproxy/main.go | 30 ++++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 016c5bd..452f1e2 100644
--- a/README.md
+++ b/README.md
@@ -94,8 +94,11 @@ you should see a 500px square coder octocat.
 By default, the imageproxy command does not cache responses, but caching can be
 enabled using the `-cache` flag.  It supports the following values:
 
- - `memory` - uses an in-memory cache.  (This can exhaust your system's
-   available memory and is not recommended for production systems)
+ - `memory` - uses an in-memory LRU cache.  By default, this is limited to
+   100mb. To customize the size of the cache or the max age for cached items,
+   use the format `memory:size:age` where size is measured in mb and age is a
+   duration.  For example, `memory:200:4h` will create a 200mb cache that will
+   cache items no longer than 4 hours.
  - directory on local disk (e.g. `/tmp/imageproxy`) - will cache images
    on disk
  - s3 URL (e.g. `s3://region/bucket-name/optional-path-prefix`) - will cache
diff --git a/cmd/imageproxy/main.go b/cmd/imageproxy/main.go
index 914ab6d..4bd02bd 100644
--- a/cmd/imageproxy/main.go
+++ b/cmd/imageproxy/main.go
@@ -23,12 +23,14 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/PaulARoy/azurestoragecache"
+	"github.com/die-net/lrucache"
 	"github.com/diegomarangoni/gcscache"
 	"github.com/garyburd/redigo/redis"
-	"github.com/gregjones/httpcache"
 	"github.com/gregjones/httpcache/diskcache"
 	rediscache "github.com/gregjones/httpcache/redis"
 	"github.com/peterbourgon/diskv"
@@ -36,6 +38,8 @@ import (
 	"willnorris.com/go/imageproxy/internal/s3cache"
 )
 
+const defaultMemorySize = 100
+
 var addr = flag.String("addr", "localhost:8080", "TCP address to listen on")
 var whitelist = flag.String("whitelist", "", "comma separated list of allowed remote hosts")
 var referrers = flag.String("referrers", "", "comma separated list of allowed referring hosts")
@@ -102,7 +106,7 @@ func parseCache() (imageproxy.Cache, error) {
 	}
 
 	if *cache == "memory" {
-		return httpcache.NewMemoryCache(), nil
+		*cache = fmt.Sprintf("memory:%d", defaultMemorySize)
 	}
 
 	u, err := url.Parse(*cache)
@@ -115,6 +119,8 @@ func parseCache() (imageproxy.Cache, error) {
 		return azurestoragecache.New("", "", u.Host)
 	case "gcs":
 		return gcscache.New(u.String()), nil
+	case "memory":
+		return lruCache(u.Opaque)
 	case "redis":
 		conn, err := redis.DialURL(u.String(), redis.DialPassword(os.Getenv("REDIS_PASSWORD")))
 		if err != nil {
@@ -130,6 +136,26 @@ func parseCache() (imageproxy.Cache, error) {
 	}
 }
 
+// lruCache creates an LRU Cache with the specified options of the form
+// "maxSize:maxAge".  maxSize is specified in megabytes, maxAge is a duration.
+func lruCache(options string) (*lrucache.LruCache, error) {
+	parts := strings.SplitN(options, ":", 2)
+	size, err := strconv.ParseInt(parts[0], 10, 64)
+	if err != nil {
+		return nil, err
+	}
+
+	var age time.Duration
+	if len(parts) > 1 {
+		age, err = time.ParseDuration(parts[1])
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return lrucache.New(size*1e6, int64(age.Seconds())), nil
+}
+
 func diskCache(path string) *diskcache.Cache {
 	d := diskv.New(diskv.Options{
 		BasePath: path,