From e70b758404955192e7c9d5e7f09466f9ad11499e Mon Sep 17 00:00:00 2001 From: Pysio Date: Wed, 15 Jan 2025 19:19:30 +0800 Subject: [PATCH] feat: add Cache :> --- .gitignore | 2 + handlers/handlers.go | 126 ++++++++++++++++++++++++++++++++++++++++++- main.go | 12 +++++ utils/cache.go | 33 ++++++++++++ 4 files changed, 171 insertions(+), 2 deletions(-) create mode 100644 utils/cache.go diff --git a/.gitignore b/.gitignore index a1ea275..9a23c5b 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,5 @@ temp/ # Build output build/ dist/ + +cache/ \ No newline at end of file diff --git a/handlers/handlers.go b/handlers/handlers.go index 5fc0945..4f27585 100644 --- a/handlers/handlers.go +++ b/handlers/handlers.go @@ -314,15 +314,46 @@ func RandomImage(c *gin.Context) { } var result models.Image - err = models.ImagesCollection.FindOne(context.Background(), bson.M{}, options.FindOne().SetSkip(int64(time.Now().UnixNano()%count))).Decode(&result) + err = models.ImagesCollection.FindOne( + context.Background(), + bson.M{}, + options.FindOne().SetProjection(bson.M{"hash": 1}).SetSkip(int64(time.Now().UnixNano()%count)), + ).Decode(&result) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + var imageData []byte + // 检查本地缓存 + if utils.ImageExistsInCache(result.Hash) { + imageData, err = utils.LoadImageFromCache(result.Hash) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to load image from cache"}) + return + } + } else { + // 从数据库获取完整图片数据 + err = models.ImagesCollection.FindOne( + context.Background(), + bson.M{"hash": result.Hash}, + ).Decode(&result) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + imageData = result.Data + + // 保存到本地缓存 + if err := utils.SaveImageToCache(result.Hash, result.Data); err != nil { + // 仅记录错误，不影响返回 + c.Error(err) + } + } + c.Header("Content-Type", "image/webp") c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s.webp"`, result.Hash)) - c.Data(http.StatusOK, "image/webp", result.Data) + c.Data(http.StatusOK, "image/webp", imageData) } type lowerCount struct { @@ -529,6 +560,19 @@ func DeleteImage(c *gin.Context) { func GetImage(c *gin.Context) { hash := c.Param("hash") + + // 先检查缓存 + if utils.ImageExistsInCache(hash) { + imageData, err := utils.LoadImageFromCache(hash) + if err == nil { + c.Header("Content-Type", "image/webp") + c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s.webp"`, hash)) + c.Data(http.StatusOK, "image/webp", imageData) + return + } + } + + // 缓存不存在或读取失败，从数据库获取 var image models.Image err := models.ImagesCollection.FindOne(context.Background(), bson.M{"hash": hash}).Decode(&image) if err != nil { @@ -540,6 +584,11 @@ func GetImage(c *gin.Context) { return } + // 保存到缓存 + if err := utils.SaveImageToCache(image.Hash, image.Data); err != nil { + c.Error(err) + } + c.Header("Content-Type", image.ContentType) c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s.webp"`, image.Hash)) c.Data(http.StatusOK, image.ContentType, image.Data) @@ -547,6 +596,29 @@ func GetImage(c *gin.Context) { func GetImageByHash(c *gin.Context) { hash := c.Param("hash") + + // 先检查缓存 + if utils.ImageExistsInCache(hash) { + imageData, err := utils.LoadImageFromCache(hash) + if err == nil { + c.Header("Content-Type", "image/webp") + c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s.webp"`, hash)) + c.Header("Cache-Control", "public, max-age=31536000, immutable") + c.Header("ETag", fmt.Sprintf(`"%s"`, hash)) + c.Header("CDN-Cache-Control", "max-age=31536000") + c.Header("Surrogate-Control", "max-age=31536000") + + if c.GetHeader("If-None-Match") == fmt.Sprintf(`"%s"`, hash) { + c.Status(http.StatusNotModified) + return + } + + c.Data(http.StatusOK, "image/webp", imageData) + return + } + } + + // 缓存不存在或读取失败，从数据库获取 var image models.Image err := models.ImagesCollection.FindOne(context.Background(), bson.M{"hash": hash}).Decode(&image) if err != nil { @@ -558,6 +630,11 @@ func GetImageByHash(c *gin.Context) { return } + // 保存到缓存 + if err := utils.SaveImageToCache(image.Hash, image.Data); err != nil { + c.Error(err) + } + c.Header("Content-Type", image.ContentType) c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s.webp"`, image.Hash)) c.Header("Cache-Control", "public, max-age=31536000, immutable") @@ -573,6 +650,51 @@ func GetImageByHash(c *gin.Context) { c.Data(http.StatusOK, image.ContentType, image.Data) } +// 新增刷新缓存的处理函数 +func RefreshCache(c *gin.Context) { + cursor, err := models.ImagesCollection.Find(context.Background(), bson.M{}) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to query images"}) + return + } + defer cursor.Close(context.Background()) + + var stats struct { + Total int `json:"total"` + Cached int `json:"cached"` + Failed int `json:"failed"` + Skipped int `json:"skipped"` + } + + for cursor.Next(context.Background()) { + var image models.Image + if err := cursor.Decode(&image); err != nil { + stats.Failed++ + continue + } + + stats.Total++ + + // 如果已经在缓存中，跳过 + if utils.ImageExistsInCache(image.Hash) { + stats.Skipped++ + continue + } + + // 保存到缓存 + if err := utils.SaveImageToCache(image.Hash, image.Data); err != nil { + stats.Failed++ + } else { + stats.Cached++ + } + } + + c.JSON(http.StatusOK, gin.H{ + "message": "Cache refresh completed", + "stats": stats, + }) +} + func Egg(c *gin.Context) { c.String(http.StatusOK, "Oops!") } diff --git a/main.go b/main.go index 3837670..288135c 100644 --- a/main.go +++ b/main.go @@ -74,6 +74,11 @@ func logEnvironmentStatus() { } func main() { + // 初始化缓存目录 + if err := utils.InitCache(); err != nil { + log.Fatalf("Failed to initialize cache: %v", err) + } + // 加载环境变量 loadEnv() @@ -115,6 +120,13 @@ func main() { r.GET("/50x", handlers.ServerError) + // 添加新的路由 + adminGroup := r.Group("/admin") + adminGroup.Use(middleware.VerifyAdminToken()) + { + adminGroup.POST("/refcache", handlers.RefreshCache) + } + // 启动服务器 r.Run(":5000") } diff --git a/utils/cache.go b/utils/cache.go new file mode 100644 index 0000000..c1af2c2 --- /dev/null +++ b/utils/cache.go @@ -0,0 +1,33 @@ +package utils + +import ( + "os" + "path/filepath" +) + +const CacheDir = "./cache/images" + +func InitCache() error { + // 确保缓存目录存在 + if err := os.MkdirAll(CacheDir, 0755); err != nil { + return err + } + return nil +} + +func GetImagePath(hash string) string { + return filepath.Join(CacheDir, hash+".webp") +} + +func SaveImageToCache(hash string, data []byte) error { + return os.WriteFile(GetImagePath(hash), data, 0644) +} + +func LoadImageFromCache(hash string) ([]byte, error) { + return os.ReadFile(GetImagePath(hash)) +} + +func ImageExistsInCache(hash string) bool { + _, err := os.Stat(GetImagePath(hash)) + return err == nil +}