From 576359798f3159032936b5cf6c727e2f38d5086a Mon Sep 17 00:00:00 2001
From: nisan <info@nisancoskun.com>
Date: Thu, 1 Aug 2024 22:29:02 +0300
Subject: [PATCH] feat: add cache support

---
 main.go | 94 ++++++++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 86 insertions(+), 8 deletions(-)

diff --git a/main.go b/main.go
index bbdb46a..5ee7a3e 100644
--- a/main.go
+++ b/main.go
@@ -11,6 +11,7 @@ import (
 	"io/fs"
 	"log"
 	"net/http"
+	"sync"
 
 	"strconv"
 	"strings"
@@ -19,6 +20,54 @@ import (
 	"golang.org/x/image/draw"
 )
 
+type cacheItem struct {
+	value      []byte
+	lastAccess int64
+}
+type AvatarCache struct {
+	cache map[string]*cacheItem
+	mu    *sync.RWMutex
+}
+
+func (ac *AvatarCache) Clear() {
+	ac.mu.Lock()
+	for k, v := range ac.cache {
+		if time.Now().Unix()-v.lastAccess > int64(100) {
+			delete(ac.cache, k)
+		}
+	}
+	ac.mu.Unlock()
+}
+
+// insert to cache
+func (ac *AvatarCache) Insert(id string, size int, value []byte) {
+	ac.mu.Lock()
+	defer ac.mu.Unlock()
+	ac.cache[id+strconv.Itoa(size)] = &cacheItem{value: value, lastAccess: time.Now().Unix()}
+	// remove oldest item if cache is full
+	if len(ac.cache) > 1000 {
+		for k := range ac.cache {
+			delete(ac.cache, k)
+			break
+		}
+	}
+}
+
+// get from cache
+func (ac *AvatarCache) Get(id string, size int) []byte {
+	ac.mu.Lock()
+	defer ac.mu.Unlock()
+	item, ok := ac.cache[id+strconv.Itoa(size)]
+	if !ok {
+		return nil
+	}
+	item.lastAccess = time.Now().Unix()
+	return item.value
+}
+
+// create new cache
+var avatarCache = AvatarCache{cache: make(map[string]*cacheItem), mu: &sync.RWMutex{}}
+
 //go:embed assets/*
 var files embed.FS
 
@@ -52,15 +101,31 @@ func serveAvatar(w http.ResponseWriter, r *http.Request) {
 
 	// get the id and size from the request
 	id := r.PathValue("id")
-	size, err := strconv.Atoi(r.PathValue("size"))
+	sizeStr := r.PathValue("size")
+	if sizeStr == "" {
+		sizeStr = r.URL.Query().Get("s")
+	}
+	size, err := strconv.Atoi(sizeStr)
 	if err != nil {
 		size = 400
 	}
-	// min size is 32, max size is 1000
+	// min size is 32, max size is 2048
 	if size < 32 {
 		size = 32
-	} else if size > 1000 {
-		size = 1000
+	} else if size > 2048 {
+		size = 2048
+	}
+	if avatarCache.Get(id, size) != nil {
+		// get from cache
+		w.Header().Set("content-type", "image/jpeg")
+		w.WriteHeader(http.StatusFound)
+		byteSize, err := w.Write(avatarCache.Get(id, size))
+		if err != nil {
+			log.Println(err.Error())
+		}
+		log.Printf("[id:%s][size:%d] [%s] [%s] [%d bytes] [%v] [cache hit]\n",
+			id, size, r.RemoteAddr, r.Proto, byteSize, time.Since(startTime))
+		return
 	}
 
 	img := image.NewRGBA(emptyRect)
@@ -74,26 +139,38 @@ func serveAvatar(w http.ResponseWriter, r *http.Request) {
 	if size != 400 {
 		// resize the image
 		dst := image.NewRGBA(image.Rect(0, 0, size, size))
-		draw.ApproxBiLinear.Scale(dst, dst.Bounds(), img, emptyRect, draw.Over, nil)
+		draw.CatmullRom.Scale(dst, dst.Bounds(), img, emptyRect, draw.Over, nil)
 		img = dst
 	}
 
 	// write the image to the response
 	w.Header().Set("content-type", "image/jpeg")
 	w.WriteHeader(http.StatusCreated)
-	err = jpeg.Encode(w, img, nil)
+	buf := new(bytes.Buffer)
+	err = jpeg.Encode(buf, img, nil)
 	if err != nil {
 		log.Println(err.Error())
 	}
-	log.Printf("[id:%s][size:%d] served to [%s] [%s] [%v]\n", id, size, r.RemoteAddr, r.Proto, time.Since(startTime))
+	byteSize, err := w.Write(buf.Bytes())
+	if err != nil {
+		log.Println(err.Error())
+	}
+	avatarCache.Insert(id, size, buf.Bytes())
+	log.Printf("[id:%s][size:%d] [%s] [%s] [%d bytes] [%v]\n",
+		id, size, r.RemoteAddr, r.Proto, byteSize, time.Since(startTime))
 }
 
 func main() {
-	log.Println("assets loaded")
+	go func() {
+		for range time.Tick(time.Minute) {
+			avatarCache.Clear()
+		}
+	}()
 	mux := http.NewServeMux()
 	mux.HandleFunc("GET /{id}", serveAvatar)
 	mux.HandleFunc("GET /{size}/{id}", serveAvatar)
 	http.ListenAndServe("0.0.0.0:3382", mux)
+	log.Println("Server started at 0.0.0.0:3382")
 }
 
 func init() {
@@ -140,4 +217,5 @@ func init() {
 			}
 		}
 	}
+	log.Printf("%d different combinations loaded\n", len(BaseRects)*len(EYES)*len(NOSE)*len(MOUTH))
 }