From cac19b2c93ca2847b49fc37e63261936a426b3e3 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Mon, 29 Jul 2024 13:31:48 +0800
Subject: [PATCH] :thread: Replace cache with thread-safe sync.Map

---
 pkg/internal/services/analyzer.go    |  2 +-
 pkg/internal/services/attachments.go | 48 +++++++++++++++++-----------
 pkg/internal/services/uploader.go    |  4 +--
 3 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/pkg/internal/services/analyzer.go b/pkg/internal/services/analyzer.go
index abce957..2c45db4 100644
--- a/pkg/internal/services/analyzer.go
+++ b/pkg/internal/services/analyzer.go
@@ -93,7 +93,7 @@ func AnalyzeAttachment(file models.Attachment) error {
 	if linked && err != nil {
 		return fmt.Errorf("unable to link file record: %v", err)
 	} else if !linked {
-		metadataCache[file.ID] = file
+		metadataCache.Store(file.ID, file)
 		if err := tx.Save(&file).Error; err != nil {
 			tx.Rollback()
 			return fmt.Errorf("unable to save file record: %v", err)
diff --git a/pkg/internal/services/attachments.go b/pkg/internal/services/attachments.go
index 51bd62c..e74e621 100644
--- a/pkg/internal/services/attachments.go
+++ b/pkg/internal/services/attachments.go
@@ -6,21 +6,23 @@ import (
 	"mime/multipart"
 	"net/http"
 	"path/filepath"
+	"sync"
 
 	"git.solsynth.dev/hydrogen/paperclip/pkg/internal/database"
 
 	"git.solsynth.dev/hydrogen/paperclip/pkg/internal/models"
 	"github.com/google/uuid"
+	"github.com/rs/zerolog/log"
 	"gorm.io/gorm"
 )
 
 const metadataCacheLimit = 512
 
-var metadataCache = make(map[uint]models.Attachment)
+var metadataCache sync.Map
 
 func GetAttachmentByID(id uint) (models.Attachment, error) {
-	if val, ok := metadataCache[id]; ok {
-		return val, nil
+	if val, ok := metadataCache.Load(id); ok {
+		return val.(models.Attachment), nil
 	}
 
 	var attachment models.Attachment
@@ -29,10 +31,8 @@ func GetAttachmentByID(id uint) (models.Attachment, error) {
 	}).Preload("Account").First(&attachment).Error; err != nil {
 		return attachment, err
 	} else {
-		if len(metadataCache) > metadataCacheLimit {
-			clear(metadataCache)
-		}
-		metadataCache[id] = attachment
+		MaintainAttachmentCache()
+		metadataCache.Store(id, attachment)
 	}
 
 	return attachment, nil
@@ -80,10 +80,8 @@ func NewAttachmentMetadata(tx *gorm.DB, user *models.Account, file *multipart.Fi
 	if err := tx.Save(&attachment).Error; err != nil {
 		return attachment, fmt.Errorf("failed to save attachment record: %v", err)
 	} else {
-		if len(metadataCache) > metadataCacheLimit {
-			clear(metadataCache)
-		}
-		metadataCache[attachment.ID] = attachment
+		MaintainAttachmentCache()
+		metadataCache.Store(attachment.ID, attachment)
 	}
 
 	return attachment, nil
@@ -108,8 +106,8 @@ func TryLinkAttachment(tx *gorm.DB, og models.Attachment, hash string) (bool, er
 		return true, err
 	}
 
-	metadataCache[prev.ID] = prev
-	metadataCache[og.ID] = og
+	metadataCache.Store(prev.ID, prev)
+	metadataCache.Store(og.ID, og)
 
 	return true, nil
 }
@@ -118,10 +116,8 @@ func UpdateAttachment(item models.Attachment) (models.Attachment, error) {
 	if err := database.C.Save(&item).Error; err != nil {
 		return item, err
 	} else {
-		if len(metadataCache) > metadataCacheLimit {
-			clear(metadataCache)
-		}
-		metadataCache[item.ID] = item
+		MaintainAttachmentCache()
+		metadataCache.Store(item.ID, item)
 	}
 
 	return item, nil
@@ -148,7 +144,7 @@ func DeleteAttachment(item models.Attachment) error {
 		tx.Rollback()
 		return err
 	} else {
-		delete(metadataCache, item.ID)
+		metadataCache.Delete(item.ID)
 	}
 
 	tx.Commit()
@@ -159,3 +155,19 @@ func DeleteAttachment(item models.Attachment) error {
 
 	return nil
 }
+
+func MaintainAttachmentCache() {
+	var keySet []uint
+	metadataCache.Range(func(k any, v any) bool {
+		keySet = append(keySet, k.(uint))
+		return true
+	})
+	if len(keySet) > metadataCacheLimit {
+		go func() {
+			log.Debug().Int("count", len(keySet)).Msg("Cleaning attachment metadata cache...")
+			for _, k := range keySet {
+				metadataCache.Delete(k)
+			}
+		}()
+	}
+}
diff --git a/pkg/internal/services/uploader.go b/pkg/internal/services/uploader.go
index 400f296..cda4362 100644
--- a/pkg/internal/services/uploader.go
+++ b/pkg/internal/services/uploader.go
@@ -81,7 +81,7 @@ func ReUploadFileToPermanent(meta models.Attachment) error {
 		}
 
 		database.C.Save(&meta)
-		metadataCache[meta.ID] = meta
+		metadataCache.Store(meta.ID, meta)
 		return nil
 	case models.DestinationTypeS3:
 		var destConfigured models.S3Destination
@@ -105,7 +105,7 @@ func ReUploadFileToPermanent(meta models.Attachment) error {
 		}
 
 		database.C.Save(&meta)
-		metadataCache[meta.ID] = meta
+		metadataCache.Store(meta.ID, meta)
 		return nil
 	default:
 		return fmt.Errorf("invalid destination: unsupported protocol %s", dest.Type)