mediorum: in memory image cache (#8544)

AudiusProject · May 17, 2024 · bebe484 · bebe484
1 parent 3072759
commit bebe484
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 6 deletions.
diff --git a/mediorum/go.mod b/mediorum/go.mod
@@ -4,7 +4,7 @@ go 1.21.0
 
 require (
 	github.com/disintegration/imaging v1.6.2
-	github.com/erni27/imcache v1.1.0
+	github.com/erni27/imcache v1.2.0
 	github.com/ethereum/go-ethereum v1.11.4
 	github.com/georgysavva/scany/v2 v2.0.0
 	github.com/gowebpki/jcs v1.0.0

diff --git a/mediorum/go.sum b/mediorum/go.sum
@@ -133,6 +133,8 @@ github.com/envoyproxy/go-control-plane v0.9.10-0.20210907150352-cf90f659a021/go.
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/erni27/imcache v1.1.0 h1:Atak8HV/vsRtQXO0WajZtriT63gAII1qKHD8MgORH4I=
 github.com/erni27/imcache v1.1.0/go.mod h1:KNUCBr1U9nOFTyUEC9CsMKZE33NboYTPavsQsuEufoA=
+github.com/erni27/imcache v1.2.0 h1:EbHHhwzJPcAYK//cVDq0use9tep1z9/kenIQhMarGnk=
+github.com/erni27/imcache v1.2.0/go.mod h1:KNUCBr1U9nOFTyUEC9CsMKZE33NboYTPavsQsuEufoA=
 github.com/ethereum/go-ethereum v1.11.4 h1:KG81SnUHXWk8LJB3mBcHg/E2yLvXoiPmRMCIRxgx3cE=
 github.com/ethereum/go-ethereum v1.11.4/go.mod h1:it7x0DWnTDMfVFdXcU6Ti4KEFQynLHVRarcSlPr0HBo=
 github.com/fjl/memsize v0.0.0-20190710130421-bcb5799ab5e5 h1:FtmdgXiUlNeRsoNMFlKLDt+S+6hbjVMEW6RGQ7aUf7c=

diff --git a/mediorum/server/serve_image.go b/mediorum/server/serve_image.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"mime"
@@ -9,6 +10,7 @@ import (
 	"time"
 
 	"github.com/AudiusProject/audius-protocol/mediorum/cidutil"
+	"github.com/erni27/imcache"
 
 	"github.com/labstack/echo/v4"
 	"gocloud.dev/blob"
@@ -22,14 +24,31 @@ func (ss *MediorumServer) serveImage(c echo.Context) error {
 	jobID := c.Param("jobID")
 	variant := c.Param("variant")
 	isOriginalJpg := variant == "original.jpg"
+	cacheKey := jobID + variant
 
-	// helper function... only sets cache-control header on success
-	serveSuccessWithReader := func(blob *blob.Reader) error {
+	serveSuccessWithBytes := func(blobData []byte, modTime time.Time) error {
 		c.Response().Header().Set(echo.HeaderCacheControl, "public, max-age=2592000, immutable")
-		http.ServeContent(c.Response(), c.Request(), jobID+variant, blob.ModTime(), blob)
+		http.ServeContent(c.Response(), c.Request(), cacheKey, modTime, bytes.NewReader(blobData))
 		return nil
 	}
 
+	// helper function... only sets cache-control header on success
+	serveSuccessWithReader := func(blob *blob.Reader) error {
+		blobData, err := io.ReadAll(blob)
+		if err != nil {
+			return err
+		}
+		blob.Close()
+		ss.imageCache.Set(cacheKey, blobData, imcache.WithNoExpiration())
+		return serveSuccessWithBytes(blobData, blob.ModTime())
+	}
+
+	// use cache if possible
+	if blobData, ok := ss.imageCache.Get(cacheKey); ok {
+		c.Response().Header().Set("x-image-cache-hit", "true")
+		return serveSuccessWithBytes(blobData, ss.StartedAt)
+	}
+
 	serveSuccess := func(blobPath string) error {
 		if blob, err := ss.bucket.NewReader(ctx, blobPath, nil); err == nil {
 			return serveSuccessWithReader(blob)

diff --git a/mediorum/server/serve_image_test.go b/mediorum/server/serve_image_test.go
@@ -33,6 +33,7 @@ func TestServeImage(t *testing.T) {
 		assert.Equal(t, 200, resp.StatusCode)
 		assert.NotEmpty(t, resp.Header.Get("x-fetch-ok"))
 		assert.NotEmpty(t, resp.Header.Get("x-resize-ok"))
+		assert.Empty(t, resp.Header.Get("x-image-cache-hit"))
 	}
 
 	// the second time it should have the variant on disk
@@ -41,6 +42,7 @@ func TestServeImage(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Equal(t, 200, resp.StatusCode)
 		assert.Empty(t, resp.Header.Get("x-resize-ok"))
+		assert.Equal(t, "true", resp.Header.Get("x-image-cache-hit"))
 	}
 
 	// it should also have the orig

diff --git a/mediorum/server/server.go b/mediorum/server/server.go
@@ -103,6 +103,7 @@ type MediorumServer struct {
 	unreachablePeers      []string
 	redirectCache         *imcache.Cache[string, string]
 	uploadOrigCidCache    *imcache.Cache[string, string]
+	imageCache            *imcache.Cache[string, []byte]
 	failsPeerReachability bool
 
 	StartedAt time.Time
@@ -266,8 +267,9 @@ func New(config MediorumConfig) (*MediorumServer, error) {
 		isAudiusdManaged: isAudiusdManaged,
 
 		peerHealths:        map[string]*PeerHealth{},
-		redirectCache:      imcache.New(imcache.WithMaxEntriesOption[string, string](50_000)),
-		uploadOrigCidCache: imcache.New(imcache.WithMaxEntriesOption[string, string](50_000)),
+		redirectCache:      imcache.New(imcache.WithMaxEntriesLimitOption[string, string](50_000, imcache.EvictionPolicyLRU)),
+		uploadOrigCidCache: imcache.New(imcache.WithMaxEntriesLimitOption[string, string](50_000, imcache.EvictionPolicyLRU)),
+		imageCache:         imcache.New(imcache.WithMaxEntriesLimitOption[string, []byte](10_000, imcache.EvictionPolicyLRU)),
 
 		StartedAt: time.Now().UTC(),
 		Config:    config,