Skip to content

Commit 44c5fac

Browse files
Calcium-IonCalcium-Ion
committed
refactor(ratio): replace maps with RWMap for improved concurrency handling
1 parent 7a146a1 commit 44c5fac

File tree

5 files changed

+115
-360
lines changed

5 files changed

+115
-360
lines changed

‎common/topup-ratio.go‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,29 +2,37 @@ package common
22

33
import (
44
"encoding/json"
5+
"sync"
56
)
67

7-
var TopupGroupRatio = map[string]float64{
8+
var topupGroupRatio = map[string]float64{
89
"default": 1,
910
"vip": 1,
1011
"svip": 1,
1112
}
13+
var topupGroupRatioMutex sync.RWMutex
1214

1315
func TopupGroupRatio2JSONString() string {
14-
jsonBytes, err := json.Marshal(TopupGroupRatio)
16+
topupGroupRatioMutex.RLock()
17+
defer topupGroupRatioMutex.RUnlock()
18+
jsonBytes, err := json.Marshal(topupGroupRatio)
1519
if err != nil {
16-
SysError("error marshalling model ratio: " + err.Error())
20+
SysError("error marshalling topup group ratio: " + err.Error())
1721
}
1822
return string(jsonBytes)
1923
}
2024

2125
func UpdateTopupGroupRatioByJSONString(jsonStr string) error {
22-
TopupGroupRatio = make(map[string]float64)
23-
return json.Unmarshal([]byte(jsonStr), &TopupGroupRatio)
26+
topupGroupRatioMutex.Lock()
27+
defer topupGroupRatioMutex.Unlock()
28+
topupGroupRatio = make(map[string]float64)
29+
return json.Unmarshal([]byte(jsonStr), &topupGroupRatio)
2430
}
2531

2632
func GetTopupGroupRatio(name string) float64 {
27-
ratio, ok := TopupGroupRatio[name]
33+
topupGroupRatioMutex.RLock()
34+
defer topupGroupRatioMutex.RUnlock()
35+
ratio, ok := topupGroupRatio[name]
2836
if !ok {
2937
SysError("topup group ratio not found: " + name)
3038
return 1

‎setting/ratio_setting/cache_ratio.go‎

Lines changed: 13 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
package ratio_setting
22

33
import (
4-
"encoding/json"
5-
"sync"
6-
7-
"github.com/QuantumNous/new-api/common"
4+
"github.com/QuantumNous/new-api/types"
85
)
96

107
var defaultCacheRatio = map[string]float64{
@@ -98,102 +95,55 @@ var defaultCreateCacheRatio = map[string]float64{
9895

9996
//var defaultCreateCacheRatio = map[string]float64{}
10097

101-
var cacheRatioMap map[string]float64
102-
var cacheRatioMapMutex sync.RWMutex
103-
104-
var createCacheRatioMap map[string]float64
105-
var createCacheRatioMapMutex sync.RWMutex
98+
var cacheRatioMap = types.NewRWMap[string, float64]()
99+
var createCacheRatioMap = types.NewRWMap[string, float64]()
106100

107-
// GetCacheRatioMap returns the cache ratio map
101+
// GetCacheRatioMap returns a copy of the cache ratio map
108102
func GetCacheRatioMap() map[string]float64 {
109-
cacheRatioMapMutex.RLock()
110-
defer cacheRatioMapMutex.RUnlock()
111-
return cacheRatioMap
103+
return cacheRatioMap.ReadAll()
112104
}
113105

114106
// CacheRatio2JSONString converts the cache ratio map to a JSON string
115107
func CacheRatio2JSONString() string {
116-
cacheRatioMapMutex.RLock()
117-
defer cacheRatioMapMutex.RUnlock()
118-
jsonBytes, err := json.Marshal(cacheRatioMap)
119-
if err != nil {
120-
common.SysLog("error marshalling cache ratio: " + err.Error())
121-
}
122-
return string(jsonBytes)
108+
return cacheRatioMap.MarshalJSONString()
123109
}
124110

125111
// CreateCacheRatio2JSONString converts the create cache ratio map to a JSON string
126112
func CreateCacheRatio2JSONString() string {
127-
createCacheRatioMapMutex.RLock()
128-
defer createCacheRatioMapMutex.RUnlock()
129-
jsonBytes, err := json.Marshal(createCacheRatioMap)
130-
if err != nil {
131-
common.SysLog("error marshalling create cache ratio: " + err.Error())
132-
}
133-
return string(jsonBytes)
113+
return createCacheRatioMap.MarshalJSONString()
134114
}
135115

136116
// UpdateCacheRatioByJSONString updates the cache ratio map from a JSON string
137117
func UpdateCacheRatioByJSONString(jsonStr string) error {
138-
cacheRatioMapMutex.Lock()
139-
defer cacheRatioMapMutex.Unlock()
140-
cacheRatioMap = make(map[string]float64)
141-
err := json.Unmarshal([]byte(jsonStr), &cacheRatioMap)
142-
if err == nil {
143-
InvalidateExposedDataCache()
144-
}
145-
return err
118+
return types.LoadFromJsonStringWithCallback(cacheRatioMap, jsonStr, InvalidateExposedDataCache)
146119
}
147120

148121
// UpdateCreateCacheRatioByJSONString updates the create cache ratio map from a JSON string
149122
func UpdateCreateCacheRatioByJSONString(jsonStr string) error {
150-
createCacheRatioMapMutex.Lock()
151-
defer createCacheRatioMapMutex.Unlock()
152-
createCacheRatioMap = make(map[string]float64)
153-
err := json.Unmarshal([]byte(jsonStr), &createCacheRatioMap)
154-
if err == nil {
155-
InvalidateExposedDataCache()
156-
}
157-
return err
123+
return types.LoadFromJsonStringWithCallback(createCacheRatioMap, jsonStr, InvalidateExposedDataCache)
158124
}
159125

160126
// GetCacheRatio returns the cache ratio for a model
161127
func GetCacheRatio(name string) (float64, bool) {
162-
cacheRatioMapMutex.RLock()
163-
defer cacheRatioMapMutex.RUnlock()
164-
ratio, ok := cacheRatioMap[name]
128+
ratio, ok := cacheRatioMap.Get(name)
165129
if !ok {
166130
return 1, false // Default to 1 if not found
167131
}
168132
return ratio, true
169133
}
170134

171135
func GetCreateCacheRatio(name string) (float64, bool) {
172-
createCacheRatioMapMutex.RLock()
173-
defer createCacheRatioMapMutex.RUnlock()
174-
ratio, ok := createCacheRatioMap[name]
136+
ratio, ok := createCacheRatioMap.Get(name)
175137
if !ok {
176138
return 1.25, false // Default to 1.25 if not found
177139
}
178140
return ratio, true
179141
}
180142

181143
func GetCacheRatioCopy() map[string]float64 {
182-
cacheRatioMapMutex.RLock()
183-
defer cacheRatioMapMutex.RUnlock()
184-
copyMap := make(map[string]float64, len(cacheRatioMap))
185-
for k, v := range cacheRatioMap {
186-
copyMap[k] = v
187-
}
188-
return copyMap
144+
return cacheRatioMap.ReadAll()
189145
}
190146

191147
func GetCreateCacheRatioCopy() map[string]float64 {
192-
createCacheRatioMapMutex.RLock()
193-
defer createCacheRatioMapMutex.RUnlock()
194-
copyMap := make(map[string]float64, len(createCacheRatioMap))
195-
for k, v := range createCacheRatioMap {
196-
copyMap[k] = v
197-
}
198-
return copyMap
148+
return createCacheRatioMap.ReadAll()
199149
}

‎setting/ratio_setting/group_ratio.go‎

Lines changed: 25 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -3,29 +3,27 @@ package ratio_setting
33
import (
44
"encoding/json"
55
"errors"
6-
"sync"
76

87
"github.com/QuantumNous/new-api/common"
98
"github.com/QuantumNous/new-api/setting/config"
109
"github.com/QuantumNous/new-api/types"
1110
)
1211

13-
var groupRatio = map[string]float64{
12+
var defaultGroupRatio = map[string]float64{
1413
"default": 1,
1514
"vip": 1,
1615
"svip": 1,
1716
}
1817

19-
var groupRatioMutex sync.RWMutex
18+
var groupRatioMap = types.NewRWMap[string, float64]()
2019

21-
var (
22-
GroupGroupRatio = map[string]map[string]float64{
23-
"vip": {
24-
"edit_this": 0.9,
25-
},
26-
}
27-
groupGroupRatioMutex sync.RWMutex
28-
)
20+
var defaultGroupGroupRatio = map[string]map[string]float64{
21+
"vip": {
22+
"edit_this": 0.9,
23+
},
24+
}
25+
26+
var groupGroupRatioMap = types.NewRWMap[string, map[string]float64]()
2927

3028
var defaultGroupSpecialUsableGroup = map[string]map[string]string{
3129
"vip": {
@@ -35,9 +33,9 @@ var defaultGroupSpecialUsableGroup = map[string]map[string]string{
3533
}
3634

3735
type GroupRatioSetting struct {
38-
GroupRatio map[string]float64 `json:"group_ratio"`
39-
GroupGroupRatio map[string]map[string]float64 `json:"group_group_ratio"`
40-
GroupSpecialUsableGroup *types.RWMap[string, map[string]string] `json:"group_special_usable_group"`
36+
GroupRatio *types.RWMap[string, float64] `json:"group_ratio"`
37+
GroupGroupRatio *types.RWMap[string, map[string]float64] `json:"group_group_ratio"`
38+
GroupSpecialUsableGroup *types.RWMap[string, map[string]string] `json:"group_special_usable_group"`
4139
}
4240

4341
var groupRatioSetting GroupRatioSetting
@@ -46,10 +44,13 @@ func init() {
4644
groupSpecialUsableGroup := types.NewRWMap[string, map[string]string]()
4745
groupSpecialUsableGroup.AddAll(defaultGroupSpecialUsableGroup)
4846

47+
groupRatioMap.AddAll(defaultGroupRatio)
48+
groupGroupRatioMap.AddAll(defaultGroupGroupRatio)
49+
4950
groupRatioSetting = GroupRatioSetting{
5051
GroupSpecialUsableGroup: groupSpecialUsableGroup,
51-
GroupRatio: groupRatio,
52-
GroupGroupRatio: GroupGroupRatio,
52+
GroupRatio: groupRatioMap,
53+
GroupGroupRatio: groupGroupRatioMap,
5354
}
5455

5556
config.GlobalConfig.Register("group_ratio_setting", &groupRatioSetting)
@@ -64,48 +65,24 @@ func GetGroupRatioSetting() *GroupRatioSetting {
6465
}
6566

6667
func GetGroupRatioCopy() map[string]float64 {
67-
groupRatioMutex.RLock()
68-
defer groupRatioMutex.RUnlock()
69-
70-
groupRatioCopy := make(map[string]float64)
71-
for k, v := range groupRatio {
72-
groupRatioCopy[k] = v
73-
}
74-
return groupRatioCopy
68+
return groupRatioMap.ReadAll()
7569
}
7670

7771
func ContainsGroupRatio(name string) bool {
78-
groupRatioMutex.RLock()
79-
defer groupRatioMutex.RUnlock()
80-
81-
_, ok := groupRatio[name]
72+
_, ok := groupRatioMap.Get(name)
8273
return ok
8374
}
8475

8576
func GroupRatio2JSONString() string {
86-
groupRatioMutex.RLock()
87-
defer groupRatioMutex.RUnlock()
88-
89-
jsonBytes, err := json.Marshal(groupRatio)
90-
if err != nil {
91-
common.SysLog("error marshalling model ratio: " + err.Error())
92-
}
93-
return string(jsonBytes)
77+
return groupRatioMap.MarshalJSONString()
9478
}
9579

9680
func UpdateGroupRatioByJSONString(jsonStr string) error {
97-
groupRatioMutex.Lock()
98-
defer groupRatioMutex.Unlock()
99-
100-
groupRatio = make(map[string]float64)
101-
return json.Unmarshal([]byte(jsonStr), &groupRatio)
81+
return types.LoadFromJsonString(groupRatioMap, jsonStr)
10282
}
10383

10484
func GetGroupRatio(name string) float64 {
105-
groupRatioMutex.RLock()
106-
defer groupRatioMutex.RUnlock()
107-
108-
ratio, ok := groupRatio[name]
85+
ratio, ok := groupRatioMap.Get(name)
10986
if !ok {
11087
common.SysLog("group ratio not found: " + name)
11188
return 1
@@ -114,10 +91,7 @@ func GetGroupRatio(name string) float64 {
11491
}
11592

11693
func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) {
117-
groupGroupRatioMutex.RLock()
118-
defer groupGroupRatioMutex.RUnlock()
119-
120-
gp, ok := GroupGroupRatio[userGroup]
94+
gp, ok := groupGroupRatioMap.Get(userGroup)
12195
if !ok {
12296
return -1, false
12397
}
@@ -129,22 +103,11 @@ func GetGroupGroupRatio(userGroup, usingGroup string) (float64, bool) {
129103
}
130104

131105
func GroupGroupRatio2JSONString() string {
132-
groupGroupRatioMutex.RLock()
133-
defer groupGroupRatioMutex.RUnlock()
134-
135-
jsonBytes, err := json.Marshal(GroupGroupRatio)
136-
if err != nil {
137-
common.SysLog("error marshalling group-group ratio: " + err.Error())
138-
}
139-
return string(jsonBytes)
106+
return groupGroupRatioMap.MarshalJSONString()
140107
}
141108

142109
func UpdateGroupGroupRatioByJSONString(jsonStr string) error {
143-
groupGroupRatioMutex.Lock()
144-
defer groupGroupRatioMutex.Unlock()
145-
146-
GroupGroupRatio = make(map[string]map[string]float64)
147-
return json.Unmarshal([]byte(jsonStr), &GroupGroupRatio)
110+
return types.LoadFromJsonString(groupGroupRatioMap, jsonStr)
148111
}
149112

150113
func CheckGroupRatio(jsonStr string) error {

0 commit comments

Comments
 (0)