fix: add read/write locks for ModelRatio and GroupRatio to prevent concurrent map read/write issues (#2067)
Some checks are pending
CI / Unit tests (push) Waiting to run
CI / commit_lint (push) Waiting to run

This commit is contained in:
longkeyy 2025-02-09 11:02:45 +08:00 committed by GitHub
parent 07808122a6
commit 3e3b8230ac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 15 additions and 2 deletions

View File

@ -3,8 +3,10 @@ package ratio
import (
"encoding/json"
"github.com/songquanpeng/one-api/common/logger"
"sync"
)
var groupRatioLock sync.RWMutex
var GroupRatio = map[string]float64{
"default": 1,
"vip": 1,
@ -20,11 +22,15 @@ func GroupRatio2JSONString() string {
}
func UpdateGroupRatioByJSONString(jsonStr string) error {
groupRatioLock.Lock()
defer groupRatioLock.Unlock()
GroupRatio = make(map[string]float64)
return json.Unmarshal([]byte(jsonStr), &GroupRatio)
}
func GetGroupRatio(name string) float64 {
groupRatioLock.RLock()
defer groupRatioLock.RUnlock()
ratio, ok := GroupRatio[name]
if !ok {
logger.SysError("group ratio not found: " + name)

View File

@ -4,6 +4,7 @@ import (
"encoding/json"
"fmt"
"strings"
"sync"
"github.com/songquanpeng/one-api/common/logger"
)
@ -15,6 +16,8 @@ const (
RMB = USD / USD2RMB
)
var modelRatioLock sync.RWMutex
// ModelRatio
// https://platform.openai.com/docs/models/model-endpoint-compatibility
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf
@ -88,11 +91,11 @@ var ModelRatio = map[string]float64{
"claude-2.1": 8.0 / 1000 * USD,
"claude-3-haiku-20240307": 0.25 / 1000 * USD,
"claude-3-5-haiku-20241022": 1.0 / 1000 * USD,
"claude-3-5-haiku-latest": 1.0 / 1000 * USD,
"claude-3-5-haiku-latest": 1.0 / 1000 * USD,
"claude-3-sonnet-20240229": 3.0 / 1000 * USD,
"claude-3-5-sonnet-20240620": 3.0 / 1000 * USD,
"claude-3-5-sonnet-20241022": 3.0 / 1000 * USD,
"claude-3-5-sonnet-latest" : 3.0 / 1000 * USD,
"claude-3-5-sonnet-latest": 3.0 / 1000 * USD,
"claude-3-opus-20240229": 15.0 / 1000 * USD,
// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlrk4akp7
"ERNIE-4.0-8K": 0.120 * RMB,
@ -417,11 +420,15 @@ func ModelRatio2JSONString() string {
}
func UpdateModelRatioByJSONString(jsonStr string) error {
modelRatioLock.Lock()
defer modelRatioLock.Unlock()
ModelRatio = make(map[string]float64)
return json.Unmarshal([]byte(jsonStr), &ModelRatio)
}
func GetModelRatio(name string, channelType int) float64 {
modelRatioLock.RLock()
defer modelRatioLock.RUnlock()
if strings.HasPrefix(name, "qwen-") && strings.HasSuffix(name, "-internet") {
name = strings.TrimSuffix(name, "-internet")
}