Skip to content

Commit 3e3b823

Browse files
authored
fix: add read/write locks for ModelRatio and GroupRatio to prevent concurrent map read/write issues (#2067)
1 parent 0780812 commit 3e3b823

File tree

2 files changed

+15
-2
lines changed

2 files changed

+15
-2
lines changed

relay/billing/ratio/group.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@ package ratio
33
import (
44
"encoding/json"
55
"github.com/songquanpeng/one-api/common/logger"
6+
"sync"
67
)
78

9+
var groupRatioLock sync.RWMutex
810
var GroupRatio = map[string]float64{
911
"default": 1,
1012
"vip": 1,
@@ -20,11 +22,15 @@ func GroupRatio2JSONString() string {
2022
}
2123

2224
func UpdateGroupRatioByJSONString(jsonStr string) error {
25+
groupRatioLock.Lock()
26+
defer groupRatioLock.Unlock()
2327
GroupRatio = make(map[string]float64)
2428
return json.Unmarshal([]byte(jsonStr), &GroupRatio)
2529
}
2630

2731
func GetGroupRatio(name string) float64 {
32+
groupRatioLock.RLock()
33+
defer groupRatioLock.RUnlock()
2834
ratio, ok := GroupRatio[name]
2935
if !ok {
3036
logger.SysError("group ratio not found: " + name)

relay/billing/ratio/model.go

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"encoding/json"
55
"fmt"
66
"strings"
7+
"sync"
78

89
"github.com/songquanpeng/one-api/common/logger"
910
)
@@ -15,6 +16,8 @@ const (
1516
RMB = USD / USD2RMB
1617
)
1718

19+
var modelRatioLock sync.RWMutex
20+
1821
// ModelRatio
1922
// https://platform.openai.com/docs/models/model-endpoint-compatibility
2023
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf
@@ -88,11 +91,11 @@ var ModelRatio = map[string]float64{
8891
"claude-2.1": 8.0 / 1000 * USD,
8992
"claude-3-haiku-20240307": 0.25 / 1000 * USD,
9093
"claude-3-5-haiku-20241022": 1.0 / 1000 * USD,
91-
"claude-3-5-haiku-latest": 1.0 / 1000 * USD,
94+
"claude-3-5-haiku-latest": 1.0 / 1000 * USD,
9295
"claude-3-sonnet-20240229": 3.0 / 1000 * USD,
9396
"claude-3-5-sonnet-20240620": 3.0 / 1000 * USD,
9497
"claude-3-5-sonnet-20241022": 3.0 / 1000 * USD,
95-
"claude-3-5-sonnet-latest" : 3.0 / 1000 * USD,
98+
"claude-3-5-sonnet-latest": 3.0 / 1000 * USD,
9699
"claude-3-opus-20240229": 15.0 / 1000 * USD,
97100
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlrk4akp7
98101
"ERNIE-4.0-8K": 0.120 * RMB,
@@ -417,11 +420,15 @@ func ModelRatio2JSONString() string {
417420
}
418421

419422
func UpdateModelRatioByJSONString(jsonStr string) error {
423+
modelRatioLock.Lock()
424+
defer modelRatioLock.Unlock()
420425
ModelRatio = make(map[string]float64)
421426
return json.Unmarshal([]byte(jsonStr), &ModelRatio)
422427
}
423428

424429
func GetModelRatio(name string, channelType int) float64 {
430+
modelRatioLock.RLock()
431+
defer modelRatioLock.RUnlock()
425432
if strings.HasPrefix(name, "qwen-") && strings.HasSuffix(name, "-internet") {
426433
name = strings.TrimSuffix(name, "-internet")
427434
}

0 commit comments

Comments
 (0)