refactor(ratio): replace maps with RWMap for improved concurrency handling

common/topup-ratio.go

@@ -2,29 +2,37 @@ package common
 
 import (
 	"encoding/json"
+	"sync"
 )
 
-var TopupGroupRatio = map[string]float64{
+var topupGroupRatio = map[string]float64{
 	"default": 1,
 	"vip":     1,
 	"svip":    1,
 }
+var topupGroupRatioMutex sync.RWMutex
 
 func TopupGroupRatio2JSONString() string {
-	jsonBytes, err := json.Marshal(TopupGroupRatio)
+	topupGroupRatioMutex.RLock()
+	defer topupGroupRatioMutex.RUnlock()
+	jsonBytes, err := json.Marshal(topupGroupRatio)
 	if err != nil {
-		SysError("error marshalling model ratio: " + err.Error())
+		SysError("error marshalling topup group ratio: " + err.Error())
 	}
 	return string(jsonBytes)
 }
 
 func UpdateTopupGroupRatioByJSONString(jsonStr string) error {
-	TopupGroupRatio = make(map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &TopupGroupRatio)
+	topupGroupRatioMutex.Lock()
+	defer topupGroupRatioMutex.Unlock()
+	topupGroupRatio = make(map[string]float64)
+	return json.Unmarshal([]byte(jsonStr), &topupGroupRatio)
 }
 
 func GetTopupGroupRatio(name string) float64 {
-	ratio, ok := TopupGroupRatio[name]
+	topupGroupRatioMutex.RLock()
+	defer topupGroupRatioMutex.RUnlock()
+	ratio, ok := topupGroupRatio[name]
 	if !ok {
 		SysError("topup group ratio not found: " + name)
 		return 1

setting/ratio_setting/cache_ratio.go

 package ratio_setting
 
 import (
-	"encoding/json"
-	"sync"
-
-	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/types"
 )
 
 var defaultCacheRatio = map[string]float64{
@@ -98,70 +95,37 @@ var defaultCreateCacheRatio = map[string]float64{
 
 //var defaultCreateCacheRatio = map[string]float64{}
 
-var cacheRatioMap map[string]float64
-var cacheRatioMapMutex sync.RWMutex
-
-var createCacheRatioMap map[string]float64
-var createCacheRatioMapMutex sync.RWMutex
+var cacheRatioMap = types.NewRWMap[string, float64]()
+var createCacheRatioMap = types.NewRWMap[string, float64]()
 
-// GetCacheRatioMap returns the cache ratio map
+// GetCacheRatioMap returns a copy of the cache ratio map
 func GetCacheRatioMap() map[string]float64 {
-	cacheRatioMapMutex.RLock()
-	defer cacheRatioMapMutex.RUnlock()
-	return cacheRatioMap
+	return cacheRatioMap.ReadAll()
 }
 
 // CacheRatio2JSONString converts the cache ratio map to a JSON string
 func CacheRatio2JSONString() string {
-	cacheRatioMapMutex.RLock()
-	defer cacheRatioMapMutex.RUnlock()
-	jsonBytes, err := json.Marshal(cacheRatioMap)
-	if err != nil {
-		common.SysLog("error marshalling cache ratio: " + err.Error())
-	}
-	return string(jsonBytes)
+	return cacheRatioMap.MarshalJSONString()
 }
 
 // CreateCacheRatio2JSONString converts the create cache ratio map to a JSON string
 func CreateCacheRatio2JSONString() string {
-	createCacheRatioMapMutex.RLock()
-	defer createCacheRatioMapMutex.RUnlock()
-	jsonBytes, err := json.Marshal(createCacheRatioMap)
-	if err != nil {
-		common.SysLog("error marshalling create cache ratio: " + err.Error())
-	}
-	return string(jsonBytes)
+	return createCacheRatioMap.MarshalJSONString()
 }
 
 // UpdateCacheRatioByJSONString updates the cache ratio map from a JSON string
 func UpdateCacheRatioByJSONString(jsonStr string) error {
-	cacheRatioMapMutex.Lock()
-	defer cacheRatioMapMutex.Unlock()
-	cacheRatioMap = make(map[string]float64)
-	err := json.Unmarshal([]byte(jsonStr), &cacheRatioMap)
-	if err == nil {
-		InvalidateExposedDataCache()
-	}
-	return err
+	return types.LoadFromJsonStringWithCallback(cacheRatioMap, jsonStr, InvalidateExposedDataCache)
 }
 
 // UpdateCreateCacheRatioByJSONString updates the create cache ratio map from a JSON string
 func UpdateCreateCacheRatioByJSONString(jsonStr string) error {
-	createCacheRatioMapMutex.Lock()
-	defer createCacheRatioMapMutex.Unlock()
-	createCacheRatioMap = make(map[string]float64)
-	err := json.Unmarshal([]byte(jsonStr), &createCacheRatioMap)
-	if err == nil {
-		InvalidateExposedDataCache()
-	}
-	return err
+	return types.LoadFromJsonStringWithCallback(createCacheRatioMap, jsonStr, InvalidateExposedDataCache)
 }
 
 // GetCacheRatio returns the cache ratio for a model
 func GetCacheRatio(name string) (float64, bool) {
-	cacheRatioMapMutex.RLock()
-	defer cacheRatioMapMutex.RUnlock()
-	ratio, ok := cacheRatioMap[name]
+	ratio, ok := cacheRatioMap.Get(name)
 	if !ok {
 		return 1, false // Default to 1 if not found
 	}
@@ -169,9 +133,7 @@ func GetCacheRatio(name string) (float64, bool) {
 }
 
 func GetCreateCacheRatio(name string) (float64, bool) {
-	createCacheRatioMapMutex.RLock()
-	defer createCacheRatioMapMutex.RUnlock()
-	ratio, ok := createCacheRatioMap[name]
+	ratio, ok := createCacheRatioMap.Get(name)
 	if !ok {
 		return 1.25, false // Default to 1.25 if not found
 	}
@@ -179,21 +141,9 @@ func GetCreateCacheRatio(name string) (float64, bool) {
 }
 
 func GetCacheRatioCopy() map[string]float64 {
-	cacheRatioMapMutex.RLock()
-	defer cacheRatioMapMutex.RUnlock()
-	copyMap := make(map[string]float64, len(cacheRatioMap))
-	for k, v := range cacheRatioMap {
-		copyMap[k] = v
-	}
-	return copyMap
+	return cacheRatioMap.ReadAll()
 }
 
 func GetCreateCacheRatioCopy() map[string]float64 {
-	createCacheRatioMapMutex.RLock()
-	defer createCacheRatioMapMutex.RUnlock()
-	copyMap := make(map[string]float64, len(createCacheRatioMap))
-	for k, v := range createCacheRatioMap {
-		copyMap[k] = v
-	}
-	return copyMap
+	return createCacheRatioMap.ReadAll()
 }

setting/ratio_setting/group_ratio.go

@@ -3,29 +3,27 @@ package ratio_setting
 import (
 	"encoding/json"
 	"errors"
-	"sync"
 
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/setting/config"
 	"github.com/QuantumNous/new-api/types"
 )
 
-var groupRatio = map[string]float64{
+var defaultGroupRatio = map[string]float64{
 	"default": 1,
 	"vip":     1,
 	"svip":    1,
 }
 
-var groupRatioMutex sync.RWMutex
+var groupRatioMap = types.NewRWMap[string, float64]()
 
-var (
-	GroupGroupRatio = map[string]map[string]float64{
-		"vip": {
-			"edit_this": 0.9,
-		},
-	}
-	groupGroupRatioMutex sync.RWMutex
-)
+var defaultGroupGroupRatio = map[string]map[string]float64{
+	"vip": {
+		"edit_this": 0.9,
+	},
+}
+
+var groupGroupRatioMap = types.NewRWMap[string, map[string]float64]()
 
 var defaultGroupSpecialUsableGroup = map[string]map[string]string{
 	"vip": {
@@ -35,9 +33,9 @@ var defaultGroupSpecialUsableGroup = map[string]map[string]string{
 }
 
 type GroupRatioSetting struct {
-	GroupRatio              map[string]float64                      `json:"group_ratio"`
-	GroupGroupRatio         map[string]map[string]float64           `json:"group_group_ratio"`
-	GroupSpecialUsableGroup *types.RWMap[string, map[string]string] `json:"group_special_usable_group"`
+	GroupRatio              *types.RWMap[string, float64]            `json:"group_ratio"`
+	GroupGroupRatio         *types.RWMap[string, map[string]float64] `json:"group_group_ratio"`
+	GroupSpecialUsableGroup *types.RWMap[string, map[string]string]  `json:"group_special_usable_group"`
 }
 
 var groupRatioSetting GroupRatioSetting
@@ -46,10 +44,13 @@ func init() {
 	groupSpecialUsableGroup := types.NewRWMap[string, map[string]string]()
 	groupSpecialUsableGroup.AddAll(defaultGroupSpecialUsableGroup)
 
+	groupRatioMap.AddAll(defaultGroupRatio)
+	groupGroupRatioMap.AddAll(defaultGroupGroupRatio)
+
 	groupRatioSetting = GroupRatioSetting{
 		GroupSpecialUsableGroup: groupSpecialUsableGroup,
-		GroupRatio:              groupRatio,
-		GroupGroupRatio:         GroupGroupRatio,
+		GroupRatio:              groupRatioMap,
+		GroupGroupRatio:         groupGroupRatioMap,
 	}
 
 	config.GlobalConfig.Register("group_ratio_setting", &groupRatioSetting)
@@ -64,48 +65,24 @@ func GetGroupRatioSetting() *GroupRatioSetting {
 }
 
 func GetGroupRatioCopy() map[string]float64 {
-	groupRatioMutex.RLock()
-	defer groupRatioMutex.RUnlock()
-
-	groupRatioCopy := make(map[string]float64)
-	for k, v := range groupRatio {
-		groupRatioCopy[k] = v
-	}
-	return groupRatioCopy
+	return groupRatioMap.ReadAll()
 }
 
 func ContainsGroupRatio(name string) bool {
-	groupRatioMutex.RLock()
-	defer groupRatioMutex.RUnlock()
-
-	_, ok := groupRatio[name]
+	_, ok := groupRatioMap.Get(name)
 	return ok
 }
 
 func GroupRatio2JSONString() string {
-	groupRatioMutex.RLock()
-	defer groupRatioMutex.RUnlock()
-
-	jsonBytes, err := json.Marshal(groupRatio)
-	if err != nil {
-		common.SysLog("error marshalling model ratio: " + err.Error())
-	}
-	return string(jsonBytes)
+	return groupRatioMap.MarshalJSONString()
 }
 
 func UpdateGroupRatioByJSONString(jsonStr string) error {
-	groupRatioMutex.Lock()
-	defer groupRatioMutex.Unlock()
-
-	groupRatio = make(map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &groupRatio)
+	return types.LoadFromJsonString(groupRatioMap, jsonStr)
 }
 
 func GetGroupRatio(name string) float64 {
-	groupRatioMutex.RLock()
-	defer groupRatioMutex.RUnlock()
-
-	ratio, ok := groupRatio[name]
+	ratio, ok := groupRatioMap.Get(name)
 	if !ok {
 		common.SysLog("group ratio not found: " + name)
 		return 1
@@ -114,10 +91,7 @@ func GetGroupRatio(name string) float64 {
 }
 
 func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) {
-	groupGroupRatioMutex.RLock()
-	defer groupGroupRatioMutex.RUnlock()
-
-	gp, ok := GroupGroupRatio[userGroup]
+	gp, ok := groupGroupRatioMap.Get(userGroup)
 	if !ok {
 		return -1, false
 	}
@@ -129,22 +103,11 @@ func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) {
 }
 
 func GroupGroupRatio2JSONString() string {
-	groupGroupRatioMutex.RLock()
-	defer groupGroupRatioMutex.RUnlock()
-
-	jsonBytes, err := json.Marshal(GroupGroupRatio)
-	if err != nil {
-		common.SysLog("error marshalling group-group ratio: " + err.Error())
-	}
-	return string(jsonBytes)
+	return groupGroupRatioMap.MarshalJSONString()
 }
 
 func UpdateGroupGroupRatioByJSONString(jsonStr string) error {
-	groupGroupRatioMutex.Lock()
-	defer groupGroupRatioMutex.Unlock()
-
-	GroupGroupRatio = make(map[string]map[string]float64)
-	return json.Unmarshal([]byte(jsonStr), &GroupGroupRatio)
+	return types.LoadFromJsonString(groupGroupRatioMap, jsonStr)
 }
 
 func CheckGroupRatio(jsonStr string) error {

types/rw_map.go

@@ -80,3 +80,24 @@ func LoadFromJsonString[K comparable, V any](m *RWMap[K, V], jsonStr string) err
 	m.data = make(map[K]V)
 	return common.Unmarshal([]byte(jsonStr), &m.data)
 }
+
+// LoadFromJsonStringWithCallback loads a JSON string into the RWMap and calls the callback on success.
+func LoadFromJsonStringWithCallback[K comparable, V any](m *RWMap[K, V], jsonStr string, onSuccess func()) error {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.data = make(map[K]V)
+	err := common.Unmarshal([]byte(jsonStr), &m.data)
+	if err == nil && onSuccess != nil {
+		onSuccess()
+	}
+	return err
+}
+
+// MarshalJSONString returns the JSON string representation of the RWMap.
+func (m *RWMap[K, V]) MarshalJSONString() string {
+	bytes, err := m.MarshalJSON()
+	if err != nil {
+		return "{}"
+	}
+	return string(bytes)
+}