fix: add job limits for thumbnails to prevent out of memory errors

This commit is contained in:
Louis Dalibard 2024-06-10 21:55:03 +02:00
parent db4aae2e5b
commit 76f1dcf83e
3 changed files with 41 additions and 3 deletions

View File

@ -8,8 +8,9 @@ import (
) )
type ConfigS struct { type ConfigS struct {
Host string `json:"host"` Host string `json:"host"`
ServeDirs map[string]string `json:"servedirs"` ThumbnailJobLimit int `json:"thumbnailjoblimit"`
ServeDirs map[string]string `json:"servedirs"`
} }
var Config ConfigS var Config ConfigS

View File

@ -1,5 +1,6 @@
{ {
"host": ":3125", "host": ":3125",
"thumbnailjoblimit": "5",
"servedirs": { "servedirs": {
"leech": "/home/ontake/Dev/go/leech" "leech": "/home/ontake/Dev/go/leech"
} }

View File

@ -27,35 +27,69 @@ var thumbnailSize = 24
var thumbnailCache = map[string][]byte{} var thumbnailCache = map[string][]byte{}
var thumbnailCacheMutex = &sync.RWMutex{} 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 { func IsSupportedFileType(completePath string) bool {
fileExt := filepath.Ext(completePath) fileExt := filepath.Ext(completePath)
return slices.Contains(SupportedFileTypes, fileExt) 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) { func GetThumbnail(c *fiber.Ctx, completePath string) {
WaitForAvailable()
c.Set(fiber.HeaderContentType, "image") c.Set(fiber.HeaderContentType, "image")
thumbnailCacheMutex.RLock() thumbnailCacheMutex.RLock()
bytesThumb, ok := thumbnailCache[completePath] bytesThumb, ok := thumbnailCache[completePath]
thumbnailCacheMutex.RUnlock() thumbnailCacheMutex.RUnlock()
if ok { if ok {
c.Write(bytesThumb) c.Write(bytesThumb)
Free()
return return
} }
fileExt := filepath.Ext(completePath) fileExt := filepath.Ext(completePath)
if !slices.Contains(SupportedFileTypes, fileExt) { if !slices.Contains(SupportedFileTypes, fileExt) {
Free()
return return
} }
f, err := os.Open(completePath) f, err := os.Open(completePath)
if err != nil { if err != nil {
Free()
return return
} }
defer f.Close()
img, _, err := image.Decode(f) img, _, err := image.Decode(f)
if err != nil { if err != nil {
Free()
return return
} }
f.Close()
// load images and make 64x64 thumbnails of them // load images and make 64x64 thumbnails of them
thumbnail := imaging.Thumbnail(img, thumbnailSize, thumbnailSize, imaging.CatmullRom) thumbnail := imaging.Thumbnail(img, thumbnailSize, thumbnailSize, imaging.CatmullRom)
// create a new blank image // create a new blank image
@ -70,10 +104,12 @@ func GetThumbnail(c *fiber.Ctx, completePath string) {
err = imaging.Encode(&buf, dst, FileTypesMap[fileExt]) err = imaging.Encode(&buf, dst, FileTypesMap[fileExt])
} }
if err != nil { if err != nil {
Free()
return return
} }
thumbnailCacheMutex.Lock() thumbnailCacheMutex.Lock()
thumbnailCache[completePath] = buf.Bytes() thumbnailCache[completePath] = buf.Bytes()
thumbnailCacheMutex.Unlock() thumbnailCacheMutex.Unlock()
c.Write(buf.Bytes()) c.Write(buf.Bytes())
Free()
} }