From 76f1dcf83e4667a1b61b4f355cddf9ad696c3768 Mon Sep 17 00:00:00 2001
From: Louis Dalibard <louis.dalibard@gmail.com>
Date: Mon, 10 Jun 2024 21:55:03 +0200
Subject: [PATCH] fix: add job limits for thumbnails to prevent out of memory
 errors

---
 config/config.go       |  5 +++--
 pkg/config.json        |  1 +
 thumbnail/thumbnail.go | 38 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/config/config.go b/config/config.go
index a0f9b09..f83e262 100644
--- a/config/config.go
+++ b/config/config.go
@@ -8,8 +8,9 @@ import (
 )
 
 type ConfigS struct {
-	Host      string            `json:"host"`
-	ServeDirs map[string]string `json:"servedirs"`
+	Host              string            `json:"host"`
+	ThumbnailJobLimit int               `json:"thumbnailjoblimit"`
+	ServeDirs         map[string]string `json:"servedirs"`
 }
 
 var Config ConfigS
diff --git a/pkg/config.json b/pkg/config.json
index eedf3f0..947f31c 100644
--- a/pkg/config.json
+++ b/pkg/config.json
@@ -1,5 +1,6 @@
 {
   "host": ":3125",
+  "thumbnailjoblimit": "5",
   "servedirs": {
     "leech": "/home/ontake/Dev/go/leech"
   }
diff --git a/thumbnail/thumbnail.go b/thumbnail/thumbnail.go
index 233fdb3..407fe92 100644
--- a/thumbnail/thumbnail.go
+++ b/thumbnail/thumbnail.go
@@ -27,35 +27,69 @@ var thumbnailSize = 24
 
 var thumbnailCache = map[string][]byte{}
 var thumbnailCacheMutex = &sync.RWMutex{}
+var jobLimit = 5
+var memLimiterMutex = &sync.RWMutex{}
+var jobCounter = 0
+var checkAgain = make(chan bool, 5)
 
 func IsSupportedFileType(completePath string) bool {
 	fileExt := filepath.Ext(completePath)
 	return slices.Contains(SupportedFileTypes, fileExt)
 }
 
+func WaitForAvailable() {
+	memLimiterMutex.RLock()
+	for jobLimit == jobCounter {
+		memLimiterMutex.RUnlock()
+		<-checkAgain
+		memLimiterMutex.RLock()
+	}
+	memLimiterMutex.RUnlock()
+	memLimiterMutex.Lock()
+	jobCounter += 1
+	memLimiterMutex.Unlock()
+}
+
+func Free() {
+	memLimiterMutex.Lock()
+	jobCounter -= 1
+	memLimiterMutex.Unlock()
+	select {
+	case checkAgain <- true:
+		return
+	default:
+		return
+	}
+}
+
 func GetThumbnail(c *fiber.Ctx, completePath string) {
+	WaitForAvailable()
 	c.Set(fiber.HeaderContentType, "image")
 	thumbnailCacheMutex.RLock()
 	bytesThumb, ok := thumbnailCache[completePath]
 	thumbnailCacheMutex.RUnlock()
 	if ok {
 		c.Write(bytesThumb)
+		Free()
 		return
 	}
 
 	fileExt := filepath.Ext(completePath)
 	if !slices.Contains(SupportedFileTypes, fileExt) {
+		Free()
 		return
 	}
 	f, err := os.Open(completePath)
 	if err != nil {
+		Free()
 		return
 	}
-	defer f.Close()
 	img, _, err := image.Decode(f)
 	if err != nil {
+		Free()
 		return
 	}
+	f.Close()
 	// load images and make 64x64 thumbnails of them
 	thumbnail := imaging.Thumbnail(img, thumbnailSize, thumbnailSize, imaging.CatmullRom)
 	// create a new blank image
@@ -70,10 +104,12 @@ func GetThumbnail(c *fiber.Ctx, completePath string) {
 		err = imaging.Encode(&buf, dst, FileTypesMap[fileExt])
 	}
 	if err != nil {
+		Free()
 		return
 	}
 	thumbnailCacheMutex.Lock()
 	thumbnailCache[completePath] = buf.Bytes()
 	thumbnailCacheMutex.Unlock()
 	c.Write(buf.Bytes())
+	Free()
 }