Feat: cache in-memory store

models/migration_test.go

@@ -1,16 +1,17 @@
 package model
 
 import (
-	"github.com/gin-gonic/gin"
+	"github.com/jinzhu/gorm"
 	"github.com/stretchr/testify/assert"
 	"testing"
 )
 
 func TestMigration(t *testing.T) {
 	asserts := assert.New(t)
-	gin.SetMode(gin.TestMode)
+	DB, _ = gorm.Open("sqlite3", ":memory:")
 
 	asserts.NotPanics(func() {
 		migration()
 	})
+	DB = mockDB
 }

models/policy.go

@@ -2,11 +2,11 @@ package model
 
 import (
 	"encoding/json"
+	"github.com/HFO4/cloudreve/pkg/cache"
 	"github.com/HFO4/cloudreve/pkg/util"
 	"github.com/jinzhu/gorm"
 	"path/filepath"
 	"strconv"
-	"sync"
 	"time"
 )
 
@@ -42,28 +42,20 @@ type PolicyOption struct {
 	RangeTransferEnabled bool     `json:"range_transfer_enabled"`
 }
 
-// 存储策略缓存，部分情况下需要频繁查询存储策略
-var policyCache = make(map[uint]Policy)
-var rw sync.RWMutex
-
 // GetPolicyByID 用ID获取存储策略
 func GetPolicyByID(ID interface{}) (Policy, error) {
 	// 尝试读取缓存
-	rw.RLock()
-	if policy, ok := policyCache[ID.(uint)]; ok {
-		rw.RUnlock()
-		return policy, nil
+	cacheKey := "policy_" + strconv.Itoa(int(ID.(uint)))
+	if policy, ok := cache.Store.Get(cacheKey); ok {
+		return policy.(Policy), nil
 	}
-	rw.RUnlock()
 
 	var policy Policy
 	result := DB.First(&policy, ID)
 
 	// 写入缓存
 	if result.Error == nil {
-		rw.Lock()
-		policyCache[policy.ID] = policy
-		rw.Unlock()
+		_ = cache.Store.Set(cacheKey, policy)
 	}
 
 	return policy, result.Error

models/policy_test.go

@@ -12,18 +12,33 @@ import (
 func TestGetPolicyByID(t *testing.T) {
 	asserts := assert.New(t)
 
-	rows := sqlmock.NewRows([]string{"name", "type", "options"}).
-		AddRow("默认存储策略", "local", "{\"op_name\":\"123\"}")
-	mock.ExpectQuery("^SELECT \\* FROM `(.+)` WHERE `(.+)`\\.`deleted_at` IS NULL AND \\(\\(`policies`.`id` = 1\\)\\)(.+)$").WillReturnRows(rows)
-	policy, err := GetPolicyByID(uint(1))
-	asserts.NoError(err)
-	asserts.Equal("默认存储策略", policy.Name)
-	asserts.Equal("123", policy.OptionsSerialized.OPName)
+	// 缓存未命中
+	{
+		rows := sqlmock.NewRows([]string{"name", "type", "options"}).
+			AddRow("默认存储策略", "local", "{\"op_name\":\"123\"}")
+		mock.ExpectQuery("^SELECT(.+)").WillReturnRows(rows)
+		policy, err := GetPolicyByID(uint(22))
+		asserts.NoError(err)
+		asserts.NoError(mock.ExpectationsWereMet())
+		asserts.Equal("默认存储策略", policy.Name)
+		asserts.Equal("123", policy.OptionsSerialized.OPName)
+
+		rows = sqlmock.NewRows([]string{"name", "type", "options"})
+		mock.ExpectQuery("^SELECT(.+)").WillReturnRows(rows)
+		policy, err = GetPolicyByID(uint(23))
+		asserts.NoError(mock.ExpectationsWereMet())
+		asserts.Error(err)
+	}
+
+	// 命中
+	{
+		policy, err := GetPolicyByID(uint(22))
+		asserts.NoError(err)
+		asserts.Equal("默认存储策略", policy.Name)
+		asserts.Equal("123", policy.OptionsSerialized.OPName)
+
+	}
 
-	rows = sqlmock.NewRows([]string{"name", "type", "options"})
-	mock.ExpectQuery("^SELECT \\* FROM `(.+)` WHERE `(.+)`\\.`deleted_at` IS NULL AND \\(\\(`policies`.`id` = 1\\)\\)(.+)$").WillReturnRows(rows)
-	policy, err = GetPolicyByID(uint(1))
-	asserts.Error(err)
 }
 
 func TestPolicy_BeforeSave(t *testing.T) {

models/setting_test.go

@@ -9,6 +9,7 @@ import (
 )
 
 var mock sqlmock.Sqlmock
+var mockDB *gorm.DB
 
 // TestMain 初始化数据库Mock
 func TestMain(m *testing.M) {
@@ -19,6 +20,7 @@ func TestMain(m *testing.M) {
 		panic("An error was not expected when opening a stub database connection")
 	}
 	DB, _ = gorm.Open("mysql", db)
+	mockDB = DB
 	defer db.Close()
 	m.Run()
 }