Skip to content

phsl / new-api

Go to a project
Commit aed8c07c by CaIon
Options

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

parent 41d478da
Showing with 69 additions and 127 deletions
common/topup-ratio.go
...@@ -2,29 +2,37 @@ package common ...@@ -2,29 +2,37 @@ package common
import ( import (
"encoding/json" "encoding/json"
"sync"
) )
var TopupGroupRatio = map[string]float64{ var topupGroupRatio = map[string]float64{
"default": 1, "default": 1,
"vip": 1, "vip": 1,
"svip": 1, "svip": 1,
} }
var topupGroupRatioMutex sync.RWMutex
func TopupGroupRatio2JSONString() string { func TopupGroupRatio2JSONString() string {
jsonBytes, err := json.Marshal(TopupGroupRatio) topupGroupRatioMutex.RLock()
defer topupGroupRatioMutex.RUnlock()
jsonBytes, err := json.Marshal(topupGroupRatio)
if err != nil { if err != nil {
SysError("error marshalling model ratio: " + err.Error()) SysError("error marshalling topup group ratio: " + err.Error())
} }
return string(jsonBytes) return string(jsonBytes)
} }
func UpdateTopupGroupRatioByJSONString(jsonStr string) error { func UpdateTopupGroupRatioByJSONString(jsonStr string) error {
TopupGroupRatio = make(map[string]float64) topupGroupRatioMutex.Lock()
return json.Unmarshal([]byte(jsonStr), &TopupGroupRatio) defer topupGroupRatioMutex.Unlock()
topupGroupRatio = make(map[string]float64)
return json.Unmarshal([]byte(jsonStr), &topupGroupRatio)
} }
func GetTopupGroupRatio(name string) float64 { func GetTopupGroupRatio(name string) float64 {
ratio, ok := TopupGroupRatio[name] topupGroupRatioMutex.RLock()
defer topupGroupRatioMutex.RUnlock()
ratio, ok := topupGroupRatio[name]
if !ok { if !ok {
SysError("topup group ratio not found: " + name) SysError("topup group ratio not found: " + name)
return 1 return 1
......
setting/ratio_setting/cache_ratio.go
package ratio_setting package ratio_setting
import ( import (
"encoding/json" "github.com/QuantumNous/new-api/types"
"sync"
"github.com/QuantumNous/new-api/common"
) )
var defaultCacheRatio = map[string]float64{ var defaultCacheRatio = map[string]float64{
...@@ -98,70 +95,37 @@ var defaultCreateCacheRatio = map[string]float64{ ...@@ -98,70 +95,37 @@ var defaultCreateCacheRatio = map[string]float64{
//var defaultCreateCacheRatio = map[string]float64{} //var defaultCreateCacheRatio = map[string]float64{}
var cacheRatioMap map[string]float64 var cacheRatioMap = types.NewRWMap[string, float64]()
var cacheRatioMapMutex sync.RWMutex var createCacheRatioMap = types.NewRWMap[string, float64]()
var createCacheRatioMap map[string]float64
var createCacheRatioMapMutex sync.RWMutex
// GetCacheRatioMap returns the cache ratio map // GetCacheRatioMap returns a copy of the cache ratio map
func GetCacheRatioMap() map[string]float64 { func GetCacheRatioMap() map[string]float64 {
cacheRatioMapMutex.RLock() return cacheRatioMap.ReadAll()
defer cacheRatioMapMutex.RUnlock()
return cacheRatioMap
} }
// CacheRatio2JSONString converts the cache ratio map to a JSON string // CacheRatio2JSONString converts the cache ratio map to a JSON string
func CacheRatio2JSONString() string { func CacheRatio2JSONString() string {
cacheRatioMapMutex.RLock() return cacheRatioMap.MarshalJSONString()
defer cacheRatioMapMutex.RUnlock()
jsonBytes, err := json.Marshal(cacheRatioMap)
if err != nil {
common.SysLog("error marshalling cache ratio: " + err.Error())
}
return string(jsonBytes)
} }
// CreateCacheRatio2JSONString converts the create cache ratio map to a JSON string // CreateCacheRatio2JSONString converts the create cache ratio map to a JSON string
func CreateCacheRatio2JSONString() string { func CreateCacheRatio2JSONString() string {
createCacheRatioMapMutex.RLock() return createCacheRatioMap.MarshalJSONString()
defer createCacheRatioMapMutex.RUnlock()
jsonBytes, err := json.Marshal(createCacheRatioMap)
if err != nil {
common.SysLog("error marshalling create cache ratio: " + err.Error())
}
return string(jsonBytes)
} }
// UpdateCacheRatioByJSONString updates the cache ratio map from a JSON string // UpdateCacheRatioByJSONString updates the cache ratio map from a JSON string
func UpdateCacheRatioByJSONString(jsonStr string) error { func UpdateCacheRatioByJSONString(jsonStr string) error {
cacheRatioMapMutex.Lock() return types.LoadFromJsonStringWithCallback(cacheRatioMap, jsonStr, InvalidateExposedDataCache)
defer cacheRatioMapMutex.Unlock()
cacheRatioMap = make(map[string]float64)
err := json.Unmarshal([]byte(jsonStr), &cacheRatioMap)
if err == nil {
InvalidateExposedDataCache()
}
return err
} }
// UpdateCreateCacheRatioByJSONString updates the create cache ratio map from a JSON string // UpdateCreateCacheRatioByJSONString updates the create cache ratio map from a JSON string
func UpdateCreateCacheRatioByJSONString(jsonStr string) error { func UpdateCreateCacheRatioByJSONString(jsonStr string) error {
createCacheRatioMapMutex.Lock() return types.LoadFromJsonStringWithCallback(createCacheRatioMap, jsonStr, InvalidateExposedDataCache)
defer createCacheRatioMapMutex.Unlock()
createCacheRatioMap = make(map[string]float64)
err := json.Unmarshal([]byte(jsonStr), &createCacheRatioMap)
if err == nil {
InvalidateExposedDataCache()
}
return err
} }
// GetCacheRatio returns the cache ratio for a model // GetCacheRatio returns the cache ratio for a model
func GetCacheRatio(name string) (float64, bool) { func GetCacheRatio(name string) (float64, bool) {
cacheRatioMapMutex.RLock() ratio, ok := cacheRatioMap.Get(name)
defer cacheRatioMapMutex.RUnlock()
ratio, ok := cacheRatioMap[name]
if !ok { if !ok {
return 1, false // Default to 1 if not found return 1, false // Default to 1 if not found
} }
...@@ -169,9 +133,7 @@ func GetCacheRatio(name string) (float64, bool) { ...@@ -169,9 +133,7 @@ func GetCacheRatio(name string) (float64, bool) {
} }
func GetCreateCacheRatio(name string) (float64, bool) { func GetCreateCacheRatio(name string) (float64, bool) {
createCacheRatioMapMutex.RLock() ratio, ok := createCacheRatioMap.Get(name)
defer createCacheRatioMapMutex.RUnlock()
ratio, ok := createCacheRatioMap[name]
if !ok { if !ok {
return 1.25, false // Default to 1.25 if not found return 1.25, false // Default to 1.25 if not found
} }
...@@ -179,21 +141,9 @@ func GetCreateCacheRatio(name string) (float64, bool) { ...@@ -179,21 +141,9 @@ func GetCreateCacheRatio(name string) (float64, bool) {
} }
func GetCacheRatioCopy() map[string]float64 { func GetCacheRatioCopy() map[string]float64 {
cacheRatioMapMutex.RLock() return cacheRatioMap.ReadAll()
defer cacheRatioMapMutex.RUnlock()
copyMap := make(map[string]float64, len(cacheRatioMap))
for k, v := range cacheRatioMap {
copyMap[k] = v
}
return copyMap
} }
func GetCreateCacheRatioCopy() map[string]float64 { func GetCreateCacheRatioCopy() map[string]float64 {
createCacheRatioMapMutex.RLock() return createCacheRatioMap.ReadAll()
defer createCacheRatioMapMutex.RUnlock()
copyMap := make(map[string]float64, len(createCacheRatioMap))
for k, v := range createCacheRatioMap {
copyMap[k] = v
}
return copyMap
} }
setting/ratio_setting/group_ratio.go
...@@ -3,29 +3,27 @@ package ratio_setting ...@@ -3,29 +3,27 @@ package ratio_setting
import ( import (
"encoding/json" "encoding/json"
"errors" "errors"
"sync"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/setting/config" "github.com/QuantumNous/new-api/setting/config"
"github.com/QuantumNous/new-api/types" "github.com/QuantumNous/new-api/types"
) )
var groupRatio = map[string]float64{ var defaultGroupRatio = map[string]float64{
"default": 1, "default": 1,
"vip": 1, "vip": 1,
"svip": 1, "svip": 1,
} }
var groupRatioMutex sync.RWMutex var groupRatioMap = types.NewRWMap[string, float64]()
var ( var defaultGroupGroupRatio = map[string]map[string]float64{
GroupGroupRatio = map[string]map[string]float64{
"vip": { "vip": {
"edit_this": 0.9, "edit_this": 0.9,
}, },
} }
groupGroupRatioMutex sync.RWMutex
) var groupGroupRatioMap = types.NewRWMap[string, map[string]float64]()
var defaultGroupSpecialUsableGroup = map[string]map[string]string{ var defaultGroupSpecialUsableGroup = map[string]map[string]string{
"vip": { "vip": {
...@@ -35,8 +33,8 @@ var defaultGroupSpecialUsableGroup = map[string]map[string]string{ ...@@ -35,8 +33,8 @@ var defaultGroupSpecialUsableGroup = map[string]map[string]string{
} }
type GroupRatioSetting struct { type GroupRatioSetting struct {
GroupRatio map[string]float64 `json:"group_ratio"` GroupRatio *types.RWMap[string, float64] `json:"group_ratio"`
GroupGroupRatio map[string]map[string]float64 `json:"group_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"` GroupSpecialUsableGroup *types.RWMap[string, map[string]string] `json:"group_special_usable_group"`
} }
...@@ -46,10 +44,13 @@ func init() { ...@@ -46,10 +44,13 @@ func init() {
groupSpecialUsableGroup := types.NewRWMap[string, map[string]string]() groupSpecialUsableGroup := types.NewRWMap[string, map[string]string]()
groupSpecialUsableGroup.AddAll(defaultGroupSpecialUsableGroup) groupSpecialUsableGroup.AddAll(defaultGroupSpecialUsableGroup)
groupRatioMap.AddAll(defaultGroupRatio)
groupGroupRatioMap.AddAll(defaultGroupGroupRatio)
groupRatioSetting = GroupRatioSetting{ groupRatioSetting = GroupRatioSetting{
GroupSpecialUsableGroup: groupSpecialUsableGroup, GroupSpecialUsableGroup: groupSpecialUsableGroup,
GroupRatio: groupRatio, GroupRatio: groupRatioMap,
GroupGroupRatio: GroupGroupRatio, GroupGroupRatio: groupGroupRatioMap,
} }
config.GlobalConfig.Register("group_ratio_setting", &groupRatioSetting) config.GlobalConfig.Register("group_ratio_setting", &groupRatioSetting)
...@@ -64,48 +65,24 @@ func GetGroupRatioSetting() *GroupRatioSetting { ...@@ -64,48 +65,24 @@ func GetGroupRatioSetting() *GroupRatioSetting {
} }
func GetGroupRatioCopy() map[string]float64 { func GetGroupRatioCopy() map[string]float64 {
groupRatioMutex.RLock() return groupRatioMap.ReadAll()
defer groupRatioMutex.RUnlock()
groupRatioCopy := make(map[string]float64)
for k, v := range groupRatio {
groupRatioCopy[k] = v
}
return groupRatioCopy
} }
func ContainsGroupRatio(name string) bool { func ContainsGroupRatio(name string) bool {
groupRatioMutex.RLock() _, ok := groupRatioMap.Get(name)
defer groupRatioMutex.RUnlock()
_, ok := groupRatio[name]
return ok return ok
} }
func GroupRatio2JSONString() string { func GroupRatio2JSONString() string {
groupRatioMutex.RLock() return groupRatioMap.MarshalJSONString()
defer groupRatioMutex.RUnlock()
jsonBytes, err := json.Marshal(groupRatio)
if err != nil {
common.SysLog("error marshalling model ratio: " + err.Error())
}
return string(jsonBytes)
} }
func UpdateGroupRatioByJSONString(jsonStr string) error { func UpdateGroupRatioByJSONString(jsonStr string) error {
groupRatioMutex.Lock() return types.LoadFromJsonString(groupRatioMap, jsonStr)
defer groupRatioMutex.Unlock()
groupRatio = make(map[string]float64)
return json.Unmarshal([]byte(jsonStr), &groupRatio)
} }
func GetGroupRatio(name string) float64 { func GetGroupRatio(name string) float64 {
groupRatioMutex.RLock() ratio, ok := groupRatioMap.Get(name)
defer groupRatioMutex.RUnlock()
ratio, ok := groupRatio[name]
if !ok { if !ok {
common.SysLog("group ratio not found: " + name) common.SysLog("group ratio not found: " + name)
return 1 return 1
...@@ -114,10 +91,7 @@ func GetGroupRatio(name string) float64 { ...@@ -114,10 +91,7 @@ func GetGroupRatio(name string) float64 {
} }
func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) { func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) {
groupGroupRatioMutex.RLock() gp, ok := groupGroupRatioMap.Get(userGroup)
defer groupGroupRatioMutex.RUnlock()
gp, ok := GroupGroupRatio[userGroup]
if !ok { if !ok {
return -1, false return -1, false
} }
...@@ -129,22 +103,11 @@ func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) { ...@@ -129,22 +103,11 @@ func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) {
} }
func GroupGroupRatio2JSONString() string { func GroupGroupRatio2JSONString() string {
groupGroupRatioMutex.RLock() return groupGroupRatioMap.MarshalJSONString()
defer groupGroupRatioMutex.RUnlock()
jsonBytes, err := json.Marshal(GroupGroupRatio)
if err != nil {
common.SysLog("error marshalling group-group ratio: " + err.Error())
}
return string(jsonBytes)
} }
func UpdateGroupGroupRatioByJSONString(jsonStr string) error { func UpdateGroupGroupRatioByJSONString(jsonStr string) error {
groupGroupRatioMutex.Lock() return types.LoadFromJsonString(groupGroupRatioMap, jsonStr)
defer groupGroupRatioMutex.Unlock()
GroupGroupRatio = make(map[string]map[string]float64)
return json.Unmarshal([]byte(jsonStr), &GroupGroupRatio)
} }
func CheckGroupRatio(jsonStr string) error { 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 ...@@ -80,3 +80,24 @@ func LoadFromJsonString[K comparable, V any](m *RWMap[K, V], jsonStr string) err
m.data = make(map[K]V) m.data = make(map[K]V)
return common.Unmarshal([]byte(jsonStr), &m.data) 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)
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment